Update app.py

app.py

@@ -65,7 +65,6 @@
 
 
 
-
 import gradio as gr
 import torch
 import numpy as np
@@ -73,94 +72,107 @@ from PIL import Image
 import os
 import legacy
 import torch_utils
+import jwt
 import requests
-import io
-import base64
-
-# Load the pre-trained StyleGAN model
+import tempfile
+from fastapi import Request
+from gradio.routes import app as fastapi_app
+from starlette.requests import Request as StarletteRequest
+from starlette.middleware.cors import CORSMiddleware
+
+# Allow frontend access
+fastapi_app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],  # You can restrict to your frontend domain
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# Load 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)
 
+# Helper to decode JWT from cookie
+def get_user_id_from_cookie(cookie_str):
+    try:
+        if 'access_token=' in cookie_str:
+            token = cookie_str.split('access_token=')[1].split(';')[0]
+            decoded = jwt.decode(token, 'your_jwt_secret', algorithms=['HS256'])
+            return decoded.get('user_id')
+    except Exception as e:
+        print("JWT decode error:", e)
+    return None
+
+# Style mixing
 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
 
-def style_mixing_interface(image1, image2, mix_value):
+# Save image to backend
+def upload_to_backend(img: Image.Image, user_id: str):
+    with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as tmp_file:
+        img.save(tmp_file.name)
+        with open(tmp_file.name, "rb") as f:
+            response = requests.post(
+                f"http://localhost:3000/customisation/upload/{user_id}",
+                files={"file": f}
+            )
+    os.remove(tmp_file.name)
+    return response.json() if response.ok else {"error": "Upload failed", "details": response.text}
+
+# Main Gradio function
+def style_mixing_interface(image1, image2, mix_value, request: StarletteRequest = None):
     if image1 is None or image2 is None:
-        return None, None
+        return None
+
+    # Get user_id from JWT cookie
+    cookie_header = request.headers.get('cookie', '')
+    user_id = get_user_id_from_cookie(cookie_header)
+    if not user_id:
+        return "❌ Invalid or missing JWT. Please log in again."
 
     selected_layers = list(range(mix_value + 1))
     mixed_img = mix_styles(image1, image2, selected_layers)
 
-    # Convert to base64
-    buffer = io.BytesIO()
-    mixed_img.save(buffer, format="PNG")
-    img_bytes = buffer.getvalue()
-    img_base64 = base64.b64encode(img_bytes).decode("utf-8")
+    # Upload image to backend
+    upload_response = upload_to_backend(mixed_img, user_id)
 
-    return mixed_img, img_base64
-
-def send_to_backend(base64_img):
-    try:
-        response = requests.post(
-            "http://localhost:3000/customisation/save",  # Change if using different port/route
-            json={"image": base64_img},
-            timeout=10
-        )
-        if response.status_code == 200:
-            return "✅ Saved to database!"
-        else:
-            return f"❌ Failed to save: {response.status_code} - {response.text}"
-    except Exception as e:
-        return f"⚠️ Error: {str(e)}"
+    print("Upload response:", upload_response)
+    return mixed_img
 
 # Gradio UI
-with gr.Blocks(title="Style Mixing for Clothing Design") as iface:
-    gr.Markdown("## Style Mixing for Clothing Design\nUpload two projected clothing images and select how many early layers to mix.")
-    
-    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")
-    base64_output = gr.Textbox(visible=False)
-    
-    download_button = gr.Button("Download & Save to Database")
-    save_status = gr.Textbox(label="Save Status", interactive=False)
-
-    def mix_and_return(image1, image2, mix_value):
-        return style_mixing_interface(image1, image2, mix_value)
-
-    mix_slider.change(mix_and_return, inputs=[image1_input, image2_input, mix_slider], outputs=[output_image, base64_output])
-
-    download_button.click(fn=send_to_backend, inputs=[base64_output], outputs=[save_status])
+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)
+    ],
+    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."
+)
 
 iface.launch()
+