import gradio as gr
import random
import numpy as np
from diffusers import DiffusionPipeline
import torch

# ডিভাইস সেটআপ
device = "cuda" if torch.cuda.is_available() else "cpu"
model_repo_id = "stabilityai/sdxl-turbo"

# মডেল লোড করুন
pipe = DiffusionPipeline.from_pretrained(
    model_repo_id,
    torch_dtype=torch.float16 if device == "cuda" else torch.float32,
    variant="fp16" if device == "cuda" else None,
    use_safetensors=True
)
pipe = pipe.to(device)

# সেটিংস
MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = 1024
DEFAULT_NUM_IMAGES = 4

def generate_images(
    prompt,
    negative_prompt="",
    seed=0,
    randomize_seed=True,
    width=768,
    height=768,
    guidance_scale=0.0,
    num_inference_steps=2,
    num_images=DEFAULT_NUM_IMAGES,
    progress=gr.Progress(track_tqdm=True)
):
    # র‍্যান্ডম সিড জেনারেট করুন
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    
    # জেনারেটর তৈরি করুন
    generator = torch.Generator(device=device).manual_seed(seed)
    
    # একাধিক ইমেজ জেনারেট করুন
    images = []
    for _ in range(num_images):
        image = pipe(
            prompt=prompt,
            negative_prompt=negative_prompt,
            guidance_scale=guidance_scale,
            num_inference_steps=num_inference_steps,
            width=width,
            height=height,
            generator=generator,
        ).images[0]
        images.append(image)
    
    return images, seed

# উদাহরণ প্রম্পট
examples = [
    "A cyberpunk cityscape at night, neon lights, rain",
    "Majestic waterfall in a tropical forest, cinematic lighting",
    "An astronaut floating in space, earth in background",
    "A futuristic robot playing guitar in a jazz club",
    "A magical forest with glowing mushrooms and fireflies"
]

# UI থিম
theme = gr.themes.Glass(
    primary_hue="purple",
    secondary_hue="blue",
    font=[gr.themes.GoogleFont("Poppins")]
)

# UI তৈরি
with gr.Blocks(theme=theme, css="footer {visibility: hidden}") as demo:
    gr.Markdown("""
    # 🚀 SDXL-Turbo Image Generator
    **Create 4 stunning AI images in seconds!** No API key required.
    """)
    
    with gr.Row():
        with gr.Column(scale=3):
            prompt = gr.Textbox(
                label="Prompt",
                placeholder="Describe your image...",
                lines=2,
                max_lines=3
            )
            negative_prompt = gr.Textbox(
                label="Negative Prompt",
                placeholder="What to exclude from the image...",
                lines=1
            )
        with gr.Column(scale=1):
            generate_btn = gr.Button("✨ Generate 4 Images", variant="primary")
    
    gallery = gr.Gallery(
        label="Generated Images",
        columns=4,
        height="auto",
        preview=True,
        object_fit="contain"
    )
    
    with gr.Accordion("Advanced Settings", open=False):
        with gr.Row():
            seed = gr.Slider(
                label="Seed",
                minimum=0,
                maximum=MAX_SEED,
                step=1,
                value=0
            )
            randomize_seed = gr.Checkbox(
                label="Random Seed",
                value=True
            )
        with gr.Row():
            width = gr.Slider(
                label="Width",
                minimum=256,
                maximum=MAX_IMAGE_SIZE,
                step=32,
                value=768
            )
            height = gr.Slider(
                label="Height",
                minimum=256,
                maximum=MAX_IMAGE_SIZE,
                step=32,
                value=768
            )
        with gr.Row():
            guidance_scale = gr.Slider(
                label="Guidance Scale",
                minimum=0.0,
                maximum=10.0,
                step=0.1,
                value=0.0
            )
            num_inference_steps = gr.Slider(
                label="Inference Steps",
                minimum=1,
                maximum=50,
                step=1,
                value=2
            )
            num_images = gr.Slider(
                label="Number of Images",
                minimum=1,
                maximum=10,
                step=1,
                value=DEFAULT_NUM_IMAGES
            )
    
    gr.Examples(
        examples=examples,
        inputs=prompt,
        label="Example Prompts"
    )
    
    # ইভেন্ট হ্যান্ডলার
    generate_btn.click(
        fn=generate_images,
        inputs=[
            prompt,
            negative_prompt,
            seed,
            randomize_seed,
            width,
            height,
            guidance_scale,
            num_inference_steps,
            num_images
        ],
        outputs=[gallery, seed]
    )
    
    gr.Markdown("""
    <div style="text-align: center; margin-top: 20px; color: #666;">
        Powered by SDXL-Turbo • No API key required • Runs on Hugging Face Spaces
    </div>
    """)

# অ্যাপ্লিকেশন চালু করুন
if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)