import gradio as gr
import torch
import ftfy
from uuid import uuid4
from diffusers import WanPipeline, AutoencoderKLWan
from diffusers.schedulers.scheduling_unipc_multistep import UniPCMultistepScheduler
from diffusers.utils import export_to_video

device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Running on {device}...")

# Load model
model_id = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
vae = AutoencoderKLWan.from_pretrained(model_id, subfolder="vae", torch_dtype=torch.float32)
scheduler = UniPCMultistepScheduler(
    prediction_type='flow_prediction',
    use_flow_sigmas=True,
    num_train_timesteps=1000,
    flow_shift=5.0
)
pipe = WanPipeline.from_pretrained(model_id, vae=vae, torch_dtype=torch.bfloat16)
pipe.scheduler = scheduler
pipe.to(device)

print("Model loaded successfully.")

def make_divisible_by_16(x):
    return int(x) - int(x) % 16

def generate_video(prompt, negative_prompt="", height=480, width=832, num_frames=81, guidance_scale=5.0):
    try:
        print(f"Generating video with prompt: {prompt}")

        if not prompt:
            raise ValueError("Prompt must be provided.")

        # Validate and adjust height/width
        height = make_divisible_by_16(int(height))
        width = make_divisible_by_16(int(width))
        num_frames = int(num_frames)
        guidance_scale = float(guidance_scale)

        output = pipe(
            prompt=prompt,
            negative_prompt=negative_prompt,
            height=height,
            width=width,
            num_frames=num_frames,
            guidance_scale=guidance_scale,
        ).frames[0]

        output_path = f"{uuid4()}.mp4"
        export_to_video(output, output_path, fps=16)

        print(f"Video generated: {output_path}")
        return output_path

    except Exception as e:
        print(f"Error during video generation: {e}")
        return None

iface = gr.Interface(
    fn=generate_video,
    inputs=[
        gr.Textbox(label="Prompt", placeholder="Describe your scene..."),
        gr.Textbox(label="Negative Prompt", value=""),
        gr.Number(label="Height", value=480),
        gr.Number(label="Width", value=832),
        gr.Number(label="Number of Frames", value=81),
        gr.Number(label="Guidance Scale", value=5.0),
    ],
    outputs=gr.File(label="Generated Video")
)

# Launch Gradio app in API mode
try:
    iface.launch(share=True)
except Exception as e:
    print(f"Error launching Gradio app: {e}")