Spaces:
Running
on
Zero
Running
on
Zero
File size: 6,246 Bytes
6a87547 dae6484 6a87547 702754c ad0cea8 702754c 579892d 6a87547 b832af5 6a87547 702754c ad0cea8 e26c748 6a87547 ad0cea8 e26c748 6a87547 ad0cea8 6a87547 ad0cea8 6a87547 ad0cea8 6a87547 ad0cea8 6a87547 ad0cea8 6a87547 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 |
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 * 0.8) # 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=8)
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() |