Update app.py

app.py

@@ -6,6 +6,7 @@ import numpy as np
 from torchvision import transforms
 import cv2
 from transformers import AutoImageProcessor, SegformerForSemanticSegmentation
+from improved_viton import ImprovedUNetGenerator, load_checkpoint
 
 # ----------------- Device -----------------
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -84,21 +85,53 @@ def generate_agnostic(image: Image.Image, segmentation):
         agnostic_np[segmentation_resized == label] = [128, 128, 128]
     return Image.fromarray(agnostic_np)
 
-def load_model(model_type):
-    model = UNetGenerator().to(device)
+# def load_model(model_type):
+#     model = UNetGenerator().to(device)
+#     if model_type == "UNet":
+#         checkpoint = torch.load("viton_unet_full_checkpoint.pth", map_location=device)
+#     elif model_type == "GAN":
+#         checkpoint = torch.load("viton_gan_full_checkpoint.pth", map_location=device)
+#     elif model_type == "Diffusion":
+#         checkpoint = torch.load("viton_diffusion_full_checkpoint.pth", map_location=device)
+#     else:
+#         raise ValueError("Invalid model type")
+
+#     model.load_state_dict(checkpoint['model_state_dict'])
+#     model.eval()
+#     return model
+
+import torch
+from improved_viton import ImprovedUNetGenerator
+
+def load_model(model_type, device):
+    # Initialize generator architecture
+    # model = ImprovedUNetGenerator(in_channels=6, out_channels=3).to(device)
+
+    # Load appropriate checkpoint
     if model_type == "UNet":
+        model = UNetGenerator().to(device)
         checkpoint = torch.load("viton_unet_full_checkpoint.pth", map_location=device)
+        state_dict = checkpoint.get('model_G_state_dict') or checkpoint.get('model_state_dict')
     elif model_type == "GAN":
+        model = ImprovedUNetGenerator(in_channels=6, out_channels=3).to(device)
         checkpoint = torch.load("viton_gan_full_checkpoint.pth", map_location=device)
+        state_dict = checkpoint.get('model_G_state_dict')
     elif model_type == "Diffusion":
+        model = ImprovedUNetGenerator(in_channels=6, out_channels=3).to(device)
         checkpoint = torch.load("viton_diffusion_full_checkpoint.pth", map_location=device)
+        state_dict = checkpoint.get('model_G_state_dict') or checkpoint.get('model_state_dict')
     else:
-        raise ValueError("Invalid model type")
+        raise ValueError(f"Invalid model type: {model_type}")
 
-    model.load_state_dict(checkpoint['model_state_dict'])
+    # Load the state dict into the model
+    if state_dict is None:
+        raise KeyError(f"No model weights found in the checkpoint for {model_type}")
+    
+    model.load_state_dict(state_dict)
     model.eval()
     return model
 
+
 def generate_tryon_output(agnostic_img, cloth_img, model):
     agnostic_tensor = img_transform(agnostic_img).unsqueeze(0).to(device)
     cloth_tensor = img_transform(cloth_img).unsqueeze(0).to(device)