import gradio as gr
import torch
from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler

# ============================================
# मॉडल लोडिंग (CPU के लिए ऑप्टिमाइज़्ड)
# ============================================

model_id = "OFA-Sys/small-stable-diffusion-v0"

print("मॉडल लोड हो रहा है... कृपया प्रतीक्षा करें।")

# Scheduler सेट करें (कम स्टेप्स में बेहतर क्वालिटी के लिए)
scheduler = DPMSolverMultistepScheduler.from_pretrained(
    model_id, 
    subfolder="scheduler"
)

# Pipeline लोड करें
pipe = StableDiffusionPipeline.from_pretrained(
    model_id,
    torch_dtype=torch.float32,      # CPU के लिए float32 ज़रूरी है
    scheduler=scheduler,
    safety_checker=None,             # मेमोरी बचाने के लिए
    requires_safety_checker=False    # Warnings हटाने के लिए
)

# CPU मेमोरी ऑप्टिमाइज़ेशन
try:
    pipe.enable_attention_slicing()
    pipe.enable_vae_slicing()
    print("मेमोरी ऑप्टिमाइज़ेशन सक्रिय।")
except Exception as e:
    print(f"ऑप्टिमाइज़ेशन में समस्या (नज़रअंदाज़ करें): {e}")

print("मॉडल तैयार है!")

# ============================================
# इमेज जनरेशन फंक्शन
# ============================================

def generate_image(prompt, negative_prompt, steps, guidance_scale, seed):
    """
    प्रॉम्प्ट के आधार पर इमेज जनरेट करता है।
    """
    # प्रॉम्प्ट खाली होने पर डिफ़ॉल्ट
    if not prompt or prompt.strip() == "":
        prompt = "a beautiful landscape, highly detailed, 4k"
    
    # सीड सेट करें
    if seed < 0:
        generator = None
    else:
        generator = torch.Generator().manual_seed(int(seed))
    
    try:
        # इमेज जनरेट करें
        result = pipe(
            prompt=prompt,
            negative_prompt=negative_prompt if negative_prompt else None,
            num_inference_steps=int(steps),
            guidance_scale=float(guidance_scale),
            generator=generator,
            width=512,
            height=512
        )
        return result.images[0]
    
    except Exception as e:
        print(f"इमेज जनरेशन में एरर: {e}")
        return None

# ============================================
# Gradio वेब इंटरफेस
# ============================================

with gr.Blocks(title="Small Stable Diffusion - CPU") as demo:
    
    # हेडर
    gr.Markdown("""
    # 🎨 Small Stable Diffusion V0 (CPU डेमो)
    
    यह **OFA-Sys/small-stable-diffusion-v0** मॉडल का उपयोग करता है।  
    यह मॉडल स्टैंडर्ड स्टेबल डिफ्यूजन से लगभग **आधे साइज़** का है, इसलिए CPU पर भी चल सकता है।
    
    ⏱️ **ध्यान दें**: CPU पर एक इमेज जनरेट होने में **35-45 सेकंड** लग सकते हैं।
    """)
    
    # मुख्य लेआउट
    with gr.Row():
        
        # बायां कॉलम - इनपुट
        with gr.Column(scale=1):
            prompt = gr.Textbox(
                label="📝 प्रॉम्प्ट",
                placeholder="जैसी इमेज चाहिए वैसा लिखें... (जैसे: a cat sitting on a chair, oil painting)",
                lines=3,
                value="a beautiful sunset over mountains, highly detailed, 4k"
            )
            
            negative_prompt = gr.Textbox(
                label="🚫 नेगेटिव प्रॉम्प्ट",
                placeholder="जो चीज़ें नहीं चाहिए... (जैसे: blurry, low quality, ugly)",
                lines=2,
                value="blurry, low quality, distorted, ugly, bad anatomy"
            )
            
            with gr.Row():
                steps = gr.Slider(
                    minimum=5,
                    maximum=30,
                    value=15,
                    step=1,
                    label="🔢 स्टेप्स",
                    info="ज़्यादा स्टेप्स = बेहतर क्वालिटी, लेकिन ज़्यादा समय"
                )
                
                guidance_scale = gr.Slider(
                    minimum=1.0,
                    maximum=15.0,
                    value=7.5,
                    step=0.5,
                    label="🎯 गाइडेंस स्केल",
                    info="प्रॉम्प्ट को कितनी सख्ती से फॉलो करना है"
                )
            
            seed = gr.Number(
                label="🎲 सीड",
                value=-1,
                precision=0,
                info="समान परिणाम के लिए कोई नंबर डालें। -1 = रैंडम"
            )
            
            generate_btn = gr.Button(
                "✨ इमेज जनरेट करें",
                variant="primary",
                size="lg"
            )
        
        # दायां कॉलम - आउटपुट
        with gr.Column(scale=1):
            output_image = gr.Image(
                label="🖼️ जनरेटेड इमेज",
                type="pil",
                height=512,
                width=512
            )
    
    # उदाहरण प्रॉम्प्ट्स
    gr.Examples(
        examples=[
            ["a red sports car on a mountain road, cinematic lighting", "blurry, low quality", 15, 7.5, 42],
            ["a cute puppy playing in a garden, watercolor painting", "ugly, deformed", 15, 7.5, 123],
            ["an astronaut riding a horse on mars, digital art", "blurry, bad anatomy", 15, 7.5, 999],
            ["a cozy cabin in snowy forest at night, warm lights", "low resolution, blurry", 15, 7.5, 777],
        ],
        inputs=[prompt, negative_prompt, steps, guidance_scale, seed],
        label="उदाहरण प्रॉम्प्ट्स - क्लिक करके ट्राई करें"
    )
    
    # बटन को फंक्शन से जोड़ें
    generate_btn.click(
        fn=generate_image,
        inputs=[prompt, negative_prompt, steps, guidance_scale, seed],
        outputs=output_image
    )
    
    # फ़ूटर
    gr.Markdown("""
    ---
    💡 **टिप्स**:
    - बेहतर क्वालिटी के लिए स्टेप्स 20-25 रखें
    - गाइडेंस स्केल 7-9 के बीच सबसे अच्छा काम करता है
    - एक ही सीड और प्रॉम्प्ट से हमेशा एक जैसी इमेज मिलेगी
    """)

# ============================================
# ऐप लॉन्च करें
# ============================================

if __name__ == "__main__":
    demo.queue(max_size=5).launch()