wehub-resource-sync
383 lines
14 KiB
Python
Executable File
383 lines
14 KiB
Python
Executable File
# Copyright 2026, MiLM Plus, Xiaomi Inc.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
import os
|
|
import argparse
|
|
import numpy as np
|
|
import torch
|
|
from diffusers import FlowMatchEulerDiscreteScheduler
|
|
from omegaconf import OmegaConf
|
|
from PIL import Image
|
|
from transformers import AutoTokenizer
|
|
import scipy
|
|
import cv2
|
|
from glob import glob
|
|
import torch.distributed as dist
|
|
from videox_fun.dist import set_multi_gpus_devices
|
|
from videox_fun.models import AutoencoderKLWan, WanT5EncoderModel, VaceWanModel
|
|
from videox_fun.pipeline import SVORPipeline
|
|
from videox_fun.utils.fp8_optimization import (
|
|
convert_model_weight_to_float8,
|
|
replace_parameters_by_name,
|
|
convert_weight_dtype_wrapper,
|
|
)
|
|
from videox_fun.utils.lora_utils import merge_lora
|
|
from videox_fun.utils.utils import save_videos_grid, filter_kwargs
|
|
|
|
|
|
def load_patch_safetensors(path):
|
|
list_tensors = glob(path + "/*.safetensors")
|
|
all = {}
|
|
for x in list_tensors:
|
|
from safetensors.torch import load_file
|
|
|
|
tmp = load_file(x)
|
|
all.update(tmp)
|
|
return all
|
|
|
|
|
|
def parse_args():
|
|
parser = argparse.ArgumentParser(description="WanFun Video Editing Script")
|
|
|
|
# GPU and memory configuration
|
|
parser.add_argument(
|
|
"--gpu_memory_mode",
|
|
type=str,
|
|
default="model_full_load",
|
|
choices=["model_full_load", "model_cpu_offload", "model_cpu_offload_and_qfloat8", "sequential_cpu_offload"],
|
|
help="GPU memory optimization mode",
|
|
)
|
|
parser.add_argument("--ulysses_degree", type=int, default=1, help="Ulysses degree for multi-GPU configuration")
|
|
parser.add_argument("--ring_degree", type=int, default=1, help="Ring degree for multi-GPU configuration")
|
|
|
|
# Model paths
|
|
parser.add_argument(
|
|
"--config_path", type=str, default="config/wan2.1/wan_civitai.yaml", help="Path to model configuration file"
|
|
)
|
|
parser.add_argument("--model_name", type=str, default="models/Wan2.1-VACE-1.3B", help="Path to pretrained model")
|
|
|
|
# Generation parameters
|
|
parser.add_argument("--sample_size", type=str, default="720,1280", help="Output size as 'height,width'")
|
|
parser.add_argument("--video_length", type=int, default=81, help="Length of generated video in frames")
|
|
parser.add_argument("--fps", type=int, default=16, help="Frames per second for output video")
|
|
parser.add_argument(
|
|
"--weight_dtype",
|
|
type=str,
|
|
default="bfloat16",
|
|
choices=["float16", "bfloat16"],
|
|
help="Data type for model weights",
|
|
)
|
|
|
|
# Prompt and generation settings
|
|
parser.add_argument(
|
|
"--prompt",
|
|
type=str,
|
|
default="Remove the target and fill the content appropriately",
|
|
help="Text prompt for generation",
|
|
)
|
|
parser.add_argument(
|
|
"--negative_prompt",
|
|
type=str,
|
|
default="色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走",
|
|
help="Negative text prompt",
|
|
)
|
|
parser.add_argument("--guidance_scale", type=float, default=6.0, help="Guidance scale for generation")
|
|
parser.add_argument("--seed", type=int, default=43, help="Random seed for reproducibility")
|
|
parser.add_argument("--context_scale", type=float, default=1.0, help="Context scale for vace control")
|
|
parser.add_argument("--dilation", type=int, default=6, help="Dilation for inp mask (only for inpaint mode)")
|
|
|
|
# Parameters for SVOR
|
|
parser.add_argument(
|
|
"--lora_path",
|
|
type=str,
|
|
default=["models/remove_model_stage1.safetensors", "models/remove_model_stage2.safetensors"],
|
|
nargs="+",
|
|
help="Optional path to LoRA checkpoint",
|
|
)
|
|
parser.add_argument(
|
|
"--lora_weight", type=float, default=[1.0, 1.0], nargs="+", help="Weight for LoRA model if used"
|
|
)
|
|
parser.add_argument("--input_video", type=str, default=None, required=True, help="Path to input video for editing")
|
|
parser.add_argument(
|
|
"--input_mask_video", type=str, default=None, required=True, help="Path to mask video for editing"
|
|
)
|
|
parser.add_argument("--num_inference_steps", type=int, default=20, help="Number of inference steps")
|
|
parser.add_argument("--save_dir", type=str, default="samples/SVOR", help="Directory to save generated videos")
|
|
|
|
return parser.parse_args()
|
|
|
|
|
|
def process_video(
|
|
input_video_path,
|
|
input_mask_video_path,
|
|
video_length,
|
|
sample_size,
|
|
dilation=0,
|
|
):
|
|
"""Process input video and mask for editing"""
|
|
if input_video_path is not None:
|
|
cap = cv2.VideoCapture(input_video_path)
|
|
frames = []
|
|
while cap.isOpened():
|
|
ret, frame = cap.read()
|
|
if not ret:
|
|
break
|
|
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
|
frames.append(Image.fromarray(frame))
|
|
cap.release()
|
|
|
|
frames = frames[:video_length]
|
|
if len(frames) < video_length:
|
|
frames += [frames[-1]] * (video_length - len(frames))
|
|
|
|
resized_frames = [frame.resize([sample_size[1], sample_size[0]]) for frame in frames]
|
|
|
|
input_video = (
|
|
torch.stack([torch.from_numpy(np.array(frame)).permute(2, 0, 1) for frame in resized_frames])
|
|
.permute(1, 0, 2, 3)
|
|
.unsqueeze(0)
|
|
) # [1, C, T, H, W]
|
|
|
|
else:
|
|
input_video = torch.zeros((1, 3, video_length, sample_size[0], sample_size[1])).float()
|
|
if input_mask_video_path is not None:
|
|
mask_cap = cv2.VideoCapture(input_mask_video_path)
|
|
mask_frames = []
|
|
while mask_cap.isOpened():
|
|
ret, frame = mask_cap.read()
|
|
if not ret:
|
|
break
|
|
if len(frame.shape) == 3:
|
|
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
|
|
_, mask = cv2.threshold(frame, 127, 255, cv2.THRESH_BINARY)
|
|
if dilation > 0:
|
|
mask_np = (mask > 0).astype(np.uint8)
|
|
mask = scipy.ndimage.binary_dilation(mask_np, iterations=dilation).astype(np.uint8) * 255
|
|
mask_frames.append(mask)
|
|
mask_cap.release()
|
|
|
|
mask_frames = mask_frames[:video_length]
|
|
if len(mask_frames) < video_length:
|
|
mask_frames += [mask_frames[-1]] * (video_length - len(mask_frames))
|
|
|
|
resized_masks = [Image.fromarray(mask).resize([sample_size[1], sample_size[0]]) for mask in mask_frames]
|
|
|
|
input_video_mask = (
|
|
torch.stack([torch.from_numpy(np.array(mask)) for mask in resized_masks]).unsqueeze(0).unsqueeze(0) / 255.0
|
|
) # [1, 1, T, H, W]
|
|
|
|
else:
|
|
input_video_mask = torch.ones((1, 1, video_length, sample_size[0], sample_size[1])).float()
|
|
|
|
if input_video_path is not None and input_video is not None:
|
|
input_video = input_video * (torch.tile(input_video_mask, [1, 3, 1, 1, 1]) < 0.5) + (128.0) * (
|
|
torch.tile(input_video_mask, [1, 3, 1, 1, 1]) >= 0.5
|
|
)
|
|
input_video = input_video.div_(127.5).sub_(1.0)
|
|
|
|
return input_video, input_video_mask
|
|
|
|
|
|
def process_single_task(
|
|
pipeline,
|
|
args,
|
|
input_video_path,
|
|
input_mask_video_path,
|
|
prompt,
|
|
):
|
|
"""Process a single video editing task"""
|
|
if input_video_path is not None:
|
|
# Get video resolution
|
|
cap = cv2.VideoCapture(input_video_path)
|
|
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
|
|
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
|
cap.release()
|
|
|
|
# Calculate aspect ratio preserving size
|
|
aspect_ratio = height / width
|
|
max_area = int(args.sample_size.split(",")[0]) * int(args.sample_size.split(",")[1])
|
|
new_height = round(np.sqrt(max_area * aspect_ratio))
|
|
new_height = (new_height + 16 - 1) // 16 * 16
|
|
new_width = round(np.sqrt(max_area / aspect_ratio))
|
|
new_width = (new_width + 16 - 1) // 16 * 16
|
|
sample_size = [new_height, new_width]
|
|
else:
|
|
sample_size = [int(args.sample_size.split(",")[0]), int(args.sample_size.split(",")[1])]
|
|
|
|
generator = torch.Generator(device=pipeline.device).manual_seed(args.seed)
|
|
|
|
with torch.no_grad():
|
|
video_length = (
|
|
int(
|
|
(args.video_length - 1)
|
|
// pipeline.vae.config.temporal_compression_ratio
|
|
* pipeline.vae.config.temporal_compression_ratio
|
|
)
|
|
+ 1
|
|
if args.video_length != 1
|
|
else 1
|
|
)
|
|
|
|
# Process video and mask
|
|
(
|
|
input_video,
|
|
input_video_mask,
|
|
) = process_video(
|
|
input_video_path,
|
|
input_mask_video_path,
|
|
video_length=video_length,
|
|
sample_size=sample_size,
|
|
dilation=args.dilation,
|
|
)
|
|
|
|
# Generate edited video
|
|
sample = pipeline(
|
|
prompt,
|
|
negative_prompt=args.negative_prompt,
|
|
height=sample_size[0],
|
|
width=sample_size[1],
|
|
generator=generator,
|
|
guidance_scale=args.guidance_scale,
|
|
num_inference_steps=args.num_inference_steps,
|
|
video=input_video,
|
|
mask_video=input_video_mask,
|
|
context_scale=args.context_scale,
|
|
).videos
|
|
|
|
return sample, video_length
|
|
|
|
|
|
def save_results(sample, args, video_length, fps, task_name=None):
|
|
"""Save the generated results"""
|
|
if not os.path.exists(args.save_dir):
|
|
os.makedirs(args.save_dir, exist_ok=True)
|
|
|
|
prefix = task_name
|
|
if video_length == 1:
|
|
video_path = os.path.join(args.save_dir, prefix + ".png")
|
|
image = sample[0, :, 0]
|
|
image = image.transpose(0, 1).transpose(1, 2)
|
|
image = (image * 255).numpy().astype(np.uint8)
|
|
image = Image.fromarray(image)
|
|
image.save(video_path)
|
|
else:
|
|
video_path = os.path.join(args.save_dir, prefix + ".mp4")
|
|
save_videos_grid(sample, video_path, fps=fps)
|
|
|
|
|
|
def main():
|
|
args = parse_args()
|
|
|
|
# Validate arguments
|
|
if args.input_video is None and args.input_mask_video is None:
|
|
raise ValueError("Must provide either --input_video and --input_mask_video")
|
|
|
|
# Convert weight dtype
|
|
weight_dtype = torch.bfloat16 if args.weight_dtype == "bfloat16" else torch.float16
|
|
device = set_multi_gpus_devices(args.ulysses_degree, args.ring_degree)
|
|
config = OmegaConf.load(args.config_path)
|
|
|
|
# Initialize transformer
|
|
transformer = VaceWanModel.from_pretrained(
|
|
os.path.join(
|
|
args.model_name, config["transformer_additional_kwargs"].get("transformer_subpath", "transformer")
|
|
),
|
|
transformer_additional_kwargs=OmegaConf.to_container(config["transformer_additional_kwargs"]),
|
|
low_cpu_mem_usage=True,
|
|
torch_dtype=weight_dtype,
|
|
)
|
|
|
|
# Get Vae
|
|
vae = AutoencoderKLWan.from_pretrained(
|
|
os.path.join(args.model_name, config["vae_kwargs"].get("vae_subpath", "vae")),
|
|
additional_kwargs=OmegaConf.to_container(config["vae_kwargs"]),
|
|
).to(weight_dtype)
|
|
|
|
# Get Tokenizer
|
|
tokenizer = AutoTokenizer.from_pretrained(
|
|
os.path.join(args.model_name, config["text_encoder_kwargs"].get("tokenizer_subpath", "tokenizer")),
|
|
)
|
|
|
|
# Get Text encoder
|
|
text_encoder = WanT5EncoderModel.from_pretrained(
|
|
os.path.join(args.model_name, config["text_encoder_kwargs"].get("text_encoder_subpath", "text_encoder")),
|
|
additional_kwargs=OmegaConf.to_container(config["text_encoder_kwargs"]),
|
|
).to(weight_dtype)
|
|
text_encoder = text_encoder.eval()
|
|
|
|
# Get Scheduler
|
|
Choosen_Scheduler = {
|
|
"Flow": FlowMatchEulerDiscreteScheduler,
|
|
}["Flow"]
|
|
scheduler = Choosen_Scheduler(
|
|
**filter_kwargs(Choosen_Scheduler, OmegaConf.to_container(config["scheduler_kwargs"]))
|
|
)
|
|
|
|
# Get Pipeline
|
|
pipeline = SVORPipeline(
|
|
transformer=transformer,
|
|
vae=vae,
|
|
tokenizer=tokenizer,
|
|
text_encoder=text_encoder,
|
|
scheduler=scheduler,
|
|
)
|
|
|
|
if args.ulysses_degree > 1 or args.ring_degree > 1:
|
|
transformer.enable_multi_gpus_inference()
|
|
|
|
if args.gpu_memory_mode == "sequential_cpu_offload":
|
|
replace_parameters_by_name(
|
|
transformer,
|
|
[
|
|
"modulation",
|
|
],
|
|
device=device,
|
|
)
|
|
transformer.freqs = transformer.freqs.to(device=device)
|
|
pipeline.enable_sequential_cpu_offload(device=device)
|
|
elif args.gpu_memory_mode == "model_cpu_offload_and_qfloat8":
|
|
convert_model_weight_to_float8(
|
|
transformer,
|
|
exclude_module_name=[
|
|
"modulation",
|
|
],
|
|
)
|
|
convert_weight_dtype_wrapper(transformer, weight_dtype)
|
|
pipeline.enable_model_cpu_offload(device=device)
|
|
elif args.gpu_memory_mode == "model_cpu_offload":
|
|
pipeline.enable_model_cpu_offload(device=device)
|
|
else:
|
|
pipeline.to(device=device)
|
|
|
|
if args.lora_path is not None:
|
|
if len(args.lora_weight) != len(args.lora_path):
|
|
args.lora_weight = [args.lora_weight[0]] * len(args.lora_path)
|
|
for lora_path, lora_weight in zip(args.lora_path, args.lora_weight):
|
|
print(f"[INFO] Loading LoRA: {lora_path}, weight: {lora_weight}")
|
|
pipeline = merge_lora(pipeline, lora_path, lora_weight)
|
|
|
|
if not os.path.exists(args.save_dir):
|
|
os.makedirs(args.save_dir, exist_ok=True)
|
|
|
|
# Single task processing
|
|
sample, video_length = process_single_task(pipeline, args, args.input_video, args.input_mask_video, args.prompt)
|
|
|
|
video_basename = os.path.splitext(os.path.basename(args.input_video))[0]
|
|
if not dist.is_initialized() or dist.get_rank() == 0:
|
|
save_results(sample, args, video_length, args.fps, task_name=video_basename)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|