app.py

# import gradio as gr
# import torch
# import numpy as np
# from PIL import Image
# import os
# import legacy
# import torch_utils

# # 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

# # Load StyleGAN Generator
# with open(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):
#     # 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):
#     if image1 is None or image2 is None:
#         return None
#     selected_layers = list(range(mix_value + 1))
#     return mix_styles(image1, image2, selected_layers)

# # 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)
#     ],
#     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()







# import gradio as gr
# import torch
# import numpy as np
# from PIL import Image
# import os
# import legacy
# import torch_utils
# import requests
# import io
# import warnings

# warnings.filterwarnings("ignore")
# device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# model_path = 'dress_model.pkl'
# with open(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://5a4d-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(inputs=None, outputs=[user_id_state])
#     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()










import gradio as gr
import torch
import numpy as np
from PIL import Image
import os
import legacy
import torch_utils
import requests
import io
import warnings
import gdown

warnings.filterwarnings("ignore")

# -------- CONFIGURATION --------
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Google Drive model setup
file_id = "12_fsSQgUfOCAPQaDtq2QPCLE74qTQEwt"
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)
    proj_file = os.path.join(proj_dir, "projected_w.npz")

    if not os.path.exists(proj_file):
        print(f"Projection for {image_name} not found. Running projector.py...")
        os.makedirs(proj_dir, exist_ok=True)
        subprocess.run([
            "python", "projector.py",
            f"--network={NETWORK_PKL}",
            f"--target={image_path}",
            f"--outdir={proj_dir}"
        ], 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)

    latent_vector_1 = np.load(proj_file1)['w']
    latent_vector_2 = np.load(proj_file2)['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



# Handles style mixing + output buffer
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

# Upload to NestJS backend
def send_to_backend(image_buffer, user_id):
    if not user_id:
        return "❌ user_id not found."
    if image_buffer is None:
        return "⚠️ No image generated. Please mix styles first."

    try:
        # Convert BytesIO to raw bytes before sending
        file_bytes = image_buffer.getvalue()

        files = {'file': ('generated_image.png', file_bytes, 'image/png')}

        # Update with your actual ngrok or server URL
        url = f"  https://68be601de1e4.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 interface
# with gr.Blocks(title="Style Mixing for Clothing Designs") as iface:
#     user_id_state = gr.State()

#     @iface.load(inputs=None, outputs=[user_id_state])
#     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", 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_button = gr.Button("Mix Styles")
#     mix_button.click(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()
import gradio as gr

with gr.Blocks(title="Style Mixing for Clothing Designs") as iface:
    user_id_state = gr.State()
    image_buffer = gr.State()

    # Load user ID from URL
    @iface.load(inputs=None, outputs=[user_id_state])
    def on_load(request: gr.Request):
        user_id = request.query_params.get('user_id', '')
        return user_id

    # Header
    gr.Markdown("## 🎨 Style Mixing for Clothing Designs")
    gr.Markdown("Upload two projected clothing images and blend their styles using the slider below.")

    # Upload Inputs
    with gr.Group():
        with gr.Row():
            image1_input = gr.Image(label="👗 First Clothing Image", type="filepath",height=256, width=256)
            image2_input = gr.Image(label="👗 Second Clothing Image", type="filepath",height=256, width=256)
        mix_slider = gr.Slider(label="🧪 Style Mixing Intensity", minimum=0, maximum=9, step=1, value=5, info="0 = Mostly Left | 9 = Mostly Right")

    # Output & Actions
    with gr.Group():
        with gr.Row():
            output_image = gr.Image(label="🧵 Mixed Clothing Design",height=256, width=256)
        with gr.Row():
            mix_button = gr.Button("✨ Mix Styles")
            save_button = gr.Button("💾 Save to Database")
        save_status = gr.Textbox(label="Status", interactive=False)

    # Functions
    def mix_and_store(image1, image2, mix_value):
        result_image, buffer = style_mixing_interface(image1, image2, mix_value)
        return result_image, buffer

    # Button Logic
    mix_button.click(
        fn=mix_and_store,
        inputs=[image1_input, image2_input, mix_slider],
        outputs=[output_image, image_buffer]
    )

    save_button.click(
        fn=send_to_backend,
        inputs=[image_buffer, user_id_state],
        outputs=[save_status]
    )

iface.launch()