pyramid-flow / app.py
multimodalart's picture
Update app.py
ad0cea8 verified
raw
history blame
6.25 kB
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
import uuid
import subprocess
subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
# 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(image, 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
if(image):
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,
)
else:
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 = f"{str(uuid.uuid4())}_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():
with gr.Accordion("Image to Video (optional)", open=False):
i2v_image = gr.Image(type="pil", label="Input Image")
t2v_prompt = gr.Textbox(label="Prompt")
with gr.Accordion("Advanced settings", open=False):
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=[i2v_image, 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_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()