Spaces:
Runtime error
Runtime error
File size: 3,934 Bytes
2d41ed4 7e9dbec 2d41ed4 7e9dbec 6731de0 a0a7f49 6731de0 8a407e4 6731de0 2d41ed4 6731de0 2d41ed4 6731de0 2d41ed4 6731de0 2d41ed4 6731de0 2d41ed4 6731de0 2d41ed4 6731de0 2d41ed4 6731de0 2d41ed4 |
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 |
import os
import shlex
import subprocess
subprocess.run(shlex.split('pip install flash-attn --no-build-isolation'), env=os.environ | {'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"})
import gradio as gr
import torch
from PIL import Image
from pyramid_dit import PyramidDiTForVideoGeneration
from diffusers.utils import load_image, export_to_video
from huggingface_hub import snapshot_download
import os
# Download and load the model
model_path = os.path.join(os.getcwd(), 'pyramid-flow-sd3')
if not os.path.exists(model_path):
snapshot_download("rain1011/pyramid-flow-sd3", local_dir=model_path, local_dir_use_symlinks=False, repo_type='model')
torch.cuda.set_device(0)
model_dtype, torch_dtype = 'bf16', torch.bfloat16
model = PyramidDiTForVideoGeneration(
model_path,
model_dtype,
model_variant='diffusion_transformer_768p',
)
model.vae.to("cuda")
model.dit.to("cuda")
model.text_encoder.to("cuda")
model.vae.enable_tiling()
def generate_video(prompt, height, width, duration, guidance_scale, video_guidance_scale):
temp = 16 if duration == "5s" else 31
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=height,
width=width,
temp=temp,
guidance_scale=guidance_scale,
video_guidance_scale=video_guidance_scale,
output_type="pil",
)
output_path = "generated_video.mp4"
export_to_video(frames, output_path, fps=24)
return output_path
def generate_video_from_image(image, prompt, video_guidance_scale):
with torch.no_grad(), torch.cuda.amp.autocast(enabled=True, dtype=torch_dtype):
frames = model.generate_i2v(
prompt=prompt,
input_image=image,
num_inference_steps=[10, 10, 10],
temp=16,
video_guidance_scale=video_guidance_scale,
output_type="pil",
)
output_path = "generated_video_from_image.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():
txt_prompt = gr.Textbox(label="Prompt")
txt_height = gr.Slider(384, 768, value=768, step=384, label="Height")
txt_width = gr.Slider(640, 1280, value=1280, step=640, label="Width")
txt_duration = gr.Radio(["5s", "10s"], value="5s", label="Duration")
txt_guidance_scale = gr.Slider(1, 15, value=9, step=0.1, label="Guidance Scale")
txt_video_guidance_scale = gr.Slider(1, 15, value=5, step=0.1, label="Video Guidance Scale")
txt_generate = gr.Button("Generate Video")
with gr.Column():
txt_output = gr.Video(label="Generated Video")
with gr.Tab("Image-to-Video"):
with gr.Row():
with gr.Column():
img_input = gr.Image(type="pil", label="Input Image")
img_prompt = gr.Textbox(label="Prompt (optional)")
img_video_guidance_scale = gr.Slider(1, 15, value=4, step=0.1, label="Video Guidance Scale")
img_generate = gr.Button("Generate Video")
with gr.Column():
img_output = gr.Video(label="Generated Video")
txt_generate.click(generate_video,
inputs=[txt_prompt, txt_height, txt_width, txt_duration, txt_guidance_scale, txt_video_guidance_scale],
outputs=txt_output)
img_generate.click(generate_video_from_image,
inputs=[img_input, img_prompt, img_video_guidance_scale],
outputs=img_output)
demo.launch()
|