#!/usr/bin/env python3
import torch
from diffusers import DiffusionPipeline, DPMSolverMultistepScheduler
import numpy as np
import gc
from diffusers.utils import export_to_video
from PIL import Image

pipe = DiffusionPipeline.from_pretrained("cerspense/zeroscope_v2_576w", torch_dtype=torch.float16)
pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
pipe.enable_model_cpu_offload()
pipe.enable_xformers_memory_efficient_attention()
pipe.enable_vae_slicing()

prompt = "spiderman running in the desert"
video_frames = pipe(prompt, num_inference_steps=2, height=320, width=576, num_frames=24, output_type="latent".frames
# video_path = export_to_video(video_frames, output_video_path="/home/patrick/videos/video_576_spiderman_24.mp4")

pipe.to("cpu")
del pipe
gc.collect()
torch.cuda.empty_cache()

import ipdb; ipdb.set_trace()

pipe = DiffusionPipeline.from_pretrained("cerspense/zeroscope_v2_XL", torch_dtype=torch.float16)
pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
pipe.enable_model_cpu_offload()
pipe.enable_xformers_memory_efficient_attention()
pipe.enable_vae_slicing()

video = [Image.fromarray(frame).resize((1024, 576)) for frame in video_frames]

video_frames = pipe(prompt, video=video, num_inference_steps=2, strength=0.6).frames
video_path = export_to_video(video_frames, output_video_path="/home/patrick/videos/video_1024_spiderman_24.mp4")