Update app.py

app.py

@@ -189,14 +189,15 @@ import torch_utils
 import requests
 import io
 import warnings
+import gdown
 
 warnings.filterwarnings("ignore")
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+# device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
-# Load pretrained StyleGAN model
-model_path = 'https://drive.google.com/file/d/12_fsSQgUfOCAPQaDtq2QPCLE74qTQEwt/view?usp=sharing'
-with open(model_path, 'rb') as f:
-    G = legacy.load_network_pkl(f)['G_ema'].to(device)
+# # Load pretrained StyleGAN model
+# model_path = 'https://drive.google.com/file/d/12_fsSQgUfOCAPQaDtq2QPCLE74qTQEwt/view?usp=sharing'
+# with open(model_path, 'rb') as f:
+#     G = legacy.load_network_pkl(f)['G_ema'].to(device)
 
 # # Style mixing between two projected images
 # def mix_styles(image1_path, image2_path, styles_to_mix):
@@ -219,6 +220,28 @@ with open(model_path, 'rb') as f:
 #     mixed_image = Image.fromarray(image[0], 'RGB')
 #     return mixed_image
 # ----------- UTIL: ENSURE PROJECTION ----------
+
+
+# -------- CONFIGURATION --------
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# Google Drive model setup
+file_id = "12_fsSQgUfOCAPQaDtq2QPCLE74qTQEwt"  # from your shared link
+output_path = "dress_model.pkl"
+
+# Download the model if it's not present
+if not os.path.exists(output_path):
+    print("Downloading StyleGAN2 model from Google Drive...")
+    gdown.download(f"https://drive.google.com/uc?id={file_id}", output_path, quiet=False)
+
+# Load the model
+with open(output_path, 'rb') as f:
+    G = legacy.load_network_pkl(f)['G_ema'].to(device)
+
+# Save model path for projector.py
+NETWORK_PKL = output_path
+
+# -------- ENSURE PROJECTION --------
 def ensure_projection(image_path):
     image_name = os.path.splitext(os.path.basename(image_path))[0]
     proj_dir = os.path.join("projection_results", image_name)
@@ -235,6 +258,7 @@ def ensure_projection(image_path):
         ], check=True)
     return proj_file
 
+# -------- STYLE MIXING --------
 def mix_styles(image1_path, image2_path, styles_to_mix):
     proj_file1 = ensure_projection(image1_path)
     proj_file2 = ensure_projection(image2_path)
@@ -256,6 +280,7 @@ def mix_styles(image1_path, image2_path, styles_to_mix):
     return mixed_image
 
 
+
 # Handles style mixing + output buffer
 def style_mixing_interface(image1, image2, mix_value):
     if image1 is None or image2 is None: