# app.py - Hyper-SD LoRA روی CPU — فوق سریع (۵–۱۸ ثانیه!)
import torch
from diffusers import StableDiffusionPipeline, DDIMScheduler
from huggingface_hub import hf_hub_download
import gradio as gr

# برای CPU فقط
device = "cpu"
torch_dtype = torch.float32

# لود base model (SD 1.5 برای سرعت بیشتر روی CPU — کیفیت عالی!)
pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    torch_dtype=torch_dtype,
    use_safetensors=True,
    safety_checker=None,  # برای سرعت بیشتر
)

# دانلود و لود LoRA Hyper-SD (۱ گام!)
ckpt_name = "Hyper-SD15-1step-lora.safetensors"
pipe.load_lora_weights(hf_hub_download("ByteDance/Hyper-SD", ckpt_name))

# fuse LoRA (خیلی مهمه برای سرعت!)
pipe.fuse_lora(lora_scale=1.0)

# scheduler مناسب برای Hyper-SD (DDIM برای ۱–۴ گام)
pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
pipe.scheduler.set_timesteps(1)  # فقط ۱ گام!

# بهینه‌سازی‌های CPU
pipe.enable_attention_slicing()  # حافظه کم
pipe.enable_vae_slicing()       # VAE سریع‌تر
pipe.unet.to(memory_format=torch.channels_last)

pipe.to(device)

def generate(prompt, negative="", steps=1, seed=-1):
    generator = torch.Generator(device=device).manual_seed(int(seed)) if seed != -1 else None
    
    image = pipe(
        prompt=prompt,
        negative_prompt=negative,
        num_inference_steps=steps,  # فقط ۱–۴!
        guidance_scale=0.0,         # برای Hyper-SD معمولاً ۰ (بدون CFG)
        generator=generator,
        height=512,                 # برای CPU سریع‌تر (بعداً upscale کن)
        width=512,
    ).images[0]
    
    return image

# رابط گرادیو
with gr.Blocks(title="Hyper-SD روی CPU — ۵–۱۸ ثانیه!") as demo:
    gr.Markdown("# Hyper-SD LoRA — سریع‌ترین مدل ۲۰۲۵ روی CPU 🚀")
    gr.Markdown("**۱ گام** → **۵–۹ ثانیه** | کیفیت عالی با LoRA ByteDance")

    with gr.Row():
        with gr.Column(scale=2):
            prompt = gr.Textbox(
                label="پرامپت",
                lines=3,
                placeholder="مثلاً: یک اژدهای غول‌پیکر در آسمان طوفانی، هنر دیجیتال",
                value="گربه فضانورد روی مریخ، سبک واقع‌گرایانه، کیفیت بالا"
            )
            negative = gr.Textbox(
                label="نگاتیو پرامپت",
                lines=2,
                placeholder="تار، بدشکل، کم‌جزئیات",
                value="تار، زشت، نویزدار"
            )

            with gr.Row():
                steps = gr.Slider(1, 4, value=1, step=1, label="گام‌ها (۱ = سریع‌ترین)")
                seed = gr.Number(value=-1, label="Seed (-1 = رندوم)")

            btn = gr.Button("تولید تصویر ⚡", variant="primary")

        with gr.Column(scale=1):
            output = gr.Image(label="نتیجه", type="pil")

    gr.Examples(
        examples=[
            ["شهر سایبرپانک در شب با باران و نئون، کیفیت بالا", "تار، تاریک"],
            ["پرتره دختر انیمه با موهای آبی، در جنگل جادویی", ""],
            ["ماشین اسپورت در جاده کوهستانی، غروب آفتاب", "تار، بدشکل"],
        ],
        inputs=[prompt, negative]
    )

    btn.click(fn=generate, inputs=[prompt, negative, steps, seed], outputs=output)

if __name__ == "__main__":
    demo.launch()