digital-garbageandClaude Sonnet 5.5
7022982aa9
Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
110 lines
4.6 KiB
Python
110 lines
4.6 KiB
Python
"""The picture inside a sampling-time latent, whatever the model packs around it."""
|
|
|
|
|
|
def video_of(x, named=None):
|
|
"""-> (video [B, C, T, H, W], rebuild(new_video) -> x), or None.
|
|
|
|
A model module that packs streams answers `video_stream`; a plain 5-D
|
|
latent is its own picture. Anything else is not guessed at. `named` is the
|
|
model call's arguments (see model_args) -- a packed result is only
|
|
unpackable with the `latent_shapes` that travel with it.
|
|
"""
|
|
from . import registry
|
|
split = registry.current().ask("video_stream", x, named or {})
|
|
if split is not None:
|
|
return split
|
|
if getattr(x, "dim", None) and not getattr(x, "is_nested", False) and x.dim() == 5:
|
|
return x, lambda video: video
|
|
return None
|
|
|
|
|
|
def av_video_stream(x, named=None):
|
|
"""The picture of an audio+video model whose latent is [video, audio], in this order, in both the
|
|
shapes that reach a feature:
|
|
|
|
* the model call's result (APPLY_MODEL) is PACKED: comfy's `_apply_model` flattens the parts into
|
|
one [B, 1, N] tensor with `pack_latents`, and the shapes to undo that ride along as
|
|
`latent_shapes`;
|
|
* elsewhere (callbacks, the finished latent) it is a two-part NestedTensor.
|
|
|
|
-> (video [B, C, T, H, W], rebuild(new_video) -> x), or None when `x` is neither.
|
|
"""
|
|
import math
|
|
if getattr(x, "is_nested", False):
|
|
if len(x.tensors) != 2 or x.tensors[0].dim() != 5:
|
|
return None
|
|
from comfy.nested_tensor import NestedTensor
|
|
rest = x.tensors[1:]
|
|
return x.tensors[0], lambda video: NestedTensor([video, *rest])
|
|
shapes = (named or {}).get("latent_shapes")
|
|
if not hasattr(x, "dim") or x.dim() != 3 or not shapes or len(shapes) != 2 \
|
|
or len(shapes[0]) != 5:
|
|
return None
|
|
n = math.prod(shapes[0][1:])
|
|
if x.shape[-1] != n + math.prod(shapes[1][1:]):
|
|
return None
|
|
b = x.shape[0]
|
|
video = x[..., :n].reshape(b, *shapes[0][1:])
|
|
rest = x[..., n:]
|
|
|
|
def rebuild(new_video):
|
|
import torch
|
|
return torch.cat([new_video.reshape(b, 1, n).to(rest.dtype), rest], dim=-1)
|
|
|
|
return video, rebuild
|
|
|
|
|
|
MODEL_ARGS = ("c_concat", "c_crossattn", "control", "transformer_options")
|
|
|
|
|
|
def with_options(args, kwargs, transformer_options):
|
|
"""The same call with other transformer_options: positional where comfy put it (4th),
|
|
by name otherwise. -> (args, kwargs)."""
|
|
if len(args) > MODEL_ARGS.index("transformer_options"):
|
|
args = list(args)
|
|
args[MODEL_ARGS.index("transformer_options")] = transformer_options
|
|
return tuple(args), kwargs
|
|
return args, {**kwargs, "transformer_options": transformer_options}
|
|
|
|
|
|
def with_context(args, kwargs, context):
|
|
"""The same call with another conditioning (c_crossattn): positional where comfy put it (2nd), by name otherwise."""
|
|
if len(args) > MODEL_ARGS.index("c_crossattn"):
|
|
args = list(args)
|
|
args[MODEL_ARGS.index("c_crossattn")] = context
|
|
return tuple(args), kwargs
|
|
return args, {**kwargs, "c_crossattn": context}
|
|
|
|
|
|
#: Set on a call that already carries its own chosen prompt option, so no other feature swaps the prompt again.
|
|
OPTION_PINNED = "funui_option_pinned"
|
|
|
|
|
|
def option_context(base, raw, like, cache, name):
|
|
"""A whole conditioning `raw` (as the encoder made it), put through the model's own text preparation (H3 refines the
|
|
prompt once per run, so the sampler's conditioning is no longer the encoder's output), on the device and dtype of
|
|
`like`. Kept in `cache` under `name` and the shape, device and dtype, so each option is prepared once per run."""
|
|
import torch
|
|
key = (name, like.device, like.dtype, tuple(like.shape))
|
|
if key not in cache:
|
|
prepare = getattr(getattr(base, "diffusion_model", None), "preprocess_text_embeds", None)
|
|
with torch.no_grad():
|
|
if prepare is not None:
|
|
raw = prepare(raw.to(device=like.device, dtype=base.get_dtype_inference()))
|
|
cache[key] = raw.to(device=like.device, dtype=like.dtype)
|
|
return cache[key]
|
|
|
|
|
|
def model_args(args, kwargs) -> dict:
|
|
"""An APPLY_MODEL wrapper's call, by name.
|
|
|
|
ComfyUI hands c_concat, c_crossattn, control and transformer_options on
|
|
POSITIONALLY (model_base.apply_model), and the model's extra conds by name.
|
|
Reading `kwargs.get("transformer_options")` alone finds nothing on a real
|
|
run -- a step gate built on it never opens, silently. A wrapper that calls
|
|
on with changed values may pass everything by name; this reads both shapes.
|
|
"""
|
|
named = dict(zip(MODEL_ARGS, args))
|
|
named.update(kwargs)
|
|
return named
|