app.py

import os
import torch
import gradio as gr
from PIL import Image, ImageOps

from huggingface_hub import snapshot_download
from pyramid_dit import PyramidDiTForVideoGeneration
from diffusers.utils import export_to_video

import spaces 

# Constants
MODEL_PATH = "pyramid-flow-model"
MODEL_REPO = "rain1011/pyramid-flow-sd3"
MODEL_VARIANT = "diffusion_transformer_768p"
MODEL_DTYPE = "bf16"

def center_crop(image, target_width, target_height):
 width, height = image.size
 aspect_ratio_target = target_width / target_height
 aspect_ratio_image = width / height

 if aspect_ratio_image > aspect_ratio_target:
 # Crop the width (left and right)
 new_width = int(height * aspect_ratio_target)
 left = (width - new_width) // 2
 right = left + new_width
 top, bottom = 0, height
 else:
 # Crop the height (top and bottom)
 new_height = int(width / aspect_ratio_target)
 top = (height - new_height) // 2
 bottom = top + new_height
 left, right = 0, width

 image = image.crop((left, top, right, bottom))
 return image

# Download and load the model
def load_model():
 if not os.path.exists(MODEL_PATH):
 snapshot_download(MODEL_REPO, local_dir=MODEL_PATH, local_dir_use_symlinks=False, repo_type='model')
 
 model = PyramidDiTForVideoGeneration(
 MODEL_PATH,
 MODEL_DTYPE,
 model_variant=MODEL_VARIANT,
 )
 
 model.vae.to("cuda")
 model.dit.to("cuda")
 model.text_encoder.to("cuda")
 model.vae.enable_tiling()
 
 return model

# Global model variable
model = load_model()

# Text-to-video generation function
@spaces.GPU(duration=240)
def generate_video(prompt, duration, guidance_scale, video_guidance_scale):
 temp = int(duration * 2.4) # Convert seconds to temp value (assuming 24 FPS)
 torch_dtype = torch.bfloat16 if MODEL_DTYPE == "bf16" else torch.float32
 
 with torch.no_grad(), torch.cuda.amp.autocast(enabled=True, dtype=torch_dtype):
 frames = model.generate(
 prompt=prompt,
 num_inference_steps=[20, 20, 20],
 video_num_inference_steps=[10, 10, 10],
 height=768,
 width=1280,
 temp=temp,
 guidance_scale=guidance_scale,
 video_guidance_scale=video_guidance_scale,
 output_type="pil",
 save_memory=True,
 )
 
 output_path = "output_video.mp4"
 export_to_video(frames, output_path, fps=24)
 return output_path

# Image-to-video generation function
@spaces.GPU(duration=240)
def generate_video_from_image(image, prompt, duration, video_guidance_scale):
 temp = int(duration * 2.4) # Convert seconds to temp value (assuming 24 FPS)
 torch_dtype = torch.bfloat16 if MODEL_DTYPE == "bf16" else torch.float32
 
 target_size = (1280, 720)
 cropped_image = center_crop(image, 1280, 720)
 resized_image = cropped_image.resize((1280, 720))
 
 with torch.no_grad(), torch.cuda.amp.autocast(enabled=True, dtype=torch_dtype):
 frames = model.generate_i2v(
 prompt=prompt,
 input_image=resized_image,
 num_inference_steps=[10, 10, 10],
 temp=temp,
 guidance_scale=7.0,
 video_guidance_scale=video_guidance_scale,
 output_type="pil",
 save_memory=True,
 )
 
 output_path = "output_video_i2v.mp4"
 export_to_video(frames, output_path, fps=24)
 return output_path

# Gradio interface
with gr.Blocks() as demo:
 gr.Markdown("# Pyramid Flow Video Generation Demo")
 
 with gr.Tab("Text-to-Video"):
 with gr.Row():
 with gr.Column():
 t2v_prompt = gr.Textbox(label="Prompt")
 t2v_duration = gr.Slider(minimum=1, maximum=10, value=5, step=1, label="Duration (seconds)")
 t2v_guidance_scale = gr.Slider(minimum=1, maximum=15, value=9, step=0.1, label="Guidance Scale")
 t2v_video_guidance_scale = gr.Slider(minimum=1, maximum=15, value=5, step=0.1, label="Video Guidance Scale")
 t2v_generate_btn = gr.Button("Generate Video")
 with gr.Column():
 t2v_output = gr.Video(label="Generated Video")
 
 t2v_generate_btn.click(
 generate_video,
 inputs=[t2v_prompt, t2v_duration, t2v_guidance_scale, t2v_video_guidance_scale],
 outputs=t2v_output
 )
 
 with gr.Tab("Image-to-Video"):
 with gr.Row():
 with gr.Column():
 i2v_image = gr.Image(type="pil", label="Input Image")
 i2v_prompt = gr.Textbox(label="Prompt")
 i2v_duration = gr.Slider(minimum=1, maximum=10, value=5, step=1, label="Duration (seconds)")
 i2v_video_guidance_scale = gr.Slider(minimum=1, maximum=15, value=4, step=0.1, label="Video Guidance Scale")
 i2v_generate_btn = gr.Button("Generate Video")
 with gr.Column():
 i2v_output = gr.Video(label="Generated Video")
 
 i2v_generate_btn.click(
 generate_video_from_image,
 inputs=[i2v_image, i2v_prompt, i2v_duration, i2v_video_guidance_scale],
 outputs=i2v_output
 )

demo.launch()