import gradio as gr
from diffusers import DiffusionPipeline
from diffusers.schedulers import DPMSolverMultistepScheduler
import torch
from diffusers import DiffusionPipeline, DPMSolverMultistepScheduler
from diffusers.utils import export_to_video
from base64 import b64encode
import torch

device = "cpu"  # Force CPU usage


# Load pipeline (outside the function for efficiency)
pipe = DiffusionPipeline.from_pretrained("damo-vilab/text-to-video-ms-1.7b", torch_dtype=torch.float16, variant="fp16")
pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
pipe.enable_model_cpu_offload()
pipe.enable_vae_slicing()

def Generate_video(prompt, video_duration_seconds):
    num_frames = video_duration_seconds * 10
    video_frames = pipe(prompt=prompt, negative_prompt="low quality",
                   num_inference_steps=25, num_frames=num_frames, device=device).frames
    video_path = export_to_video(video_frames)  # Assuming you have this function defined
    return video_path

# Create Gradio interface
iface = gr.Interface(
    fn=Generate_video,
    inputs=[
        gr.Textbox(lines=5, label="Prompt"),
        gr.Number(label="Video Duration (seconds)", value=3),
    ],
    outputs=gr.Video(label="Generated Video"),
)

# Launch the app
iface.launch(debug=True)