Upload inference.py

inference.py

@@ -43,14 +43,20 @@ 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(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")
+    tokenizer = CLIPTokenizer.from_pretrained(model_id, subfolder="tokenizer")
+    text_encoder = CLIPTextModel.from_pretrained(model_id, subfolder="text_encoder").to(device)
+    feature_extractor = CLIPImageProcessor.from_pretrained(model_id, subfolder="feature_extractor")
     
     # 2) Chargez votre UNet non‑conditionné et wrappez‑le
-    base_unet = UNet2DModel.from_pretrained(f"{model_id}/unet").to(device)
+    base_unet = UNet2DModel.from_pretrained(model_id, subfolder="unet").to(device)
     wrapped_unet = UNetNoCondWrapper(base_unet).to(device)
     
     # 3) Construisez la pipeline manuellement