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

device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.float16 if device == "cuda" else torch.float32
variant = "fp16" if device == "cuda" else None

pipe = DiffusionPipeline.from_pretrained(
    "stabilityai/sdxl-turbo",
    torch_dtype=dtype,
    variant=variant
).to(device)

def infer(color_prompt, phone_type_prompt, design_prompt):
    prompt = (
        f"A single vertical {color_prompt} colored {phone_type_prompt} back cover featuring a bold {design_prompt} design on the front, hanging on the plain wall. The soft light and shadows, creating a striking contrast against the minimal background, evoking modern sophistication."
    )
    image = pipe(prompt).images[0]
    message = "Design generated successfully!"
    return image, message

def save_design(image):
    if image is None:
        return "No image to save. Please generate a design first."
    file_path = "saved_design.png"
    image.save(file_path)
    return f"Design saved as {file_path}!"

with gr.Blocks() as interface:
    gr.Markdown("# **AI Phone Cover Designer**")
    with gr.Row():
        with gr.Column(scale=1):
            color_prompt = gr.Textbox(label="Color")
            phone_type_prompt = gr.Textbox(label="Mobile Type")
            design_prompt = gr.Textbox(label="Design Details")
            generate_button = gr.Button("Generate Design")
            save_button = gr.Button("Save Design")
        with gr.Column(scale=1):
            output_image = gr.Image(label="Generated Design")
            output_message = gr.Textbox(label="Status", interactive=False)

    generate_button.click(
        infer,
        inputs=[color_prompt, phone_type_prompt, design_prompt],
        outputs=[output_image, output_message],
    )
    save_button.click(
        save_design,
        inputs=[output_image],
        outputs=output_message,
    )

interface.launch(debug=True)