import gradio as gr
from diffusers import StableDiffusionPipeline
import torch

# ✅ Choose device (GPU if available, else CPU)
device = "cuda" if torch.cuda.is_available() else "cpu"

# ✅ Load the model with correct settings
pipe = StableDiffusionPipeline.from_pretrained(
    "stabilityai/stable-diffusion-2-1",
    torch_dtype=torch.float16 if device == "cuda" else torch.float32,
    revision="fp16" if device == "cuda" else "main"
).to(device)

# ✅ Generate image or placeholder for 3D
def generate_muse(prompt, output_type):
    if output_type == "Image":
        image = pipe(prompt).images[0]
        return image
    else:
        return "⚠️ 3D Model generation coming soon. Please select 'Image'."

# ✅ Gradio UI layout
with gr.Blocks() as demo:
    gr.Markdown("## 🎨 Welcome to **MUSE**")
    gr.Markdown("Turn your architectural, artistic, or design ideas into stunning visual concepts.\n\n🧠 Just describe your **muse**, pick a format, and let the AI do the rest.")

    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(
                label="🖋️ Describe your muse idea",
                placeholder="e.g. a floating bamboo pavilion shaped like a jellyfish with fabric sails",
                lines=3
            )
            output_type = gr.Radio(["Image", "3D Model"], label="📌 Choose output type", value="Image")
            generate_button = gr.Button("✨ Generate Muse")

        with gr.Column():
            image_output = gr.Image(label="🎨 Your Muse Output")

    generate_button.click(fn=generate_muse, inputs=[prompt, output_type], outputs=image_output)

demo.launch()