xinyun/ComfyUI
1
0
Fork 0
mirror of https://git.datalinker.icu/comfyanonymous/ComfyUI synced 2026-09-02 09:37:06 +08:00
Code Issues Packages Projects Releases Wiki Activity

Batch image/trajectory

Browse Source
This commit is contained in:
Eugene Fairley 2025-07-24 15:17:37 -07:00
parent f97f88f58b
commit 99c6fcaa1c
2 changed files with 44 additions and 78 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

20
comfy/utils.py
View File

@ -698,6 +698,26 @@ def resize_to_batch_size(tensor, batch_size):
return output return output
def resize_list_to_batch_size(l, batch_size):
in_batch_size = len(l)
if in_batch_size == batch_size or in_batch_size == 0:
return l
if batch_size <= 1:
return l[:batch_size]
output = []
if batch_size < in_batch_size:
scale = (in_batch_size - 1) / (batch_size - 1)
for i in range(batch_size):
output.append(l[min(round(i * scale), in_batch_size - 1)])
else:
scale = in_batch_size / batch_size
for i in range(batch_size):
output.append(l[min(math.floor((i + 0.5) * scale), in_batch_size - 1)])
return output
def convert_sd_to(state_dict, dtype): def convert_sd_to(state_dict, dtype):
keys = list(state_dict.keys()) keys = list(state_dict.keys())
for k in keys: for k in keys:

102
comfy_extras/nodes_wan.py
View File

@ -415,36 +415,6 @@ def parse_json_tracks(tracks):
tracks_data = [] tracks_data = []
return tracks_data return tracks_data
def tracks_to_tensor(tracks_data, length, width, height, batch_size=1):
"""Convert parsed track data to tensor format (B, T, N, 4)"""
if not tracks_data:
# Return empty tracks if no data
return torch.zeros((batch_size, length, 1, 4))
num_tracks = len(tracks_data)
tracks_tensor = torch.zeros((batch_size, length, num_tracks, 4))
for batch_idx in range(batch_size):
for track_idx, track in enumerate(tracks_data):
for frame_idx in range(min(length, len(track))):
point = track[frame_idx]
if isinstance(point, dict):
x = point.get('x', 0)
y = point.get('y', 0)
# Normalize coordinates to [-1, 1] range
x_norm = (x / width) * 2 - 1
y_norm = (y / height) * 2 - 1
visible = point.get('visible', 1)
tracks_tensor[batch_idx, frame_idx, track_idx] = torch.tensor([
track_idx, # track_id
x_norm, # x coordinate
y_norm, # y coordinate
visible # visibility
])
return tracks_tensor
def process_tracks(tracks_np: np.ndarray, frame_size: Tuple[int, int], num_frames, quant_multi: int = 8, **kwargs): def process_tracks(tracks_np: np.ndarray, frame_size: Tuple[int, int], num_frames, quant_multi: int = 8, **kwargs):
# tracks: shape [t, h, w, 3] => samples align with 24 fps, model trained with 16 fps. # tracks: shape [t, h, w, 3] => samples align with 24 fps, model trained with 16 fps.
# frame_size: tuple (W, H) # frame_size: tuple (W, H)
@ -577,7 +547,7 @@ def _patch_motion_single(
grid_mode = "bilinear" grid_mode = "bilinear"
point_feature = torch.nn.functional.grid_sample( point_feature = torch.nn.functional.grid_sample(
vid[vae_divide[0]:].permute(1, 0, 2, 3)[:1], vid.permute(1, 0, 2, 3)[:1],
tracks_n[:, :1].type(vid.dtype), tracks_n[:, :1].type(vid.dtype),
mode=grid_mode, mode=grid_mode,
padding_mode="zeros", padding_mode="zeros",
@ -589,9 +559,9 @@ def _patch_motion_single(
out_weight = vert_weight.sum(-1) # T - 1, H, W out_weight = vert_weight.sum(-1) # T - 1, H, W
# out feature -> already soft weighted # out feature -> already soft weighted
mix_feature = out_feature + vid[vae_divide[0]:, 1:] * (1 - out_weight.clamp(0, 1)) mix_feature = out_feature + vid[:, 1:] * (1 - out_weight.clamp(0, 1))
out_feature_full = torch.cat([vid[vae_divide[0]:, :1], mix_feature], dim=1) # C, T, H, W out_feature_full = torch.cat([vid[:, :1], mix_feature], dim=1) # C, T, H, W
out_mask_full = torch.cat([torch.ones_like(out_weight[:1]), out_weight], dim=0) # T, H, W out_mask_full = torch.cat([torch.ones_like(out_weight[:1]), out_weight], dim=0) # T, H, W
return out_mask_full[None].expand(vae_divide[0], -1, -1, -1), out_feature_full return out_mask_full[None].expand(vae_divide[0], -1, -1, -1), out_feature_full
@ -613,7 +583,7 @@ def patch_motion(
for b in range(B): for b in range(B):
mask, feature = _patch_motion_single( mask, feature = _patch_motion_single(
tracks[b], # (T, N, 4) tracks[b], # (T, N, 4)
vid, # (C, T, H, W) vid[b], # (C, T, H, W)
temperature, temperature,
vae_divide, vae_divide,
topk topk
@ -656,7 +626,6 @@ class WanTrackToVideo:
def encode(self, positive, negative, vae, tracks, width, height, length, batch_size, def encode(self, positive, negative, vae, tracks, width, height, length, batch_size,
temperature, topk, start_image=None, clip_vision_output=None): temperature, topk, start_image=None, clip_vision_output=None):
# Parse tracks from JSON
tracks_data = parse_json_tracks(tracks) tracks_data = parse_json_tracks(tracks)
if not tracks_data: if not tracks_data:
@ -664,7 +633,7 @@ class WanTrackToVideo:
latent = torch.zeros([batch_size, 16, ((length - 1) // 4) + 1, height // 8, width // 8], latent = torch.zeros([batch_size, 16, ((length - 1) // 4) + 1, height // 8, width // 8],
device=comfy.model_management.intermediate_device()) device=comfy.model_management.intermediate_device())
# Convert tracks to tensor format
if isinstance(tracks_data[0][0], dict): if isinstance(tracks_data[0][0], dict):
tracks_data = [tracks_data] tracks_data = [tracks_data]
@ -679,64 +648,41 @@ class WanTrackToVideo:
processed_tracks.append(process_tracks(tracks_np, (width, height), length - 1).unsqueeze(0)) processed_tracks.append(process_tracks(tracks_np, (width, height), length - 1).unsqueeze(0))
if start_image is not None: if start_image is not None:
start_image = comfy.utils.common_upscale(start_image[:length].movedim(-1, 1), width, height, "bilinear", "center").movedim(1, -1) start_image = comfy.utils.common_upscale(start_image[:batch_size].movedim(-1, 1), width, height, "bilinear", "center").movedim(1, -1)
videos = torch.ones((start_image.shape[0], length, height, width, start_image.shape[-1]), device=start_image.device, dtype=start_image.dtype) * 0.5
for i in range(start_image.shape[0]):
videos[i, 0] = start_image[i]
lat_h = height // 8 latent_videos = []
lat_w = width // 8 videos = comfy.utils.resize_to_batch_size(videos, batch_size)
for i in range(batch_size):
latent_videos += [vae.encode(videos[i, :, :, :, :3])]
y = torch.cat(latent_videos, dim=0)
msk = torch.ones(1, length, lat_h, lat_w, device=start_image.device)
msk[:, 1:] = 0
# repeat first frame 4 times
msk = torch.concat([
torch.repeat_interleave(msk[:, 0:1], repeats=4, dim=1), msk[:, 1:]
],
dim=1)
# Reshape mask into groups of 4 frames
msk = msk.view(1, msk.shape[1] // 4, 4, lat_h, lat_w)
# first batch
msk = msk.transpose(1, 2)
dummy_frames = torch.ones(3, length - 1, height, width) * .5
start_image = start_image.permute(3,0,1,2) # C, T, H, W
res = torch.concat([
start_image.to(start_image.device),
dummy_frames
],
dim=1).to(start_image.device)
res = res.permute(1,2,3,0)[:, :, :, :3] # T, H, W, C
y = vae.encode(res)
# Scale latent since patch_motion is non-linear # Scale latent since patch_motion is non-linear
y = comfy.latent_formats.Wan21().process_in(y) y = comfy.latent_formats.Wan21().process_in(y)
c = torch.cat((msk, y), dim=1)
processed_tracks = comfy.utils.resize_list_to_batch_size(processed_tracks, batch_size)
res = patch_motion( res = patch_motion(
processed_tracks, c[0], temperature=temperature, topk=topk, vae_divide=(4, 16) processed_tracks, y, temperature=temperature, topk=topk, vae_divide=(4, 16)
) )
msk, y = res mask, concat_latent_image = res
y = comfy.latent_formats.Wan21().process_out(y) concat_latent_image = comfy.latent_formats.Wan21().process_out(concat_latent_image)
mask = -mask + 1.0 # Invert mask to match expected format
msk = -msk + 1.0 # Invert mask to match expected format
positive = node_helpers.conditioning_set_values(positive, positive = node_helpers.conditioning_set_values(positive,
{"tracks": processed_tracks, {"tracks": processed_tracks,
"concat_mask": msk, "concat_mask": mask,
"concat_latent_image": y, "concat_latent_image": concat_latent_image,
"ati_temperature": temperature, "ati_temperature": temperature,
"ati_topk": topk}) "ati_topk": topk})
negative = node_helpers.conditioning_set_values(negative, negative = node_helpers.conditioning_set_values(negative,
{"tracks": processed_tracks, {"tracks": processed_tracks,
"concat_mask": msk, "concat_mask": mask,
"concat_latent_image": y, "concat_latent_image": concat_latent_image,
"ati_temperature": temperature, "ati_temperature": temperature,
"ati_topk": topk}) "ati_topk": topk})
# Handle clip vision output if provided
if clip_vision_output is not None: if clip_vision_output is not None:
positive = node_helpers.conditioning_set_values(positive, {"clip_vision_output": clip_vision_output}) positive = node_helpers.conditioning_set_values(positive, {"clip_vision_output": clip_vision_output})
negative = node_helpers.conditioning_set_values(negative, {"clip_vision_output": clip_vision_output}) negative = node_helpers.conditioning_set_values(negative, {"clip_vision_output": clip_vision_output})