Update app.py

app.py

@@ -64,7 +64,6 @@
 # iface.launch()
 
 
-
 import gradio as gr
 import torch
 import numpy as np
@@ -74,91 +73,71 @@ import legacy
 import torch_utils
 import jwt
 import requests
-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
+# Load the pre-trained StyleGAN model
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-model_path = 'dress_model.pkl'
+model_path = 'dress_model.pkl'  # Place your .pkl in the same directory or update path
 
+# 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
+# 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
 
-# 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):
+# Function to handle style mixing via Gradio UI
+def style_mixing_interface(image1, image2, mix_value, cookie):
     if image1 is None or image2 is 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."
+    # 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."
 
     selected_layers = list(range(mix_value + 1))
-    mixed_img = mix_styles(image1, image2, selected_layers)
+    mixed_image = mix_styles(image1, image2, selected_layers)
 
-    # Upload image to backend
-    upload_response = upload_to_backend(mixed_img, 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)
 
-    print("Upload response:", upload_response)
-    return mixed_img
+        if response.status_code == 200:
+            return "Image uploaded successfully!"
+        else:
+            return f"Failed to upload image: {response.text}"
+    else:
+        return "User ID not found in token."
 
 # Gradio UI
 iface = gr.Interface(
@@ -166,13 +145,16 @@ iface = gr.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.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."
+    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
 iface.launch()
 
+