Update app.py

app.py

@@ -63,7 +63,6 @@
 
 # iface.launch()
 
-
 import gradio as gr
 import torch
 import numpy as np
@@ -71,90 +70,113 @@ from PIL import Image
 import os
 import legacy
 import torch_utils
-import jwt
 import requests
+import io
+from http.cookies import SimpleCookie
+import jwt  # PyJWT
+import warnings
+
+# Suppress deprecated torch warnings
+warnings.filterwarnings("ignore")
 
-# Load the pre-trained StyleGAN model
+# --- Load the pre-trained StyleGAN model ---
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-model_path = 'dress_model.pkl'  # Place your .pkl in the same directory or update path
+model_path = 'dress_model.pkl'
 
-# Load StyleGAN Generator
 with open(model_path, 'rb') as f:
     G = legacy.load_network_pkl(f)['G_ema'].to(device)
 
-# Function to mix styles of two clothing images
 def mix_styles(image1_path, image2_path, styles_to_mix):
-    # Extract image names (without extensions)
     image1_name = os.path.splitext(os.path.basename(image1_path))[0]
     image2_name = os.path.splitext(os.path.basename(image2_path))[0]
 
-    # Load latent vectors from .npz
     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']
 
-    # Convert to torch tensors
     latent_1_tensor = torch.from_numpy(latent_vector_1).to(device)
     latent_2_tensor = torch.from_numpy(latent_vector_2).to(device)
 
-    # Mix layers
     mixed_latent = latent_1_tensor.clone()
     mixed_latent[:, styles_to_mix] = latent_2_tensor[:, styles_to_mix]
 
-    # Generate image
     with torch.no_grad():
         image = G.synthesis(mixed_latent, noise_mode='const')
 
-    # Convert to image
     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
 
-# Function to handle style mixing via Gradio UI
-def style_mixing_interface(image1, image2, mix_value, cookie):
+# --- Step 1: Mixing and preparing the image ---
+def style_mixing_interface(image1, image2, mix_value):
     if image1 is None or image2 is None:
-        return 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
+
+# --- Step 2: Extract user_id from JWT cookie and upload image ---
+def send_to_backend(image_buffer, request: gr.Request):
+    cookie_header = request.headers.get('cookie', '')
+    cookies = SimpleCookie()
+    cookies.load(cookie_header)
+
+    jwt_token_cookie = cookies.get("access_token")
+    if not jwt_token_cookie:
+        return "❌ JWT token not found in cookies."
+
+    token = jwt_token_cookie.value
 
-    # Extract user_id from the JWT token passed via cookies (assuming JWT token is passed as 'cookie' in the request)
     try:
-        decoded_token = jwt.decode(cookie, options={"verify_exp": False})  # Decode token without verifying expiration
-        user_id = decoded_token.get("user_id", None)
-    except jwt.ExpiredSignatureError:
-        return "Session expired, please log in again."
-    except jwt.InvalidTokenError:
-        return "Invalid token, please log in again."
+        # Decode JWT (no verification assuming it's a trusted local token for this case)
+        payload = jwt.decode(token, options={"verify_signature": False})
+        user_id = payload.get("user_id")
+        if not user_id:
+            return "❌ user_id not found in JWT."
 
-    selected_layers = list(range(mix_value + 1))
-    mixed_image = mix_styles(image1, image2, selected_layers)
+        # Prepare image upload
+        files = {'file': ('generated_image.png', image_buffer, 'image/png')}
+        url = f"http://localhost:3000/customisation/upload/{user_id}"
 
-    # Call backend API to save the image
-    if user_id:
-        upload_url = f"http://localhost:3000/customisation/upload/{user_id}"
-        files = {'file': ('mixed_image.png', mixed_image.tobytes(), 'image/png')}
-        response = requests.post(upload_url, files=files)
+        response = requests.post(url, files=files)
 
-        if response.status_code == 200:
-            return "Image uploaded successfully!"
+        if response.status_code == 201:
+            return "✅ Image uploaded and saved to database!"
         else:
-            return f"Failed to upload image: {response.text}"
-    else:
-        return "User ID not found in token."
-
-# Gradio UI
-iface = gr.Interface(
-    fn=style_mixing_interface,
-    inputs=[
-        gr.Image(label="First Clothing Image", type="filepath"),
-        gr.Image(label="Second Clothing Image", type="filepath"),
-        gr.Slider(label="Style Mixing Strength (Layers 0 to N)", minimum=0, maximum=9, step=1, value=5),
-        gr.Textbox(label="JWT Token (as cookie)", type="text")  # You may pass JWT token here for testing purposes
-    ],
-    outputs=gr.Image(label="Mixed Clothing Design"),
-    live=True,
-    title="Style Mixing for Clothing Design",
-    description="Upload two projected images and choose how many early layers to mix. The resulting image will be saved after mixing."
-)
-
-# Launch the Gradio interface directly
+            return f"❌ Upload failed: {response.status_code} - {response.text}"
+
+    except jwt.DecodeError:
+        return "❌ Invalid JWT token."
+    except Exception as e:
+        return f"⚠️ Error: {str(e)}"
+
+# --- Gradio Interface ---
+with gr.Blocks(title="Style Mixing for Clothing Design") as iface:
+    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)
+
+    output_image = gr.Image(label="Mixed Clothing Design")
+    image_buffer = gr.State()
+
+    save_button = gr.Button("Download & Save to Database")
+    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], outputs=[save_status])
+
 iface.launch()
 
 
+