# Install Gradio and dependencies

# --- Imports ---
from diffusers import StableDiffusionPipeline
import torch
import gradio as gr
import uuid

# --- Configuration ---
class CGF:
    device = "cuda" if torch.cuda.is_available() else "cpu"
    seed = 42
    generator = torch.Generator(device=device).manual_seed(seed)
    image_gen_model_id = "my_text_to_image_saved_sd_model"
    image_gen_steps = 50
    image_gen_guidance_scale = 10
    image_gen_size = (768, 768)

# --- Load Model ---
pipe = StableDiffusionPipeline.from_pretrained(
    CGF.image_gen_model_id,
    torch_dtype=torch.float16,
    use_safetensors=True,
    revision="fp16",
).to(CGF.device)

# --- Generation Function ---
def generate_image(prompt):
    result = pipe(
        prompt,
        num_inference_steps=CGF.image_gen_steps,
        guidance_scale=CGF.image_gen_guidance_scale,
        generator=CGF.generator,
    ).images[0]

    filename = f"{prompt.replace(' ', '_')}_{uuid.uuid4().hex[:6]}.png"
    result.save(filename)
    return result, filename

# --- Gradio UI ---
examples = [
    "A majestic mountain landscape during sunrise",
    "A futuristic cyberpunk city at night with neon lights",
    "A photorealistic portrait of a happy Indian woman in a green saree",
    "A fantasy dragon flying over a castle at sunset",
    "A cat astronaut floating in space",
]

with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("## 🎨 Text-to-Image Generator with Stable Diffusion 2.1")
    gr.Markdown("Type your prompt or select an example below to generate high-quality images.")

    with gr.Row():
        prompt = gr.Textbox(label="Enter your prompt", placeholder="e.g. A girl sitting on a hill in a summer breeze 🌼")
        generate_btn = gr.Button("🚀 Generate Image")

    with gr.Row():
        with gr.Column():
            gallery = gr.Image(label="Generated Image", type="filepath", interactive=False)
            download = gr.File(label="Download Image")

        with gr.Column():
            gr.Examples(examples=examples, inputs=prompt)

    generate_btn.click(fn=generate_image, inputs=prompt, outputs=[gallery, download])
# Launch Gradio app with a public link

# --- Launch App ---
demo.launch(share=True)