import gradio as gr
import torch
from diffusers import DiffusionPipeline
import os

# التحقق من وجود GPU (ضروري جداً لهذا النموذج)
device = "cuda" if torch.cuda.is_available() else "cpu"

# تحميل النموذج (استخدمنا الإصدار الأساسي كمثال)
# ملاحظة: قد يحتاج النموذج لمساحة تخزين كبيرة (أكثر من 10GB)
model_id = "robbyant/lingbot-world-base-cam"
pipe = DiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16 if device=="cuda" else torch.float32)
pipe.to(device)

def generate_world(input_image, prompt):
    if input_image is None:
        return None
    
    # عملية التوليد (Inference)
    # ملاحظة: البارامترات قد تختلف قليلاً حسب التحديث الأخير للنموذج
    output_video = pipe(prompt=prompt, image=input_image, num_frames=16).frames[0]
    
    # حفظ الفيديو مؤقتاً لعرضه
    output_path = "output_world.mp4"
    # كود حفظ الفيديو هنا (يعتمد على التنسيق الراجع من النموذج)
    return output_path

# بناء الواجهة
with gr.Blocks() as demo:
    gr.Markdown("# Lingbot-World Online (Genie 3 Alternative)")
    with gr.Row():
        with gr.Column():
            input_img = gr.Image(label="الصورة المرجعية (Initial Frame)")
            prompt_text = gr.Textbox(label="الوصف (Prompt)", placeholder="مثلاً: تحرك الكاميرا لليمين...")
            run_btn = gr.Button("توليد العالم")
        with gr.Column():
            output_vid = gr.Video(label="العالم المتولد")

    run_btn.click(fn=generate_world, inputs=[input_img, prompt_text], outputs=output_vid)

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