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 = 'pyramid_flow_model' 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()