#!/usr/bin/env python
"""
This is the full application script for VideoPainter.
It first checks for and (if necessary) installs missing dependencies.
When installing the custom packages (diffusers and app),
it uses the flag --no-build-isolation so that the installed torch is seen.
If the custom diffusers package fails to provide the expected submodules,
the script will force-install the official diffusers package.
"""
import os
import sys
import subprocess
import warnings
import time
import importlib
warnings.filterwarnings("ignore")
# Set Gradio temp directory via environment variable
GRADIO_TEMP_DIR = "./tmp_gradio"
os.makedirs(GRADIO_TEMP_DIR, exist_ok=True)
os.makedirs(f"{GRADIO_TEMP_DIR}/track", exist_ok=True)
os.makedirs(f"{GRADIO_TEMP_DIR}/inpaint", exist_ok=True)
os.environ["GRADIO_TEMP_DIR"] = GRADIO_TEMP_DIR
def install_package(package_spec):
print(f"Installing {package_spec} ...")
try:
subprocess.check_call([sys.executable, "-m", "pip", "install", package_spec])
print(f"Successfully installed {package_spec}")
return True
except Exception as e:
print(f"Failed to install {package_spec}: {e}")
return False
print("Checking for PyTorch ...")
try:
import torch
print("PyTorch is already installed.")
except ImportError:
print("PyTorch not found, installing...")
if not install_package("torch>=2.0.0 torchvision>=0.15.0"):
print("Failed to install PyTorch, which is required.")
sys.exit(1)
# First, install wheel package which is needed for bdist_wheel command
install_package("wheel")
# Install ninja for faster builds
install_package("ninja")
# Check and install other critical dependencies
critical_dependencies = [
("hydra", "hydra-core>=1.3.2"),
("omegaconf", "omegaconf>=2.3.0"),
("decord", "decord>=0.6.0"),
("diffusers", "diffusers>=0.24.0"), # Will be replaced with custom one
("transformers", "transformers>=4.35.0"),
("gradio", "gradio>=4.0.0"),
("numpy", "numpy>=1.24.0"),
("cv2", "opencv-python>=4.8.0"),
("PIL", "Pillow>=10.0.0"),
("scipy", "scipy>=1.11.0"),
("einops", "einops>=0.7.0"),
("onnxruntime", "onnxruntime>=1.16.0"),
("timm", "timm>=0.9.0"),
("safetensors", "safetensors>=0.4.0"),
("moviepy", "moviepy>=1.0.3"),
("imageio", "imageio>=2.30.0"),
("tqdm", "tqdm>=4.64.0"),
("openai", "openai>=1.5.0"),
("psutil", "psutil>=5.9.0")
]
for mod_name, pkg_spec in critical_dependencies:
try:
if mod_name == "PIL":
from PIL import Image
elif mod_name == "cv2":
import cv2
else:
__import__(mod_name)
print(f"{mod_name} is already installed.")
except ImportError:
print(f"{mod_name} not found, installing {pkg_spec} ...")
install_package(pkg_spec)
print("Setting up environment...")
# Clone the VideoPainter repository if not present
if not os.path.exists("VideoPainter"):
print("Cloning VideoPainter repository...")
os.system("git clone https://github.com/TencentARC/VideoPainter.git")
# Add necessary paths to sys.path
sys.path.append(os.path.join(os.getcwd(), "VideoPainter"))
sys.path.append(os.path.join(os.getcwd(), "VideoPainter/app"))
sys.path.append(os.path.join(os.getcwd(), "app"))
sys.path.append(".")
# Ensure custom diffusers is importable
if os.path.exists("VideoPainter/diffusers"):
print("Installing custom diffusers...")
# First, remove any existing diffusers installation
subprocess.call([sys.executable, "-m", "pip", "uninstall", "-y", "diffusers"])
# Copy the files directly into the site-packages directory instead of using pip install -e
import site
site_packages = site.getsitepackages()[0]
diffusers_src = os.path.join(os.getcwd(), "VideoPainter/diffusers/src/diffusers")
diffusers_dst = os.path.join(site_packages, "diffusers")
print(f"Copying diffusers from {diffusers_src} to {diffusers_dst}")
if not os.path.exists(diffusers_dst):
os.makedirs(diffusers_dst, exist_ok=True)
# Copy diffusers files directly
os.system(f"cp -r {diffusers_src}/* {diffusers_dst}/")
# Also add VideoPainter/diffusers/src to sys.path
sys.path.append(os.path.join(os.getcwd(), "VideoPainter/diffusers/src"))
# Verify the custom model is available
try:
# Force reload diffusers to pick up the new files
if "diffusers" in sys.modules:
del sys.modules["diffusers"]
import diffusers
print(f"Diffusers version: {diffusers.__version__}")
print(f"Available modules in diffusers: {dir(diffusers)}")
# Check if models directory exists in custom diffusers
models_dir = os.path.join(diffusers_dst, "models")
if os.path.exists(models_dir):
print(f"Models in diffusers: {os.listdir(models_dir)}")
except Exception as e:
print(f"Error verifying diffusers installation: {e}")
# Copy the app directory if needed
if not os.path.exists("app"):
os.makedirs("app", exist_ok=True)
print("Copying VideoPainter/app to local app directory...")
os.system("cp -r VideoPainter/app/* app/")
# Don't try to install app package, just add to path
print("Adding app directory to Python path...")
app_path = os.path.join(os.getcwd(), "app")
sys.path.insert(0, app_path)
# Insert the VideoPainter path at the beginning of sys.path to ensure it takes precedence
sys.path.insert(0, os.path.join(os.getcwd(), "VideoPainter"))
print("Importing standard modules and dependencies ...")
try:
import gradio as gr
import cv2
import numpy as np
import scipy
import torchvision
from PIL import Image
from huggingface_hub import snapshot_download
from decord import VideoReader
except ImportError as e:
print(f"Error importing basic modules: {e}")
sys.exit(1)
# Import specialized modules with better error handling
try:
# Import our custom modules
from sam2.build_sam import build_sam2_video_predictor
# Force reload of diffusers after direct copy
if "diffusers" in sys.modules:
del sys.modules["diffusers"]
# Now import diffusers with explicit path to the files we need
sys.path.insert(0, os.path.join(os.getcwd(), "VideoPainter/app"))
# Import utils after setting up correct paths
from utils import load_model, generate_frames
print("All modules imported successfully!")
except ImportError as e:
print(f"Error importing specialized modules: {e}")
print("Paths:", sys.path)
# Try to diagnose and fix the specific issue
if "CogvideoXBranchModel" in str(e):
print("Trying to fix missing CogvideoXBranchModel...")
# Check if the model file exists in the repository
branch_model_file = "VideoPainter/diffusers/src/diffusers/models/cogvideox_branch.py"
if os.path.exists(branch_model_file):
print(f"Found branch model file at {branch_model_file}")
# Manually import the module
import sys
sys.path.insert(0, os.path.join(os.getcwd(), "VideoPainter/diffusers/src"))
# Add the import to __init__.py if not already there
init_file = os.path.join(site_packages, "diffusers/__init__.py")
with open(init_file, 'r') as f:
init_content = f.read()
if "CogvideoXBranchModel" not in init_content:
print("Adding CogvideoXBranchModel to diffusers/__init__.py")
with open(init_file, 'a') as f:
f.write("\nfrom .models.cogvideox_branch import CogvideoXBranchModel\n")
# Force reload diffusers
if "diffusers" in sys.modules:
del sys.modules["diffusers"]
# Try importing again
from utils import load_model, generate_frames
print("Fixed CogvideoXBranchModel import issue!")
else:
print(f"Could not find {branch_model_file}")
sys.exit(1)
else:
sys.exit(1)
###############################
# Begin Application Code (VideoPainter demo)
###############################
def download_models():
print("Downloading models from Hugging Face Hub...")
models = {
"CogVideoX-5b-I2V": "THUDM/CogVideoX-5b-I2V",
"VideoPainter": "TencentARC/VideoPainter"
}
model_paths = {}
os.makedirs("ckpt", exist_ok=True)
for name, repo_id in models.items():
print(f"Downloading {name} from {repo_id}...")
path = snapshot_download(repo_id=repo_id)
model_paths[name] = path
print(f"Downloaded {name} to {path}")
try:
flux_path = snapshot_download(repo_id="black-forest-labs/FLUX.1-Fill-dev")
model_paths["FLUX"] = flux_path
except Exception as e:
print(f"Failed to download FLUX model: {e}")
model_paths["FLUX"] = None
os.makedirs("ckpt/Grounded-SAM-2", exist_ok=True)
sam2_path = "ckpt/Grounded-SAM-2/sam2_hiera_large.pt"
if not os.path.exists(sam2_path):
print(f"Downloading SAM2 to {sam2_path}...")
os.system(f"wget -O {sam2_path} https://huggingface.co/spaces/sam2/sam2/resolve/main/sam2_hiera_large.pt")
model_paths["SAM2"] = sam2_path
return model_paths
print("Initializing application environment...")
if not os.path.exists("app"):
print("Setting up app folder from VideoPainter repository ...")
os.system("git clone https://github.com/TencentARC/VideoPainter.git")
os.makedirs("app", exist_ok=True)
os.system("cp -r VideoPainter/app/* app/")
os.system("pip install --no-build-isolation -e VideoPainter/diffusers")
os.chdir("app")
os.system("pip install --no-build-isolation -e .")
os.chdir("..")
sys.path.append("app")
sys.path.append(".")
# Import project modules (again, to be safe)
try:
from decord import VideoReader
from sam2.build_sam import build_sam2_video_predictor
from utils import load_model, generate_frames
except ImportError as e:
print(f"Failed to import specialized modules: {e}")
sys.exit(1)
# Set up OpenRouter / OpenAI (for caption generation)
try:
from openai import OpenAI
vlm_model = OpenAI(
api_key=os.getenv("OPENROUTER_API_KEY", ""),
base_url="https://openrouter.ai/api/v1"
)
print("OpenRouter client initialized successfully")
except Exception as e:
print(f"OpenRouter API not available: {e}")
class DummyModel:
def __getattr__(self, name):
return self
def __call__(self, *args, **kwargs):
return self
def create(self, *args, **kwargs):
class DummyResponse:
choices = [type('obj', (object,), {'message': type('obj', (object,), {'content': "OpenRouter API not available. Using default prompt."})})]
return DummyResponse()
vlm_model = DummyModel()
###############################
# Download models and initialize predictors
###############################
model_paths = download_models()
base_model_path = model_paths["CogVideoX-5b-I2V"]
videopainter_path = model_paths["VideoPainter"]
inpainting_branch = os.path.join(videopainter_path, "checkpoints/branch")
id_adapter = os.path.join(videopainter_path, "VideoPainterID/checkpoints")
img_inpainting_model = model_paths.get("FLUX")
sam2_checkpoint = "ckpt/Grounded-SAM-2/sam2_hiera_large.pt"
model_cfg = "sam2_hiera_l.yaml"
try:
predictor = build_sam2_video_predictor(model_cfg, sam2_checkpoint)
print("Build SAM2 predictor done!")
validation_pipeline, validation_pipeline_img = load_model(
model_path=base_model_path,
inpainting_branch=inpainting_branch,
id_adapter=id_adapter,
img_inpainting_model=img_inpainting_model
)
print("Load model done!")
except Exception as e:
print(f"Error initializing models: {e}")
sys.exit(1)
###############################
# Helper functions & state definitions
###############################
EXAMPLES = [
[
"https://huggingface.co/spaces/TencentARC/VideoPainter/resolve/main/examples/ferry.mp4",
"A white ferry with red and blue accents, named 'COLONIA', cruises on a calm river...",
"White and red passenger ferry boat labeled 'COLONIA 6' with multiple windows, life buoys, and upper deck seating.",
"Positive",
"Inpaint",
"",
42,
6.0,
16,
[[[320, 240]], [1]],
],
[
"https://huggingface.co/spaces/TencentARC/VideoPainter/resolve/main/examples/street.mp4",
"A bustling city street at night illuminated by festive lights, a red double-decker bus...",
"The rear of a black car with illuminated red tail lights and a visible license plate.",
"Positive",
"Inpaint",
"",
42,
6.0,
16,
[[[200, 400]], [1]],
],
]
class StatusMessage:
INFO = "Info"
WARNING = "Warning"
ERROR = "Error"
SUCCESS = "Success"
def create_status(message, status_type=StatusMessage.INFO):
timestamp = time.strftime("%H:%M:%S")
return [("", ""), (f"[{timestamp}]: {message}\n", status_type)]
def update_status(previous_status, new_message, status_type=StatusMessage.INFO):
timestamp = time.strftime("%H:%M:%S")
history = previous_status[-3:]
history.append((f"[{timestamp}]: {new_message}\n", status_type))
return [("", "")] + history
def init_state(offload_video_to_cpu=False, offload_state_to_cpu=False):
inference_state = {}
inference_state["images"] = torch.zeros([1, 3, 100, 100])
inference_state["num_frames"] = 1
inference_state["offload_video_to_cpu"] = offload_video_to_cpu
inference_state["offload_state_to_cpu"] = offload_state_to_cpu
inference_state["video_height"] = 100
inference_state["video_width"] = 100
inference_state["device"] = torch.device("cuda")
inference_state["storage_device"] = torch.device("cpu") if offload_state_to_cpu else torch.device("cuda")
inference_state["point_inputs_per_obj"] = {}
inference_state["mask_inputs_per_obj"] = {}
inference_state["cached_features"] = {}
inference_state["constants"] = {}
inference_state["obj_id_to_idx"] = OrderedDict()
inference_state["obj_idx_to_id"] = OrderedDict()
inference_state["obj_ids"] = []
inference_state["output_dict"] = {"cond_frame_outputs": {}, "non_cond_frame_outputs": {}}
inference_state["output_dict_per_obj"] = {}
inference_state["temp_output_dict_per_obj"] = {}
inference_state["consolidated_frame_inds"] = {"cond_frame_outputs": set(), "non_cond_frame_outputs": set()}
inference_state["tracking_has_started"] = False
inference_state["frames_already_tracked"] = {}
inference_state = gr.State(inference_state)
return inference_state
# (All additional helper functions such as get_frames_from_video, sam_refine, vos_tracking_video,
# inpaint_video, generate_video_from_frames, process_example, reset_all, etc. are defined below.)
# For brevity, they are included here in full as in your original code.
def get_frames_from_video(video_input, video_state):
video_path = video_input
frames = []
user_name = time.time()
vr = VideoReader(video_path)
original_fps = vr.get_avg_fps()
if original_fps > 8:
total_frames = len(vr)
sample_interval = max(1, int(original_fps / 8))
frame_indices = list(range(0, total_frames, sample_interval))
frames = vr.get_batch(frame_indices).asnumpy()
else:
frames = vr.get_batch(list(range(len(vr)))).asnumpy()
frames = frames[:49]
resized_frames = [cv2.resize(frame, (720, 480)) for frame in frames]
frames = np.array(resized_frames)
init_start = time.time()
inference_state = predictor.init_state(images=frames, offload_video_to_cpu=True, async_loading_frames=True)
init_time = time.time() - init_start
print(f"Inference state initialization took {init_time:.2f}s")
fps = 8
image_size = (frames[0].shape[0], frames[0].shape[1])
video_state = {
"user_name": user_name,
"video_name": os.path.split(video_path)[-1],
"origin_images": frames,
"painted_images": frames.copy(),
"masks": [np.zeros((frames[0].shape[0], frames[0].shape[1]), np.uint8)] * len(frames),
"logits": [None] * len(frames),
"select_frame_number": 0,
"fps": fps,
"ann_obj_id": 0
}
video_info = f"Video Name: {video_state['video_name']}, FPS: {video_state['fps']}, Total Frames: {len(frames)}, Image Size: {image_size}"
video_input_path = generate_video_from_frames(frames, output_path=f"{GRADIO_TEMP_DIR}/inpaint/original_{video_state['video_name']}", fps=fps)
return (gr.update(visible=True), gr.update(visible=True), inference_state, video_state, video_info,
video_state["origin_images"][0], gr.update(visible=False, maximum=len(frames), value=1, interactive=True),
gr.update(visible=False, maximum=len(frames), value=len(frames), interactive=True), gr.update(visible=True, interactive=True),
gr.update(visible=True, interactive=True), gr.update(visible=True, interactive=True), gr.update(visible=True),
gr.update(visible=True, interactive=False), create_status("Upload video complete. Ready to select targets.", StatusMessage.SUCCESS), video_input_path)
def select_template(image_selection_slider, video_state, interactive_state, previous_status):
image_selection_slider -= 1
video_state["select_frame_number"] = image_selection_slider
return video_state["painted_images"][image_selection_slider], video_state, interactive_state, update_status(previous_status, f"Set tracking start at frame {image_selection_slider}.", StatusMessage.INFO)
def get_end_number(track_pause_number_slider, video_state, interactive_state, previous_status):
interactive_state["track_end_number"] = track_pause_number_slider
return video_state["painted_images"][track_pause_number_slider], interactive_state, update_status(previous_status, f"Set tracking finish at frame {track_pause_number_slider}.", StatusMessage.INFO)
def sam_refine(inference_state, video_state, point_prompt, click_state, interactive_state, evt, previous_status):
ann_obj_id = 0
ann_frame_idx = video_state["select_frame_number"]
if point_prompt == "Positive":
coordinate = f"[[{evt.index[0]},{evt.index[1]},1]]"
interactive_state["positive_click_times"] += 1
else:
coordinate = f"[[{evt.index[0]},{evt.index[1]},0]]"
interactive_state["negative_click_times"] += 1
print(f"sam_refine, point_prompt: {point_prompt}, click_state: {click_state}")
prompt = {"prompt_type":["click"], "input_point": click_state[0], "input_label": click_state[1], "multimask_output": "True"}
points = np.array(prompt["input_point"])
labels = np.array(prompt["input_label"])
height, width = video_state["origin_images"][0].shape[0:2]
for i in range(len(points)):
points[i, 0] = int(points[i, 0])
points[i, 1] = int(points[i, 1])
print(f"sam_refine points: {points}, labels: {labels}")
frame_idx, obj_ids, mask = predictor.add_new_points(inference_state=inference_state, frame_idx=ann_frame_idx, obj_id=ann_obj_id, points=points, labels=labels)
mask_ = mask.cpu().squeeze().detach().numpy()
mask_[mask_ <= 0] = 0
mask_[mask_ > 0] = 1
org_image = video_state["origin_images"][video_state["select_frame_number"]]
mask_ = cv2.resize(mask_, (width, height))
mask_ = mask_[:, :, None]
mask_[mask_ > 0.5] = 1
mask_[mask_ <= 0.5] = 0
color = 63 * np.ones((height, width, 3)) * np.array([[[np.random.randint(5), np.random.randint(5), np.random.randint(5)]]])
painted_image = np.uint8((1 - 0.5 * mask_) * org_image + 0.5 * mask_ * color)
video_state["masks"][video_state["select_frame_number"]] = mask_
video_state["painted_images"][video_state["select_frame_number"]] = painted_image
return painted_image, video_state, interactive_state, update_status(previous_status, "Segmentation updated. Add more points or continue tracking.", StatusMessage.SUCCESS)
def clear_click(inference_state, video_state, click_state, previous_status):
predictor.reset_state(inference_state)
click_state = [[], []]
template_frame = video_state["origin_images"][video_state["select_frame_number"]]
return inference_state, template_frame, click_state, update_status(previous_status, "Click history cleared.", StatusMessage.INFO)
def vos_tracking_video(inference_state, video_state, interactive_state, previous_status):
height, width = video_state["origin_images"][0].shape[0:2]
masks = []
for out_frame_idx, out_obj_ids, out_mask_logits in predictor.propagate_in_video(inference_state):
mask = np.zeros([480, 720, 1])
for i in range(len(out_mask_logits)):
out_mask = out_mask_logits[i].cpu().squeeze().detach().numpy()
out_mask[out_mask > 0] = 1
out_mask[out_mask <= 0] = 0
out_mask = out_mask[:, :, None]
mask += out_mask
mask = cv2.resize(mask, (width, height))
mask = mask[:, :, None]
mask[mask > 0.5] = 1
mask[mask < 1] = 0
mask = scipy.ndimage.binary_dilation(mask, iterations=6)
masks.append(mask)
masks = np.array(masks)
if interactive_state.get("track_end_number") is not None:
video_state["masks"][video_state["select_frame_number"]:interactive_state["track_end_number"]] = masks
org_images = video_state["origin_images"][video_state["select_frame_number"]:interactive_state["track_end_number"]]
color = 255 * np.ones((1, org_images.shape[-3], org_images.shape[-2], 3)) * np.array([[[[0, 1, 1]]]])
painted_images = np.uint8((1 - 0.5 * masks) * org_images + 0.5 * masks * color)
video_state["painted_images"][video_state["select_frame_number"]:interactive_state["track_end_number"]] = painted_images
else:
video_state["masks"] = masks
org_images = video_state["origin_images"]
color = 255 * np.ones((1, org_images.shape[-3], org_images.shape[-2], 3)) * np.array([[[[0, 1, 1]]]])
painted_images = np.uint8((1 - 0.5 * masks) * org_images + 0.5 * masks * color)
video_state["painted_images"] = painted_images
video_output = generate_video_from_frames(video_state["painted_images"], output_path=f"{GRADIO_TEMP_DIR}/track/{video_state['video_name']}", fps=video_state["fps"])
interactive_state["inference_times"] += 1
print(f"vos_tracking_video output: {video_output}")
return inference_state, video_output, video_state, interactive_state, update_status(previous_status, "Tracking complete.", StatusMessage.SUCCESS), gr.Button.update(interactive=True), gr.Button.update(interactive=True), gr.Button.update(interactive=True), gr.Button.update(interactive=True)
def inpaint_video(video_state, video_caption, target_region_frame1_caption, interactive_state, previous_status, seed_param, cfg_scale, dilate_size):
seed = int(seed_param) if int(seed_param) >= 0 else np.random.randint(0, 2**32 - 1)
validation_images = video_state["origin_images"]
validation_masks = video_state["masks"]
validation_masks = [np.squeeze(mask) for mask in validation_masks]
validation_masks = [(mask > 0).astype(np.uint8) * 255 for mask in validation_masks]
validation_masks = [np.stack([m, m, m], axis=-1) for m in validation_masks]
validation_images = [Image.fromarray(np.uint8(img)).convert('RGB') for img in validation_images]
validation_masks = [Image.fromarray(np.uint8(mask)).convert('RGB') for mask in validation_masks]
validation_images = [img.resize((720, 480)) for img in validation_images]
validation_masks = [mask.resize((720, 480)) for mask in validation_masks]
print("Inpainting: video_caption=", video_caption)
images = generate_frames(
images=validation_images,
masks=validation_masks,
pipe=validation_pipeline,
pipe_img_inpainting=validation_pipeline_img,
prompt=str(video_caption),
image_inpainting_prompt=str(target_region_frame1_caption),
seed=seed,
cfg_scale=float(cfg_scale),
dilate_size=int(dilate_size)
)
images = (images * 255).astype(np.uint8)
video_output = generate_video_from_frames(images, output_path=f"{GRADIO_TEMP_DIR}/inpaint/{video_state['video_name']}", fps=8)
print(f"Inpaint_video output: {video_output}")
return video_output, update_status(previous_status, "Inpainting complete.", StatusMessage.SUCCESS)
def generate_video_from_frames(frames, output_path, fps=8):
frames_tensor = torch.from_numpy(np.asarray(frames)).to(torch.uint8)
if not os.path.exists(os.path.dirname(output_path)):
os.makedirs(os.path.dirname(output_path))
torchvision.io.write_video(output_path, frames_tensor, fps=fps, video_codec="libx264")
return output_path
def process_example(video_input, video_caption, target_region_frame1_caption, prompt, click_state):
if video_input is None or video_input == "":
return (gr.update(value=""), gr.update(value=""), init_state(),
{"user_name": "", "video_name": "", "origin_images": None, "painted_images": None, "masks": None, "inpaint_masks": None, "logits": None, "select_frame_number": 0, "fps": 8, "ann_obj_id": 0},
"", None,
gr.update(value=1, visible=False, interactive=False),
gr.update(value=1, visible=False, interactive=False),
gr.update(value="Positive", interactive=False),
gr.update(visible=True, interactive=False),
gr.update(visible=True, interactive=False),
gr.update(value=None),
gr.update(visible=True, interactive=False),
create_status("Reset complete. Ready for new input.", StatusMessage.INFO),
gr.update(value=None))
video_state = gr.State({
"user_name": "",
"video_name": "",
"origin_images": None,
"painted_images": None,
"masks": None,
"inpaint_masks": None,
"logits": None,
"select_frame_number": 0,
"fps": 8,
"ann_obj_id": 0
})
results = get_frames_from_video(video_input, video_state)
if click_state[0] and click_state[1]:
print("Example detected, executing sam_refine")
(video_caption, target_region_frame1_caption, inference_state, video_state, video_info, template_frame, image_selection_slider, track_pause_number_slider, point_prompt, clear_button, tracking_button, video_output, inpaint_button, run_status, video_input) = results
class MockEvent:
def __init__(self, points, point_idx=0):
self.index = points[point_idx]
for i_click in range(len(click_state[0])):
evt = MockEvent(click_state[0], i_click)
prompt_type = "Positive" if click_state[1][i_click] == 1 else "Negative"
template_frame, video_state, interactive_state, run_status = sam_refine(inference_state, video_state, prompt_type, click_state, {"inference_times": 0, "negative_click_times": 0, "positive_click_times": 0, "mask_save": False, "multi_mask": {"mask_names": [], "masks": []}, "track_end_number": None}, evt, run_status)
return (video_caption, target_region_frame1_caption, inference_state, video_state, video_info, template_frame, image_selection_slider, track_pause_number_slider, point_prompt, clear_button, tracking_button, video_output, inpaint_button, run_status, video_input)
return results
def reset_all():
return (gr.update(value=None), gr.update(value=""), gr.update(value=""), init_state(),
{"user_name": "", "video_name": "", "origin_images": None, "painted_images": None, "masks": None, "inpaint_masks": None, "logits": None, "select_frame_number": 0, "fps": 8, "ann_obj_id": 0},
{"inference_times": 0, "negative_click_times": 0, "positive_click_times": 0, "mask_save": False, "multi_mask": {"mask_names": [], "masks": []}, "track_end_number": None},
[[], []], None, gr.update(visible=True, interactive=True), "",
gr.update(value=1, visible=False, interactive=False), gr.update(value=1, visible=False, interactive=False),
gr.update(value="Positive", interactive=False), gr.Button.update(interactive=False),
gr.Button.update(interactive=False), gr.Button.update(interactive=False),
gr.Button.update(interactive=False), gr.Button.update(interactive=False),
gr.Button.update(interactive=False), gr.Number.update(value=42),
gr.Slider.update(value=6.0), gr.Slider.update(value=16),
create_status("Reset complete. Ready for new input.", StatusMessage.INFO))
###############################
# Build Gradio Interface
###############################
title = """VideoPainter
"""
with gr.Blocks() as iface:
gr.HTML("""
""")
click_state = gr.State([[], []])
interactive_state = gr.State({
"inference_times": 0,
"negative_click_times": 0,
"positive_click_times": 0,
"mask_save": False,
"multi_mask": {"mask_names": [], "masks": []},
"track_end_number": None,
})
video_state = gr.State({
"user_name": "",
"video_name": "",
"origin_images": None,
"painted_images": None,
"masks": None,
"inpaint_masks": None,
"logits": None,
"select_frame_number": 0,
"fps": 8,
"ann_obj_id": 0
})
inference_state = init_state()
with gr.Row():
with gr.Column():
with gr.Row():
video_input = gr.Video(label="Original Video", visible=True)
with gr.Row():
with gr.Column(scale=3):
template_frame = gr.Image(type="pil", interactive=True, elem_id="template_frame", visible=True)
with gr.Column(scale=1):
with gr.Accordion("Segmentation Point Prompt", open=True):
point_prompt = gr.Radio(choices=["Positive", "Negative"], value="Positive", label="Point Type", interactive=False, visible=True)
clear_button_click = gr.Button(value="Clear clicks", interactive=False, visible=True)
gr.Markdown("✨ Positive: Include target region.
✨ Negative: Exclude target region.")
image_selection_slider = gr.Slider(minimum=1, maximum=100, step=1, value=1, label="Track start frame", visible=False, interactive=False)
track_pause_number_slider = gr.Slider(minimum=1, maximum=100, step=1, value=1, label="Track end frame", visible=False, interactive=False)
video_output = gr.Video(label="Generated Video", visible=True)
with gr.Row():
tracking_video_predict_button = gr.Button(value="Tracking", interactive=False, visible=True)
inpaint_video_predict_button = gr.Button(value="Inpainting", interactive=False, visible=True)
reset_button = gr.Button(value="Reset All", interactive=True, visible=True)
with gr.Column():
with gr.Accordion("Global Video Caption", open=True):
video_caption = gr.Textbox(label="Global Video Caption", placeholder="Input global video caption...", interactive=True, visible=True, max_lines=5, show_copy_button=True)
with gr.Row():
gr.Markdown("✨ Enhance prompt using GPT-4o (optional).")
enhance_button = gr.Button("✨ Enhance Prompt(Optional)", interactive=False)
with gr.Accordion("Target Object Caption", open=True):
target_region_frame1_caption = gr.Textbox(label="Target Object Caption", placeholder="Input target object caption...", interactive=True, visible=True, max_lines=5, show_copy_button=True)
with gr.Row():
gr.Markdown("✨ Generate target caption (optional).")
enhance_target_region_frame1_button = gr.Button("✨ Target Prompt Generation (Optional)", interactive=False)
with gr.Accordion("Editing Instruction", open=False):
gr.Markdown("✨ Modify captions based on your instruction using GPT-4o.")
with gr.Row():
editing_instruction = gr.Textbox(label="Editing Instruction", placeholder="Input editing instruction...", interactive=True, visible=True, max_lines=5, show_copy_button=True)
enhance_editing_instruction_button = gr.Button("✨ Modify Caption(For Editing)", interactive=False)
with gr.Accordion("Advanced Sampling Settings", open=False):
cfg_scale = gr.Slider(value=6.0, label="Classifier-Free Guidance Scale", minimum=1, maximum=10, step=0.1, interactive=True)
seed_param = gr.Number(label="Inference Seed (>=0)", interactive=True, value=42)
dilate_size = gr.Slider(value=16, label="Mask Dilate Size", minimum=0, maximum=32, step=1, interactive=True)
video_info = gr.Textbox(label="Video Info", visible=True, interactive=False)
model_type = gr.Textbox(label="Type", placeholder="Model type...", interactive=True, visible=False)
notes_accordion = gr.Accordion("Notes", open=False)
with notes_accordion:
gr.HTML("🧐 Reminder: VideoPainter may produce unexpected outputs. Adjust settings if needed.
")
run_status = gr.HighlightedText(value=[("", "")], visible=True, label="Operation Status", show_label=True,
color_map={"Success": "green", "Error": "red", "Warning": "orange", "Info": "blue"})
with gr.Row():
examples = gr.Examples(label="Quick Examples", examples=EXAMPLES,
inputs=[video_input, video_caption, target_region_frame1_caption, point_prompt, model_type, editing_instruction, seed_param, cfg_scale, dilate_size, click_state],
examples_per_page=20, cache_examples=False)
video_input.change(fn=process_example, inputs=[video_input, video_caption, target_region_frame1_caption, point_prompt, click_state],
outputs=[video_caption, target_region_frame1_caption, inference_state, video_state, video_info,
template_frame, image_selection_slider, track_pause_number_slider, point_prompt, clear_button_click,
tracking_video_predict_button, video_output, inpaint_video_predict_button, run_status, video_input])
image_selection_slider.release(fn=select_template, inputs=[image_selection_slider, video_state, interactive_state, run_status],
outputs=[template_frame, video_state, interactive_state, run_status])
track_pause_number_slider.release(fn=get_end_number, inputs=[track_pause_number_slider, video_state, interactive_state, run_status],
outputs=[template_frame, interactive_state, run_status])
template_frame.select(fn=sam_refine, inputs=[inference_state, video_state, point_prompt, click_state, interactive_state, run_status],
outputs=[template_frame, video_state, interactive_state, run_status])
tracking_video_predict_button.click(fn=vos_tracking_video, inputs=[inference_state, video_state, interactive_state, run_status],
outputs=[inference_state, video_output, video_state, interactive_state, run_status,
inpaint_video_predict_button, enhance_button, enhance_target_region_frame1_button, enhance_editing_instruction_button, notes_accordion])
inpaint_video_predict_button.click(fn=inpaint_video, inputs=[video_state, video_caption, target_region_frame1_caption, interactive_state, run_status, seed_param, cfg_scale, dilate_size],
outputs=[video_output, run_status], api_name=False, show_progress="full")
def enhance_prompt_func(video_caption):
return video_caption # Replace with your convert_prompt() if available
def enhance_target_region_frame1_prompt_func(target_region_frame1_caption, video_state):
return target_region_frame1_caption # Replace with your convert_prompt_target_region_frame1() if available
def enhance_editing_instruction_prompt_func(editing_instruction, video_caption, target_region_frame1_caption, video_state):
return video_caption, target_region_frame1_caption # Replace with your convert_prompt_editing_instruction() if available
enhance_button.click(enhance_prompt_func, inputs=[video_caption], outputs=[video_caption])
enhance_target_region_frame1_button.click(enhance_target_region_frame1_prompt_func, inputs=[target_region_frame1_caption, video_state], outputs=[target_region_frame1_caption])
enhance_editing_instruction_button.click(enhance_editing_instruction_prompt_func, inputs=[editing_instruction, video_caption, target_region_frame1_caption, video_state],
outputs=[video_caption, target_region_frame1_caption])
video_input.clear(fn=lambda: (gr.update(visible=True), gr.update(visible=True), init_state(),
{"user_name": "", "video_name": "", "origin_images": None, "painted_images": None, "masks": None, "inpaint_masks": None, "logits": None, "select_frame_number": 0, "fps": 8, "ann_obj_id": 0},
{"inference_times": 0, "negative_click_times": 0, "positive_click_times": 0, "mask_save": False, "multi_mask": {"mask_names": [], "masks": []}, "track_end_number": 0},
[[], []], None, None,
gr.update(visible=True), gr.update(visible=True), gr.update(visible=True), gr.update(visible=True),
gr.update(visible=True), gr.update(visible=True), gr.update(visible=True), gr.update(visible=True, value=[]),
gr.update(visible=True), gr.update(visible=True), gr.update(visible=True),
gr.Button.update(interactive=False), gr.Button.update(interactive=False), gr.Button.update(interactive=False)),
outputs=[video_caption, target_region_frame1_caption, inference_state, video_state, interactive_state, click_state, video_output, template_frame, tracking_video_predict_button, image_selection_slider, track_pause_number_slider, point_prompt, clear_button_click, template_frame, tracking_video_predict_button, video_output, inpaint_video_predict_button, run_status], queue=False, show_progress=False)
clear_button_click.click(fn=clear_click, inputs=[inference_state, video_state, click_state, run_status],
outputs=[inference_state, template_frame, click_state, run_status])
reset_button.click(fn=reset_all, inputs=[], outputs=[video_input, video_caption, target_region_frame1_caption, inference_state, video_state, interactive_state, click_state, video_output, template_frame, video_info, image_selection_slider, track_pause_number_slider, point_prompt, clear_button_click, tracking_video_predict_button, inpaint_video_predict_button, enhance_button, enhance_target_region_frame1_button, enhance_editing_instruction_button, seed_param, cfg_scale, dilate_size, run_status])
iface.queue().launch(share=False)