Spaces:
Running
Running
| import gradio as gr | |
| import os | |
| import subprocess | |
| import torch | |
| is_shared_ui = True if "fffiloni/DimensionX" in os.environ['SPACE_ID'] else False | |
| is_gpu_associated = torch.cuda.is_available() | |
| import gc | |
| from diffusers import AutoencoderKLCogVideoX, CogVideoXImageToVideoPipeline, CogVideoXTransformer3DModel | |
| from diffusers.utils import export_to_video, load_image | |
| from transformers import T5EncoderModel, T5Tokenizer | |
| from datetime import datetime | |
| import random | |
| from moviepy.editor import VideoFileClip | |
| import ffmpeg | |
| from huggingface_hub import hf_hub_download | |
| # Ensure 'checkpoint' directory exists | |
| os.makedirs("checkpoints", exist_ok=True) | |
| if not is_shared_ui and is_gpu_associated: | |
| # Download LoRA weights | |
| hf_hub_download( | |
| repo_id="wenqsun/DimensionX", | |
| filename="orbit_left_lora_weights.safetensors", | |
| local_dir="checkpoints" | |
| ) | |
| hf_hub_download( | |
| repo_id="wenqsun/DimensionX", | |
| filename="orbit_up_lora_weights.safetensors", | |
| local_dir="checkpoints" | |
| ) | |
| # Load models in the global scope | |
| model_id = "THUDM/CogVideoX-5b-I2V" | |
| transformer = CogVideoXTransformer3DModel.from_pretrained(model_id, subfolder="transformer", torch_dtype=torch.float16).to("cpu") | |
| text_encoder = T5EncoderModel.from_pretrained(model_id, subfolder="text_encoder", torch_dtype=torch.float16).to("cpu") | |
| vae = AutoencoderKLCogVideoX.from_pretrained(model_id, subfolder="vae", torch_dtype=torch.float16).to("cpu") | |
| tokenizer = T5Tokenizer.from_pretrained(model_id, subfolder="tokenizer") | |
| pipe = CogVideoXImageToVideoPipeline.from_pretrained(model_id, tokenizer=tokenizer, text_encoder=text_encoder, transformer=transformer, vae=vae, torch_dtype=torch.float16) | |
| # Add this near the top after imports | |
| os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True' | |
| def calculate_resize_dimensions(width, height, max_width=1024): | |
| """Calculate new dimensions maintaining aspect ratio""" | |
| if width <= max_width: | |
| return width, height | |
| aspect_ratio = height / width | |
| new_width = max_width | |
| new_height = int(max_width * aspect_ratio) | |
| # Make height even number for video encoding | |
| new_height = new_height - (new_height % 2) | |
| return new_width, new_height | |
| def infer(image_path, prompt, orbit_type, progress=gr.Progress(track_tqdm=True)): | |
| # Move everything to CPU initially | |
| pipe.to("cpu") | |
| torch.cuda.empty_cache() | |
| # Load and get original image dimensions | |
| image = load_image(image_path) | |
| original_width, original_height = image.size | |
| print(f"IMAGE INPUT SIZE: {original_width} x {original_height}") | |
| # Calculate target dimensions maintaining aspect ratio | |
| target_width, target_height = calculate_resize_dimensions(original_width, original_height) | |
| print(f"TARGET SIZE: {target_width} x {target_height}") | |
| lora_path = "checkpoints/" | |
| weight_name = "orbit_left_lora_weights.safetensors" if orbit_type == "Left" else "orbit_up_lora_weights.safetensors" | |
| lora_rank = 256 | |
| adapter_timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") | |
| # Load LoRA weights on CPU | |
| pipe.load_lora_weights(lora_path, weight_name=weight_name, adapter_name=f"adapter_{adapter_timestamp}") | |
| pipe.fuse_lora(lora_scale=1 / lora_rank) | |
| try: | |
| # Move to GPU just before inference | |
| pipe.to("cuda") | |
| torch.cuda.empty_cache() | |
| prompt = f"{prompt}. High quality, ultrarealistic detail and breath-taking movie-like camera shot." | |
| seed = random.randint(0, 2**8 - 1) | |
| with torch.inference_mode(): | |
| video = pipe( | |
| image, | |
| prompt, | |
| num_inference_steps=50, | |
| guidance_scale=7.0, | |
| use_dynamic_cfg=True, | |
| generator=torch.Generator(device="cpu").manual_seed(seed) | |
| ) | |
| finally: | |
| # Ensure cleanup happens even if inference fails | |
| pipe.to("cpu") | |
| pipe.unfuse_lora() | |
| pipe.unload_lora_weights() | |
| torch.cuda.empty_cache() | |
| gc.collect() | |
| # Generate initial output video | |
| timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") | |
| temp_path = f"output_{timestamp}_temp.mp4" | |
| final_path = f"output_{timestamp}.mp4" | |
| # First export the original video | |
| export_to_video(video.frames[0], temp_path, fps=8) | |
| try: | |
| # Use ffmpeg via subprocess | |
| cmd = [ | |
| 'ffmpeg', | |
| '-i', temp_path, | |
| '-vf', f'scale={target_width}:{target_height}', | |
| '-c:v', 'libx264', | |
| '-preset', 'medium', | |
| '-crf', '23', | |
| '-y', # Overwrite output file if it exists | |
| final_path | |
| ] | |
| subprocess.run(cmd, check=True, capture_output=True) | |
| except subprocess.CalledProcessError as e: | |
| print(f"FFmpeg error: {e.stderr.decode()}") | |
| raise e | |
| finally: | |
| if os.path.exists(temp_path): | |
| os.remove(temp_path) | |
| return final_path | |
| # Set up Gradio U | |
| css = """ | |
| div#warning-duplicate { | |
| background-color: #ebf5ff; | |
| padding: 0 16px 16px; | |
| margin: 0px 0; | |
| color: #030303!important; | |
| } | |
| div#warning-duplicate > .gr-prose > h2, div#warning-duplicate > .gr-prose > p { | |
| color: #0f4592!important; | |
| } | |
| div#warning-duplicate strong { | |
| color: #0f4592; | |
| } | |
| p.actions { | |
| display: flex; | |
| align-items: center; | |
| margin: 20px 0; | |
| } | |
| div#warning-duplicate .actions a { | |
| display: inline-block; | |
| margin-right: 10px; | |
| } | |
| div#warning-setgpu { | |
| background-color: #fff4eb; | |
| padding: 0 16px 16px; | |
| margin: 0px 0; | |
| color: #030303!important; | |
| } | |
| div#warning-setgpu > .gr-prose > h2, div#warning-setgpu > .gr-prose > p { | |
| color: #92220f!important; | |
| } | |
| div#warning-setgpu a, div#warning-setgpu b { | |
| color: #91230f; | |
| } | |
| div#warning-setgpu p.actions > a { | |
| display: inline-block; | |
| background: #1f1f23; | |
| border-radius: 40px; | |
| padding: 6px 24px; | |
| color: antiquewhite; | |
| text-decoration: none; | |
| font-weight: 600; | |
| font-size: 1.2em; | |
| } | |
| div#warning-ready { | |
| background-color: #ecfdf5; | |
| padding: 0 16px 16px; | |
| margin: 0px 0; | |
| color: #030303!important; | |
| } | |
| div#warning-ready > .gr-prose > h2, div#warning-ready > .gr-prose > p { | |
| color: #057857!important; | |
| } | |
| .custom-color { | |
| color: #030303 !important; | |
| } | |
| """ | |
| with gr.Blocks(css=css, analytics_enabled=False) as demo: | |
| with gr.Column(elem_id="col-container"): | |
| gr.Markdown("# DimensionX") | |
| gr.Markdown("### Create Any 3D and 4D Scenes from a Single Image with Controllable Video Diffusion") | |
| gr.HTML(""" | |
| <div style="display:flex;column-gap:4px;"> | |
| <a href="https://github.com/wenqsun/DimensionX"> | |
| <img src='https://img.shields.io/badge/GitHub-Repo-blue'> | |
| </a> | |
| <a href="https://chenshuo20.github.io/DimensionX/"> | |
| <img src='https://img.shields.io/badge/Project-Page-green'> | |
| </a> | |
| <a href="https://arxiv.org/abs/2411.04928"> | |
| <img src='https://img.shields.io/badge/ArXiv-Paper-red'> | |
| </a> | |
| <a href="https://huggingface.co/spaces/fffiloni/DimensionX?duplicate=true"> | |
| <img src="https://huggingface.co/datasets/huggingface/badges/resolve/main/duplicate-this-space-sm.svg" alt="Duplicate this Space"> | |
| </a> | |
| <a href="https://huggingface.co/fffiloni"> | |
| <img src="https://huggingface.co/datasets/huggingface/badges/resolve/main/follow-me-on-HF-sm-dark.svg" alt="Follow me on HF"> | |
| </a> | |
| </div> | |
| """) | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| if is_shared_ui: | |
| top_description = gr.HTML(f''' | |
| <div class="gr-prose"> | |
| <h2 class="custom-color"><svg xmlns="http://www.w3.org/2000/svg" width="18px" height="18px" style="margin-right: 0px;display: inline-block;"fill="none"><path fill="#fff" d="M7 13.2a6.3 6.3 0 0 0 4.4-10.7A6.3 6.3 0 0 0 .6 6.9 6.3 6.3 0 0 0 7 13.2Z"/><path fill="#fff" fill-rule="evenodd" d="M7 0a6.9 6.9 0 0 1 4.8 11.8A6.9 6.9 0 0 1 0 7 6.9 6.9 0 0 1 7 0Zm0 0v.7V0ZM0 7h.6H0Zm7 6.8v-.6.6ZM13.7 7h-.6.6ZM9.1 1.7c-.7-.3-1.4-.4-2.2-.4a5.6 5.6 0 0 0-4 1.6 5.6 5.6 0 0 0-1.6 4 5.6 5.6 0 0 0 1.6 4 5.6 5.6 0 0 0 4 1.7 5.6 5.6 0 0 0 4-1.7 5.6 5.6 0 0 0 1.7-4 5.6 5.6 0 0 0-1.7-4c-.5-.5-1.1-.9-1.8-1.2Z" clip-rule="evenodd"/><path fill="#000" fill-rule="evenodd" d="M7 2.9a.8.8 0 1 1 0 1.5A.8.8 0 0 1 7 3ZM5.8 5.7c0-.4.3-.6.6-.6h.7c.3 0 .6.2.6.6v3.7h.5a.6.6 0 0 1 0 1.3H6a.6.6 0 0 1 0-1.3h.4v-3a.6.6 0 0 1-.6-.7Z" clip-rule="evenodd"/></svg> | |
| Attention: this Space need to be duplicated to work</h2> | |
| <p class="main-message custom-color"> | |
| To make it work, <strong>duplicate the Space</strong> and run it on your own profile using a <strong>private</strong> GPU (L40s recommended).<br /> | |
| A L40s costs <strong>US$1.80/h</strong>. | |
| </p> | |
| <p class="actions custom-color"> | |
| <a href="https://huggingface.co/spaces/{os.environ['SPACE_ID']}?duplicate=true"> | |
| <img src="https://huggingface.co/datasets/huggingface/badges/resolve/main/duplicate-this-space-lg-dark.svg" alt="Duplicate this Space" /> | |
| </a> | |
| to start experimenting with this demo | |
| </p> | |
| </div> | |
| ''', elem_id="warning-duplicate") | |
| else: | |
| if(is_gpu_associated): | |
| top_description = gr.HTML(f''' | |
| <div class="gr-prose"> | |
| <h2 class="custom-color"><svg xmlns="http://www.w3.org/2000/svg" width="18px" height="18px" style="margin-right: 0px;display: inline-block;"fill="none"><path fill="#fff" d="M7 13.2a6.3 6.3 0 0 0 4.4-10.7A6.3 6.3 0 0 0 .6 6.9 6.3 6.3 0 0 0 7 13.2Z"/><path fill="#fff" fill-rule="evenodd" d="M7 0a6.9 6.9 0 0 1 4.8 11.8A6.9 6.9 0 0 1 0 7 6.9 6.9 0 0 1 7 0Zm0 0v.7V0ZM0 7h.6H0Zm7 6.8v-.6.6ZM13.7 7h-.6.6ZM9.1 1.7c-.7-.3-1.4-.4-2.2-.4a5.6 5.6 0 0 0-4 1.6 5.6 5.6 0 0 0-1.6 4 5.6 5.6 0 0 0 1.6 4 5.6 5.6 0 0 0 4 1.7 5.6 5.6 0 0 0 4-1.7 5.6 5.6 0 0 0 1.7-4 5.6 5.6 0 0 0-1.7-4c-.5-.5-1.1-.9-1.8-1.2Z" clip-rule="evenodd"/><path fill="#000" fill-rule="evenodd" d="M7 2.9a.8.8 0 1 1 0 1.5A.8.8 0 0 1 7 3ZM5.8 5.7c0-.4.3-.6.6-.6h.7c.3 0 .6.2.6.6v3.7h.5a.6.6 0 0 1 0 1.3H6a.6.6 0 0 1 0-1.3h.4v-3a.6.6 0 0 1-.6-.7Z" clip-rule="evenodd"/></svg> | |
| You have successfully associated a GPU to this Space 🎉</h2> | |
| <p class="custom-color"> | |
| You will be billed by the minute from when you activated the GPU until when it is turned off. | |
| </p> | |
| </div> | |
| ''', elem_id="warning-ready") | |
| else: | |
| top_description = gr.HTML(f''' | |
| <div class="gr-prose"> | |
| <h2 class="custom-color"><svg xmlns="http://www.w3.org/2000/svg" width="18px" height="18px" style="margin-right: 0px;display: inline-block;"fill="none"><path fill="#fff" d="M7 13.2a6.3 6.3 0 0 0 4.4-10.7A6.3 6.3 0 0 0 .6 6.9 6.3 6.3 0 0 0 7 13.2Z"/><path fill="#fff" fill-rule="evenodd" d="M7 0a6.9 6.9 0 0 1 4.8 11.8A6.9 6.9 0 0 1 0 7 6.9 6.9 0 0 1 7 0Zm0 0v.7V0ZM0 7h.6H0Zm7 6.8v-.6.6ZM13.7 7h-.6.6ZM9.1 1.7c-.7-.3-1.4-.4-2.2-.4a5.6 5.6 0 0 0-4 1.6 5.6 5.6 0 0 0-1.6 4 5.6 5.6 0 0 0 1.6 4 5.6 5.6 0 0 0 4 1.7 5.6 5.6 0 0 0 4-1.7 5.6 5.6 0 0 0 1.7-4 5.6 5.6 0 0 0-1.7-4c-.5-.5-1.1-.9-1.8-1.2Z" clip-rule="evenodd"/><path fill="#000" fill-rule="evenodd" d="M7 2.9a.8.8 0 1 1 0 1.5A.8.8 0 0 1 7 3ZM5.8 5.7c0-.4.3-.6.6-.6h.7c.3 0 .6.2.6.6v3.7h.5a.6.6 0 0 1 0 1.3H6a.6.6 0 0 1 0-1.3h.4v-3a.6.6 0 0 1-.6-.7Z" clip-rule="evenodd"/></svg> | |
| You have successfully duplicated the MimicMotion Space 🎉</h2> | |
| <p class="custom-color">There's only one step left before you can properly play with this demo: <a href="https://huggingface.co/spaces/{os.environ['SPACE_ID']}/settings" style="text-decoration: underline" target="_blank">attribute a GPU</b> to it (via the Settings tab)</a> and run the app below. | |
| You will be billed by the minute from when you activate the GPU until when it is turned off.</p> | |
| <p class="actions custom-color"> | |
| <a href="https://huggingface.co/spaces/{os.environ['SPACE_ID']}/settings">🔥 Set recommended GPU</a> | |
| </p> | |
| </div> | |
| ''', elem_id="warning-setgpu") | |
| image_in = gr.Image(label="Image Input", type="filepath") | |
| prompt = gr.Textbox(label="Prompt") | |
| orbit_type = gr.Radio(label="Orbit type", choices=["Left", "Up"], value="Left", interactive=True) | |
| submit_btn = gr.Button("Submit", interactive=False if is_shared_ui else True) | |
| with gr.Column(scale=2): | |
| video_out = gr.Video(label="Video output") | |
| examples = gr.Examples( | |
| examples = [ | |
| [ | |
| "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg", | |
| "An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space realised in the background.", | |
| "Left", | |
| "./examples/output_astronaut_left.mp4" | |
| ], | |
| [ | |
| "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg", | |
| "An astronaut hatching from an egg, on the surface of the moon, the darkness and depth of space realised in the background.", | |
| "Up", | |
| "./examples/output_astronaut_up.mp4" | |
| ] | |
| ], | |
| inputs=[image_in, prompt, orbit_type, video_out] | |
| ) | |
| submit_btn.click( | |
| fn=infer, | |
| inputs=[image_in, prompt, orbit_type], | |
| outputs=[video_out] | |
| ) | |
| demo.queue().launch(show_error=True, show_api=False, ssr_mode=False) |