add app

app.py

@@ -0,0 +1,266 @@
+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)