import gradio as gr
import torch
from diffusers import SanaVideoPipeline
from diffusers.utils import export_to_video
import tempfile

# إعداد النموذج: نسخة 0.6B هي الأخف في العالم حالياً
model_id = "NVLabs/Sana-1.1-0.6B" # أو النسخة المتاحة في بيئتك

print("جاري تحميل نموذج Sana الذكي...")
# تحميل النموذج بوضعية توفير الطاقة القصوى
pipe = SanaVideoPipeline.from_pretrained(
    model_id, 
    torch_dtype=torch.float32,
    use_safetensors=True
)

# تفعيل خاصية نقل الأجزاء غير المستخدمة من الذاكرة
pipe.enable_model_cpu_offload()

def generate_video(prompt):
    if not prompt:
        return None
    
    # Sana يتميز بقدرته على إنتاج جودة عالية بـ 4-8 خطوات فقط!
    output = pipe(
        prompt=prompt,
        height=320, 
        width=320,
        num_frames=8,
        num_inference_steps=4, # السرعة هنا!
        guidance_scale=5.0,
    ).frames[0]
    
    # حفظ النتيجة
    temp_file = tempfile.NamedTemporaryFile(suffix='.mp4', delete=False)
    export_to_video(output, temp_file.name, fps=8)
    return temp_file.name

# واجهة المستخدم
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🚀 Sana Video: Ultra-Light CPU Gen")
    gr.Markdown("هذا النموذج يعمل بتقنية Linear Attention لضمان السرعة على المساحات المجانية.")
    
    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(label="وصف الفيديو", placeholder="A cinematic shot of a golden retriever puppy playing in the snow...")
            btn = gr.Button("توليد الفيديو (دقيقتين تقريباً)", variant="primary")
        with gr.Column():
            result = gr.Video(label="النتيجة")
    
    btn.click(fn=generate_video, inputs=prompt, outputs=result)

demo.launch()