import torch
import gradio as gr
from diffusers import DiffusionPipeline
from PIL import Image
from io import BytesIO
import os
import requests
import time
from tqdm import tqdm

# Load Stable Diffusion XL base model with LoRA weights
pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0")
pipe.load_lora_weights("networks/TShirtDesignRedmondV2-Tshirtdesign-TshirtDesignAF.safetensors")
pipe = pipe.to("cuda")

def infer(color_prompt, dress_type_prompt, design_prompt):
    prompt = (
        f"A single {color_prompt} colored {dress_type_prompt} featuring a bold {design_prompt} design printed on the {dress_type_prompt},"
        " hanging on a plain wall. The soft light and shadows create a striking contrast against the minimal background, evoking modern sophistication."
    )
    
    print("Generating image locally with prompt:", prompt)
    try:
        image = pipe(prompt).images[0]
        return image
    except Exception as e:
        print("Local generation failed.", str(e))
        return None

# Gradio Interface
iface = gr.Interface(
    fn=infer,
    inputs=[
        gr.Textbox(lines=1, placeholder="Color"),
        gr.Textbox(lines=1, placeholder="Dress Type"),
        gr.Textbox(lines=2, placeholder="Design"),
    ],
    outputs="image",
    title="AI-Generated T-Shirt Designs",
    description="Generate custom t-shirt designs using AI!",
    examples=[["Red", "T-shirt", "Minimalistic logo"]]
)

print("Launching Gradio interface...")
iface.launch()