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 # Constants MODEL_PATH = "pyramid-flow-model" MODEL_REPO = "rain1011/pyramid-flow-sd3" MODEL_VARIANT = "diffusion_transformer_768p" MODEL_DTYPE = "bf16" # 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 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 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) image = ImageOps.fit(image, target_size, method=Image.LANCZOS, centering=(0.5, 0.5)) 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=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()