946 lines
48 KiB
Python
946 lines
48 KiB
Python
import logging
|
|
import math
|
|
import types
|
|
import torch
|
|
import torch.nn as nn
|
|
import numpy as np
|
|
import base64
|
|
import io as pyio
|
|
from PIL import Image
|
|
import comfy.utils
|
|
import folder_paths
|
|
import node_helpers
|
|
import comfy.model_management
|
|
import comfy.ldm.modules.attention
|
|
import json
|
|
from comfy_api.latest import io, ComfyExtension
|
|
from typing_extensions import override
|
|
|
|
log = logging.getLogger(__name__)
|
|
|
|
# ==============================================================================
|
|
# 1. CORE UTILS & MATH
|
|
# ==============================================================================
|
|
|
|
def build_temporal_cost(q_token_idx, Lq, Lk, device, dtype, tokens_per_frame):
|
|
offset = torch.zeros(Lq, Lk, device=device, dtype=dtype)
|
|
query_frames = torch.arange(Lq, device=device, dtype=torch.long) // tokens_per_frame
|
|
for seg in q_token_idx:
|
|
local = seg["local_token_idx"].to(device=device)
|
|
d = (query_frames.float()[:, None] - seg["midpoint"]).abs()
|
|
strength = seg.get("strength", 1.0)
|
|
cost = strength * (torch.relu(d - seg["window"]) ** 2) / (2 * seg["sigma"] ** 2)
|
|
offset[:, local] = cost.to(offset.dtype)
|
|
return offset
|
|
|
|
def build_temporal_cost_scaled(q_token_idx, Lq, Lk, device, dtype, latent_frames):
|
|
offset = torch.zeros(Lq, Lk, device=device, dtype=dtype)
|
|
query_frames = torch.arange(Lq, device=device, dtype=torch.float32) * latent_frames / Lq
|
|
for seg in q_token_idx:
|
|
local = seg["local_token_idx"].to(device=device)
|
|
d = (query_frames[:, None] - seg["midpoint"]).abs()
|
|
sigma_a = seg.get("sigma_audio", seg["sigma"])
|
|
window_a = seg.get("window_audio", seg["window"])
|
|
strength_a = seg.get("strength_audio", 1.0)
|
|
cost = strength_a * (torch.relu(d - window_a) ** 2) / (2 * sigma_a ** 2)
|
|
offset[:, local] = cost.to(offset.dtype)
|
|
return offset
|
|
|
|
def create_mask_fn(q_token_idx, fallback_tokens_per_frame, latent_frames):
|
|
cache = {}
|
|
max_token_idx = max(int(seg["local_token_idx"].max().item()) for seg in q_token_idx) + 1 if q_token_idx else 0
|
|
def mask_fn(q, k, transformer_options):
|
|
Lq, Lk = q.shape[1], k.shape[1]
|
|
if Lq == Lk: return None
|
|
cond_or_uncond = transformer_options.get("cond_or_uncond", [])
|
|
if 1 in cond_or_uncond and 0 not in cond_or_uncond: return None
|
|
grid_sizes = transformer_options.get("grid_sizes", None)
|
|
video_tpf = int(grid_sizes[1]) * int(grid_sizes[2]) if grid_sizes is not None else fallback_tokens_per_frame
|
|
video_lq = latent_frames * video_tpf
|
|
if Lk == video_lq or Lk < max_token_idx: return None
|
|
mode = "video" if Lq == video_lq else "scaled"
|
|
key = (Lq, Lk, mode, q.device)
|
|
if key not in cache:
|
|
if mode == "video":
|
|
cost = build_temporal_cost(q_token_idx, Lq, Lk, q.device, q.dtype, video_tpf)
|
|
else:
|
|
cost = build_temporal_cost_scaled(q_token_idx, Lq, Lk, q.device, q.dtype, latent_frames)
|
|
cache[key] = -cost
|
|
|
|
nonzero = (cost > 0).sum().item()
|
|
total = Lq * Lk
|
|
log.info(f"[RikanPromptRelay] Built penalty matrix ({mode}): Lq={Lq}, Lk={Lk}, nonzero={nonzero}/{total}")
|
|
|
|
return cache[key].to(q.dtype)
|
|
return mask_fn
|
|
|
|
def build_segments(token_ranges, segment_lengths, epsilon=1e-3):
|
|
sigma = 1.0 / math.log(1.0 / epsilon) if 0 < epsilon < 1 else 0.1448
|
|
q_token_idx = []
|
|
frame_cursor = 0
|
|
for (tok_start, tok_end), L in zip(token_ranges, segment_lengths):
|
|
if L <= 0:
|
|
frame_cursor += L
|
|
continue
|
|
midpoint = (2 * frame_cursor + L) // 2
|
|
base_window = max(L // 2 - 2, 0)
|
|
q_token_idx.append({
|
|
"local_token_idx": torch.arange(tok_start, tok_end),
|
|
"midpoint": midpoint,
|
|
"window": float(base_window),
|
|
"sigma": sigma,
|
|
"strength": 1.0,
|
|
"window_audio": float(base_window),
|
|
"sigma_audio": sigma,
|
|
"strength_audio": 1.0,
|
|
})
|
|
frame_cursor += L
|
|
return q_token_idx
|
|
|
|
def get_raw_tokenizer(clip):
|
|
tokenizer_wrapper = clip.tokenizer
|
|
for attr_name in dir(tokenizer_wrapper):
|
|
if attr_name.startswith("_"): continue
|
|
inner = getattr(tokenizer_wrapper, attr_name, None)
|
|
if inner is not None and hasattr(inner, "tokenizer"):
|
|
return inner.tokenizer
|
|
raise RuntimeError("Could not find raw tokenizer.")
|
|
|
|
def map_token_indices(raw_tokenizer, global_prompt, local_prompts):
|
|
prefixed_locals = [" " + lp for lp in local_prompts]
|
|
has_eos = getattr(raw_tokenizer, "add_eos", False)
|
|
eos_adj = 1 if has_eos else 0
|
|
prev_len = len(raw_tokenizer(global_prompt)["input_ids"]) - eos_adj
|
|
token_ranges = []
|
|
built = global_prompt
|
|
for plp in prefixed_locals:
|
|
built += plp
|
|
cur_len = len(raw_tokenizer(built)["input_ids"]) - eos_adj
|
|
token_ranges.append((prev_len, cur_len))
|
|
prev_len = cur_len
|
|
return global_prompt + "".join(prefixed_locals), token_ranges
|
|
|
|
def distribute_segment_lengths(num_segments, latent_frames, specified_lengths=None):
|
|
if specified_lengths:
|
|
lengths = specified_lengths
|
|
else:
|
|
step = -(-latent_frames // num_segments)
|
|
lengths = [step] * num_segments
|
|
effective, cursor = [], 0
|
|
for L in lengths:
|
|
end = min(cursor + L, latent_frames)
|
|
effective.append(max(end - cursor, 0))
|
|
cursor = end
|
|
return effective
|
|
|
|
def _convert_to_latent_lengths(pixel_lengths, temporal_stride, latent_frames):
|
|
if not pixel_lengths: return []
|
|
total_pixel = sum(pixel_lengths)
|
|
if total_pixel <= 0: return [1] * len(pixel_lengths)
|
|
target_total = min(latent_frames, max(1, round(total_pixel / temporal_stride)))
|
|
if target_total >= latent_frames - 1: target_total = latent_frames
|
|
exact = [p * target_total / total_pixel for p in pixel_lengths]
|
|
result = [int(e) for e in exact]
|
|
diff = target_total - sum(result)
|
|
if diff > 0:
|
|
order = sorted(range(len(exact)), key=lambda i: -(exact[i] - int(exact[i])))
|
|
for k in range(diff): result[order[k % len(order)]] += 1
|
|
return result
|
|
|
|
# ==============================================================================
|
|
# 2. MODEL PATCHING & LORA GATE HELPERS
|
|
# ==============================================================================
|
|
|
|
def detect_model_type(model):
|
|
diff_model = model.model.diffusion_model
|
|
if hasattr(diff_model, "patch_size") and not hasattr(diff_model, "patchifier"):
|
|
return "wan", tuple(diff_model.patch_size), 4
|
|
if hasattr(diff_model, "patchifier"):
|
|
return "ltx", (1, 1, 1), int(diff_model.vae_scale_factors[0])
|
|
raise ValueError("Unsupported model type.")
|
|
|
|
class _CrossAttnPatch:
|
|
def __init__(self, impl, mask_fn):
|
|
self.impl, self.mask_fn = impl, mask_fn
|
|
def __get__(self, obj, objtype=None):
|
|
def wrapped(self_module, *args, **kwargs):
|
|
return self.impl(self_module, self.mask_fn, *args, **kwargs)
|
|
return types.MethodType(wrapped, obj)
|
|
|
|
def apply_patches(model_clone, arch, mask_fn):
|
|
diffusion_model = model_clone.get_model_object("diffusion_model")
|
|
if arch == "wan":
|
|
from comfy.ldm.wan.model import WanI2VCrossAttention
|
|
for idx, block in enumerate(diffusion_model.blocks):
|
|
impl = _wan_i2v_forward if isinstance(block.cross_attn, WanI2VCrossAttention) else _wan_t2v_forward
|
|
model_clone.add_object_patch(f"diffusion_model.blocks.{idx}.cross_attn.forward", _CrossAttnPatch(impl, mask_fn).__get__(block.cross_attn))
|
|
elif arch == "ltx":
|
|
for idx, block in enumerate(diffusion_model.transformer_blocks):
|
|
for attr in ("attn2", "audio_attn2"):
|
|
module = getattr(block, attr, None)
|
|
if module: model_clone.add_object_patch(f"diffusion_model.transformer_blocks.{idx}.{attr}.forward", _CrossAttnPatch(_ltx_forward, mask_fn).__get__(module))
|
|
|
|
def _wan_t2v_forward(self, mask_fn, x, context, transformer_options={}, **kwargs):
|
|
q, k, v = self.norm_q(self.q(x)), self.norm_k(self.k(context)), self.v(context)
|
|
mask = mask_fn(q, k, transformer_options)
|
|
return self.o(comfy.ldm.modules.attention.attention_pytorch(q, k, v, self.num_heads, mask=mask, transformer_options=transformer_options) if mask is not None else comfy.ldm.modules.attention.optimized_attention(q, k, v, self.num_heads, transformer_options=transformer_options))
|
|
|
|
def _wan_i2v_forward(self, mask_fn, x, context, context_img_len, transformer_options={}, **kwargs):
|
|
q = self.norm_q(self.q(x))
|
|
img_x = comfy.ldm.modules.attention.optimized_attention(q, self.norm_k_img(self.k_img(context[:, :context_img_len])), self.v_img(context[:, :context_img_len]), self.num_heads, transformer_options=transformer_options)
|
|
mask = mask_fn(q, self.norm_k(self.k(context[:, context_img_len:])), transformer_options)
|
|
attn_txt = comfy.ldm.modules.attention.attention_pytorch(q, self.norm_k(self.k(context[:, context_img_len:])), self.v(context[:, context_img_len:]), self.num_heads, mask=mask, transformer_options=transformer_options) if mask is not None else comfy.ldm.modules.attention.optimized_attention(q, self.norm_k(self.k(context[:, context_img_len:])), self.v(context[:, context_img_len:]), self.num_heads, transformer_options=transformer_options)
|
|
return self.o(attn_txt + img_x)
|
|
|
|
def _ltx_forward(self, mask_fn, x, context=None, mask=None, pe=None, k_pe=None, transformer_options={}):
|
|
from comfy.ldm.lightricks.model import apply_rotary_emb
|
|
is_self = context is None
|
|
c = x if is_self else context
|
|
q, k, v = self.q_norm(self.to_q(x)), self.k_norm(self.to_k(c)), self.to_v(c)
|
|
if pe is not None: q, k = apply_rotary_emb(q, pe), apply_rotary_emb(k, pe if k_pe is None else k_pe)
|
|
if not is_self:
|
|
tm = mask_fn(q, k, transformer_options)
|
|
mask = tm if mask is None else mask + tm
|
|
out = comfy.ldm.modules.attention.optimized_attention(q, k, v, self.heads, mask=mask, transformer_options=transformer_options)
|
|
if self.to_gate_logits: out = out * (2.0 * torch.sigmoid(self.to_gate_logits(x))).unsqueeze(-1)
|
|
return self.to_out(out)
|
|
|
|
# ── LORA GATE INTERNAL CLASSES ───────────────────────────────────────────────
|
|
|
|
class _TemporalGatedLoRA:
|
|
def __init__(self, original, lora_down, lora_up, scale, strength, gate_weights, tokens_per_frame):
|
|
self.original = original
|
|
self._lora_down = lora_down.cpu()
|
|
self._lora_up = lora_up.cpu()
|
|
self.scale = scale
|
|
self.strength = strength
|
|
self._gate_weights = gate_weights.cpu()
|
|
self.tokens_per_frame = tokens_per_frame
|
|
|
|
def forward(self, x, *args, **kwargs):
|
|
base = self.original(x, *args, **kwargs)
|
|
lora_down = self._lora_down.to(device=x.device, dtype=x.dtype)
|
|
lora_up = self._lora_up.to(device=x.device, dtype=x.dtype)
|
|
|
|
delta = (x @ lora_down.T) @ lora_up.T
|
|
delta = delta * (self.scale * self.strength)
|
|
|
|
if x.dim() < 2: return base + delta
|
|
|
|
seq_len = x.shape[-2]
|
|
n_frames = len(self._gate_weights)
|
|
n_video_tokens = n_frames * self.tokens_per_frame
|
|
|
|
if seq_len >= n_video_tokens and n_video_tokens > 0:
|
|
gate_weights = self._gate_weights.to(device=x.device, dtype=x.dtype)
|
|
frame_idx = torch.arange(n_video_tokens, device=x.device) // self.tokens_per_frame
|
|
frame_idx = frame_idx.clamp(0, n_frames - 1)
|
|
video_gates = gate_weights[frame_idx]
|
|
|
|
if seq_len > n_video_tokens:
|
|
extra = torch.zeros(seq_len - n_video_tokens, device=x.device, dtype=x.dtype)
|
|
token_gates = torch.cat([video_gates, extra])
|
|
else:
|
|
token_gates = video_gates
|
|
|
|
view_shape = (1,) * (x.dim() - 2) + (seq_len, 1)
|
|
delta = delta * token_gates.view(view_shape)
|
|
|
|
return base + delta
|
|
|
|
def _build_gate_weights(seg, latent_frames):
|
|
midpoint, window, sigma = float(seg["midpoint"]), float(seg["window"]), float(seg["sigma"])
|
|
weights = []
|
|
for frame in range(latent_frames):
|
|
d = abs(frame - midpoint)
|
|
cost = (max(d - window, 0.0) ** 2) / (2.0 * sigma ** 2)
|
|
weights.append(math.exp(-cost))
|
|
return torch.tensor(weights, dtype=torch.float32)
|
|
|
|
def _get_module(diffusion_model, local_path):
|
|
m = diffusion_model
|
|
for part in local_path.split("."): m = getattr(m, part)
|
|
return m
|
|
|
|
def _apply_gated_lora(model_clone, lora_path, segment, strength, tokens_per_frame, latent_frames):
|
|
try:
|
|
try: import comfy.loras as _lora_mod
|
|
except ImportError: import comfy.lora as _lora_mod
|
|
except ImportError:
|
|
log.error("[Rikan MultiLoraGate] comfy.loras is not available.")
|
|
return
|
|
|
|
lora_file = comfy.utils.load_torch_file(lora_path, safe_load=True)
|
|
key_map = {}
|
|
try:
|
|
_lora_mod.model_lora_keys_unet(model_clone.model, key_map)
|
|
except Exception as e:
|
|
log.error("[Rikan MultiLoraGate] Failed to create LoRA key map: %s", e)
|
|
return
|
|
|
|
try:
|
|
patches = _lora_mod.load_lora(lora_file, key_map)
|
|
except Exception as e:
|
|
log.error("[Rikan MultiLoraGate] Failed to parse LoRA file: %s", e)
|
|
return
|
|
|
|
gate_weights = _build_gate_weights(segment, latent_frames)
|
|
diffusion_model = model_clone.get_model_object("diffusion_model")
|
|
|
|
applied = 0
|
|
skipped_not_weight = 0
|
|
skipped_not_diffusion = 0
|
|
skipped_bad_data = 0
|
|
skipped_not_linear = 0
|
|
_logged_patch_sample = False
|
|
|
|
for param_key, patch_data in patches.items():
|
|
if not _logged_patch_sample:
|
|
log.info("[Rikan MultiLoraGate] patch_data sample — key=%s attrs=%s", param_key, [a for a in dir(patch_data) if not a.startswith("_")])
|
|
log.info("[Rikan MultiLoraGate] patch_data.weights=%s multiplier=%s", repr(patch_data.weights), repr(patch_data.multiplier))
|
|
_logged_patch_sample = True
|
|
|
|
if not param_key.endswith(".weight"):
|
|
skipped_not_weight += 1
|
|
continue
|
|
|
|
module_path = param_key[:-7]
|
|
if not module_path.startswith("diffusion_model."):
|
|
skipped_not_diffusion += 1
|
|
continue
|
|
|
|
local_path = module_path[len("diffusion_model."):]
|
|
|
|
try:
|
|
lora_up = patch_data.weights[0]
|
|
lora_down = patch_data.weights[1]
|
|
alpha = patch_data.multiplier
|
|
except (TypeError, IndexError, AttributeError):
|
|
skipped_bad_data += 1
|
|
continue
|
|
|
|
if lora_up is None or lora_down is None:
|
|
skipped_bad_data += 1
|
|
continue
|
|
|
|
if lora_down.dim() != 2 or lora_up.dim() != 2:
|
|
skipped_not_linear += 1
|
|
continue
|
|
|
|
rank = lora_down.shape[0]
|
|
scale = float(alpha) / rank if alpha is not None else 1.0
|
|
|
|
target = model_clone.object_patches.get(module_path)
|
|
if target is None:
|
|
try:
|
|
target = _get_module(diffusion_model, local_path)
|
|
except AttributeError:
|
|
skipped_not_diffusion += 1
|
|
continue
|
|
|
|
wrapped = _TemporalGatedLoRA(target.forward, lora_down, lora_up, scale, strength, gate_weights, tokens_per_frame)
|
|
model_clone.add_object_patch(module_path + ".forward", wrapped.forward)
|
|
applied += 1
|
|
|
|
log.info(
|
|
"[Rikan MultiLoraGate] Segment midpoint=%.1f window=%.1f — applied %d, not_weight=%d, not_diffusion=%d, bad_data=%d, not_linear=%d",
|
|
float(segment["midpoint"]), float(segment["window"]), applied,
|
|
skipped_not_weight, skipped_not_diffusion, skipped_bad_data, skipped_not_linear,
|
|
)
|
|
|
|
# ==============================================================================
|
|
# 3. NODES
|
|
# ==============================================================================
|
|
|
|
class RikanPromptRelayEncodeTimeline(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(
|
|
node_id="RikanPromptRelayEncodeTimeline",
|
|
display_name="Rikan Prompt Relay Encode (Timeline)",
|
|
category="Rikannodes",
|
|
inputs=[
|
|
io.Model.Input("model"),
|
|
io.Clip.Input("clip"),
|
|
io.Latent.Input("latent"),
|
|
io.String.Input("global_prompt", multiline=True, default=""),
|
|
io.Int.Input("max_frames", default=129, min=1, max=10000),
|
|
io.String.Input("timeline_data", default=""),
|
|
io.String.Input("local_prompts", multiline=True, default=""),
|
|
io.String.Input("segment_lengths", default=""),
|
|
io.Float.Input("epsilon", default=1e-3, min=1e-6, max=0.99, step=1e-4),
|
|
io.Float.Input("fps", default=24.0, min=0.1, max=240.0, optional=True),
|
|
io.Float.Input("seconds", default=5.375, min=0.1, max=1000.0, step=0.1, optional=True),
|
|
io.Combo.Input("time_units", options=["frames", "seconds"], default="frames", optional=True),
|
|
],
|
|
outputs=[
|
|
io.Model.Output(display_name="model"),
|
|
io.Conditioning.Output(display_name="positive"),
|
|
io.Float.Output(display_name="fps_out"),
|
|
io.Int.Output(display_name="max_frames_out"),
|
|
],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, model, clip, latent, global_prompt, max_frames, timeline_data, local_prompts, segment_lengths, epsilon, fps=24.0, seconds=5.375, time_units="frames") -> io.NodeOutput:
|
|
locals_list = [p.strip() for p in local_prompts.split("|") if p.strip()]
|
|
if not locals_list: raise ValueError("At least one local prompt is required.")
|
|
|
|
arch, patch_size, temporal_stride = detect_model_type(model)
|
|
samples = latent["samples"]
|
|
latent_frames = samples.shape[2]
|
|
tokens_per_frame = (samples.shape[3] // patch_size[1]) * (samples.shape[4] // patch_size[2])
|
|
|
|
parsed_lengths = None
|
|
if segment_lengths.strip():
|
|
pixel_lengths = [int(x.strip()) for x in segment_lengths.split(",") if x.strip()]
|
|
parsed_lengths = _convert_to_latent_lengths(pixel_lengths, temporal_stride, latent_frames)
|
|
|
|
full_prompt, token_ranges = map_token_indices(get_raw_tokenizer(clip), global_prompt, locals_list)
|
|
|
|
log.info("[RikanPromptRelay] Global: tokens [0:%d] (%d tokens)", token_ranges[0][0], token_ranges[0][0])
|
|
for i, (s, e) in enumerate(token_ranges):
|
|
log.info("[RikanPromptRelay] Segment %d: tokens [%d:%d] (%d tokens)", i, s, e, e - s)
|
|
|
|
conditioning = clip.encode_from_tokens_scheduled(clip.tokenize(full_prompt))
|
|
effective_lengths = distribute_segment_lengths(len(locals_list), latent_frames, parsed_lengths)
|
|
|
|
log.info(
|
|
"[RikanPromptRelay] Latent: %d frames, %d tokens/frame, segments: %s",
|
|
latent_frames, tokens_per_frame, effective_lengths,
|
|
)
|
|
|
|
q_token_idx = build_segments(token_ranges, effective_lengths, epsilon)
|
|
mask_fn = create_mask_fn(q_token_idx, tokens_per_frame, latent_frames)
|
|
|
|
patched = model.clone()
|
|
apply_patches(patched, arch, mask_fn)
|
|
|
|
for key in ["rikan_pr_segments", "pr_segments"]: patched.model_options[key] = q_token_idx
|
|
for key in ["rikan_pr_latent_frames", "pr_latent_frames"]: patched.model_options[key] = latent_frames
|
|
for key in ["rikan_pr_tokens_per_frame", "pr_tokens_per_frame"]: patched.model_options[key] = tokens_per_frame
|
|
|
|
print(f"\n[Timeline] SUCCESS: Injected {len(q_token_idx)} segments.\n")
|
|
return io.NodeOutput(patched, conditioning, float(fps), int(max_frames))
|
|
|
|
class RikanPromptRelayMultiLoraGate(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(
|
|
node_id="RikanPromptRelayMultiLoraGate",
|
|
display_name="Rikan Prompt Relay Multi LoRA Gate",
|
|
category="Rikannodes",
|
|
inputs=[
|
|
io.Model.Input("model"),
|
|
io.String.Input("lora_data", default="[]", multiline=True),
|
|
],
|
|
outputs=[io.Model.Output(display_name="model")],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, model, lora_data) -> io.NodeOutput:
|
|
segments = model.model_options.get("rikan_pr_segments")
|
|
latent_frames = model.model_options.get("rikan_pr_latent_frames")
|
|
tokens_per_frame = model.model_options.get("rikan_pr_tokens_per_frame")
|
|
|
|
if segments is None:
|
|
log.warning("[Rikan MultiLoraGate] Prompt Relay metadata not found. Ensure it is connected to the Timeline node.")
|
|
return io.NodeOutput(model)
|
|
|
|
try:
|
|
loras_to_apply = json.loads(lora_data)
|
|
except Exception as e:
|
|
log.error(f"[Rikan MultiLoraGate] Failed to process JSON data: {e}")
|
|
return io.NodeOutput(model)
|
|
|
|
patched = model.clone()
|
|
total_segments = len(segments)
|
|
|
|
for lora in loras_to_apply:
|
|
if not lora.get("enable", False) or lora.get("name") == "None":
|
|
continue
|
|
|
|
lora_name = lora.get("name")
|
|
seg_idx = int(lora.get("segment", 0))
|
|
strength = float(lora.get("modelStr", 1.0))
|
|
|
|
if seg_idx < 0 or seg_idx >= total_segments:
|
|
log.warning(f"[Rikan MultiLoraGate] SKIP: LoRA '{lora_name}' attempts to use Segment {seg_idx}, "
|
|
f"but the Timeline only has {total_segments} segment(s) (Available indices 0 to {total_segments-1}).")
|
|
continue
|
|
|
|
lora_path = folder_paths.get_full_path("loras", lora_name)
|
|
if lora_path is None:
|
|
log.warning(f"[Rikan MultiLoraGate] LoRA file not found: {lora_name}")
|
|
continue
|
|
|
|
_apply_gated_lora(patched, lora_path, segments[seg_idx], strength, tokens_per_frame, latent_frames)
|
|
|
|
patched.model_options["rikan_pr_segments"] = segments
|
|
patched.model_options["rikan_pr_latent_frames"] = latent_frames
|
|
patched.model_options["rikan_pr_tokens_per_frame"] = tokens_per_frame
|
|
|
|
return io.NodeOutput(patched)
|
|
|
|
|
|
class RikanI2VPainterTiledVAE(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(
|
|
node_id="rikan-i2vpainter-tiled-vae",
|
|
display_name="Rikan I2V Painter (Tiled VAE)",
|
|
category="Rikannodes",
|
|
inputs=[
|
|
io.Conditioning.Input("positive"),
|
|
io.Conditioning.Input("negative"),
|
|
io.Vae.Input("vae"),
|
|
io.Int.Input("width", default=832, min=16, max=4096, step=16),
|
|
io.Int.Input("height", default=480, min=16, max=4096, step=16),
|
|
io.Int.Input("length", default=81, min=1, max=4096, step=4),
|
|
io.Int.Input("batch_size", default=1, min=1, max=4096),
|
|
io.Int.Input("tile_size", default=512, min=64, max=4096, step=64),
|
|
io.Int.Input("overlap", default=64, min=0, max=4096, step=32),
|
|
io.Int.Input("temporal_size", default=64, min=8, max=4096, step=4),
|
|
io.Int.Input("temporal_overlap", default=24, min=4, max=4096, step=4),
|
|
io.Float.Input("motion_amplitude", default=1.3, min=1.0, max=2.0, step=0.05),
|
|
io.Boolean.Input("color_protect", default=True),
|
|
io.Float.Input("correct_strength", default=0.01, min=0.0, max=0.3, step=0.01),
|
|
io.ClipVisionOutput.Input("clip_vision", optional=True),
|
|
io.Image.Input("start_image", optional=True),
|
|
],
|
|
outputs=[
|
|
io.Conditioning.Output(display_name="high_positive"),
|
|
io.Conditioning.Output(display_name="high_negative"),
|
|
io.Conditioning.Output(display_name="low_positive"),
|
|
io.Conditioning.Output(display_name="low_negative"),
|
|
io.Latent.Output(display_name="latent"),
|
|
]
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, positive, negative, vae, width, height, length, batch_size,
|
|
tile_size=512, overlap=64, temporal_size=64, temporal_overlap=8,
|
|
motion_amplitude=1.3, color_protect=True, correct_strength=0.01,
|
|
start_image=None, clip_vision=None) -> io.NodeOutput:
|
|
|
|
latent = torch.zeros([batch_size, 16, ((length - 1) // 4) + 1, height // 8, width // 8],
|
|
device=comfy.model_management.intermediate_device())
|
|
|
|
positive_original = positive
|
|
negative_original = negative
|
|
|
|
if start_image is not None:
|
|
start_image = start_image[:1]
|
|
start_image = comfy.utils.common_upscale(start_image.movedim(-1, 1), width, height, "bilinear", "center").movedim(1, -1)
|
|
image = torch.ones((length, height, width, 3), device=start_image.device, dtype=start_image.dtype) * 0.5
|
|
image[0] = start_image[0]
|
|
|
|
concat_latent_image = vae.encode_tiled(image, tile_x=tile_size, tile_y=tile_size, overlap=overlap, tile_t=temporal_size, overlap_t=temporal_overlap)
|
|
mask = torch.ones((1, 1, latent.shape[2], concat_latent_image.shape[-2], concat_latent_image.shape[-1]), device=start_image.device)
|
|
mask[:, :, 0] = 0.0
|
|
|
|
concat_latent_image_original = concat_latent_image.clone()
|
|
|
|
if motion_amplitude > 1.0:
|
|
base_latent = concat_latent_image[:, :, 0:1]
|
|
gray_latent = concat_latent_image[:, :, 1:]
|
|
|
|
diff = gray_latent - base_latent
|
|
diff_mean = diff.mean(dim=(1, 3, 4), keepdim=True)
|
|
diff_centered = diff - diff_mean
|
|
|
|
scaled_latent = base_latent + diff_centered * motion_amplitude + diff_mean
|
|
scaled_latent = torch.clamp(scaled_latent, -6, 6)
|
|
concat_latent_image = torch.cat([base_latent, scaled_latent], dim=2)
|
|
|
|
post_enhanced = concat_latent_image.clone()
|
|
|
|
if color_protect and correct_strength > 0:
|
|
orig_mean = concat_latent_image_original.mean(dim=(2, 3, 4))
|
|
enhanced_mean = post_enhanced.mean(dim=(2, 3, 4))
|
|
|
|
mean_drift = torch.abs(enhanced_mean - orig_mean) / (torch.abs(orig_mean) + 1e-6)
|
|
problem_channels = mean_drift > 0.18
|
|
|
|
if problem_channels.any():
|
|
drift_amount = enhanced_mean - orig_mean
|
|
correction = drift_amount * problem_channels.float() * correct_strength * 0.03
|
|
|
|
for b in range(batch_size):
|
|
for c in range(16):
|
|
if correction[b, c].abs() > 0:
|
|
post_enhanced[b, c] = torch.where(
|
|
post_enhanced[b, c] > 0,
|
|
post_enhanced[b, c] - correction[b, c],
|
|
post_enhanced[b, c]
|
|
)
|
|
|
|
orig_brightness = concat_latent_image_original.mean()
|
|
enhanced_brightness = post_enhanced.mean()
|
|
|
|
if enhanced_brightness < orig_brightness * 0.92:
|
|
brightness_boost = min(orig_brightness / (enhanced_brightness + 1e-6), 1.05)
|
|
post_enhanced = torch.where(
|
|
post_enhanced < 0.5,
|
|
post_enhanced * brightness_boost,
|
|
post_enhanced
|
|
)
|
|
|
|
concat_latent_image = torch.clamp(post_enhanced, -6, 6)
|
|
|
|
positive = node_helpers.conditioning_set_values(positive, {"concat_latent_image": concat_latent_image, "concat_mask": mask})
|
|
negative = node_helpers.conditioning_set_values(negative, {"concat_latent_image": concat_latent_image, "concat_mask": mask})
|
|
|
|
positive_original = node_helpers.conditioning_set_values(positive_original, {"concat_latent_image": concat_latent_image_original, "concat_mask": mask})
|
|
negative_original = node_helpers.conditioning_set_values(negative_original, {"concat_latent_image": concat_latent_image_original, "concat_mask": mask})
|
|
|
|
ref_latent = vae.encode_tiled(start_image[:,:,:,:3], tile_x=tile_size, tile_y=tile_size, overlap=overlap, tile_t=temporal_size, overlap_t=temporal_overlap)
|
|
positive = node_helpers.conditioning_set_values(positive, {"reference_latents": [ref_latent]}, append=True)
|
|
negative = node_helpers.conditioning_set_values(negative, {"reference_latents": [torch.zeros_like(ref_latent)]}, append=True)
|
|
|
|
positive_original = node_helpers.conditioning_set_values(positive_original, {"reference_latents": [ref_latent]}, append=True)
|
|
negative_original = node_helpers.conditioning_set_values(negative_original, {"reference_latents": [torch.zeros_like(ref_latent)]}, append=True)
|
|
|
|
if clip_vision is not None:
|
|
positive = node_helpers.conditioning_set_values(positive, {"clip_vision_output": clip_vision})
|
|
negative = node_helpers.conditioning_set_values(negative, {"clip_vision_output": clip_vision})
|
|
|
|
positive_original = node_helpers.conditioning_set_values(positive_original, {"clip_vision_output": clip_vision})
|
|
negative_original = node_helpers.conditioning_set_values(negative_original, {"clip_vision_output": clip_vision})
|
|
|
|
out_latent = {"samples": latent}
|
|
return io.NodeOutput(positive, negative, positive_original, negative_original, out_latent)
|
|
|
|
|
|
class RikanI2VPainter(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(
|
|
node_id="rikan-i2vpainter",
|
|
display_name="Rikan I2V Painter",
|
|
category="Rikannodes",
|
|
inputs=[
|
|
io.Conditioning.Input("positive"),
|
|
io.Conditioning.Input("negative"),
|
|
io.Vae.Input("vae"),
|
|
io.Int.Input("width", default=832, min=16, max=4096, step=16),
|
|
io.Int.Input("height", default=480, min=16, max=4096, step=16),
|
|
io.Int.Input("length", default=81, min=1, max=4096, step=4),
|
|
io.Int.Input("batch_size", default=1, min=1, max=4096),
|
|
io.Float.Input("motion_amplitude", default=1.3, min=1.0, max=2.0, step=0.05),
|
|
io.Boolean.Input("color_protect", default=True),
|
|
io.Float.Input("correct_strength", default=0.01, min=0.0, max=0.3, step=0.01),
|
|
io.ClipVisionOutput.Input("clip_vision", optional=True),
|
|
io.Image.Input("start_image", optional=True),
|
|
],
|
|
outputs=[
|
|
io.Conditioning.Output(display_name="high_positive"),
|
|
io.Conditioning.Output(display_name="high_negative"),
|
|
io.Conditioning.Output(display_name="low_positive"),
|
|
io.Conditioning.Output(display_name="low_negative"),
|
|
io.Latent.Output(display_name="latent"),
|
|
]
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, positive, negative, vae, width, height, length, batch_size,
|
|
motion_amplitude=1.3, color_protect=True, correct_strength=0.01,
|
|
start_image=None, clip_vision=None) -> io.NodeOutput:
|
|
|
|
latent = torch.zeros([batch_size, 16, ((length - 1) // 4) + 1, height // 8, width // 8],
|
|
device=comfy.model_management.intermediate_device())
|
|
|
|
positive_original = positive
|
|
negative_original = negative
|
|
|
|
if start_image is not None:
|
|
start_image = start_image[:1]
|
|
start_image = comfy.utils.common_upscale(start_image.movedim(-1, 1), width, height, "bilinear", "center").movedim(1, -1)
|
|
image = torch.ones((length, height, width, 3), device=start_image.device, dtype=start_image.dtype) * 0.5
|
|
image[0] = start_image[0]
|
|
|
|
concat_latent_image = vae.encode(image[:, :, :, :3])
|
|
mask = torch.ones((1, 1, latent.shape[2], concat_latent_image.shape[-2], concat_latent_image.shape[-1]), device=start_image.device)
|
|
mask[:, :, 0] = 0.0
|
|
|
|
concat_latent_image_original = concat_latent_image.clone()
|
|
|
|
if motion_amplitude > 1.0:
|
|
base_latent = concat_latent_image[:, :, 0:1]
|
|
gray_latent = concat_latent_image[:, :, 1:]
|
|
|
|
diff = gray_latent - base_latent
|
|
diff_mean = diff.mean(dim=(1, 3, 4), keepdim=True)
|
|
diff_centered = diff - diff_mean
|
|
|
|
scaled_latent = base_latent + diff_centered * motion_amplitude + diff_mean
|
|
scaled_latent = torch.clamp(scaled_latent, -6, 6)
|
|
concat_latent_image = torch.cat([base_latent, scaled_latent], dim=2)
|
|
|
|
post_enhanced = concat_latent_image.clone()
|
|
|
|
if color_protect and correct_strength > 0:
|
|
orig_mean = concat_latent_image_original.mean(dim=(2, 3, 4))
|
|
enhanced_mean = post_enhanced.mean(dim=(2, 3, 4))
|
|
|
|
mean_drift = torch.abs(enhanced_mean - orig_mean) / (torch.abs(orig_mean) + 1e-6)
|
|
problem_channels = mean_drift > 0.18
|
|
|
|
if problem_channels.any():
|
|
drift_amount = enhanced_mean - orig_mean
|
|
correction = drift_amount * problem_channels.float() * correct_strength * 0.03
|
|
|
|
for b in range(batch_size):
|
|
for c in range(16):
|
|
if correction[b, c].abs() > 0:
|
|
post_enhanced[b, c] = torch.where(
|
|
post_enhanced[b, c] > 0,
|
|
post_enhanced[b, c] - correction[b, c],
|
|
post_enhanced[b, c]
|
|
)
|
|
|
|
orig_brightness = concat_latent_image_original.mean()
|
|
enhanced_brightness = post_enhanced.mean()
|
|
|
|
if enhanced_brightness < orig_brightness * 0.92:
|
|
brightness_boost = min(orig_brightness / (enhanced_brightness + 1e-6), 1.05)
|
|
post_enhanced = torch.where(
|
|
post_enhanced < 0.5,
|
|
post_enhanced * brightness_boost,
|
|
post_enhanced
|
|
)
|
|
|
|
concat_latent_image = torch.clamp(post_enhanced, -6, 6)
|
|
|
|
positive = node_helpers.conditioning_set_values(positive, {"concat_latent_image": concat_latent_image, "concat_mask": mask})
|
|
negative = node_helpers.conditioning_set_values(negative, {"concat_latent_image": concat_latent_image, "concat_mask": mask})
|
|
|
|
positive_original = node_helpers.conditioning_set_values(positive_original, {"concat_latent_image": concat_latent_image_original, "concat_mask": mask})
|
|
negative_original = node_helpers.conditioning_set_values(negative_original, {"concat_latent_image": concat_latent_image_original, "concat_mask": mask})
|
|
|
|
ref_latent = vae.encode(start_image[:,:,:,:3])
|
|
positive = node_helpers.conditioning_set_values(positive, {"reference_latents": [ref_latent]}, append=True)
|
|
negative = node_helpers.conditioning_set_values(negative, {"reference_latents": [torch.zeros_like(ref_latent)]}, append=True)
|
|
|
|
positive_original = node_helpers.conditioning_set_values(positive_original, {"reference_latents": [ref_latent]}, append=True)
|
|
negative_original = node_helpers.conditioning_set_values(negative_original, {"reference_latents": [torch.zeros_like(ref_latent)]}, append=True)
|
|
|
|
if clip_vision is not None:
|
|
positive = node_helpers.conditioning_set_values(positive, {"clip_vision_output": clip_vision})
|
|
negative = node_helpers.conditioning_set_values(negative, {"clip_vision_output": clip_vision})
|
|
|
|
positive_original = node_helpers.conditioning_set_values(positive_original, {"clip_vision_output": clip_vision})
|
|
negative_original = node_helpers.conditioning_set_values(negative_original, {"clip_vision_output": clip_vision})
|
|
|
|
out_latent = {"samples": latent}
|
|
return io.NodeOutput(positive, negative, positive_original, negative_original, out_latent)
|
|
|
|
|
|
class RikanHiddenBase64ImageLoader(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(
|
|
node_id="RikanHiddenBase64ImageLoader",
|
|
display_name="Rikan Hidden Base64 Image Loader",
|
|
category="Rikannodes",
|
|
inputs=[io.String.Input("base64_data", multiline=True, default="")],
|
|
outputs=[io.Image.Output(display_name="image"), io.String.Output(display_name="raw_base64")]
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, base64_data) -> io.NodeOutput:
|
|
if not base64_data: return io.NodeOutput(torch.zeros((1, 64, 64, 3)), "")
|
|
try:
|
|
if "," in base64_data: base64_data = base64_data.split(",")[1]
|
|
img = Image.open(pyio.BytesIO(base64.b64decode(base64_data))).convert("RGB")
|
|
img = torch.from_numpy(np.array(img).astype(np.float32) / 255.0)[None,]
|
|
return io.NodeOutput(img, base64_data)
|
|
except: return io.NodeOutput(torch.zeros((1, 64, 64, 3)), "")
|
|
|
|
# ── NEW: WAN SPATIO-TEMPORAL TILED VAE DECODE ───────────────────────────────
|
|
|
|
class RikanWanSpatioTemporalTiledVAEDecode(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(
|
|
node_id="RikanWanSpatioTemporalTiledVAEDecode",
|
|
display_name="Rikan Wan Spatio-Temporal Tiled VAE Decode",
|
|
category="Rikannodes",
|
|
inputs=[
|
|
io.Vae.Input("vae"),
|
|
io.Latent.Input("latents"),
|
|
io.Int.Input("spatial_tiles", default=4, min=1, max=8),
|
|
io.Int.Input("spatial_overlap", default=4, min=0, max=8),
|
|
io.Int.Input("temporal_tile_length", default=16, min=2, max=1000),
|
|
io.Int.Input("temporal_overlap", default=4, min=0, max=8),
|
|
],
|
|
outputs=[
|
|
io.Image.Output(display_name="image"),
|
|
]
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, vae, latents, spatial_tiles=2, spatial_overlap=1, temporal_tile_length=16, temporal_overlap=1) -> io.NodeOutput:
|
|
samples = latents["samples"]
|
|
batch, channels, frames, height, width = samples.shape
|
|
|
|
# Menganalisis faktor skala secara dinamis (mengadaptasi berbagai jenis model)
|
|
dummy_latent = torch.zeros((1, channels, 2, 2, 2), device=samples.device, dtype=samples.dtype)
|
|
dummy_out = vae.decode(dummy_latent)
|
|
|
|
# PERBAIKAN: Membaca shape dari belakang agar aman dari perbedaan 4D vs 5D tensor
|
|
f_out = dummy_out.shape[-4]
|
|
h_out = dummy_out.shape[-3]
|
|
w_out = dummy_out.shape[-2]
|
|
|
|
time_scale_factor = (f_out - 1) // (2 - 1)
|
|
height_scale_factor = h_out // 2
|
|
width_scale_factor = w_out // 2
|
|
|
|
image_frames = 1 + (frames - 1) * time_scale_factor
|
|
output_height = height * height_scale_factor
|
|
output_width = width * width_scale_factor
|
|
|
|
target_device = samples.device
|
|
target_dtype = torch.float32
|
|
|
|
def compute_chunk_boundaries(chunk_start, temp_len, temp_overlap, total_frames):
|
|
if chunk_start == 0:
|
|
return chunk_start, min(chunk_start + temp_len, total_frames)
|
|
overlap_start = max(1, chunk_start - temp_overlap - 1)
|
|
extra_frames = chunk_start - overlap_start
|
|
chunk_end = min(chunk_start + temp_len - extra_frames, total_frames)
|
|
return overlap_start, chunk_end
|
|
|
|
def decode_spatial(temporal_chunk):
|
|
chunk_frames = temporal_chunk.shape[2]
|
|
chunk_out_frames = 1 + (chunk_frames - 1) * time_scale_factor
|
|
|
|
base_tile_height = (height + (spatial_tiles - 1) * spatial_overlap) // spatial_tiles
|
|
base_tile_width = (width + (spatial_tiles - 1) * spatial_overlap) // spatial_tiles
|
|
|
|
out_chunk = torch.zeros((batch, chunk_out_frames, output_height, output_width, 3), device=target_device, dtype=target_dtype)
|
|
out_weights = torch.zeros((batch, chunk_out_frames, output_height, output_width, 1), device=target_device, dtype=target_dtype)
|
|
|
|
for v in range(spatial_tiles):
|
|
for h in range(spatial_tiles):
|
|
h_start = h * (base_tile_width - spatial_overlap)
|
|
v_start = v * (base_tile_height - spatial_overlap)
|
|
h_end = min(h_start + base_tile_width, width) if h < spatial_tiles - 1 else width
|
|
v_end = min(v_start + base_tile_height, height) if v < spatial_tiles - 1 else height
|
|
|
|
tile = temporal_chunk[:, :, :, v_start:v_end, h_start:h_end]
|
|
|
|
log.info(f"[Rikan VAE] Decoding spatial tile ({v},{h}) size: {tile.shape}")
|
|
decoded_tile = vae.decode(tile).to(target_device, target_dtype)
|
|
|
|
# PERBAIKAN: Memastikan formatnya selalu 5D [Batch, Frame, Height, Width, Channels]
|
|
if len(decoded_tile.shape) == 4:
|
|
decoded_tile = decoded_tile.view(batch, chunk_out_frames, decoded_tile.shape[-3], decoded_tile.shape[-2], decoded_tile.shape[-1])
|
|
|
|
out_h_start = v_start * height_scale_factor
|
|
out_h_end = v_end * height_scale_factor
|
|
out_w_start = h_start * width_scale_factor
|
|
out_w_end = h_end * width_scale_factor
|
|
|
|
tile_weights = torch.ones((batch, chunk_out_frames, out_h_end - out_h_start, out_w_end - out_w_start, 1), device=target_device, dtype=target_dtype)
|
|
|
|
overlap_out_h = spatial_overlap * height_scale_factor
|
|
overlap_out_w = spatial_overlap * width_scale_factor
|
|
|
|
if h > 0:
|
|
h_blend = torch.linspace(0, 1, overlap_out_w, device=target_device, dtype=target_dtype)
|
|
tile_weights[:, :, :, :overlap_out_w, :] *= h_blend.view(1, 1, 1, -1, 1)
|
|
if h < spatial_tiles - 1:
|
|
h_blend = torch.linspace(1, 0, overlap_out_w, device=target_device, dtype=target_dtype)
|
|
tile_weights[:, :, :, -overlap_out_w:, :] *= h_blend.view(1, 1, 1, -1, 1)
|
|
if v > 0:
|
|
v_blend = torch.linspace(0, 1, overlap_out_h, device=target_device, dtype=target_dtype)
|
|
tile_weights[:, :, :overlap_out_h, :, :] *= v_blend.view(1, 1, -1, 1, 1)
|
|
if v < spatial_tiles - 1:
|
|
v_blend = torch.linspace(1, 0, overlap_out_h, device=target_device, dtype=target_dtype)
|
|
tile_weights[:, :, -overlap_out_h:, :, :] *= v_blend.view(1, 1, -1, 1, 1)
|
|
|
|
out_chunk[:, :, out_h_start:out_h_end, out_w_start:out_w_end, :] += decoded_tile * tile_weights
|
|
out_weights[:, :, out_h_start:out_h_end, out_w_start:out_w_end, :] += tile_weights
|
|
|
|
out_chunk /= out_weights + 1e-8
|
|
return out_chunk
|
|
|
|
output = torch.empty((batch, image_frames, output_height, output_width, 3), device=target_device, dtype=target_dtype)
|
|
|
|
chunk_start = 0
|
|
while chunk_start < frames:
|
|
overlap_start, chunk_end = compute_chunk_boundaries(chunk_start, temporal_tile_length, temporal_overlap, frames)
|
|
|
|
temp_tile = samples[:, :, overlap_start:chunk_end]
|
|
|
|
log.info(f"[Rikan VAE] Processing temporal chunk {overlap_start}:{chunk_end} (length: {chunk_end - overlap_start})")
|
|
decoded_chunk = decode_spatial(temp_tile)
|
|
|
|
if chunk_start == 0:
|
|
output[:, :decoded_chunk.shape[1]] = decoded_chunk
|
|
else:
|
|
decoded_chunk = decoded_chunk[:, 1:]
|
|
out_t_start = 1 + overlap_start * time_scale_factor
|
|
out_t_end = out_t_start + decoded_chunk.shape[1]
|
|
|
|
overlap_frames = temporal_overlap * time_scale_factor
|
|
frame_weights = torch.linspace(0, 1, overlap_frames + 2, device=target_device, dtype=target_dtype)[1:-1]
|
|
tile_weights = frame_weights.view(1, -1, 1, 1, 1)
|
|
|
|
after_overlap = out_t_start + overlap_frames
|
|
|
|
overlap_out = decoded_chunk[:, :overlap_frames]
|
|
output[:, out_t_start:after_overlap] *= (1 - tile_weights)
|
|
output[:, out_t_start:after_overlap] += (tile_weights * overlap_out)
|
|
|
|
output[:, after_overlap:out_t_end] = decoded_chunk[:, overlap_frames:]
|
|
|
|
chunk_start = chunk_end
|
|
|
|
output = output.view(batch * image_frames, output_height, output_width, 3)
|
|
|
|
if target_device.type == "cuda":
|
|
torch.cuda.empty_cache()
|
|
|
|
return io.NodeOutput(output)
|
|
|
|
|
|
# ==============================================================================
|
|
# MAPPINGS & REGISTRATION
|
|
# ==============================================================================
|
|
|
|
class RikannodesExtension(ComfyExtension):
|
|
@override
|
|
async def get_node_list(self) -> list[type[io.ComfyNode]]:
|
|
return [
|
|
RikanPromptRelayEncodeTimeline,
|
|
RikanPromptRelayMultiLoraGate,
|
|
RikanI2VPainterTiledVAE,
|
|
RikanI2VPainter,
|
|
RikanHiddenBase64ImageLoader,
|
|
RikanWanSpatioTemporalTiledVAEDecode
|
|
]
|
|
|
|
async def comfy_entrypoint() -> RikannodesExtension:
|
|
return RikannodesExtension()
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"RikanPromptRelayEncodeTimeline": RikanPromptRelayEncodeTimeline,
|
|
"RikanPromptRelayMultiLoraGate": RikanPromptRelayMultiLoraGate,
|
|
"rikan-i2vpainter-tiled-vae": RikanI2VPainterTiledVAE,
|
|
"rikan-i2vpainter": RikanI2VPainter,
|
|
"RikanHiddenBase64ImageLoader": RikanHiddenBase64ImageLoader,
|
|
"RikanWanSpatioTemporalTiledVAEDecode": RikanWanSpatioTemporalTiledVAEDecode
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"RikanPromptRelayEncodeTimeline": "Rikan Prompt Relay Encode (Timeline)",
|
|
"RikanPromptRelayMultiLoraGate": "Rikan Prompt Relay Multi LoRA Gate",
|
|
"rikan-i2vpainter-tiled-vae": "Rikan I2V Painter (Tiled VAE)",
|
|
"rikan-i2vpainter": "Rikan I2V Painter",
|
|
"RikanHiddenBase64ImageLoader": "Rikan Hidden Base64 Image Loader",
|
|
"RikanWanSpatioTemporalTiledVAEDecode": "Rikan Wan Spatio-Temporal Tiled VAE Decode"
|
|
}
|