import torch
import gradio as gr
from diffusers import DiffusionPipeline
from PIL import Image

# Load Stable Diffusion XL base model with LoRA weights (requires PEFT backend)
pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0")
pipe.to("cpu")

# Ensure PEFT is installed and compatible
try:
    from peft import LoraModel
    pipe.load_lora_weights("networks/TShirtDesignRedmondV2-Tshirtdesign-TshirtDesignAF.safetensors")
except ImportError:
    raise ImportError("PEFT is required for loading LoRA weights. Install it using `pip install peft`." )

pipe.unet.set_default_attn_processor()
pipe.vae.set_default_attn_processor()

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()