Update app.py

app.py

@@ -1,8 +1,26 @@
 import torch
-import torchvision.transforms as transforms
 from PIL import Image
 import gradio as gr
+import pickle
+from torchvision import transforms
 from tqdm import tqdm
+
+# Load the fine-tuned model
+def load_model(model_path='fine_tuned_stylegan.pth'):
+    with open('ffhq.pkl', 'rb') as f:
+        data = pickle.load(f)
+    G = data['G_ema']
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    G = G.to(device)
+    
+    # Load the fine-tuned weights
+    G.load_state_dict(torch.load(model_path))
+    G.eval()
+    
+    return G
+
+G = load_model('fine_tuned_stylegan.pth')
+
 def optimize_latent_vector(G, target_image, num_iterations=1000):
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
     target_image = transforms.Resize((G.img_resolution, G.img_resolution))(target_image)
@@ -29,25 +47,20 @@ def optimize_latent_vector(G, target_image, num_iterations=1000):
 def generate_from_upload(uploaded_image):
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
-    # Optimize latent vector for the uploaded image
     optimized_z = optimize_latent_vector(G, uploaded_image)
 
-    # Generate variations
     num_variations = 4
     variation_strength = 0.1
     varied_z = optimized_z + torch.randn((num_variations, G.z_dim), device=device) * variation_strength
 
-    # Generate the variations
     with torch.no_grad():
         imgs = G(varied_z, c=None, truncation_psi=0.7, noise_mode='const')
 
     imgs = (imgs * 127.5 + 128).clamp(0, 255).to(torch.uint8)
     imgs = imgs.permute(0, 2, 3, 1).cpu().numpy()
 
-    # Convert the generated image tensors to PIL Images
     generated_images = [Image.fromarray(img) for img in imgs]
 
-    # Return the images separately
     return generated_images[0], generated_images[1], generated_images[2], generated_images[3]
 
 # Create the Gradio interface
@@ -60,17 +73,3 @@ iface = gr.Interface(
 
 # Launch the Gradio interface
 iface.launch(share=True, debug=True)
-
-# If you want to test it without the Gradio interface:
-"""
-# Load an image explicitly
-image_path = "path/to/your/image.jpg"
-image = Image.open(image_path)
-
-# Call the generate method explicitly
-generated_images = generate_from_upload(image)
-
-# Display the generated images
-for img in generated_images:
-    img.show()
-"""