Upload inference.py

inference.py

@@ -42,12 +42,6 @@ class UNetNoCondWrapper(nn.Module):
     
 def inference(model_id,device, img1, img2):
     
-
-    # vae = AutoencoderKL.from_pretrained(f"{model_id}/vae").to(device)
-    # scheduler = DDPMScheduler.from_pretrained(f"{model_id}/scheduler")
-    # tokenizer = CLIPTokenizer.from_pretrained(f"{model_id}/tokenizer")
-    # text_encoder = CLIPTextModel.from_pretrained(f"{model_id}/text_encoder").to(device)
-    # feature_extractor = CLIPImageProcessor.from_pretrained(f"{model_id}/feature_extractor")
     
     vae = AutoencoderKL.from_pretrained(model_id, subfolder="vae").to(device)
     scheduler = DDPMScheduler.from_pretrained(model_id, subfolder="scheduler")
@@ -69,9 +63,10 @@ def inference(model_id,device, img1, img2):
         safety_checker=None,
         feature_extractor=feature_extractor,
     )
-    pipe = pipe.to(torch.float16).to(device)
+    # pipe = pipe.to(torch.float16).to(device)
+    pipe = pipe.to(torch.float32).to(device)
     
-    generator = torch.Generator("cuda").manual_seed(0)
+    generator = torch.Generator("cpu").manual_seed(0)
     
     
     img1 = img1.resize((512, 512))