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

print("Loading Rookus T1 model...")
pipe = DiffusionPipeline.from_pretrained("stabilityai/sdxl-turbo", torch_dtype=torch.float16)
pipe.to("cuda" if torch.cuda.is_available() else "cpu")
pipe.set_progress_bar_config(disable=True)

def infer(color_prompt, dress_type_prompt, design_prompt):
    # Construct prompt
    prompt = (
        f"A high-quality digital image of a {color_prompt} {dress_type_prompt}, "
        f"featuring a {design_prompt} printed in sharp detail on the {dress_type_prompt}, "
        f"facing front, hanging on a plain wall. "
        f"The fabric has realistic texture, smooth folds, and accurate lighting. "
        f"The design is perfectly aligned, with natural shadows and highlights, "
        f"creating a photorealistic look."
    )

    print("Generating image with prompt:", prompt)

    # Generate image
    result = pipe(prompt, num_inference_steps=4, guidance_scale=0.0)  # SDXL-Turbo is optimized for fast steps
    image = result.images[0]

    return image

# Gradio Interface
iface = gr.Interface(
    fn=infer,
    inputs=[
        gr.Textbox(lines=1, placeholder="Color Prompt"),
        gr.Textbox(lines=1, placeholder="Dress Type Prompt"),
        gr.Textbox(lines=2, placeholder="Design Prompt"),
    ],
    outputs="image",
    title="Make your Brand",
    description="Generation of clothes",
    examples=[["Red", "T-shirt", "Simple design"]]
)

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