import spaces import gradio as gr import torch from huggingface_hub import snapshot_download from xora.models.autoencoders.causal_video_autoencoder import CausalVideoAutoencoder from xora.models.transformers.transformer3d import Transformer3DModel from xora.models.transformers.symmetric_patchifier import SymmetricPatchifier from xora.schedulers.rf import RectifiedFlowScheduler from xora.pipelines.pipeline_xora_video import XoraVideoPipeline from transformers import T5EncoderModel, T5Tokenizer from xora.utils.conditioning_method import ConditioningMethod from pathlib import Path import safetensors.torch import json import numpy as np import cv2 from PIL import Image import tempfile import os # Load Hugging Face token if needed hf_token = os.getenv("HF_TOKEN") # Set model download directory within Hugging Face Spaces model_path = "asset" if not os.path.exists(model_path): snapshot_download("Lightricks/Xora", local_dir=model_path, repo_type='model', token=hf_token) # Global variables to load components vae_dir = Path(model_path) / 'vae' unet_dir = Path(model_path) / 'unet' scheduler_dir = Path(model_path) / 'scheduler' device = torch.device("cuda" if torch.cuda.is_available() else "cpu") def load_vae(vae_dir): vae_ckpt_path = vae_dir / "vae_diffusion_pytorch_model.safetensors" vae_config_path = vae_dir / "config.json" with open(vae_config_path, 'r') as f: vae_config = json.load(f) vae = CausalVideoAutoencoder.from_config(vae_config) vae_state_dict = safetensors.torch.load_file(vae_ckpt_path) vae.load_state_dict(vae_state_dict) return vae.cuda().to(torch.bfloat16) def load_unet(unet_dir): unet_ckpt_path = unet_dir / "unet_diffusion_pytorch_model.safetensors" unet_config_path = unet_dir / "config.json" transformer_config = Transformer3DModel.load_config(unet_config_path) transformer = Transformer3DModel.from_config(transformer_config) unet_state_dict = safetensors.torch.load_file(unet_ckpt_path) transformer.load_state_dict(unet_state_dict, strict=True) return transformer.to(device) def load_scheduler(scheduler_dir): scheduler_config_path = scheduler_dir / "scheduler_config.json" scheduler_config = RectifiedFlowScheduler.load_config(scheduler_config_path) return RectifiedFlowScheduler.from_config(scheduler_config) # Helper function for image processing def center_crop_and_resize(frame, target_height, target_width): h, w, _ = frame.shape aspect_ratio_target = target_width / target_height aspect_ratio_frame = w / h if aspect_ratio_frame > aspect_ratio_target: new_width = int(h * aspect_ratio_target) x_start = (w - new_width) // 2 frame_cropped = frame[:, x_start:x_start + new_width] else: new_height = int(w / aspect_ratio_target) y_start = (h - new_height) // 2 frame_cropped = frame[y_start:y_start + new_height, :] frame_resized = cv2.resize(frame_cropped, (target_width, target_height)) return frame_resized def load_image_to_tensor_with_resize(image_path, target_height=512, target_width=768): image = Image.open(image_path).convert("RGB") image_np = np.array(image) frame_resized = center_crop_and_resize(image_np, target_height, target_width) frame_tensor = torch.tensor(frame_resized).permute(2, 0, 1).float() frame_tensor = (frame_tensor / 127.5) - 1.0 return frame_tensor.unsqueeze(0).unsqueeze(2) # Preset options for resolution and frame configuration preset_options = [ {"label": "704x1216, 41 frames", "height": 704, "width": 1216, "num_frames": 41}, {"label": "704x1088, 49 frames", "height": 704, "width": 1088, "num_frames": 49}, {"label": "640x1056, 57 frames", "height": 640, "width": 1056, "num_frames": 57}, {"label": "608x992, 65 frames", "height": 608, "width": 992, "num_frames": 65}, {"label": "608x896, 73 frames", "height": 608, "width": 896, "num_frames": 73}, {"label": "544x896, 81 frames", "height": 544, "width": 896, "num_frames": 81}, {"label": "544x832, 89 frames", "height": 544, "width": 832, "num_frames": 89}, {"label": "512x800, 97 frames", "height": 512, "width": 800, "num_frames": 97}, {"label": "512x768, 97 frames", "height": 512, "width": 768, "num_frames": 97}, {"label": "480x800, 105 frames", "height": 480, "width": 800, "num_frames": 105}, {"label": "480x736, 113 frames", "height": 480, "width": 736, "num_frames": 113}, {"label": "480x704, 121 frames", "height": 480, "width": 704, "num_frames": 121}, {"label": "448x704, 129 frames", "height": 448, "width": 704, "num_frames": 129}, {"label": "448x672, 137 frames", "height": 448, "width": 672, "num_frames": 137}, {"label": "416x640, 153 frames", "height": 416, "width": 640, "num_frames": 153}, {"label": "384x672, 161 frames", "height": 384, "width": 672, "num_frames": 161}, {"label": "384x640, 169 frames", "height": 384, "width": 640, "num_frames": 169}, {"label": "384x608, 177 frames", "height": 384, "width": 608, "num_frames": 177}, {"label": "384x576, 185 frames", "height": 384, "width": 576, "num_frames": 185}, {"label": "352x608, 193 frames", "height": 352, "width": 608, "num_frames": 193}, {"label": "352x576, 201 frames", "height": 352, "width": 576, "num_frames": 201}, {"label": "352x544, 209 frames", "height": 352, "width": 544, "num_frames": 209}, {"label": "352x512, 225 frames", "height": 352, "width": 512, "num_frames": 225}, {"label": "352x512, 233 frames", "height": 352, "width": 512, "num_frames": 233}, {"label": "320x544, 241 frames", "height": 320, "width": 544, "num_frames": 241}, {"label": "320x512, 249 frames", "height": 320, "width": 512, "num_frames": 249}, {"label": "320x512, 257 frames", "height": 320, "width": 512, "num_frames": 257}, {"label": "Custom", "height": None, "width": None, "num_frames": None} ] # Function to toggle visibility of sliders based on preset selection def preset_changed(preset): if preset != "Custom": selected = next(item for item in preset_options if item["label"] == preset) return ( selected["height"], selected["width"], selected["num_frames"], gr.update(visible=False), gr.update(visible=False), gr.update(visible=False) ) else: return None, None, None, gr.update(visible=True), gr.update(visible=True), gr.update(visible=True) # Load models vae = load_vae(vae_dir) unet = load_unet(unet_dir) scheduler = load_scheduler(scheduler_dir) patchifier = SymmetricPatchifier(patch_size=1) text_encoder = T5EncoderModel.from_pretrained("PixArt-alpha/PixArt-XL-2-1024-MS", subfolder="text_encoder").to(device) tokenizer = T5Tokenizer.from_pretrained("PixArt-alpha/PixArt-XL-2-1024-MS", subfolder="tokenizer") pipeline = XoraVideoPipeline( transformer=unet, patchifier=patchifier, text_encoder=text_encoder, tokenizer=tokenizer, scheduler=scheduler, vae=vae, ).to(device) # Modified function to include validation with gr.Error @spaces.GPU(duration=120) def generate_video(image_path=None, prompt="", negative_prompt="", seed=171198, num_inference_steps=40, num_images_per_prompt=1, guidance_scale=3, height=512, width=768, num_frames=121, frame_rate=25, progress=gr.Progress()): # Check prompt length and raise an error if it's too short if len(prompt.strip()) < 50: raise gr.Error("Prompt must be at least 50 characters long. Please provide more details for the best results.", duration=5) if image_path: media_items = load_image_to_tensor_with_resize(image_path, height, width).to(device) else: raise ValueError("Image path must be provided.") sample = { "prompt": prompt, 'prompt_attention_mask': None, 'negative_prompt': negative_prompt, 'negative_prompt_attention_mask': None, 'media_items': media_items, } generator = torch.Generator(device="cpu").manual_seed(seed) def gradio_progress_callback(self, step, timestep, kwargs): progress((step + 1) / num_inference_steps) images = pipeline( num_inference_steps=num_inference_steps, num_images_per_prompt=num_images_per_prompt, guidance_scale=guidance_scale, generator=generator, output_type="pt", height=height, width=width, num_frames=num_frames, frame_rate=frame_rate, **sample, is_video=True, vae_per_channel_normalize=True, conditioning_method=ConditioningMethod.FIRST_FRAME, mixed_precision=True, callback_on_step_end=gradio_progress_callback ).images output_path = tempfile.mktemp(suffix=".mp4") video_np = images.squeeze(0).permute(1, 2, 3, 0).cpu().float().numpy() video_np = (video_np * 255).astype(np.uint8) height, width = video_np.shape[1:3] out = cv2.VideoWriter(output_path, cv2.VideoWriter_fourcc(*'mp4v'), frame_rate, (width, height)) for frame in video_np[..., ::-1]: out.write(frame) out.release() return output_path # Define the Gradio interface with presets with gr.Blocks() as iface: gr.Markdown("# Video Generation with Xora") with gr.Row(): with gr.Column(): image_input = gr.Image(type="filepath", label="Image Input") prompt = gr.Textbox(label="Prompt", value="A man riding a motorcycle down a winding road, surrounded by lush, green scenery and distant mountains. The sky is clear with a few wispy clouds, and the sunlight glistens on the motorcycle as it speeds along. The rider is dressed in a black leather jacket and helmet, leaning slightly forward as the wind rustles through nearby trees. The wheels kick up dust, creating a slight trail behind the motorcycle, adding a sense of speed and excitement to the scene.") negative_prompt = gr.Textbox(label="Negative Prompt", value="worst quality, inconsistent motion...") # Preset dropdown for resolution and frame settings preset_dropdown = gr.Dropdown( choices=[p["label"] for p in preset_options], value="704x1216, 41 frames", label="Resolution Preset" ) # Advanced options section with gr.Accordion("Advanced Options", open=False): seed = gr.Slider(label="Seed", minimum=0, maximum=1000000, step=1, value=171198) inference_steps = gr.Slider(label="Inference Steps", minimum=1, maximum=100, step=1, value=40) images_per_prompt = gr.Slider(label="Images per Prompt", minimum=1, maximum=10, step=1, value=1) guidance_scale = gr.Slider(label="Guidance Scale", minimum=1.0, maximum=20.0, step=0.1, value=3.0) # Sliders to appear at the end of the advanced settings height_slider = gr.Slider(label="Height", minimum=256, maximum=1024, step=64, value=704, visible=False) width_slider = gr.Slider(label="Width", minimum=256, maximum=1024, step=64, value=1216, visible=False) num_frames_slider = gr.Slider(label="Number of Frames", minimum=1, maximum=200, step=1, value=41, visible=False) frame_rate = gr.Slider(label="Frame Rate", minimum=1, maximum=60, step=1, value=25, visible=False) generate_button = gr.Button("Generate Video") with gr.Column(): output_video = gr.Video(label="Generated Video") # Link dropdown change to update sliders visibility and values preset_dropdown.change( fn=preset_changed, inputs=[preset_dropdown], outputs=[height_slider, width_slider, num_frames_slider, height_slider, width_slider, frame_rate] ) generate_button.click( fn=generate_video, inputs=[image_input, prompt, negative_prompt, seed, inference_steps, images_per_prompt, guidance_scale, height_slider, width_slider, num_frames_slider, frame_rate], outputs=output_video ) iface.launch(share=True)