Update app.py

app.py

@@ -69,8 +69,8 @@ import torch
 import numpy as np
 from PIL import Image
 import os
-import legacy
-import torch_utils
+import legacy  # your local module
+import torch_utils  # your local module
 import requests
 import io
 import warnings
@@ -78,10 +78,12 @@ import warnings
 warnings.filterwarnings("ignore")
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
+# Load StyleGAN model
 model_path = 'dress_model.pkl'
 with open(model_path, 'rb') as f:
     G = legacy.load_network_pkl(f)['G_ema'].to(device)
 
+# Function to perform style mixing
 def mix_styles(image1_path, image2_path, styles_to_mix):
     image1_name = os.path.splitext(os.path.basename(image1_path))[0]
     image2_name = os.path.splitext(os.path.basename(image2_path))[0]
@@ -102,6 +104,7 @@ def mix_styles(image1_path, image2_path, styles_to_mix):
     mixed_image = Image.fromarray(image[0], 'RGB')
     return mixed_image
 
+# Gradio-friendly wrapper
 def style_mixing_interface(image1, image2, mix_value):
     if image1 is None or image2 is None:
         return None, None
@@ -113,6 +116,7 @@ def style_mixing_interface(image1, image2, mix_value):
     buffer.seek(0)
     return mixed_img, buffer
 
+# Upload to backend
 def send_to_backend(image_buffer, user_id):
     if not user_id:
         return "❌ user_id not found in URL."
@@ -120,7 +124,6 @@ def send_to_backend(image_buffer, user_id):
     try:
         files = {'file': ('generated_image.png', image_buffer, 'image/png')}
         url = f"https://361d-103-40-74-78.ngrok-free.app/customisation/upload/{user_id}"
-
         response = requests.post(url, files=files)
 
         if response.status_code == 201:
@@ -131,8 +134,10 @@ def send_to_backend(image_buffer, user_id):
     except Exception as e:
         return f"⚠️ Error: {str(e)}"
 
+# Gradio Interface
 with gr.Blocks(title="Style Mixing for Clothing Design") as iface:
     user_id_state = gr.State()
+    image_buffer = gr.State()
 
     gr.Markdown("## Style Mixing for Clothing Design\nUpload two projected clothing images and mix their styles.")
 
@@ -146,115 +151,30 @@ with gr.Blocks(title="Style Mixing for Clothing Design") as iface:
         output_image = gr.Image(label="Mixed Clothing Design")
         save_button = gr.Button("Download & Save to Database")
 
-    image_buffer = gr.State()
     save_status = gr.Textbox(label="Save Status", interactive=False)
 
+    # Mix and update preview image + store in memory buffer
     def mix_and_store(image1, image2, mix_value):
         result_image, buffer = style_mixing_interface(image1, image2, mix_value)
         return result_image, buffer
 
     mix_slider.change(
-        mix_and_store, 
-        inputs=[image1_input, image2_input, mix_slider], 
+        fn=mix_and_store,
+        inputs=[image1_input, image2_input, mix_slider],
         outputs=[output_image, image_buffer]
     )
 
     save_button.click(
-        send_to_backend, 
-        inputs=[image_buffer, user_id_state], 
+        fn=send_to_backend,
+        inputs=[image_buffer, user_id_state],
         outputs=[save_status]
     )
 
-    # Initialization function that extracts user_id
+    # Extract user_id from URL query
+    @iface.load
     def init_fn(request: gr.Request):
         user_id = request.query_params.get("user_id", "")
         return {user_id_state: user_id}
 
-    iface.load(fn=None, inputs=None, outputs=None, preprocess=False, queue=False, show_progress=False)
-    iface.launch(initialize=init_fn)
-(model_path, 'rb') as f:
-    G = legacy.load_network_pkl(f)['G_ema'].to(device)
-
-def mix_styles(image1_path, image2_path, styles_to_mix):
-    image1_name = os.path.splitext(os.path.basename(image1_path))[0]
-    image2_name = os.path.splitext(os.path.basename(image2_path))[0]
-
-    latent_vector_1 = np.load(os.path.join("projection_results", image1_name, "projected_w.npz"))['w']
-    latent_vector_2 = np.load(os.path.join("projection_results", image2_name, "projected_w.npz"))['w']
-
-    latent_1_tensor = torch.from_numpy(latent_vector_1).to(device)
-    latent_2_tensor = torch.from_numpy(latent_vector_2).to(device)
-
-    mixed_latent = latent_1_tensor.clone()
-    mixed_latent[:, styles_to_mix] = latent_2_tensor[:, styles_to_mix]
-
-    with torch.no_grad():
-        image = G.synthesis(mixed_latent, noise_mode='const')
-
-    image = (image.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255).to(torch.uint8).cpu().numpy()
-    mixed_image = Image.fromarray(image[0], 'RGB')
-    return mixed_image
-
-def style_mixing_interface(image1, image2, mix_value):
-    if image1 is None or image2 is None:
-        return None, None
-    selected_layers = list(range(mix_value + 1))
-    mixed_img = mix_styles(image1, image2, selected_layers)
-
-    buffer = io.BytesIO()
-    mixed_img.save(buffer, format="PNG")
-    buffer.seek(0)
-    return mixed_img, buffer
-
-def send_to_backend(image_buffer, user_id):
-    if not user_id:
-        return "❌ user_id not found."
-
-    try:
-        files = {'file': ('generated_image.png', image_buffer, 'image/png')}
-        url = f"https://361d-103-40-74-78.ngrok-free.app/customisation/upload/{user_id}"
-
-        response = requests.post(url, files=files)
-
-        if response.status_code == 201:
-            return "✅ Image uploaded and saved to database!"
-        else:
-            return f"❌ Upload failed: {response.status_code} - {response.text}"
-
-    except Exception as e:
-        return f"⚠️ Error: {str(e)}"
-
-# --- Gradio UI ---
-with gr.Blocks(title="Style Mixing for Clothing Design") as iface:
-
-    user_id_state = gr.State()
-
-    @iface.load
-    def on_load(request: gr.Request):
-        user_id = request.query_params.get('user_id', '')
-        return user_id
-
-    gr.Markdown("## Style Mixing for Clothing Design\nUpload two projected clothing images and mix their styles.")
-
-    with gr.Row():
-        image1_input = gr.Image(label="First Clothing Image", type="filepath")
-        image2_input = gr.Image(label="Second Clothing Image", type="filepath")
-
-    mix_slider = gr.Slider(label="Style Mixing Strength (Layers 0 to N)", minimum=0, maximum=9, step=1, value=5)
-
-    with gr.Row():
-        output_image = gr.Image(label="Mixed Clothing Design")
-        save_button = gr.Button("Download & Save to Database")
-
-    image_buffer = gr.State()
-    save_status = gr.Textbox(label="Save Status", interactive=False)
-
-    def mix_and_store(image1, image2, mix_value):
-        result_image, buffer = style_mixing_interface(image1, image2, mix_value)
-        return result_image, buffer
-
-    mix_slider.change(mix_and_store, inputs=[image1_input, image2_input, mix_slider], outputs=[output_image, image_buffer])
-    save_button.click(send_to_backend, inputs=[image_buffer, user_id_state], outputs=[save_status])
-
 iface.launch()