385 lines
24 KiB
Python
385 lines
24 KiB
Python
import torch
|
|
from typing import List, Tuple
|
|
import gc
|
|
|
|
# ComfyUI Imports
|
|
import comfy.samplers
|
|
import node_helpers
|
|
import comfy.clip_vision
|
|
import comfy.utils
|
|
import comfy.model_management
|
|
import comfy.model_sampling
|
|
import nodes
|
|
|
|
# ====================================================================================================
|
|
# Logging Utility
|
|
# ====================================================================================================
|
|
def _log(msg: str):
|
|
"""Simple print logger for the console."""
|
|
try:
|
|
print(f"[WanVideoLooper] {msg}", flush=True)
|
|
except Exception:
|
|
pass
|
|
|
|
# ====================================================================================================
|
|
# Core Tensor Utilities
|
|
# ====================================================================================================
|
|
def _align_to_multiple(x, multiple=16):
|
|
"""Aligns a number to the nearest larger multiple."""
|
|
return max(multiple, (int(x) // multiple) * multiple)
|
|
|
|
def _get_model_device_info(model) -> Tuple[torch.device, torch.dtype]:
|
|
"""Gets the device and dtype from a model."""
|
|
try:
|
|
model_device = next(model.model.parameters()).device
|
|
model_dtype = next(model.model.parameters()).dtype
|
|
except Exception:
|
|
model_device = comfy.model_management.intermediate_device()
|
|
model_dtype = torch.float32
|
|
return model_device, model_dtype
|
|
|
|
def _move_tensor(tensor: torch.Tensor, device, dtype):
|
|
"""Moves a tensor to a device/dtype if it's not already there."""
|
|
if tensor.device != device or tensor.dtype != dtype:
|
|
return tensor.to(device=device, dtype=dtype, non_blocking=True)
|
|
return tensor
|
|
|
|
# ====================================================================================================
|
|
# Model & Sampling Helpers (MoE Logic)
|
|
# ====================================================================================================
|
|
def _apply_sigma_shift(model, sigma_shift):
|
|
"""Applies sigma shift to the model's sampling parameters."""
|
|
model_sampler = model.get_model_object("model_sampling")
|
|
if not model_sampler:
|
|
class CustomModelSampling(comfy.model_sampling.ModelSamplingDiscreteFlow, comfy.model_sampling.CONST): pass
|
|
model_sampler = CustomModelSampling(model.model.model_config)
|
|
model_sampler.set_parameters(shift=sigma_shift, multiplier=1000)
|
|
model.add_object_patch("model_sampling", model_sampler)
|
|
return model
|
|
|
|
def _find_model_switch_index(model, scheduler_name, total_steps, boundary_point):
|
|
"""Calculates the step index to switch models based on the boundary."""
|
|
try:
|
|
model_sampler = model.get_model_object("model_sampling")
|
|
sigmas = comfy.samplers.calculate_sigmas(model_sampler, scheduler_name, total_steps)
|
|
timesteps = [model_sampler.timestep(sigma) / 1000.0 for sigma in sigmas.tolist()]
|
|
for i, t in enumerate(timesteps):
|
|
if t < boundary_point:
|
|
switch_index = max(0, i - 1)
|
|
return switch_index
|
|
return total_steps
|
|
except Exception as e:
|
|
_log(f"Warning: Error calculating model switch point, defaulting to 50/50 split. Details: {e}")
|
|
return total_steps // 2
|
|
|
|
def _run_sampler_step(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent, denoise, disable_noise=False, start_step=None, last_step=None, force_full_denoise=False):
|
|
"""Thin wrapper for the comfy KSampler node call."""
|
|
latent_dict = {"samples": latent["samples"]}
|
|
output, = nodes.common_ksampler(
|
|
model=model, seed=int(seed), steps=int(steps), cfg=float(cfg),
|
|
sampler_name=str(sampler_name), scheduler=str(scheduler),
|
|
positive=positive, negative=negative, latent=latent_dict,
|
|
denoise=float(denoise), disable_noise=disable_noise,
|
|
start_step=start_step, last_step=last_step,
|
|
force_full_denoise=force_full_denoise
|
|
)
|
|
return {"samples": output["samples"]}
|
|
|
|
# ====================================================================================================
|
|
# Conditioning & Latent Helpers
|
|
# ====================================================================================================
|
|
def _encode_seed_image_vision(clip_vision, seed_image, cropping="center"):
|
|
"""Encodes the seed image using CLIP Vision, handling potential errors."""
|
|
if clip_vision is None: return None
|
|
if not isinstance(seed_image, torch.Tensor): return None
|
|
image_tensor = seed_image.to(torch.float32).clamp(0, 1)
|
|
if image_tensor.dim() == 3: image_tensor = image_tensor.unsqueeze(0)
|
|
if image_tensor.dim() == 4 and image_tensor.shape[1] in (1, 3, 4):
|
|
image_tensor = image_tensor.permute(0, 2, 3, 1).contiguous()
|
|
if image_tensor.dim() != 4 or image_tensor.shape[-1] not in (1, 3, 4):
|
|
_log(f"Warning: Invalid image shape for CLIP Vision: {tuple(image_tensor.shape)}")
|
|
return None
|
|
vision_output = clip_vision.encode_image(image_tensor, crop=(cropping == "center"))
|
|
return vision_output
|
|
|
|
def _tokenize_and_encode(clip, positive_prompt, negative_prompt):
|
|
"""Encodes positive and negative text prompts using the provided CLIP model."""
|
|
pos_tokens = clip.tokenize(positive_prompt or "")
|
|
pos_data = clip.encode_from_tokens(pos_tokens, return_pooled=True, return_dict=True)
|
|
positive_cond = [(pos_data["cond"], {"pooled_output": pos_data.get("pooled_output")})]
|
|
neg_tokens = clip.tokenize(negative_prompt or "")
|
|
neg_data = clip.encode_from_tokens(neg_tokens, return_pooled=True, return_dict=True)
|
|
negative_cond = [(neg_data["cond"], {"pooled_output": neg_data.get("pooled_output")})]
|
|
return positive_cond, negative_cond
|
|
|
|
def _prepare_latent_window(positive_cond, negative_cond, vae, width, height, frame_count, batch_size, seed_image=None, vision_output=None):
|
|
"""Creates the latent tensor window and conditions it with seed image/CLIP Vision."""
|
|
device = comfy.model_management.intermediate_device()
|
|
latent_height = height // 8
|
|
latent_width = width // 8
|
|
time_dim = ((frame_count - 1) // 4) + 1
|
|
empty_latent = torch.zeros([batch_size, 16, time_dim, latent_height, latent_width], device=device)
|
|
|
|
if seed_image is not None:
|
|
seed_image_upscaled = comfy.utils.common_upscale(seed_image[:frame_count].movedim(-1, 1), width, height, "bilinear", "center").movedim(1, -1)
|
|
image_padded = torch.ones((frame_count, height, width, seed_image_upscaled.shape[-1]), device=seed_image_upscaled.device, dtype=seed_image_upscaled.dtype) * 0.5
|
|
image_padded[:seed_image_upscaled.shape[0]] = seed_image_upscaled
|
|
seed_image_latent = vae.encode(image_padded[:, :, :, :3])
|
|
mask_time_dim = int(seed_image_latent.size(2))
|
|
mask = torch.ones((1, 1, mask_time_dim, latent_height, latent_width), device=seed_image.device, dtype=seed_image.dtype)
|
|
zero_frames_dim = min(mask_time_dim, ((seed_image_upscaled.shape[0] - 1) // 4) + 1)
|
|
mask[:, :, :zero_frames_dim] = 0.0
|
|
positive_cond = node_helpers.conditioning_set_values(positive_cond, {"concat_latent_image": seed_image_latent, "concat_mask": mask})
|
|
negative_cond = node_helpers.conditioning_set_values(negative_cond, {"concat_latent_image": seed_image_latent, "concat_mask": mask})
|
|
|
|
if vision_output is not None:
|
|
positive_cond = node_helpers.conditioning_set_values(positive_cond, {"clip_vision_output": vision_output})
|
|
negative_cond = node_helpers.conditioning_set_values(negative_cond, {"clip_vision_output": vision_output})
|
|
|
|
return positive_cond, negative_cond, {"samples": empty_latent}
|
|
|
|
# ====================================================================================================
|
|
# Main WanVideoLooper Node Class
|
|
# ====================================================================================================
|
|
class WanVideoLooper:
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model_high": ("MODEL", {"tooltip": "The BASE high-noise model (can be pre-patched with global LoRAs)."}),
|
|
"model_low": ("MODEL", {"tooltip": "The BASE low-noise model (can be pre-patched with global LoRAs)."}),
|
|
"clip": ("CLIP", {"tooltip": "The BASE CLIP model (can be pre-patched with global LoRAs)."}),
|
|
"vae": ("VAE", {"tooltip": "The VAE model used for decoding the final image."}),
|
|
"start_image": ("IMAGE", {"tooltip": "The seed image that the first loop will start from."}),
|
|
"positive_prompts": ("*", {"forceInput": True, "tooltip": "Connect the 'prompt_list' output from the WanVideo Multi-Prompt node."}),
|
|
"negative_prompt": ("STRING", {"tooltip": "A single, global negative prompt applied to all loops."}),
|
|
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "tooltip": "The random seed. This will be used for all loops."}),
|
|
"steps": ("INT", {"default": 6, "min": 1, "max": 10000, "tooltip": "The total number of sampling steps for each loop."}),
|
|
"enable_motion_cfg": ("BOOLEAN", {"default": False, "tooltip": "Enable a separate CFG for the very first step to control motion."}),
|
|
"cfg_motion_step": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.1, "round": 0.01, "tooltip": "CFG for the first step ONLY. A value > 1.0 will engage the negative prompt for initial motion guidance."}),
|
|
"cfg_high_noise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.1, "round": 0.01, "tooltip": "CFG for the high-noise model (structure/motion). 1.0 disables guidance."}),
|
|
"cfg_low_noise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.1, "round": 0.01, "tooltip": "CFG for the low-noise model (details). 1.0 disables guidance."}),
|
|
"sampler_name": (comfy.samplers.KSampler.SAMPLERS, {"tooltip": "The sampling algorithm (e.g., dpmpp_2m_sde_gpu)."}),
|
|
"scheduler": (comfy.samplers.KSampler.SCHEDULERS, {"tooltip": "The noise scheduler (e.g., karras)."}),
|
|
"model_switch_point": ("FLOAT", {"default": 0.9, "min": 0.0, "max": 1.0, "step": 0.001, "round": 0.001, "tooltip": "Timestep to switch from high-noise to low-noise model. 0.9 (I2V) or 0.875 (T2V) recommended."}),
|
|
"sigma_shift": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step":0.01, "tooltip": "Applies sigma shift to both models. 8.0 is a good default."}),
|
|
"denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Denoise amount. 1.0 means generate a full new image."}),
|
|
|
|
"width": ("INT", {"default": 832, "min": 64, "max": 8192, "step": 16, "tooltip": "The width of the output video."}),
|
|
"height": ("INT", {"default": 480, "min": 64, "max": 8192, "step": 16, "tooltip": "The height of the output video."}),
|
|
"frame_merge": ("INT", {"default": 1, "min": 0, "max": 8, "step": 1, "tooltip": "Number of frames to overlap/merge between loops for smooth transitions."}),
|
|
"duration_sec": ("INT", {"default": 5, "min": 1, "max": 10, "step": 1, "tooltip": "Duration of each segment in seconds. 5s is recommended, max 10."}),
|
|
"enable_dry_run": ("BOOLEAN", {"default": False, "tooltip": "If enabled, skips sampling and decoding to quickly check setup and logs."}),
|
|
},
|
|
"optional": {
|
|
"clip_vision": ("CLIP_VISION", {"tooltip": "(Optional) A CLIP Vision model for guiding the start image."}),
|
|
"model_clip_sequence": ("ANY", {"forceInput": True, "tooltip": "(Optional) Connect a WanVideo Lora Sequencer to use different pre-patched models/clips per segment."}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "IMAGE")
|
|
RETURN_NAMES = ("images", "last_frame")
|
|
FUNCTION = "loop_video"
|
|
CATEGORY = "WanVideoLooper"
|
|
|
|
# --- Main Looping Function ---
|
|
def loop_video(
|
|
self,
|
|
model_high, model_low, clip, vae,
|
|
start_image, positive_prompts, negative_prompt,
|
|
seed, steps, enable_motion_cfg, cfg_motion_step, cfg_high_noise, cfg_low_noise,
|
|
sampler_name, scheduler, model_switch_point, sigma_shift, denoise,
|
|
width, height, frame_merge, duration_sec, enable_dry_run,
|
|
clip_vision=None,
|
|
model_clip_sequence=None
|
|
):
|
|
|
|
# --- 1. Validation & Initial Setup ---
|
|
if not isinstance(positive_prompts, list):
|
|
_log("Error: 'positive_prompts' is not a valid list..."); return (torch.zeros((1, 64, 64, 3)),)*2
|
|
|
|
# --- Check if model_clip_sequence is valid, ignore if bypassed ---
|
|
valid_sequence = isinstance(model_clip_sequence, list)
|
|
if not valid_sequence and model_clip_sequence is not None:
|
|
_log("Warning: Connected 'model_clip_sequence' is not a valid list (likely bypassed). Ignoring.")
|
|
model_clip_sequence = None
|
|
|
|
all_frames_collected = []
|
|
current_seed_image = start_image
|
|
merge_frame_count = int(frame_merge)
|
|
duration = max(1, int(duration_sec))
|
|
window_frame_count = (duration * 16) + 1
|
|
|
|
_log("Cloning base models and clip for looping.")
|
|
base_model_high = model_high.clone()
|
|
base_model_low = model_low.clone()
|
|
base_clip = clip.clone()
|
|
|
|
_log(f"Starting loop with {len(positive_prompts)} prompts...")
|
|
if enable_dry_run: _log("!!! DRY RUN MODE ENABLED - Sampling and decoding will be skipped !!!")
|
|
|
|
model_device, model_dtype = _get_model_device_info(base_model_high)
|
|
cpu_device = torch.device("cpu")
|
|
render_width = _align_to_multiple(width, 16)
|
|
render_height = _align_to_multiple(height, 16)
|
|
current_seed = int(seed) & 0xFFFFFFFFFFFFFFFF
|
|
total_steps = int(steps)
|
|
|
|
# --- 2. Pre-Loop Preparations (Applied Once) ---
|
|
_log(f"Applying Sigma Shift: {sigma_shift}")
|
|
base_model_high = _apply_sigma_shift(base_model_high, sigma_shift)
|
|
base_model_low = _apply_sigma_shift(base_model_low, sigma_shift)
|
|
_, global_negative_cond = _tokenize_and_encode(base_clip, "", negative_prompt)
|
|
switch_step_index = _find_model_switch_index(base_model_high, scheduler, total_steps, model_switch_point)
|
|
_log(f"Model switch point {model_switch_point} calculated split at step: {switch_step_index} / {total_steps}")
|
|
if switch_step_index >= total_steps: switch_step_index = total_steps
|
|
elif switch_step_index <= 0: switch_step_index = 0
|
|
|
|
# --- 3. Main Generation Loop (Per Prompt) ---
|
|
for i, prompt_text in enumerate(positive_prompts):
|
|
_log(f"--- Starting Loop {i+1}/{len(positive_prompts)} ---")
|
|
_log(f"Prompt: {prompt_text[:80]}...")
|
|
|
|
# --- 3a. Select/Clone Models for this Loop ---
|
|
_log(f"Segment {i+1}: Cloning fresh base models and clip.")
|
|
loop_model_high = base_model_high.clone()
|
|
loop_model_low = base_model_low.clone()
|
|
loop_clip = base_clip.clone()
|
|
|
|
segment_models_provided = False
|
|
using_provided_high, using_provided_low, using_provided_clip = False, False, False
|
|
|
|
# Use sequence only if it's a valid list (handles bypass)
|
|
if valid_sequence and i < len(model_clip_sequence):
|
|
segment_data = model_clip_sequence[i]
|
|
if segment_data and isinstance(segment_data, tuple) and len(segment_data) == 3:
|
|
mh, ml, c = segment_data
|
|
if mh is not None: loop_model_high = mh; segment_models_provided = True; using_provided_high = True
|
|
if ml is not None: loop_model_low = ml; segment_models_provided = True; using_provided_low = True
|
|
if c is not None: loop_clip = c; segment_models_provided = True; using_provided_clip = True
|
|
|
|
if segment_models_provided:
|
|
high_source = "PROVIDED (LoRA?)" if using_provided_high else "CLONED BASE"
|
|
low_source = "PROVIDED (LoRA?)" if using_provided_low else "CLONED BASE"
|
|
clip_source = "PROVIDED (LoRA?)" if using_provided_clip else "CLONED BASE"
|
|
_log(f"Segment {i+1}: Using models - High: {high_source}, Low: {low_source}, Clip: {clip_source}")
|
|
else:
|
|
_log(f"Segment {i+1}: Using CLONED BASE models and clip (no specific models provided for this segment).")
|
|
|
|
# --- 3b. Prepare Conditionings and Latent ---
|
|
global_positive_cond, _ = _tokenize_and_encode(loop_clip, prompt_text, "")
|
|
vision_output = _encode_seed_image_vision(clip_vision, current_seed_image, "center")
|
|
loop_positive_cond, loop_negative_cond, latent_on_cpu = _prepare_latent_window(
|
|
global_positive_cond, global_negative_cond, vae,
|
|
render_width, render_height, window_frame_count, 1,
|
|
current_seed_image, vision_output
|
|
)
|
|
latent_on_gpu = {"samples": _move_tensor(latent_on_cpu["samples"], model_device, model_dtype)}
|
|
final_latent = None
|
|
|
|
# --- 3c/3d. Run Samplers OR Skip for Dry Run ---
|
|
if not enable_dry_run:
|
|
_log(f"Segment {i+1}: Running HIGH noise sampler (Steps 0 to {switch_step_index}).")
|
|
if enable_motion_cfg and switch_step_index > 1:
|
|
_log(f"Segment {i+1}: Using special Motion CFG ({cfg_motion_step}) for step 0.")
|
|
initial_latent = _run_sampler_step(
|
|
loop_model_high, current_seed, total_steps, cfg_motion_step, sampler_name, scheduler,
|
|
loop_positive_cond, loop_negative_cond, latent_on_gpu, denoise,
|
|
disable_noise=False, start_step=0, last_step=1, force_full_denoise=False
|
|
)
|
|
high_noise_latent = _run_sampler_step(
|
|
loop_model_high, current_seed, total_steps, cfg_high_noise, sampler_name, scheduler,
|
|
loop_positive_cond, loop_negative_cond, initial_latent, denoise,
|
|
disable_noise=True, start_step=1, last_step=switch_step_index, force_full_denoise=False
|
|
)
|
|
else:
|
|
high_noise_latent = _run_sampler_step(
|
|
loop_model_high, current_seed, total_steps, cfg_high_noise, sampler_name, scheduler,
|
|
loop_positive_cond, loop_negative_cond, latent_on_gpu, denoise,
|
|
disable_noise=False, start_step=0, last_step=switch_step_index, force_full_denoise=False
|
|
)
|
|
|
|
_log(f"Segment {i+1}: Running LOW noise sampler (Steps {switch_step_index} to {total_steps}).")
|
|
final_latent = _run_sampler_step(
|
|
loop_model_low, current_seed, total_steps, cfg_low_noise, sampler_name, scheduler,
|
|
loop_positive_cond, loop_negative_cond, high_noise_latent, denoise,
|
|
disable_noise=True, start_step=switch_step_index, last_step=total_steps, force_full_denoise=True
|
|
)
|
|
_log(f"Segment {i+1}: Sampling finished.")
|
|
else:
|
|
_log(f"Segment {i+1}: DRY RUN - Skipping HIGH noise sampler.")
|
|
_log(f"Segment {i+1}: DRY RUN - Skipping LOW noise sampler.")
|
|
final_latent = latent_on_gpu
|
|
high_noise_latent = latent_on_gpu
|
|
|
|
# --- 3e. Decode & Normalize OR Skip for Dry Run ---
|
|
decoded_image_batch = None
|
|
if not enable_dry_run and final_latent is not None:
|
|
_log(f"Segment {i+1}: Decoding latent.")
|
|
final_latent_on_cpu = _move_tensor(final_latent["samples"], cpu_device, torch.float32)
|
|
decoded_data = vae.decode(final_latent_on_cpu)
|
|
decoded_image_batch = decoded_data["image"] if isinstance(decoded_data, dict) else decoded_data
|
|
if decoded_image_batch.dim() == 3: decoded_image_batch = decoded_image_batch.unsqueeze(0)
|
|
elif decoded_image_batch.dim() == 5:
|
|
if decoded_image_batch.shape[-1] not in (1, 3, 4): decoded_image_batch = decoded_image_batch.permute(0, 2, 3, 4, 1).contiguous()
|
|
b, t, h, w, c = decoded_image_batch.shape; decoded_image_batch = decoded_image_batch.reshape(b * t, h, w, c)
|
|
elif decoded_image_batch.dim() == 4:
|
|
if decoded_image_batch.shape[1] in (1, 3, 4) and decoded_image_batch.shape[-1] not in (1, 3, 4):
|
|
decoded_image_batch = decoded_image_batch.permute(0, 2, 3, 1).contiguous()
|
|
decoded_image_batch = decoded_image_batch.to(torch.float32).clamp(0, 1)
|
|
else:
|
|
_log(f"Segment {i+1}: DRY RUN - Skipping VAE Decode.")
|
|
num_frames_to_gen = window_frame_count if i == 0 else window_frame_count - merge_frame_count
|
|
decoded_image_batch = torch.zeros((num_frames_to_gen, render_height, render_width, 3), dtype=torch.float32)
|
|
|
|
# --- 3f. Store Frames & Update Seed ---
|
|
if decoded_image_batch.shape[0] >= merge_frame_count and merge_frame_count > 0:
|
|
current_seed_image = decoded_image_batch[-merge_frame_count:-merge_frame_count+1 if merge_frame_count > 1 else None]
|
|
elif decoded_image_batch.shape[0] > 0: current_seed_image = decoded_image_batch[-1:]
|
|
else: _log("Warning: No frames decoded/generated..."); current_seed_image = None
|
|
frames_to_save = decoded_image_batch if i == 0 else decoded_image_batch[merge_frame_count:]
|
|
all_frames_collected.append(frames_to_save.cpu())
|
|
|
|
# --- 3g. Per-Iteration Cleanup ---
|
|
_log(f"Segment {i+1}: Cleaning up loop resources.")
|
|
del loop_model_high, loop_model_low, loop_clip
|
|
_log(f"Segment {i+1}: Deleted loop model/clip variables.")
|
|
if 'high_noise_latent' in locals(): del high_noise_latent
|
|
if final_latent is not None: del final_latent
|
|
if 'latent_on_gpu' in locals(): del latent_on_gpu
|
|
if 'final_latent_on_cpu' in locals(): del final_latent_on_cpu
|
|
if 'decoded_data' in locals(): del decoded_data
|
|
if decoded_image_batch is not None: del decoded_image_batch
|
|
if 'frames_to_save' in locals(): del frames_to_save
|
|
_log(f"Segment {i+1}: Forcing garbage collection.")
|
|
gc.collect()
|
|
if torch.cuda.is_available():
|
|
_log(f"Segment {i+1}: Emptying CUDA cache.")
|
|
torch.cuda.memory.empty_cache()
|
|
_log(f"--- Finished Loop {i+1} ---")
|
|
|
|
# --- 4. Finalize & Return ---
|
|
if not all_frames_collected:
|
|
_log("Error: No frames were collected."); return (torch.zeros((1, 64, 64, 3)),)*2
|
|
_log("Concatenating final video batch.")
|
|
num_channels = all_frames_collected[0].shape[-1]
|
|
consistent_batches = []
|
|
for batch in all_frames_collected:
|
|
if batch.shape[-1] == num_channels:
|
|
consistent_batches.append(batch)
|
|
else:
|
|
_log(f"Warning: Channel mismatch found (Expected {num_channels}, Got {batch.shape[-1]}). Adjusting batch.")
|
|
if batch.shape[-1] == 1 and num_channels == 3: batch = batch.repeat(1, 1, 1, 3)
|
|
elif batch.shape[-1] == 4 and num_channels == 3: batch = batch[:, :, :, :3]
|
|
elif batch.shape[-1] == 3 and num_channels == 4:
|
|
alpha = torch.ones_like(batch[:, :, :, :1]); batch = torch.cat([batch, alpha], dim=-1)
|
|
else: batch = torch.zeros((batch.shape[0], batch.shape[1], batch.shape[2], num_channels), dtype=batch.dtype)
|
|
consistent_batches.append(batch)
|
|
|
|
final_batch = torch.cat(consistent_batches, dim=0)
|
|
last_frame = final_batch[-1:] if final_batch.shape[0] > 0 else torch.zeros((1, render_height, render_width, 3))
|
|
log_message = "DRY RUN Finished" if enable_dry_run else "WanVideoLooper finished"
|
|
_log(f"{log_message}. Total frames: {final_batch.shape[0]}")
|
|
return (final_batch, last_frame) |