Update app.py

app.py

@@ -89,9 +89,9 @@ def load_model(model_type):
     if model_type == "UNet":
         checkpoint = torch.load("viton_unet_full_checkpoint.pth", map_location=device)
     elif model_type == "GAN":
-        checkpoint = torch.load("viton_gan_generator_checkpoint.pth", map_location=device)
+        checkpoint = torch.load("viton_gan_full_checkpoint.pth", map_location=device)
     elif model_type == "Diffusion":
-        checkpoint = torch.load("viton_diffusion_generator_checkpoint.pth", map_location=device)
+        checkpoint = torch.load("viton_diffusion_full_checkpoint.pth", map_location=device)
     else:
         raise ValueError("Invalid model type")