Support Fun Control-Camera

This commit is contained in:
kijai
2025-04-29 20:41:27 +03:00
parent 4a41fb0aaf
commit 77f78ceded
5 changed files with 2146 additions and 19 deletions
+3
View File
@@ -3,15 +3,18 @@ from .recammaster.nodes import NODE_CLASS_MAPPINGS as RECAM_MASTER_NODE_CLASS_MA
from .unianimate.nodes import NODE_CLASS_MAPPINGS as UNIANIMATE_NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as UNIANIMATE_NODE_DISPLAY_NAME_MAPPINGS
from .skyreels.nodes import NODE_CLASS_MAPPINGS as SKYREELS_NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as SKYREELS_NODE_DISPLAY_NAME_MAPPINGS
from .fantasytalking.nodes import NODE_CLASS_MAPPINGS as FANTASYTALKING_NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as FANTASYTALKING_NODE_DISPLAY_NAME_MAPPINGS
from .fun_camera.nodes import NODE_CLASS_MAPPINGS as FUN_CAMERA_NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as FUN_CAMERA_NODE_DISPLAY_NAME_MAPPINGS
NODE_CLASS_MAPPINGS.update(RECAM_MASTER_NODE_CLASS_MAPPINGS)
NODE_CLASS_MAPPINGS.update(UNIANIMATE_NODE_CLASS_MAPPINGS)
NODE_CLASS_MAPPINGS.update(SKYREELS_NODE_CLASS_MAPPINGS)
NODE_CLASS_MAPPINGS.update(FANTASYTALKING_NODE_CLASS_MAPPINGS)
NODE_CLASS_MAPPINGS.update(FUN_CAMERA_NODE_CLASS_MAPPINGS)
NODE_DISPLAY_NAME_MAPPINGS.update(RECAM_MASTER_NODE_DISPLAY_NAME_MAPPINGS)
NODE_DISPLAY_NAME_MAPPINGS.update(UNIANIMATE_NODE_DISPLAY_NAME_MAPPINGS)
NODE_DISPLAY_NAME_MAPPINGS.update(SKYREELS_NODE_DISPLAY_NAME_MAPPINGS)
NODE_DISPLAY_NAME_MAPPINGS.update(FANTASYTALKING_NODE_DISPLAY_NAME_MAPPINGS)
NODE_DISPLAY_NAME_MAPPINGS.update(FUN_CAMERA_NODE_DISPLAY_NAME_MAPPINGS)
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
File diff suppressed because it is too large Load Diff
+186
View File
@@ -0,0 +1,186 @@
import numpy as np
import os
import torch
from einops import rearrange
script_directory = os.path.dirname(os.path.abspath(__file__))
class Camera(object):
"""Copied from https://github.com/hehao13/CameraCtrl/blob/main/inference.py
"""
def __init__(self, entry):
fx, fy, cx, cy = entry[1:5]
self.fx = fx
self.fy = fy
self.cx = cx
self.cy = cy
w2c_mat = np.array(entry[7:]).reshape(3, 4)
w2c_mat_4x4 = np.eye(4)
w2c_mat_4x4[:3, :] = w2c_mat
self.w2c_mat = w2c_mat_4x4
self.c2w_mat = np.linalg.inv(w2c_mat_4x4)
def custom_meshgrid(*args):
"""Copied from https://github.com/hehao13/CameraCtrl/blob/main/inference.py
"""
# ref: https://pytorch.org/docs/stable/generated/torch.meshgrid.html?highlight=meshgrid#torch.meshgrid
return torch.meshgrid(*args)
def get_relative_pose(cam_params):
"""Copied from https://github.com/hehao13/CameraCtrl/blob/main/inference.py
"""
abs_w2cs = [cam_param.w2c_mat for cam_param in cam_params]
abs_c2ws = [cam_param.c2w_mat for cam_param in cam_params]
cam_to_origin = 0
target_cam_c2w = np.array([
[1, 0, 0, 0],
[0, 1, 0, -cam_to_origin],
[0, 0, 1, 0],
[0, 0, 0, 1]
])
abs2rel = target_cam_c2w @ abs_w2cs[0]
ret_poses = [target_cam_c2w, ] + [abs2rel @ abs_c2w for abs_c2w in abs_c2ws[1:]]
ret_poses = np.array(ret_poses, dtype=np.float32)
return ret_poses
def ray_condition(K, c2w, H, W, device):
"""Copied from https://github.com/hehao13/CameraCtrl/blob/main/inference.py
"""
# c2w: B, V, 4, 4
# K: B, V, 4
B = K.shape[0]
j, i = custom_meshgrid(
torch.linspace(0, H - 1, H, device=device, dtype=c2w.dtype),
torch.linspace(0, W - 1, W, device=device, dtype=c2w.dtype),
)
i = i.reshape([1, 1, H * W]).expand([B, 1, H * W]) + 0.5 # [B, HxW]
j = j.reshape([1, 1, H * W]).expand([B, 1, H * W]) + 0.5 # [B, HxW]
fx, fy, cx, cy = K.chunk(4, dim=-1) # B,V, 1
zs = torch.ones_like(i) # [B, HxW]
xs = (i - cx) / fx * zs
ys = (j - cy) / fy * zs
zs = zs.expand_as(ys)
directions = torch.stack((xs, ys, zs), dim=-1) # B, V, HW, 3
directions = directions / directions.norm(dim=-1, keepdim=True) # B, V, HW, 3
rays_d = directions @ c2w[..., :3, :3].transpose(-1, -2) # B, V, 3, HW
rays_o = c2w[..., :3, 3] # B, V, 3
rays_o = rays_o[:, :, None].expand_as(rays_d) # B, V, 3, HW
# c2w @ dirctions
rays_dxo = torch.cross(rays_o, rays_d)
plucker = torch.cat([rays_dxo, rays_d], dim=-1)
plucker = plucker.reshape(B, c2w.shape[1], H, W, 6) # B, V, H, W, 6
# plucker = plucker.permute(0, 1, 4, 2, 3)
return plucker
def process_poses(poses, width=672, height=384, original_pose_width=1280, original_pose_height=720, device='cpu', return_poses=False):
"""Modified from https://github.com/hehao13/CameraCtrl/blob/main/inference.py
"""
cam_params = [[float(x) for x in pose] for pose in poses]
if return_poses:
return cam_params
else:
cam_params = [Camera(cam_param) for cam_param in cam_params]
sample_wh_ratio = width / height
pose_wh_ratio = original_pose_width / original_pose_height # Assuming placeholder ratios, change as needed
if pose_wh_ratio > sample_wh_ratio:
resized_ori_w = height * pose_wh_ratio
for cam_param in cam_params:
cam_param.fx = resized_ori_w * cam_param.fx / width
else:
resized_ori_h = width / pose_wh_ratio
for cam_param in cam_params:
cam_param.fy = resized_ori_h * cam_param.fy / height
intrinsic = np.asarray([[cam_param.fx * width,
cam_param.fy * height,
cam_param.cx * width,
cam_param.cy * height]
for cam_param in cam_params], dtype=np.float32)
K = torch.as_tensor(intrinsic)[None] # [1, 1, 4]
c2ws = get_relative_pose(cam_params) # Assuming this function is defined elsewhere
c2ws = torch.as_tensor(c2ws)[None] # [1, n_frame, 4, 4]
plucker_embedding = ray_condition(K, c2ws, height, width, device=device)[0].permute(0, 3, 1, 2).contiguous() # V, 6, H, W
plucker_embedding = plucker_embedding[None]
plucker_embedding = rearrange(plucker_embedding, "b f c h w -> b f h w c")[0]
return plucker_embedding
class WanVideoFunCameraEmbeds:
@classmethod
def INPUT_TYPES(s):
return {"required": {
"poses": ("CAMERACTRL_POSES", ),
"width": ("INT", {"default": 832, "min": 64, "max": 2048, "step": 8, "tooltip": "Width of the image to encode"}),
"height": ("INT", {"default": 480, "min": 64, "max": 29048, "step": 8, "tooltip": "Height of the image to encode"}),
"strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Strength of the camera motion"}),
"start_percent": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Start percent of the steps to apply camera motion"}),
"end_percent": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "End percent of the steps to apply camera motion"}),
},
# "optional": {
# "fun_ref_image": ("LATENT", {"tooltip": "Reference latent for the Fun 1.1 -model"}),
# }
}
RETURN_TYPES = ("WANVIDIMAGE_EMBEDS",)
RETURN_NAMES = ("image_embeds",)
FUNCTION = "process"
CATEGORY = "WanVideoWrapper"
def process(self, poses, width, height, strength, start_percent, end_percent, fun_ref_image=None):
num_frames = len(poses)
control_camera_video = process_poses(poses, width, height)
control_camera_video = control_camera_video.permute([3, 0, 1, 2]).unsqueeze(0)
print("control_camera_video.shape", control_camera_video.shape)
# Rearrange dimensions
# Concatenate and transpose dimensions
control_camera_latents = torch.concat(
[
torch.repeat_interleave(control_camera_video[:, :, 0:1], repeats=4, dim=2),
control_camera_video[:, :, 1:]
], dim=2
).transpose(1, 2)
# Reshape, transpose, and view into desired shape
b, f, c, h, w = control_camera_latents.shape
control_camera_latents = control_camera_latents.contiguous().view(b, f // 4, 4, c, h, w).transpose(2, 3)
control_camera_latents = control_camera_latents.contiguous().view(b, f // 4, c * 4, h, w).transpose(1, 2)
print("control_camera_latents.shape", control_camera_latents.shape)
vae_stride = (4, 8, 8)
target_shape = (16, (num_frames - 1) // vae_stride[0] + 1,
height // vae_stride[1],
width // vae_stride[2])
embeds = {
"target_shape": target_shape,
"num_frames": num_frames,
"control_embeds": {
"control_camera_latents": control_camera_latents * strength,
"control_camera_start_percent": start_percent,
"control_camera_end_percent": end_percent,
"fun_ref_image": fun_ref_image["samples"][:,:, 0] if fun_ref_image is not None else None,
}
}
return (embeds,)
NODE_CLASS_MAPPINGS = {
"WanVideoFunCameraEmbeds": WanVideoFunCameraEmbeds,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"WanVideoFunCameraEmbeds": "WanVideo FunCamera Embeds",
}
+40 -17
View File
@@ -559,6 +559,9 @@ class WanVideoModelLoader:
model_type = "i2v"
elif in_channels == 16:
model_type = "t2v"
elif "control_adapter.conv.weight" in sd:
model_type = "t2v"
num_heads = 40 if dim == 5120 else 12
num_layers = 40 if dim == 5120 else 30
@@ -625,8 +628,7 @@ class WanVideoModelLoader:
"add_ref_conv": True if "ref_conv.weight" in sd else False,
"in_dim_ref_conv": sd["ref_conv.weight"].shape[1] if "ref_conv.weight" in sd else None,
"add_control_adapter": True if "control_adapter.conv.weight" in sd else False,
"in_dim_control_adapter": sd["control_adapter.conv.weight"].shape[1] if "control_adapter.conv.weight" in sd else None,
}
}
with init_empty_weights():
transformer = WanModel(**TRANSFORMER_CONFIG)
@@ -668,7 +670,7 @@ class WanVideoModelLoader:
dtype = torch.float8_e5m2
else:
dtype = base_dtype
params_to_keep = {"norm", "head", "bias", "time_in", "vector_in", "patch_embedding", "time_", "img_emb", "modulation", "text_embedding"}
params_to_keep = {"norm", "head", "bias", "time_in", "vector_in", "patch_embedding", "time_", "img_emb", "modulation", "text_embedding", "adapter"}
#if lora is not None:
# transformer_load_device = device
if not lora_low_mem_load:
@@ -1588,13 +1590,13 @@ class WanVideoImageToVideoEncode:
"width": ("INT", {"default": 832, "min": 64, "max": 2048, "step": 8, "tooltip": "Width of the image to encode"}),
"height": ("INT", {"default": 480, "min": 64, "max": 29048, "step": 8, "tooltip": "Height of the image to encode"}),
"num_frames": ("INT", {"default": 81, "min": 1, "max": 10000, "step": 4, "tooltip": "Number of frames to encode"}),
"clip_embeds": ("WANVIDIMAGE_CLIPEMBEDS", {"tooltip": "Clip vision encoded image"}),
"noise_aug_strength": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 10.0, "step": 0.001, "tooltip": "Strength of noise augmentation, helpful for I2V where some noise can add motion and give sharper results"}),
"start_latent_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.001, "tooltip": "Additional latent multiplier, helpful for I2V where lower values allow for more motion"}),
"end_latent_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.001, "tooltip": "Additional latent multiplier, helpful for I2V where lower values allow for more motion"}),
"force_offload": ("BOOLEAN", {"default": True}),
},
"optional": {
"clip_embeds": ("WANVIDIMAGE_CLIPEMBEDS", {"tooltip": "Clip vision encoded image"}),
"start_image": ("IMAGE", {"tooltip": "Image to encode"}),
"end_image": ("IMAGE", {"tooltip": "end frame"}),
"control_embeds": ("WANVIDIMAGE_EMBEDS", {"tooltip": "Control signal for the Fun -model"}),
@@ -1609,9 +1611,9 @@ class WanVideoImageToVideoEncode:
FUNCTION = "process"
CATEGORY = "WanVideoWrapper"
def process(self, vae, width, height, num_frames, clip_embeds, force_offload, noise_aug_strength,
def process(self, vae, width, height, num_frames, force_offload, noise_aug_strength,
start_latent_strength, end_latent_strength, start_image=None, end_image=None, control_embeds=None, fun_or_fl2v_model=False,
temporal_mask=None, extra_latents=None):
temporal_mask=None, extra_latents=None, clip_embeds=None):
device = mm.get_torch_device()
offload_device = mm.unet_offload_device()
@@ -1722,8 +1724,8 @@ class WanVideoImageToVideoEncode:
image_embeds = {
"image_embeds": y,
"clip_context": clip_embeds.get("clip_embeds", None),
"negative_clip_context": clip_embeds.get("negative_clip_embeds", None),
"clip_context": clip_embeds.get("clip_embeds", None) if clip_embeds is not None else None,
"negative_clip_context": clip_embeds.get("negative_clip_embeds", None) if clip_embeds is not None else None,
"max_seq_len": max_seq_len,
"num_frames": num_frames,
"lat_h": lat_h,
@@ -2400,9 +2402,9 @@ class WanVideoSampler:
seed_g = torch.Generator(device=torch.device("cpu"))
seed_g.manual_seed(seed)
control_latents, clip_fea, clip_fea_neg, end_image, recammaster, camera_embed, unianim_data = None, None, None, None, None, None, None
vace_data, vace_context, vace_scale = None, None, None
fun_or_fl2v_model, has_ref, drop_last, = False, False, False
control_latents = control_camera_latents = clip_fea = clip_fea_neg = end_image = recammaster = camera_embed = unianim_data = None
vace_data = vace_context = vace_scale = None
fun_or_fl2v_model = has_ref = drop_last = False
phantom_latents = None
fun_ref_image = None
@@ -2435,9 +2437,12 @@ class WanVideoSampler:
control_embeds = image_embeds.get("control_embeds", None)
if control_embeds is not None:
if transformer.in_dim != 48:
if transformer.in_dim not in [48, 32]:
raise ValueError("Control signal only works with Fun-Control model")
control_latents = control_embeds["control_images"].to(device)
control_latents = control_embeds.get("control_images", None)
control_camera_latents = control_embeds.get("control_camera_latents", None)
control_camera_start_percent = control_embeds.get("control_camera_start_percent", 0.0)
control_camera_end_percent = control_embeds.get("control_camera_end_percent", 1.0)
control_start_percent = control_embeds.get("start_percent", 0.0)
control_end_percent = control_embeds.get("end_percent", 1.0)
drop_last = image_embeds.get("drop_last", False)
@@ -2498,7 +2503,15 @@ class WanVideoSampler:
control_embeds = image_embeds.get("control_embeds", None)
if control_embeds is not None:
control_latents = control_embeds["control_images"].to(device)
control_latents = control_embeds.get("control_images", None)
if control_latents is not None:
control_latents = control_latents.to(device)
control_camera_latents = control_embeds.get("control_camera_latents", None)
control_camera_start_percent = control_embeds.get("control_camera_start_percent", 0.0)
control_camera_end_percent = control_embeds.get("control_camera_end_percent", 1.0)
if control_camera_latents is not None:
control_camera_latents = control_camera_latents.to(device)
if control_lora:
image_cond = control_latents.to(device)
if not patcher.model.is_patched:
@@ -2506,9 +2519,9 @@ class WanVideoSampler:
patcher = apply_lora(patcher, device, device, low_mem_load=False)
patcher.model.is_patched = True
else:
if transformer.in_dim != 48:
if transformer.in_dim not in [48, 32]:
raise ValueError("Control signal only works with Fun-Control model")
image_cond = torch.zeros_like(control_latents).to(device) #fun control
image_cond = torch.zeros_like(noise).to(device) #fun control
clip_fea = None
fun_ref_image = control_embeds.get("fun_ref_image", None)
control_start_percent = control_embeds.get("start_percent", 0.0)
@@ -2691,6 +2704,8 @@ class WanVideoSampler:
for name, param in transformer.named_parameters():
if "block" not in name:
param.data = param.data.to(device)
if "control_adapter" in name:
param.data = param.data.to(device)
elif block_swap_args["offload_txt_emb"] and "txt_emb" in name:
param.data = param.data.to(offload_device, non_blocking=transformer.use_non_blocking)
elif block_swap_args["offload_img_emb"] and "img_emb" in name:
@@ -2846,10 +2861,17 @@ class WanVideoSampler:
if not patcher.model.is_patched:
log.info("Loading LoRA...")
patcher = apply_lora(patcher, device, device, low_mem_load=False)
patcher.model.is_patched = True
patcher.model.is_patched = True
else:
image_cond_input = image_cond.to(z) if image_cond is not None else None
if control_camera_latents is not None:
if (control_camera_start_percent <= current_step_percentage <= control_camera_end_percent) or \
(control_end_percent > 0 and idx == 0 and current_step_percentage >= control_camera_start_percent):
control_camera_input = control_camera_latents.to(z)
else:
control_camera_input = None
if recammaster is not None:
z = torch.cat([z, recam_latents.to(z)], dim=1)
use_phantom = False
@@ -2876,6 +2898,7 @@ class WanVideoSampler:
'camera_embed': camera_embed,
'unianim_data': unianim_data,
'fun_ref': fun_ref_input if fun_ref_image is not None else None,
'fun_camera': control_camera_input if control_camera_latents is not None else None,
'audio_proj': audio_proj if fantasytalking_embeds is not None else None,
'audio_context_lens': audio_context_lens if fantasytalking_embeds is not None else None,
'audio_scale': audio_scale if fantasytalking_embeds is not None else None,
+2 -2
View File
@@ -1066,7 +1066,7 @@ class WanModel(ModelMixin, ConfigMixin):
camera_embed=None,
unianim_data=None,
fps_embeds=None,
fun_ref = None,
fun_ref=None,
fun_camera=None,
audio_proj=None,
audio_context_lens=None,
@@ -1214,7 +1214,7 @@ class WanModel(ModelMixin, ConfigMixin):
self.text_embedding.to(self.offload_device, non_blocking=self.use_non_blocking)
clip_embed = None
if clip_fea is not None:
if clip_fea is not None and hasattr(self, "img_emb"):
clip_fea = clip_fea.to(self.main_device)
if self.offload_img_emb:
self.img_emb.to(self.main_device)