import torch
import gradio as gr
from diffnext.pipelines import NOVAPipeline
from diffnext.utils import export_to_image, export_to_video

# Inizializzazione del modello
model_id = "BAAI/nova-d48w1024-osp480"
model_args = {"torch_dtype": torch.float16, "trust_remote_code": True}
pipe = NOVAPipeline.from_pretrained(model_id, **model_args)
pipe = pipe.to("cuda")

# Funzioni per generare immagine e video
def generate_image(prompt: str):
    image = pipe(prompt, max_latent_length=1).frames[0, 0]
    export_to_image(image, "output.jpg")
    return "output.jpg"

def generate_video(prompt: str):
    video = pipe(prompt, max_latent_length=9).frames[0]
    export_to_video(video, "output.mp4", fps=12)
    return "output.mp4"

def generate_video_high_quality(prompt: str):
    video = pipe(
        prompt,
        max_latent_length=9,
        num_inference_steps=128,
        num_diffusion_steps=100,
    ).frames[0]
    export_to_video(video, "output_v2.mp4", fps=12)
    return "output_v2.mp4"


)

iface3 = gr.Interface(
    fn=generate_video_high_quality,
    inputs=gr.Textbox(label="Enter Prompt"),
    outputs=gr.Video(label="High Quality Generated Video"),
    live=True
)


iface3.launch(share=True)