Spaces:
Running
on
Zero
Running
on
Zero
| import os, sys, shutil | |
| import csv | |
| import numpy as np | |
| import ffmpeg | |
| import cv2 | |
| import collections | |
| import json | |
| import math | |
| import time | |
| import imageio | |
| import random | |
| import ast | |
| import spaces | |
| import gradio as gr | |
| from omegaconf import OmegaConf | |
| from PIL import Image | |
| from segment_anything import SamPredictor, sam_model_registry | |
| os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1" | |
| import torch | |
| from torch.utils.data import DataLoader, Dataset | |
| from torchvision import transforms | |
| from diffusers import AutoencoderKLCogVideoX | |
| from transformers import T5EncoderModel | |
| from diffusers.utils import export_to_video, load_image | |
| # Import files from the local fodler | |
| root_path = os.path.abspath('.') | |
| sys.path.append(root_path) | |
| from pipelines.pipeline_cogvideox_i2v_motion_FrameINO import CogVideoXImageToVideoPipeline | |
| from architecture.cogvideox_transformer_3d import CogVideoXTransformer3DModel | |
| from data_loader.video_dataset_motion import VideoDataset_Motion | |
| from architecture.transformer_wan import WanTransformer3DModel | |
| from pipelines.pipeline_wan_i2v_motion_FrameINO import WanImageToVideoPipeline | |
| from architecture.autoencoder_kl_wan import AutoencoderKLWan | |
| MARKDOWN = \ | |
| """ | |
| <div align='center'> | |
| <h1> Frame In-N-Out </h1> \ | |
| <h2 style='font-weight: 450; font-size: 1rem; margin-bottom: 1rem;'>\ | |
| <a href='https://kiteretsu77.github.io/BoyangWang/'>Boyang Wang</a>, <a href='https://xuweiyichen.github.io/'>Xuweiyi Chen</a>, <a href='http://mgadelha.me/'>Matheus Gadelha</a>, <a href='https://sites.google.com/site/zezhoucheng/'>Zezhou Cheng</a>\ | |
| </h2> \ | |
| <div style="display: flex; flex-wrap: wrap; justify-content: center; gap: 2rem; margin-bottom: 1rem;"> | |
| <a href="https://arxiv.org/abs/2505.21491" target="_blank" | |
| style="display: inline-flex; align-items: center; padding: 0.5rem 1rem; background-color: #f0f0f0; /* 浅灰色背景 */ color: #333; /* 深色文字 */ text-decoration: none; border-radius: 9999px; font-weight: 500; transition: background-color 0.3s;"> | |
| <span style="margin-right: 0.5rem;">📄</span> | |
| <span>Paper</span> | |
| </a> | |
| <a href="https://github.com/UVA-Computer-Vision-Lab/FrameINO" target="_blank" | |
| style="display: inline-flex; align-items: center; padding: 0.5rem 1rem; background-color: #f0f0f0; color: #333; text-decoration: none; border-radius: 9999px; font-weight: 500; transition: background-color 0.3s;"> | |
| <span style="margin-right: 0.5rem;">💻</span> | |
| <span>GitHub</span> | |
| </a> | |
| <a href="https://uva-computer-vision-lab.github.io/Frame-In-N-Out" target="_blank" | |
| style="display: inline-flex; align-items: center; padding: 0.5rem 1rem; background-color: #f0f0f0; color: #333; text-decoration: none; border-radius: 9999px; font-weight: 500; transition: background-color 0.3s;"> | |
| <span style="margin-right: 0.5rem;">🤖</span> | |
| <span>Project Page</span> | |
| </a> | |
| <a href="https://huggingface.co/collections/uva-cv-lab/frame-in-n-out" target="_blank" | |
| style="display: inline-flex; align-items: center; padding: 0.5rem 1rem; background-color: #f0f0f0; color: #333; text-decoration: none; border-radius: 9999px; font-weight: 500; transition: background-color 0.3s;"> | |
| <span style="margin-right: 0.5rem;">🤗</span> | |
| <span>HF Model and Data</span> | |
| </a> | |
| </div> | |
| </div> | |
| Frame In-N-Out expands the first-frame to a larger canvas, where it allows users to assign motion trajectories to existing objects and introduce new identities that enter the scene with specified trajectories.<br> | |
| The model we used here is [<b>Wan2.2-5B</b> V1.6](https://huggingface.co/uva-cv-lab/FrameINO_Wan2.2_5B_Stage2_MotionINO_v1.6) trained on our Frame In-N-Out control mechanism. | |
| <p style="color: red;"> | |
| <b>Easiest way:</b> Choose one from <b>Examples</b> below and then simply click <b>Generate</b>. | |
| </p> | |
| ❗️❗️❗️Instruction Steps: | |
| <br> | |
| 1️⃣ Upload your <b>Input Image 🖼️ </b>. | |
| Next, set <b>Resized Height for Input Image</b> and <b>Resized Width for Input Image</b> for the size you want. | |
| <br> | |
| 2️⃣ Set <b>Top-Left Expand Height</b>, <b>Top-Left Expand Width</b>, <b>Bottom-Right Expand Height</b>, and <b>Bottom-Right Expand Width</b> for the expansion amount. | |
| <br> The Canvas Height (Resized Height + Top-Left Expand Height + Bottom-Right Expand Height) and Canvas Width (Resized Width + Top-Left Expand Width + Bottom-Right Expand Width) should be the multiplier of 32. | |
| <br> Recommend <b>Canvas Height = 704</b> and <b>Canvas Width = 1280</b> for the best performance (pre-trained model default resolution). | |
| <br> | |
| 3️⃣ Click <b>Build the Canvas</b>. | |
| 4️⃣ Provide the motion trajectory of the object by clicking on the <b>Expanded Canvas 🖼️ </b>. | |
| You can make additional trajectory for the same object by clicking <b>Add New Traj Line (Same Obj)</b>. | |
| Reset by <b>Clear All Traj</b>. | |
| <br> | |
| 5️⃣ Provide the <b>Identity Reference</b> image and its trajectory (optional). | |
| Since image is segmented by SAM first (providng center point as query), it will be nice for the inputs to be center cropped. | |
| <br> New instance trajectory can be done by clicking <b>Add New Instance (New Obj, including new ID)</b>. | |
| <br> | |
| 6️⃣ Write a detailed <b>text prompt</b>. | |
| <br> | |
| 7️⃣ Click the <b>Generate!</b> button to start the Video Generation. | |
| <br> | |
| If **Frame In-N-Out** is helpful, please help star the [GitHub Repo](https://github.com/UVA-Computer-Vision-Lab/FrameINO?tab=readme-ov-file). Thanks! | |
| """ | |
| # Color | |
| all_color_codes = [(255, 0, 0), (255, 255, 0), (0, 255, 0), (0, 255, 255), | |
| (255, 0, 255), (0, 0, 255), (128, 128, 128), (64, 224, 208), | |
| (233, 150, 122)] | |
| for _ in range(100): # Should not be over 100 colors | |
| all_color_codes.append((random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))) | |
| # Data Transforms | |
| train_transforms = transforms.Compose( | |
| [ | |
| transforms.Lambda(lambda x: x / 255.0 * 2.0 - 1.0), | |
| ] | |
| ) | |
| # Download to avoid uneasy git lfs install | |
| if not os.path.exists("__assets__"): # Check if the assets images exists; if not, download and unzip one | |
| os.system("wget -q https://github.com/user-attachments/files/23571020/assets.zip") | |
| os.system("unzip assets.zip") | |
| ######################################################## Wan2.2 5B ################################################################# | |
| # Path Setting | |
| model_code_name = "Wan" | |
| base_model_id = "Wan-AI/Wan2.2-TI2V-5B-Diffusers" | |
| transformer_ckpt_path = "uva-cv-lab/FrameINO_Wan2.2_5B_Stage2_MotionINO_v1.6" | |
| # Load Model | |
| print("Loading the model!") | |
| transformer = WanTransformer3DModel.from_pretrained(transformer_ckpt_path, torch_dtype=torch.float16) | |
| vae = AutoencoderKLWan.from_pretrained(base_model_id, subfolder="vae", torch_dtype=torch.float32) | |
| # Create the Pipeline | |
| print("Loading the pipeline!") | |
| pipe = WanImageToVideoPipeline.from_pretrained(base_model_id, transformer=transformer, vae=vae, torch_dtype=torch.bfloat16) | |
| pipe.to("cuda") | |
| # pipe.enable_model_cpu_offload() | |
| ##################################################################################################################################### | |
| ########################################################## Other Auxiliary Func ################################################################# | |
| # # Init SAM model | |
| model_type = "vit_h" #vit-h has the most number of paramter | |
| sam_pretrained_path = "pretrained/sam_vit_h_4b8939.pth" | |
| if not os.path.exists(sam_pretrained_path): | |
| os.system("wget -q https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth -P pretrained/") | |
| sam = sam_model_registry[model_type](checkpoint = sam_pretrained_path).to(device="cuda") | |
| sam_predictor = SamPredictor(sam) # There is a lot of setting here | |
| ##################################################################################################################################### | |
| # Examples Sample | |
| def get_example(): | |
| case = [ | |
| [ | |
| '__assets__/horse.jpg', | |
| 480, | |
| 736, | |
| 128, | |
| 224, | |
| 96, | |
| 320, | |
| '__assets__/sheep.png', | |
| "A brown horse with a black mane walks to the right on a wooden path in a green forest, and then a white sheep enters from the left and walks toward it. Natural daylight, realistic texture, smooth motion, cinematic focus, 4K detail.", | |
| [[[[299, 241], [390, 236], [461, 245], [521, 249], [565, 240], [612, 246], [666, 245]], [[449, 224], [488, 212], [512, 206], [531, 209], [552, 202], [581, 204], [609, 210], [657, 206], [703, 202], [716, 211]]], [[[24, 305], [104, 300], [167, 299], [219, 303], [270, 296], [295, 304]]]], | |
| ], | |
| [ | |
| '__assets__/cup.jpg', | |
| 448, | |
| 736, | |
| 256, | |
| 64, | |
| 0, | |
| 480, | |
| '__assets__/hand2.png', | |
| "A human hand reaches into the frame, gently grabbing the black metal cup with a golden character design on the front, lifting it off the table and taking it away.", | |
| [[[[565, 324], [473, 337], [386, 345], [346, 340], [339, 324], [352, 212], [328, 114], [328, 18], [348, 0]]]], | |
| ], | |
| [ | |
| '__assets__/grass.jpg', | |
| 512, | |
| 800, | |
| 64, | |
| 64, | |
| 160, | |
| 416, | |
| '__assets__/dog.png', | |
| "A fluffy, adorable puppy joyfully sprints onto the bright green grass, its fur bouncing with each step as sunlight highlights its soft coat. The scene takes place in a peaceful park filled with tall trees casting gentle shadows across the lawn. After dashing forward with enthusiasm, the puppy slows to a happy trot, continuing farther ahead into the deeper area of the park, disappearing toward the more shaded grass beneath the trees.", | |
| [[[[600, 412], [512, 394], [408, 358], [333, 336], [270, 313], [259, 260], [236, 222], [231, 180]], [[592, 392], [295, 305], [256, 217], [243, 163]]]], | |
| ], | |
| [ | |
| '__assets__/man_scene.jpg', | |
| 576, | |
| 1024, | |
| 64, | |
| 32, | |
| 64, | |
| 224, | |
| None, | |
| "A single hiker, equipped with a backpack, walks toward the right side of a rugged mountainside trail. The bright sunlight highlights the pale rocky terrain around him, while massive stone cliffs loom in the background. Sparse patches of grass and scattered boulders sit along the path, emphasizing the isolation and vastness of the mountain environment as he steadily continues his journey.", | |
| [[[[342, 247], [415, 247], [478, 262], [518, 271], [570, 275], [613, 283], [646, 308], [690, 307], [705, 325]], [[349, 227], [461, 232], [536, 254], [595, 252], [638, 269], [691, 289], [715, 291]], [[341, 283], [415, 291], [500, 316], [590, 317], [632, 354], [675, 362], [711, 372]]]], | |
| ] | |
| ] | |
| return case | |
| def on_example_click( | |
| input_image, resized_height, resized_width, | |
| top_left_height, top_left_width, bottom_right_height, bottom_right_width, | |
| identity_image, text_prompt, traj_lists, | |
| ): | |
| # Convert | |
| traj_lists = ast.literal_eval(traj_lists) | |
| # Note: No need for the rest like resized_width and resized_height, because these will be replaced in function | |
| # Sequentially build the canvas (We don't accept the empty traj_lists & traj_instance_idx returned by build_canvas) | |
| print("Build Canvas by Example!") | |
| visual_canvas, initial_visual_canvas, inference_canvas, _, _ = build_canvas(input_image, resized_height, resized_width, top_left_height, top_left_width, bottom_right_height, bottom_right_width) | |
| print("Init visual_canvas shape is", visual_canvas.shape) | |
| print("Init inference_canvas shape is", inference_canvas.shape) | |
| # Sequentially load the Trajs of all instances on the canvas | |
| print("Build Traj by Example!") | |
| visual_canvas, traj_instance_idx = fn_vis_all_instance_traj(visual_canvas, traj_lists) | |
| print("Drawn visual_canvas shape is", visual_canvas.shape) | |
| return visual_canvas, initial_visual_canvas, inference_canvas, traj_instance_idx | |
| def build_canvas(input_image_path, resized_height, resized_width, top_left_height, top_left_width, bottom_right_height, bottom_right_width): | |
| # Init | |
| canvas_color = (250, 249, 246) # This color is like white color used in painting paper | |
| # Convert the string to integer | |
| if not resized_height.isdigit(): | |
| raise gr.Error("resized_height must be integer input!") | |
| resized_height = int(resized_height) | |
| if not resized_width.isdigit(): | |
| raise gr.Error("resized_width must be integer input!") | |
| resized_width = int(resized_width) | |
| if not top_left_height.isdigit(): | |
| raise gr.Error("top_left_height must be integer input!") | |
| top_left_height = int(top_left_height) | |
| if not top_left_width.isdigit(): | |
| raise gr.Error("top_left_width must be integer input!") | |
| top_left_width = int(top_left_width) | |
| if not bottom_right_height.isdigit(): | |
| raise gr.Error("bottom_right_height must be integer input!") | |
| bottom_right_height = int(bottom_right_height) | |
| if not bottom_right_width.isdigit(): | |
| raise gr.Error("bottom_right_width must be integer input!") | |
| bottom_right_width = int(bottom_right_width) | |
| # Read the original image and preprare the placeholder | |
| first_frame_img = np.uint8(np.asarray(Image.open(input_image_path))) # NOTE: this is BGR form, be careful for the later cropping process for ID Reference | |
| print("first_frame_img shape is ", first_frame_img.shape) | |
| # Resize to a uniform resolution | |
| first_frame_img = cv2.resize(first_frame_img, (resized_width, resized_height), interpolation = cv2.INTER_AREA) | |
| print("first_frame_img is resized to", first_frame_img.shape) | |
| # Expand to Outside Region to form the Canvas | |
| expand_height = resized_height + top_left_height + bottom_right_height | |
| expand_width = resized_width + top_left_width + bottom_right_width | |
| inference_canvas = np.uint8(np.zeros((expand_height, expand_width, 3))) # Whole Black Canvas, same as other inference | |
| visual_canvas = np.full((expand_height, expand_width, 3), canvas_color, dtype=np.uint8) | |
| print("Init Visual Canvas shape is", visual_canvas.shape) | |
| print("Init Inference Canvs shape is", inference_canvas.shape) | |
| # Sanity Check | |
| if expand_height % 32 != 0: | |
| raise gr.Error("The Height of resized_height + top_left_height + bottom_right_height must be divisible by 32!") | |
| if expand_width % 32 != 0: | |
| raise gr.Error("The Width of resized_width + top_left_width + bottom_right_width must be divisible by 32!") | |
| # Draw the Region Box Region (Original Resolution) | |
| bottom_len = inference_canvas.shape[0] - bottom_right_height | |
| right_len = inference_canvas.shape[1] - bottom_right_width | |
| inference_canvas[top_left_height:bottom_len, top_left_width:right_len, :] = first_frame_img | |
| visual_canvas[top_left_height:bottom_len, top_left_width:right_len, :] = first_frame_img | |
| # Resize to the uniform height and width | |
| visual_canvas = cv2.resize(visual_canvas, (uniform_width, uniform_height), interpolation = cv2.INTER_AREA) | |
| print("Visual Canvas resized to", visual_canvas.shape) | |
| # Return the visual_canvas (for visualizaiton) and canvas map | |
| # Corresponds to: visual_canvas, initial_visual_canvas, inference_canvas, traj_instance_idx, traj_lists | |
| return visual_canvas, visual_canvas.copy(), inference_canvas, 0, [ [ [] ] ] # The last two is initialized with the trajectory instance idx and trajectory list | |
| def process_points(traj_list, num_frames=81): | |
| if len(traj_list) < 2: # First point | |
| return [traj_list[0]] * num_frames | |
| elif len(traj_list) >= num_frames: | |
| raise gr.Info("The number of trajectory points is more than the limits, we will do cropping!") | |
| skip = len(traj_list) // num_frames | |
| return traj_list[::skip][: num_frames - 1] + traj_list[-1:] | |
| else: | |
| insert_num = num_frames - len(traj_list) | |
| insert_num_dict = {} | |
| interval = len(traj_list) - 1 | |
| n = insert_num // interval | |
| m = insert_num % interval | |
| for i in range(interval): | |
| insert_num_dict[i] = n | |
| for i in range(m): | |
| insert_num_dict[i] += 1 | |
| res = [] | |
| for i in range(interval): | |
| insert_points = [] | |
| x0, y0 = traj_list[i] | |
| x1, y1 = traj_list[i + 1] | |
| delta_x = x1 - x0 | |
| delta_y = y1 - y0 | |
| for j in range(insert_num_dict[i]): | |
| x = x0 + (j + 1) / (insert_num_dict[i] + 1) * delta_x | |
| y = y0 + (j + 1) / (insert_num_dict[i] + 1) * delta_y | |
| insert_points.append([int(x), int(y)]) | |
| res += traj_list[i : i + 1] + insert_points | |
| res += traj_list[-1:] | |
| # return | |
| return res | |
| def fn_vis_realtime_traj(visual_canvas, traj_list, traj_instance_idx): # Visualize the traj on canvas | |
| # Process Points | |
| points = process_points(traj_list) | |
| # Draw straight line to connect | |
| for i in range(len(points) - 1): | |
| p = points[i] | |
| p1 = points[i + 1] | |
| cv2.line(visual_canvas, p, p1, all_color_codes[traj_instance_idx], 5) | |
| return visual_canvas | |
| def fn_vis_all_instance_traj(visual_canvas, traj_lists): # Visualize all traj from all instances on canvas | |
| for traj_instance_idx, traj_list_instance in enumerate(traj_lists): | |
| for traj_list_line in traj_list_instance: | |
| visual_canvas = fn_vis_realtime_traj(visual_canvas, traj_list_line, traj_instance_idx) | |
| return visual_canvas, traj_instance_idx # Also return the instance idx | |
| def add_traj_point( | |
| visual_canvas, | |
| traj_lists, | |
| traj_instance_idx, | |
| evt: gr.SelectData, | |
| ): # Add new Traj and then visualize | |
| # Convert | |
| traj_lists = ast.literal_eval(traj_lists) | |
| # Mark New Trajectory Key Point | |
| hotizontal, vertical = evt.index | |
| # traj_lists data structure is: (Num of Instnace, Num of Trajecotries, Num of Points, [X, Y]) | |
| traj_lists[-1][-1].append( [int(hotizontal), int(vertical)] ) | |
| # Draw new trajectory on the Canvas image | |
| visual_canvas = fn_vis_realtime_traj(visual_canvas, traj_lists[-1][-1], traj_instance_idx) | |
| # Return New Traj Marked Canvas image | |
| return visual_canvas, traj_lists | |
| def clear_traj_points(initial_visual_canvas): | |
| return initial_visual_canvas.copy(), 0, [ [ [] ] ] # 1sr One is the initial state canvas; 2nd one is the traj instance idx; 3rd one is the traj list (with the same data structure) | |
| def traj_point_update(traj_lists): | |
| # Convert | |
| traj_lists = ast.literal_eval(traj_lists) | |
| # Append on the last trajecotry line | |
| traj_lists[-1].append([]) | |
| return traj_lists | |
| def traj_instance_update(traj_instance_idx, traj_lists): | |
| # Convert | |
| traj_lists = ast.literal_eval(traj_lists) | |
| # Update one index | |
| if traj_instance_idx >= len(all_color_codes): | |
| raise gr.Error("The trajectory instance number is over the limit!") | |
| # Add one for the traj instance | |
| traj_instance_idx = traj_instance_idx + 1 | |
| # Append a new empty list to the traj lists | |
| traj_lists.append([[]]) | |
| # Reutn | |
| return traj_instance_idx, traj_lists | |
| def sample_traj_by_length(points, num_samples): | |
| # Sample points evenly from traj based on the euclidean distance | |
| pts = np.array(points, dtype=float) # shape (M, 2) | |
| # 1) 每段长度 | |
| seg = pts[1:] - pts[:-1] | |
| seg_len = np.sqrt((seg**2).sum(axis=1)) # shape (M-1,) | |
| # 2) 累积长度 | |
| cum = np.cumsum(seg_len) | |
| total_length = cum[-1] | |
| # 3) 目标等距长度位置 | |
| target = np.linspace(0, total_length, num_samples) | |
| res = [] | |
| for t in target: | |
| # 4) 找到它落在哪一段 | |
| idx = np.searchsorted(cum, t) | |
| if idx == 0: | |
| prev = 0. | |
| else: | |
| prev = cum[idx-1] | |
| # 5) 在该段内插值 | |
| ratio = (t - prev) / seg_len[idx] | |
| p = pts[idx] * ratio + pts[idx+1] * (1-ratio) # careful: direction reversed? | |
| # Actually want: start*(1-ratio) + end*ratio | |
| p = pts[idx] * (1 - ratio) + pts[idx+1] * ratio | |
| res.append(p) | |
| return np.array(res) | |
| def inference(inference_canvas, visual_canvas, text_prompt, traj_lists, main_reference_img, | |
| resized_height, resized_width, top_left_height, top_left_width, bottom_right_height, bottom_right_width): | |
| # TODO: enhance the text prompt by Qwen3-VL-32B? | |
| print("Start Inference Pipeline!") | |
| print("Check: inference_canvas shape is", inference_canvas.shape) | |
| # Convert | |
| resized_height = int(resized_height) | |
| resized_width = int(resized_width) | |
| top_left_height = int(top_left_height) | |
| top_left_width = int(top_left_width) | |
| bottom_right_height = int(bottom_right_height) | |
| bottom_right_width = int(bottom_right_width) | |
| traj_lists = ast.literal_eval(traj_lists) | |
| # Init Some Fixed Setting | |
| if model_code_name == "Wan": | |
| config_path = "config/train_wan_motion_FrameINO.yaml" | |
| dot_radius = 7 | |
| num_frames = 81 | |
| elif model_code_name == "CogVideoX": | |
| config_path = "config/train_cogvideox_i2v_motion_FrameINO.yaml" | |
| dot_radius = 6 | |
| num_frames = 49 | |
| config = OmegaConf.load(config_path) | |
| # Prepare tmp folders | |
| print() | |
| store_folder_path = "tmp_app_example_" + str(int(time.time())) | |
| if os.path.exists(store_folder_path): | |
| shutil.rmtree(store_folder_path) | |
| os.makedirs(store_folder_path) | |
| # Write the visual canvas | |
| visual_canvas_store_path = os.path.join(store_folder_path, "visual_canvas.png") | |
| cv2.imwrite( visual_canvas_store_path, cv2.cvtColor(visual_canvas, cv2.COLOR_BGR2RGB) ) | |
| # Resize the map | |
| canvas_width = resized_width + top_left_width + bottom_right_width | |
| canvas_height = resized_height + top_left_height + bottom_right_height | |
| # inference_canvas = cv2.resize(visual_canvas, (canvas_width, canvas_height), interpolation = cv2.INTER_AREA) | |
| print("Canvas Shape is", str(canvas_height) + "x" + str(canvas_width) ) | |
| # TODO: 还要去enhance这个text prompt要跟QWen的保持一致的complexity的感觉。。。 | |
| # Save the text prompt | |
| print("Text Prompt is", text_prompt) | |
| with open(os.path.join(store_folder_path, 'text_prompt.txt'), 'w') as file: | |
| file.write(text_prompt) | |
| ################################################## Motion Trajectory Condition ##################################################### | |
| # #Prepare the points in the linear way | |
| full_pred_tracks = [[] for _ in range(num_frames)] | |
| ID_tensor = None | |
| # Iterate all tracking information for all objects | |
| print("traj_lists is", traj_lists) | |
| for instance_idx, traj_list_per_object in enumerate(traj_lists): | |
| # Iterate all trajectory lines in one instance | |
| for traj_idx, single_trajectory in enumerate(traj_list_per_object): | |
| # Sanity Check | |
| if len(single_trajectory) < 2: | |
| raise gr.Error("One of the trajectory provided is too short!") | |
| # Sampled the point based on the Euclidean distance | |
| sampled_points = sample_traj_by_length(single_trajectory, num_frames) | |
| # Iterate all points | |
| temporal_idx = 0 | |
| for (raw_point_x, raw_point_y) in sampled_points: | |
| # Scale the point coordinate to the Infernece Size (Realistic Canvas size) | |
| point_x, point_y = int(raw_point_x * canvas_width / uniform_width), int(raw_point_y * canvas_height / uniform_height) # Clicking on the board is with respect to the Uniform Preset Height and Width | |
| if traj_idx == 0: # Needs to init the list in list | |
| full_pred_tracks[temporal_idx].append( [] ) | |
| full_pred_tracks[temporal_idx][-1].append( (point_x, point_y) ) # [-1] and [instance_idx] should have the same effect | |
| temporal_idx += 1 | |
| # Create the traj tensor | |
| print("visual_canvas shape is", visual_canvas.shape) | |
| print("inference_canvas shape is", inference_canvas.shape) | |
| traj_tensor, traj_imgs_np, _, img_with_traj = VideoDataset_Motion.prepare_traj_tensor( | |
| full_pred_tracks, canvas_height, canvas_width, | |
| [], dot_radius, canvas_width, canvas_height, | |
| idx=0, first_frame_img = inference_canvas | |
| ) | |
| # Store Trajectory | |
| imageio.mimsave(os.path.join(store_folder_path, "traj_video.mp4"), traj_imgs_np, fps=8) | |
| print("Save Traj video to tmp folder!") | |
| ###################################################################################################################################################### | |
| ########################################## Prepare the Identity Reference Condition ##################################################### | |
| # ID reference preparation | |
| if main_reference_img is not None: | |
| print("We have an ID reference being used!") | |
| # Fetch | |
| ref_h, ref_w, _ = main_reference_img.shape | |
| # Using breakpoint to extract the points | |
| sam_predictor.set_image(np.uint8(main_reference_img)) | |
| # Define the sample point | |
| sam_points = [(ref_w//2, ref_h//2)] # We don't need that many points to express [:len(traj_points)//2] | |
| # Reverse traj_points | |
| positive_point_cords = np.array(sam_points) | |
| positive_point_labels = np.ones(len(positive_point_cords)) | |
| # Predict the mask based on the point and bounding box designed | |
| masks, scores, logits = sam_predictor.predict( | |
| point_coords = positive_point_cords, | |
| point_labels = positive_point_labels, | |
| multimask_output = False, | |
| ) | |
| mask = masks[0] | |
| main_reference_img[mask == False] = 0 # Merge the mask the first first frame | |
| # Resize to the same resolution as the first frame | |
| scale_h = canvas_height / max(ref_h, ref_w) | |
| scale_w = canvas_width / max(ref_h, ref_w) | |
| new_h, new_w = int(ref_h * scale_h), int(ref_w * scale_w) | |
| main_reference_img = cv2.resize(main_reference_img, (new_w, new_h), interpolation = cv2.INTER_AREA) | |
| # Calculate padding amounts on all direction | |
| pad_height1 = (canvas_height - main_reference_img.shape[0]) // 2 | |
| pad_height2 = canvas_height - main_reference_img.shape[0] - pad_height1 | |
| pad_width1 = (canvas_width - main_reference_img.shape[1]) // 2 | |
| pad_width2 = canvas_width - main_reference_img.shape[1] - pad_width1 | |
| # Apply padding to same resolution as the training farmes | |
| main_reference_img = np.pad( | |
| main_reference_img, | |
| ((pad_height1, pad_height2), (pad_width1, pad_width2), (0, 0)), | |
| mode = 'constant', | |
| constant_values = 0 | |
| ) | |
| cv2.imwrite(os.path.join(store_folder_path, "ID.png"), cv2.cvtColor(main_reference_img, cv2.COLOR_BGR2RGB)) | |
| elif main_reference_img is None: | |
| # Whole Black Color placeholder | |
| main_reference_img = np.uint8(np.zeros((canvas_height, canvas_width, 3))) | |
| # Convert to tensor | |
| ID_tensor = torch.tensor(main_reference_img) | |
| ID_tensor = train_transforms(ID_tensor).permute(2, 0, 1).contiguous() | |
| if model_code_name == "Wan": # Needs to be the shape (B, C, F, H, W) | |
| ID_tensor = ID_tensor.unsqueeze(0).unsqueeze(2) | |
| ############################################################################################################################################### | |
| ############################################# Call the Inference Pipeline ########################################################## | |
| image = Image.fromarray(inference_canvas) | |
| if model_code_name == "Wan": | |
| video = pipe( | |
| image = image, | |
| prompt = text_prompt, negative_prompt = "", # Empty string as negative text prompt | |
| traj_tensor = traj_tensor, # Should be shape (F, C, H, W) | |
| ID_tensor = ID_tensor, # Should be shape (B, C, F, H, W) | |
| height = canvas_height, width = canvas_width, num_frames = num_frames, | |
| num_inference_steps = 50, # 38 is also ok | |
| guidance_scale = 5.0, | |
| ).frames[0] | |
| elif model_code_name == "CogVideoX": | |
| video = pipe( | |
| image = image, | |
| prompt = text_prompt, | |
| traj_tensor = traj_tensor, | |
| ID_tensor = ID_tensor, | |
| height = canvas_height, width = canvas_width, num_frames = len(traj_tensor), | |
| guidance_scale = 6, use_dynamic_cfg = False, | |
| num_inference_steps = 50, | |
| add_ID_reference_augment_noise = True, | |
| ).frames[0] | |
| # Store the reuslt | |
| export_to_video(video, os.path.join(store_folder_path, "generated_video_padded.mp4"), fps=8) | |
| # Save frames | |
| print("Writing as Frames") | |
| video_file_path = os.path.join(store_folder_path, "generated_video.mp4") | |
| writer = imageio.get_writer(video_file_path, fps = 8) | |
| for frame_idx, frame in enumerate(video): | |
| # Extract Unpadded version | |
| # frame = np.uint8(frame) | |
| if model_code_name == "CogVideoX": | |
| frame = np.asarray(frame) # PIL to RGB | |
| bottom_right_y = frame.shape[0] - bottom_right_height | |
| bottom_right_x = frame.shape[1] - bottom_right_width | |
| cropped_region_frame = np.uint8(frame[top_left_height: bottom_right_y, top_left_width : bottom_right_x] * 255) | |
| writer.append_data(cropped_region_frame) | |
| writer.close() | |
| ##################################################################################################################################### | |
| return gr.update(value = video_file_path, width = uniform_width, height = uniform_height) | |
| if __name__ == '__main__': | |
| # Global Setting | |
| uniform_height = 480 # Visual Canvas as 480x720 is decent | |
| uniform_width = 720 | |
| # Draw the Website | |
| block = gr.Blocks().queue(max_size=10) | |
| with block: | |
| with gr.Row(): | |
| gr.Markdown(MARKDOWN) | |
| with gr.Row(elem_classes=["container"]): | |
| with gr.Column(scale=2): | |
| # Input image | |
| input_image = gr.Image(type="filepath", label="Input Image 🖼️ ") | |
| # uploaded_files = gr.Gallery(label="Your images", visible=False, columns=5, rows=1, height=200) | |
| with gr.Column(scale=2): | |
| # Input image | |
| resized_height = gr.Textbox(label="Resized Height for Input Image") | |
| resized_width = gr.Textbox(label="Resized Width for Input Image") | |
| # gr.Number(value=unit_height, label="Fixed", interactive=False) | |
| # gr.Number(value=unit_height * 1.77777, label="Fixed", interactive=False) | |
| # Input the expansion factor | |
| top_left_height = gr.Textbox(label="Top-Left Expand Height") | |
| top_left_width = gr.Textbox(label="Top-Left Expand Width") | |
| bottom_right_height = gr.Textbox(label="Bottom-Right Expand Height") | |
| bottom_right_width = gr.Textbox(label="Bottom-Right Expand Width") | |
| # Button | |
| build_canvas_btn = gr.Button(value="Build the Canvas") | |
| with gr.Row(): | |
| with gr.Column(scale=3): | |
| with gr.Row(scale=3): | |
| visual_canvas = gr.Image(height = uniform_height, width = uniform_width, type="numpy", label='Expanded Canvas 🖼️ ') | |
| # inference_canvas = gr.Image(height = uniform_height, width = uniform_width, type="numpy") | |
| # inference_canvas = None | |
| with gr.Row(scale=1): | |
| # TODO: 还差clear traj的选择 | |
| add_point = gr.Button(value = "Add New Traj Line (Same Obj)", visible = True) # Add new trajectory for the same instance | |
| add_traj = gr.Button(value = "Add New Instance (New Obj, including new ID)", visible = True) | |
| clear_traj_button = gr.Button("Clear All Traj", visible=True) | |
| with gr.Column(scale=2): | |
| with gr.Row(scale=2): | |
| identity_image = gr.Image(type="numpy", label="Identity Reference (SAM on center point only) 🖼️ ") | |
| with gr.Row(scale=2): | |
| text_prompt = gr.Textbox(label="Text Prompt", lines=3) | |
| with gr.Row(): | |
| # Button | |
| generation_btn = gr.Button(value="Generate!") | |
| with gr.Row(): | |
| generated_video = gr.Video(value = None, label="Generated Video", show_label = True, height = uniform_height, width = uniform_width) | |
| ################################################################## Click + Select + Any Effect Area ########################################################################### | |
| # Init some states that will be supporting purposes | |
| traj_lists = gr.Textbox(label="Trajectory", visible = False) # gr.State(None) # Data Structure is: (Number of Instance, Number of Trajectories, Points) Init as [ [ [] ] ] | |
| inference_canvas = gr.Image(height = uniform_height, width = uniform_width, type="numpy", visible=False) # TODO: Check | |
| traj_instance_idx = gr.State(0) | |
| initial_visual_canvas = gr.State(None) # gr.Image(height = uniform_height, width = uniform_width, type="numpy", label='Canvas Expanded Image (Initial State)') # This is the initila visual, used to load back in clearing | |
| # Canvas Click | |
| build_canvas_btn.click( | |
| build_canvas, | |
| inputs = [input_image, resized_height, resized_width, top_left_height, top_left_width, bottom_right_height, bottom_right_width], | |
| outputs = [visual_canvas, initial_visual_canvas, inference_canvas, traj_instance_idx, traj_lists] # inference_canvas is used for inference; visual_canvas is for gradio visualization | |
| ) | |
| # Draw Trajectory for each click on the canvas | |
| visual_canvas.select( | |
| fn = add_traj_point, | |
| inputs = [visual_canvas, traj_lists, traj_instance_idx], | |
| outputs = [visual_canvas, traj_lists] | |
| ) | |
| # Add new Trajectory | |
| add_point.click( | |
| fn = traj_point_update, | |
| inputs = [traj_lists], | |
| outputs = [traj_lists], | |
| ) | |
| add_traj.click( | |
| fn = traj_instance_update, | |
| inputs = [traj_instance_idx, traj_lists], | |
| outputs = [traj_instance_idx, traj_lists], | |
| ) | |
| # Clean all the traj points | |
| clear_traj_button.click( | |
| clear_traj_points, | |
| [initial_visual_canvas], | |
| [visual_canvas, traj_instance_idx, traj_lists], | |
| ) | |
| # Inference Generation | |
| generation_btn.click( | |
| inference, | |
| inputs = [inference_canvas, visual_canvas, text_prompt, traj_lists, identity_image, resized_height, resized_width, top_left_height, top_left_width, bottom_right_height, bottom_right_width], | |
| outputs = [generated_video], | |
| ) | |
| # Load Examples | |
| with gr.Row(elem_classes=["container"]): | |
| gr.Examples( | |
| examples = get_example(), | |
| inputs = [input_image, resized_height, resized_width, top_left_height, top_left_width, bottom_right_height, bottom_right_width, identity_image, text_prompt, traj_lists], | |
| run_on_click = True, | |
| fn = on_example_click, | |
| outputs = [visual_canvas, initial_visual_canvas, inference_canvas, traj_instance_idx], | |
| ) | |
| block.launch(share=True) | |