import gradio as gr
import torch
from diffusers import DiffusionPipeline
from PIL import Image
import numpy as np
import io

# تحميل النموذج
model_id = "stabilityai/stable-video-diffusion-img2vid"
pipe = DiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16).to("cuda")

def generate_video(image):
    # تحويل الصورة إلى تنسور
    image_tensor = torch.tensor(np.array(image)).float().unsqueeze(0).permute(0, 3, 1, 2).to("cuda") / 255.0
    
    # توليد الفيديو
    video_frames = pipe(image_tensor, num_frames=25).images  # افترضنا 25 إطار كعدد افتراضي

    # حفظ الفيديو الناتج إلى ملف
    video_bytes = io.BytesIO()
    video_frames[0].save(video_bytes, format="mp4")  # احفظ الإطار الأول كفيديو للإشارة
    
    return video_bytes.getvalue()

# إعداد واجهة Gradio
interface = gr.Interface(
    fn=generate_video,
    inputs=gr.Image(type="pil", label="Upload an Image"),
    outputs=gr.File(label="Download Video"),
    title="Stable Video Diffusion - Image to Video",
    description="Upload an image to generate a video using the Stable Video Diffusion model.",
)

# تشغيل الواجهة
if __name__ == "__main__":
    interface.launch()