Update ootd/inference_ootd.py

ootd/inference_ootd.py

@@ -33,7 +33,7 @@ MODEL_PATH = "./checkpoints/ootd"
 class OOTDiffusion:
 
     def __init__(self, gpu_id):
-        self.gpu_id = 'cuda:' + str(gpu_id)
+        # self.gpu_id = 'cuda:' + str(gpu_id)
 
         vae = AutoencoderKL.from_pretrained(
             VAE_PATH,
@@ -64,12 +64,12 @@ class OOTDiffusion:
             use_safetensors=True,
             safety_checker=None,
             requires_safety_checker=False,
-        ).to(self.gpu_id)
+        )#.to(self.gpu_id)
 
         self.pipe.scheduler = UniPCMultistepScheduler.from_config(self.pipe.scheduler.config)
         
         self.auto_processor = AutoProcessor.from_pretrained(VIT_PATH)
-        self.image_encoder = CLIPVisionModelWithProjection.from_pretrained(VIT_PATH).to(self.gpu_id)
+        self.image_encoder = CLIPVisionModelWithProjection.from_pretrained(VIT_PATH)#.to(self.gpu_id)
 
         self.tokenizer = CLIPTokenizer.from_pretrained(
             MODEL_PATH,
@@ -78,7 +78,7 @@ class OOTDiffusion:
         self.text_encoder = CLIPTextModel.from_pretrained(
             MODEL_PATH,
             subfolder="text_encoder",
-        ).to(self.gpu_id)
+        )#.to(self.gpu_id)
 
 
     def tokenize_captions(self, captions, max_length):
@@ -107,14 +107,14 @@ class OOTDiffusion:
         generator = torch.manual_seed(seed)
 
         with torch.no_grad():
-            prompt_image = self.auto_processor(images=image_garm, return_tensors="pt").to(self.gpu_id)
+            prompt_image = self.auto_processor(images=image_garm, return_tensors="pt").to('cuda')
             prompt_image = self.image_encoder(prompt_image.data['pixel_values']).image_embeds
             prompt_image = prompt_image.unsqueeze(1)
             if model_type == 'hd':
-                prompt_embeds = self.text_encoder(self.tokenize_captions([""], 2).to(self.gpu_id))[0]
+                prompt_embeds = self.text_encoder(self.tokenize_captions([""], 2).to('cuda'))[0]
                 prompt_embeds[:, 1:] = prompt_image[:]
             elif model_type == 'dc':
-                prompt_embeds = self.text_encoder(self.tokenize_captions([category], 3).to(self.gpu_id))[0]
+                prompt_embeds = self.text_encoder(self.tokenize_captions([category], 3).to('cuda'))[0]
                 prompt_embeds = torch.cat([prompt_embeds, prompt_image], dim=1)
             else:
                 raise ValueError("model_type must be \'hd\' or \'dc\'!")