import gradio as gr
import torch
from diffusers import StableDiffusionPipeline

MODEL_ID = "runwayml/stable-diffusion-v1-5"

pipe = StableDiffusionPipeline.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
)

if torch.cuda.is_available():
    pipe = pipe.to("cuda")

STYLES = {
    "None": "",
    "Cinematic": "cinematic lighting, dramatic, ultra detailed",
    "Anime": "anime style, vibrant colors, sharp lines",
    "Oil Painting": "oil painting, textured brush strokes",
    "Cyberpunk": "cyberpunk, neon lights, futuristic",
    "Fantasy Art": "fantasy art, magical, epic composition"
}

def generate(prompt, style, steps, guidance):
    full_prompt = prompt
    if STYLES[style]:
        full_prompt = f"{prompt}, {STYLES[style]}"

    image = pipe(
        prompt=full_prompt,
        num_inference_steps=steps,
        guidance_scale=guidance
    ).images[0]

    return image

with gr.Blocks(theme=gr.themes.Soft(primary_hue="purple")) as demo:
    gr.Markdown(
        """
        # ✨ Instruct Image
        Create images from text instructions using AI.
        """
    )

    with gr.Row():
        with gr.Column(scale=1):
            prompt = gr.Textbox(
                label="Instruction",
                placeholder="A futuristic city at night with glowing neon lights",
                lines=3
            )

            style = gr.Dropdown(
                choices=list(STYLES.keys()),
                value="None",
                label="Style"
            )

            steps = gr.Slider(15, 50, value=30, step=1, label="Steps")
            guidance = gr.Slider(1, 15, value=7.5, step=0.5, label="Guidance")

            generate_btn = gr.Button("✨ Generate Image", variant="primary")

        with gr.Column(scale=1):
            output = gr.Image(label="Result", height=512)

    generate_btn.click(
        fn=generate,
        inputs=[prompt, style, steps, guidance],
        outputs=output
    )

    gr.Markdown(
        """
        **Model:** Stable Diffusion v1.5  
        **Note:** This demo is for creative and educational use.
        """
    )

demo.launch()