Merge remote-tracking branch 'origin/master' into fix/minimax-h3-attention-patches
This commit is contained in:
commit
debf4bc50d
|
|
@ -11,6 +11,10 @@ reviews:
|
|||
review_status: false
|
||||
review_details: true
|
||||
commit_status: true
|
||||
# Without this, a review that never happened (rate limit, internal error)
|
||||
# still posts a green "CodeRabbit" commit status, so a throttled review is
|
||||
# indistinguishable from a clean one.
|
||||
fail_commit_status: true
|
||||
collapse_walkthrough: true
|
||||
changed_files_summary: false
|
||||
sequence_diagrams: false
|
||||
|
|
|
|||
|
|
@ -305,6 +305,12 @@
|
|||
|
||||
- Follow existing node conventions: `INPUT_TYPES`, `RETURN_TYPES`, `FUNCTION`,
|
||||
`CATEGORY`, and registration through the local mapping used by that file.
|
||||
- Treat legacy combo inputs, `io.Combo`, and `io.DynamicCombo` values as
|
||||
untrusted when they affect filesystem access. Any value used as a file or
|
||||
folder name, path component, format, or extension must be validated again at
|
||||
the load/save boundary using an existing `folder_paths` resolver or
|
||||
containment helper, or a fixed allowlist/mapping. Do not rely only on the
|
||||
advertised combo options or prompt validation.
|
||||
- Keep node changes backward compatible by default. Add inputs with sensible
|
||||
defaults and avoid changing output types unless the request requires it.
|
||||
- Model implementations should add the minimal number of ComfyUI nodes required
|
||||
|
|
|
|||
93
README.md
93
README.md
|
|
@ -66,67 +66,24 @@ Supports all operating systems and GPU types (NVIDIA, AMD, Intel, Apple Silicon,
|
|||
See what ComfyUI can do with the [newer template workflows](https://comfy.org/workflows) or old [example workflows](https://comfyanonymous.github.io/ComfyUI_examples/).
|
||||
|
||||
## Features
|
||||
- Nodes/graph/flowchart interface to experiment and create complex Stable Diffusion workflows without needing to code anything.
|
||||
- NOTE: There are many more models supported than the list below, if you want to see what is supported see our templates list inside ComfyUI.
|
||||
- Image Models
|
||||
- SD1.x, SD2.x ([unCLIP](https://comfyanonymous.github.io/ComfyUI_examples/unclip/))
|
||||
- [SDXL](https://comfyanonymous.github.io/ComfyUI_examples/sdxl/), [SDXL Turbo](https://comfyanonymous.github.io/ComfyUI_examples/sdturbo/)
|
||||
- [Stable Cascade](https://comfyanonymous.github.io/ComfyUI_examples/stable_cascade/)
|
||||
- [SD3 and SD3.5](https://comfyanonymous.github.io/ComfyUI_examples/sd3/)
|
||||
- Pixart Alpha and Sigma
|
||||
- [AuraFlow](https://comfyanonymous.github.io/ComfyUI_examples/aura_flow/)
|
||||
- [HunyuanDiT](https://comfyanonymous.github.io/ComfyUI_examples/hunyuan_dit/)
|
||||
- [Flux](https://comfyanonymous.github.io/ComfyUI_examples/flux/)
|
||||
- [Lumina Image 2.0](https://comfyanonymous.github.io/ComfyUI_examples/lumina2/)
|
||||
- [HiDream](https://comfyanonymous.github.io/ComfyUI_examples/hidream/)
|
||||
- [Qwen Image](https://comfyanonymous.github.io/ComfyUI_examples/qwen_image/)
|
||||
- [Hunyuan Image 2.1](https://comfyanonymous.github.io/ComfyUI_examples/hunyuan_image/)
|
||||
- [Flux 2](https://comfyanonymous.github.io/ComfyUI_examples/flux2/)
|
||||
- [Z Image](https://comfyanonymous.github.io/ComfyUI_examples/z_image/)
|
||||
- Ernie Image
|
||||
- Image Editing Models
|
||||
- [Omnigen 2](https://comfyanonymous.github.io/ComfyUI_examples/omnigen/)
|
||||
- [Flux Kontext](https://comfyanonymous.github.io/ComfyUI_examples/flux/#flux-kontext-image-editing-model)
|
||||
- [HiDream E1.1](https://comfyanonymous.github.io/ComfyUI_examples/hidream/#hidream-e11)
|
||||
- [Qwen Image Edit](https://comfyanonymous.github.io/ComfyUI_examples/qwen_image/#edit-model)
|
||||
- Video Models
|
||||
- [Stable Video Diffusion](https://comfyanonymous.github.io/ComfyUI_examples/video/)
|
||||
- [Mochi](https://comfyanonymous.github.io/ComfyUI_examples/mochi/)
|
||||
- [LTX-Video](https://comfyanonymous.github.io/ComfyUI_examples/ltxv/)
|
||||
- [Hunyuan Video](https://comfyanonymous.github.io/ComfyUI_examples/hunyuan_video/)
|
||||
- [Wan 2.1](https://comfyanonymous.github.io/ComfyUI_examples/wan/)
|
||||
- [Wan 2.2](https://comfyanonymous.github.io/ComfyUI_examples/wan22/)
|
||||
- [Hunyuan Video 1.5](https://docs.comfy.org/tutorials/video/hunyuan/hunyuan-video-1-5)
|
||||
- Audio Models
|
||||
- [Stable Audio](https://comfyanonymous.github.io/ComfyUI_examples/audio/)
|
||||
- [ACE Step](https://comfyanonymous.github.io/ComfyUI_examples/audio/)
|
||||
- 3D Models
|
||||
- [Hunyuan3D 2.0](https://docs.comfy.org/tutorials/3d/hunyuan3D-2)
|
||||
- Asynchronous Queue system
|
||||
- Many optimizations: Only re-executes the parts of the workflow that changes between executions.
|
||||
- Smart memory management: can automatically run large models on GPUs with as low as 1GB vram with smart offloading.
|
||||
- Works even if you don't have a GPU with: ```--cpu``` (slow)
|
||||
- Can load ckpt and safetensors: All in one checkpoints or standalone diffusion models, VAEs and CLIP models.
|
||||
- Safe loading of ckpt, pt, pth, etc.. files.
|
||||
- Embeddings/Textual inversion
|
||||
- [Loras (regular, locon and loha)](https://comfyanonymous.github.io/ComfyUI_examples/lora/)
|
||||
- [Hypernetworks](https://comfyanonymous.github.io/ComfyUI_examples/hypernetworks/)
|
||||
- Loading full workflows (with seeds) from generated PNG, WebP and FLAC files.
|
||||
- Saving/Loading workflows as Json files.
|
||||
- Nodes interface can be used to create complex workflows like one for [Hires fix](https://comfyanonymous.github.io/ComfyUI_examples/2_pass_txt2img/) or much more advanced ones.
|
||||
- [Area Composition](https://comfyanonymous.github.io/ComfyUI_examples/area_composition/)
|
||||
- [Inpainting](https://comfyanonymous.github.io/ComfyUI_examples/inpaint/) with both regular and inpainting models.
|
||||
- [ControlNet and T2I-Adapter](https://comfyanonymous.github.io/ComfyUI_examples/controlnet/)
|
||||
- [Upscale Models (ESRGAN, ESRGAN variants, SwinIR, Swin2SR, etc...)](https://comfyanonymous.github.io/ComfyUI_examples/upscale_models/)
|
||||
- [GLIGEN](https://comfyanonymous.github.io/ComfyUI_examples/gligen/)
|
||||
- [Model Merging](https://comfyanonymous.github.io/ComfyUI_examples/model_merging/)
|
||||
- [LCM models and Loras](https://comfyanonymous.github.io/ComfyUI_examples/lcm/)
|
||||
- Latent previews with [TAESD](#how-to-show-high-quality-previews)
|
||||
- Works fully offline: core will never download anything unless you want to.
|
||||
- Optional API nodes to use paid models from external providers through the online [Comfy API](https://docs.comfy.org/tutorials/api-nodes/overview) disable with: `--disable-api-nodes`
|
||||
- [Config file](extra_model_paths.yaml.example) to set the search paths for models.
|
||||
- A visual node graph for building and reusing image, video, audio, 3D, and text workflows without code.
|
||||
- Reusable subgraphs, workflow templates, App Mode, and a local API for integrating workflows into applications.
|
||||
- Efficient local execution with asynchronous queueing, partial graph re-execution, smart VRAM and RAM management, model offloading, and support for quantized models.
|
||||
- Broad native model support. This is a representative list; browse the [workflow library](https://comfy.org/workflows/) for maintained, ready-to-run templates.
|
||||
- [Image generation](https://comfy.org/workflows/tag/text-to-image/): Stable Diffusion 1.5, SDXL, SD3.5, Flux.1, Flux.2, Qwen Image, Z-Image, Hunyuan Image 2.1, HiDream, Lumina Image 2.0, Chroma, Anima, LongCat Image, Ideogram 4, Krea 2, MageFlow, Microsoft Lens, PixelDiT, Kandinsky 5, and Ernie Image.
|
||||
- [Image editing](https://comfy.org/workflows/tag/image-edit/): Flux Kontext, Flux.2 Klein, Qwen Image Edit, HiDream E1.1 and O1, OmniGen2, Boogu, JoyImage Edit, MageFlow Edit, and LongCat Image Edit.
|
||||
- [Video generation](https://comfy.org/workflows/tag/video-generation/): Wan 2.1 and 2.2, LTX-Video 2 and 2.3, HunyuanVideo 1.5, Kandinsky 5 Video, CogVideoX, Cosmos Predict2, Bernini-R, SCAIL 2, and Mochi.
|
||||
- [Audio and video generation](https://comfy.org/workflows/): MiniMax H3 and LTX-AV.
|
||||
- [Audio generation](https://comfy.org/workflows/tag/text-to-audio/): ACE-Step 1.5 and Stable Audio 3.
|
||||
- [3D and vision](https://comfy.org/workflows/): Hunyuan3D 2.1, TripoSplat, SeedVR2, SUPIR, Depth Anything 3, MoGe, SAM 3 and 3.1, RT-DETRv4, and BiRefNet.
|
||||
- [Text generation](https://comfy.org/workflows/tag/text-generation/): Gemma 3 and 4, Qwen3, Qwen3.5, and Qwen3-VL, including multimodal inputs.
|
||||
- Load complete checkpoints or separate diffusion models, VAEs, text encoders, LoRAs, ControlNets, adapters, and upscalers from supported model formats.
|
||||
- Built-in tools for inpainting, outpainting, reference conditioning, masks and compositing, model merging, upscaling, frame interpolation, segmentation, depth estimation, and media processing.
|
||||
- Save and load workflows as JSON, or recover complete workflows and seeds from supported generated media.
|
||||
- Runs fully offline: core does not download anything unless you request it. Use `--disable-api-nodes` to disable the optional paid [Comfy API nodes](https://docs.comfy.org/tutorials/api-nodes/overview) and force all built-in functionality to stay offline.
|
||||
- Extend ComfyUI with custom nodes
|
||||
- Configure additional model locations with [`extra_model_paths.yaml`](extra_model_paths.yaml.example).
|
||||
|
||||
Workflow examples can be found on the [Examples page](https://comfyanonymous.github.io/ComfyUI_examples/)
|
||||
|
||||
## Release Process
|
||||
|
||||
|
|
@ -188,11 +145,19 @@ ComfyUI follows a weekly release cycle targeting Monday but this regularly chang
|
|||
|
||||
# Installing
|
||||
|
||||
## Windows and Mac
|
||||
|
||||
We highly recommend using the [desktop app](https://comfy.org/download):
|
||||
|
||||
### [Link to Download](https://comfy.org/download)
|
||||
|
||||
The desktop app is the easiest and best way to use ComfyUI for new users.
|
||||
|
||||
## Windows Portable
|
||||
|
||||
There is a portable standalone build for Windows that should work for running on Nvidia GPUs or for running on your CPU only on the [releases page](https://github.com/comfyanonymous/ComfyUI/releases).
|
||||
There is a portable standalone build for Windows that should work for running on Nvidia GPUs or for running on your CPU only. It is not recommended for regular users. Regular users should use the desktop app above.
|
||||
|
||||
### [Direct link to download](https://github.com/comfyanonymous/ComfyUI/releases/latest/download/ComfyUI_windows_portable_nvidia.7z)
|
||||
[Direct link to download (nvidia)](https://github.com/comfyanonymous/ComfyUI/releases/latest/download/ComfyUI_windows_portable_nvidia.7z)
|
||||
|
||||
Simply download, extract with [7-Zip](https://7-zip.org) or with the windows explorer on recent windows versions and run. For smaller models you normally only need to put the checkpoints (the huge ckpt/safetensors files) in: ComfyUI\models\checkpoints but many of the larger models have multiple files. Make sure to follow the instructions to know which subfolder to put them in ComfyUI\models\
|
||||
|
||||
|
|
@ -208,7 +173,7 @@ The portable above currently comes with python 3.13 and pytorch cuda 13.0. Updat
|
|||
|
||||
[Portable for Nvidia GPUs](https://github.com/comfyanonymous/ComfyUI/releases/latest/download/ComfyUI_windows_portable_nvidia.7z) (supports 20 series and above).
|
||||
|
||||
[Portable for Nvidia GPUs with pytorch cuda 12.6 and python 3.12](https://github.com/comfyanonymous/ComfyUI/releases/latest/download/ComfyUI_windows_portable_nvidia_cu126.7z) (Supports Nvidia 10 series and older GPUs).
|
||||
[Portable for Nvidia GPUs with pytorch cuda 12.6 and python 3.12](https://github.com/comfyanonymous/ComfyUI/releases/latest/download/ComfyUI_windows_portable_nvidia_cu126.7z) (Supports Nvidia 10 series and older GPUs, DO NOT USE THIS ON NEWER 20 SERIES AND ABOVE GPUS).
|
||||
|
||||
#### How do I share models between another UI and ComfyUI?
|
||||
|
||||
|
|
@ -229,7 +194,7 @@ Python 3.14 works but some custom nodes may have issues. The free threaded varia
|
|||
|
||||
Python 3.13 is very well supported. If you have trouble with some custom node dependencies on 3.13 you can try 3.12
|
||||
|
||||
torch 2.5 is minimally supported but using a newer version is extremely recommended. Some features and optimizations might only work on newer versions. We generally recommend using the latest major version of pytorch with the latest cuda version unless it is less than 2 weeks old. If your pytorch is more than 6 months old, please update it.
|
||||
torch 2.7 is minimally supported but using a newer version is extremely recommended. Using a cu130 or above version of pytorch is required on Nvidia 20 series and above. Some features and optimizations might only work on newer versions. We generally recommend using the latest major version of pytorch with the latest cuda version unless it is less than 2 weeks old. If your pytorch is more than 6 months old, please update it.
|
||||
|
||||
### Instructions:
|
||||
|
||||
|
|
|
|||
|
|
@ -433,19 +433,16 @@ class DeformableConv2d(nn.Module):
|
|||
def forward(self, x):
|
||||
offset = self.offset_conv(x)
|
||||
modulator = 2. * torch.sigmoid(self.modulator_conv(x))
|
||||
weight, bias, offload_info = comfy.ops.cast_bias_weight(self.regular_conv, x, offloadable=True)
|
||||
|
||||
x = deform_conv2d(
|
||||
input=x,
|
||||
offset=offset,
|
||||
weight=weight,
|
||||
bias=None,
|
||||
padding=self.padding,
|
||||
mask=modulator,
|
||||
stride=self.stride,
|
||||
)
|
||||
comfy.ops.uncast_bias_weight(self.regular_conv, weight, bias, offload_info)
|
||||
return x
|
||||
with comfy.ops.CastBiasWeightContext(self.regular_conv, x, offloadable=True) as (weight, _bias):
|
||||
return deform_conv2d(
|
||||
input=x,
|
||||
offset=offset,
|
||||
weight=weight,
|
||||
bias=None,
|
||||
padding=self.padding,
|
||||
mask=modulator,
|
||||
stride=self.stride,
|
||||
)
|
||||
|
||||
class BasicDecBlk(nn.Module):
|
||||
def __init__(self, in_channels=64, out_channels=64, inter_channels=64, device=None, dtype=None, operations=None):
|
||||
|
|
|
|||
|
|
@ -381,13 +381,10 @@ class ControlLoraOps:
|
|||
self.bias = None
|
||||
|
||||
def forward(self, input):
|
||||
weight, bias, offload_stream = comfy.ops.cast_bias_weight(self, input, offloadable=True)
|
||||
if self.up is not None:
|
||||
x = torch.nn.functional.linear(input, weight + (torch.mm(self.up.flatten(start_dim=1), self.down.flatten(start_dim=1))).reshape(self.weight.shape).type(input.dtype), bias)
|
||||
else:
|
||||
x = torch.nn.functional.linear(input, weight, bias)
|
||||
comfy.ops.uncast_bias_weight(self, weight, bias, offload_stream)
|
||||
return x
|
||||
with comfy.ops.CastBiasWeightContext(self, input, offloadable=True) as (weight, bias):
|
||||
if self.up is None:
|
||||
return torch.nn.functional.linear(input, weight, bias)
|
||||
return torch.nn.functional.linear(input, weight + (torch.mm(self.up.flatten(start_dim=1), self.down.flatten(start_dim=1))).reshape(self.weight.shape).type(input.dtype), bias)
|
||||
|
||||
class Conv2d(torch.nn.Module, comfy.ops.CastWeightBiasOp):
|
||||
def __init__(
|
||||
|
|
|
|||
|
|
@ -1445,7 +1445,7 @@ def res_multistep(model, x, sigmas, extra_args=None, callback=None, disable=None
|
|||
x = sigma_fn(h) * x + h * (b1 * denoised + b2 * old_denoised)
|
||||
|
||||
# Noise addition
|
||||
if sigmas[i + 1] > 0:
|
||||
if sigma_up > 0:
|
||||
x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up
|
||||
|
||||
if cfg_pp:
|
||||
|
|
|
|||
|
|
@ -16,7 +16,9 @@ from comfy.ldm.lightricks.model import (
|
|||
from comfy.ldm.lightricks.symmetric_patchifier import AudioPatchifier
|
||||
from comfy.ldm.lightricks.embeddings_connector import Embeddings1DConnector
|
||||
import comfy.ldm.common_dit
|
||||
import comfy.model_management
|
||||
import comfy.model_prefetch
|
||||
import comfy.quant_ops
|
||||
|
||||
class CompressedTimestep:
|
||||
"""Store video timestep embeddings in compressed form using per-frame indexing."""
|
||||
|
|
@ -271,7 +273,10 @@ class BasicAVTransformerBlock(nn.Module):
|
|||
if run_vx:
|
||||
# video self-attention
|
||||
vshift_msa, vscale_msa = (self.get_ada_values(self.scale_shift_table, vx.shape[0], v_timestep, slice(0, 2)))
|
||||
norm_vx = comfy.ldm.common_dit.rms_norm(vx) * (1 + vscale_msa) + vshift_msa
|
||||
if comfy.model_management.in_training:
|
||||
norm_vx = comfy.ldm.common_dit.rms_norm(vx) * (1 + vscale_msa) + vshift_msa
|
||||
else:
|
||||
norm_vx = comfy.quant_ops.ck.rms_adaln(vx, vscale_msa, vshift_msa)
|
||||
del vshift_msa, vscale_msa
|
||||
attn1_out = self.attn1(norm_vx, pe=v_pe, mask=self_attention_mask, transformer_options=transformer_options)
|
||||
del norm_vx
|
||||
|
|
@ -305,7 +310,6 @@ class BasicAVTransformerBlock(nn.Module):
|
|||
|
||||
# video - audio cross attention.
|
||||
if run_a2v or run_v2a:
|
||||
vx_norm3 = comfy.ldm.common_dit.rms_norm(vx)
|
||||
ax_norm3 = comfy.ldm.common_dit.rms_norm(ax)
|
||||
|
||||
# audio to video cross attention
|
||||
|
|
@ -315,7 +319,10 @@ class BasicAVTransformerBlock(nn.Module):
|
|||
scale_ca_video_hidden_states_a2v_v, shift_ca_video_hidden_states_a2v_v = self.get_ada_values(
|
||||
self.scale_shift_table_a2v_ca_video[:4, :], vx.shape[0], v_cross_scale_shift_timestep)[:2]
|
||||
|
||||
vx_scaled = vx_norm3 * (1 + scale_ca_video_hidden_states_a2v_v) + shift_ca_video_hidden_states_a2v_v
|
||||
if comfy.model_management.in_training:
|
||||
vx_scaled = comfy.ldm.common_dit.rms_norm(vx) * (1 + scale_ca_video_hidden_states_a2v_v) + shift_ca_video_hidden_states_a2v_v
|
||||
else:
|
||||
vx_scaled = comfy.quant_ops.ck.rms_adaln(vx, scale_ca_video_hidden_states_a2v_v, shift_ca_video_hidden_states_a2v_v)
|
||||
ax_scaled = ax_norm3 * (1 + scale_ca_audio_hidden_states_a2v) + shift_ca_audio_hidden_states_a2v
|
||||
del scale_ca_video_hidden_states_a2v_v, shift_ca_video_hidden_states_a2v_v, scale_ca_audio_hidden_states_a2v, shift_ca_audio_hidden_states_a2v
|
||||
|
||||
|
|
@ -334,7 +341,10 @@ class BasicAVTransformerBlock(nn.Module):
|
|||
self.scale_shift_table_a2v_ca_video[:4, :], vx.shape[0], v_cross_scale_shift_timestep)[2:4]
|
||||
|
||||
ax_scaled = ax_norm3 * (1 + scale_ca_audio_hidden_states_v2a) + shift_ca_audio_hidden_states_v2a
|
||||
vx_scaled = vx_norm3 * (1 + scale_ca_video_hidden_states_v2a) + shift_ca_video_hidden_states_v2a
|
||||
if comfy.model_management.in_training:
|
||||
vx_scaled = comfy.ldm.common_dit.rms_norm(vx) * (1 + scale_ca_video_hidden_states_v2a) + shift_ca_video_hidden_states_v2a
|
||||
else:
|
||||
vx_scaled = comfy.quant_ops.ck.rms_adaln(vx, scale_ca_video_hidden_states_v2a, shift_ca_video_hidden_states_v2a)
|
||||
del scale_ca_video_hidden_states_v2a, shift_ca_video_hidden_states_v2a, scale_ca_audio_hidden_states_v2a, shift_ca_audio_hidden_states_v2a
|
||||
|
||||
v2a_out = self.video_to_audio_attn(ax_scaled, context=vx_scaled, pe=a_cross_pe, k_pe=v_cross_pe, transformer_options=transformer_options)
|
||||
|
|
@ -344,12 +354,14 @@ class BasicAVTransformerBlock(nn.Module):
|
|||
ax.addcmul_(v2a_out, gate_out_v2a)
|
||||
del gate_out_v2a, v2a_out
|
||||
|
||||
del vx_norm3, ax_norm3
|
||||
|
||||
# video feedforward
|
||||
if run_vx:
|
||||
vshift_mlp, vscale_mlp = self.get_ada_values(self.scale_shift_table, vx.shape[0], v_timestep, slice(3, 5))
|
||||
vx_scaled = comfy.ldm.common_dit.rms_norm(vx) * (1 + vscale_mlp) + vshift_mlp
|
||||
if comfy.model_management.in_training:
|
||||
vx_scaled = comfy.ldm.common_dit.rms_norm(vx) * (1 + vscale_mlp) + vshift_mlp
|
||||
else:
|
||||
vx_scaled = comfy.quant_ops.ck.rms_adaln(vx, vscale_mlp, vshift_mlp)
|
||||
del vshift_mlp, vscale_mlp
|
||||
|
||||
ff_out = self.ff(vx_scaled)
|
||||
|
|
|
|||
|
|
@ -13,6 +13,7 @@ import comfy.patcher_extension
|
|||
import comfy.ldm.modules.attention
|
||||
import comfy.ldm.common_dit
|
||||
import comfy.model_management
|
||||
import comfy.ops
|
||||
import comfy.quant_ops
|
||||
|
||||
from .symmetric_patchifier import SymmetricPatchifier, latent_to_pixel_coords
|
||||
|
|
@ -321,7 +322,11 @@ class FeedForward(nn.Module):
|
|||
)
|
||||
|
||||
def forward(self, x):
|
||||
return self.net(x)
|
||||
# net = [GELU_approx(proj), Dropout, Linear]; the fused path skips the
|
||||
# Dropout, so leave it to the stock path whenever it could be active.
|
||||
if comfy.model_management.in_training:
|
||||
return self.net(x)
|
||||
return comfy.ops.linear_input_act(self.net[2], self.net[0].proj(x), "gelu_tanh")
|
||||
|
||||
def apply_rotary_emb(input_tensor, freqs_cis):
|
||||
rotation_matrix, split_pe = freqs_cis
|
||||
|
|
@ -535,7 +540,12 @@ class BasicTransformerBlock(nn.Module):
|
|||
def forward(self, x, context=None, attention_mask=None, timestep=None, pe=None, transformer_options={}, self_attention_mask=None, prompt_timestep=None):
|
||||
shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (self.scale_shift_table[None, None, :6].to(device=x.device, dtype=x.dtype) + timestep.reshape(x.shape[0], timestep.shape[1], self.scale_shift_table.shape[0], -1)[:, :, :6, :]).unbind(dim=2)
|
||||
|
||||
x += self.attn1(comfy.ldm.common_dit.rms_norm(x) * (1 + scale_msa) + shift_msa, pe=pe, mask=self_attention_mask, transformer_options=transformer_options) * gate_msa
|
||||
if comfy.model_management.in_training:
|
||||
norm_x = comfy.ldm.common_dit.rms_norm(x) * (1 + scale_msa) + shift_msa
|
||||
else:
|
||||
norm_x = comfy.quant_ops.ck.rms_adaln(x, scale_msa, shift_msa)
|
||||
|
||||
x += self.attn1(norm_x, pe=pe, mask=self_attention_mask, transformer_options=transformer_options) * gate_msa
|
||||
|
||||
if self.cross_attention_adaln:
|
||||
shift_q_mca, scale_q_mca, gate_mca = (self.scale_shift_table[None, None, 6:9].to(device=x.device, dtype=x.dtype) + timestep.reshape(x.shape[0], timestep.shape[1], self.scale_shift_table.shape[0], -1)[:, :, 6:9, :]).unbind(dim=2)
|
||||
|
|
@ -589,7 +599,10 @@ def apply_cross_attention_adaln(
|
|||
prompt_scale_shift_table[None, None].to(device=x.device, dtype=x.dtype)
|
||||
+ prompt_timestep.reshape(batch_size, prompt_timestep.shape[1], 2, -1)
|
||||
).unbind(dim=2)
|
||||
attn_input = comfy.ldm.common_dit.rms_norm(x) * (1 + q_scale) + q_shift
|
||||
if comfy.model_management.in_training:
|
||||
attn_input = comfy.ldm.common_dit.rms_norm(x) * (1 + q_scale) + q_shift
|
||||
else:
|
||||
attn_input = comfy.quant_ops.ck.rms_adaln(x, q_scale, q_shift)
|
||||
encoder_hidden_states = context * (1 + scale_kv) + shift_kv
|
||||
return attn(attn_input, context=encoder_hidden_states, mask=attention_mask, transformer_options=transformer_options) * q_gate
|
||||
|
||||
|
|
|
|||
|
|
@ -99,7 +99,7 @@ class UpSample1d(nn.Module):
|
|||
def forward(self, x):
|
||||
_, C, _ = x.shape
|
||||
x = F.pad(x, (self.pad, self.pad), mode="replicate")
|
||||
x = F.conv_transpose1d(x, self.filter.expand(C, -1, -1).to(x.dtype), stride=self.stride, groups=C).mul_(self.ratio)
|
||||
x = F.conv_transpose1d(x, comfy.ops.cast_to_input(self.filter.expand(C, -1, -1), x), stride=self.stride, groups=C).mul_(self.ratio)
|
||||
x = x[..., self.pad_left:-self.pad_right]
|
||||
return x
|
||||
|
||||
|
|
@ -115,7 +115,7 @@ class LowPassFilter1d(nn.Module):
|
|||
def forward(self, x):
|
||||
_, C, _ = x.shape
|
||||
x = F.pad(x, (self.pad_left, self.pad_right), mode="replicate")
|
||||
return F.conv1d(x, self.filter.expand(C, -1, -1).to(x.dtype), stride=self.stride, groups=C)
|
||||
return F.conv1d(x, comfy.ops.cast_to_input(self.filter.expand(C, -1, -1), x), stride=self.stride, groups=C)
|
||||
|
||||
|
||||
class DownSample1d(nn.Module):
|
||||
|
|
|
|||
|
|
@ -9,8 +9,9 @@ The packed sequence is:
|
|||
Timestep domain: the model receives the *video* sigma from the sampler and
|
||||
derives per-token timesteps t = 1 - sigma internally; the audio stream runs on
|
||||
its own shifted schedule (sigma_shift video 12.0 / audio 3.0), mapped from the
|
||||
video sigma in closed form. The audio velocity is returned scaled by the
|
||||
schedule map's derivative d(sigma_a)/d(sigma_v).
|
||||
video sigma in closed form. The sampler carries the audio latent scaled onto the
|
||||
video schedule (ModelSamplingAV); forward() undoes that scale and converts the
|
||||
velocity back, so _forward only ever sees the stream's own latent.
|
||||
"""
|
||||
|
||||
import math
|
||||
|
|
@ -38,17 +39,6 @@ def time_shift_sigma(sigma, from_shift, to_shift):
|
|||
return to_shift * base / (1.0 + (to_shift - 1.0) * base)
|
||||
|
||||
|
||||
def time_shift_slope(sigma, from_shift, to_shift):
|
||||
"""d(sigma_to)/d(sigma_from) at the same base-grid point.
|
||||
|
||||
Scaling a stream's returned velocity by this slope makes the flat ODE that
|
||||
any sampler integrates on the from-schedule equal to that stream's true ODE
|
||||
on its own schedule.
|
||||
"""
|
||||
base = sigma / (from_shift + sigma * (1.0 - from_shift))
|
||||
return (to_shift * (1.0 + (from_shift - 1.0) * base) ** 2) / (from_shift * (1.0 + (to_shift - 1.0) * base) ** 2)
|
||||
|
||||
|
||||
def patchify_video(latent, patch_size=(1, 2, 2)):
|
||||
# [B, C, T, H, W] -> [B*t*h*w, C*pt*ph*pw]
|
||||
b, c, t_full, h_full, w_full = latent.shape
|
||||
|
|
@ -526,12 +516,30 @@ class MiniMaxH3Model(nn.Module):
|
|||
return torch.cat(rows, dim=0) if rows else None
|
||||
|
||||
def forward(self, x, timestep, context, transformer_options={}, minimax_payload=None, **kwargs):
|
||||
return comfy.patcher_extension.WrapperExecutor.new_class_executor(
|
||||
# the sampler carries the audio as (sigma_v / sigma_a) * x_audio; undo it outside
|
||||
# the wrappers so they and the network see the stream's own latent and velocity
|
||||
scale = float((minimax_payload or {}).get("audio_scale", 1.0))
|
||||
audio_src = x[1]
|
||||
if scale != 1.0:
|
||||
shift_v = float(transformer_options.get("minimax_h3_sigma_shift_video", self.sigma_shift_video))
|
||||
shift_a = float(transformer_options.get("minimax_h3_sigma_shift_audio", self.sigma_shift_audio))
|
||||
sigma_v = (timestep.flatten()[0] / 1000.0).float().clamp(min=1e-6)
|
||||
sigma_a = time_shift_sigma(sigma_v, shift_v, shift_a)
|
||||
carry = (sigma_a / sigma_v).to(audio_src.dtype)
|
||||
x = [x[0], audio_src * carry]
|
||||
|
||||
out = comfy.patcher_extension.WrapperExecutor.new_class_executor(
|
||||
self._forward,
|
||||
self,
|
||||
comfy.patcher_extension.get_all_wrappers(comfy.patcher_extension.WrappersMP.DIFFUSION_MODEL, transformer_options)
|
||||
).execute(x, timestep, context, transformer_options, minimax_payload=minimax_payload, **kwargs)
|
||||
|
||||
if scale != 1.0:
|
||||
# d/d(sigma_v) of the carried variable
|
||||
out[1] = ((1.0 - scale) * (audio_src * carry)
|
||||
+ (1.0 + (scale - 1.0) * sigma_a).to(out[1].dtype) * out[1])
|
||||
return out
|
||||
|
||||
def _forward(self, x, timestep, context, transformer_options={}, minimax_payload=None, **kwargs):
|
||||
video_x, audio_x = x[0], x[1]
|
||||
orig_t, orig_h, orig_w = video_x.shape[2], video_x.shape[3], video_x.shape[4]
|
||||
|
|
@ -672,8 +680,4 @@ class MiniMaxH3Model(nn.Module):
|
|||
video_out = video_out[:, :, :orig_t, :orig_h, :orig_w]
|
||||
audio_out = unpack_audio(a)
|
||||
|
||||
# The sampler integrates the flat ODE dX/dsigma_v = (X - denoised)/sigma_v.
|
||||
# Scaling the audio velocity by d(sigma_a)/d(sigma_v) makes that ODE equal
|
||||
# to the audio stream's true ODE on its own shifted schedule.
|
||||
slope_a = time_shift_slope(sigma_v, shift_v, shift_a).to(audio_out.dtype)
|
||||
return [-video_out.to(video_x.dtype), (-slope_a) * audio_out.to(audio_x.dtype)]
|
||||
return [-video_out.to(video_x.dtype), -audio_out.to(audio_x.dtype)]
|
||||
|
|
|
|||
|
|
@ -199,11 +199,11 @@ class RotaryEmbeddingND(nn.Module):
|
|||
|
||||
class FeedForward(nn.Module):
|
||||
# Gated SiLU FFN.
|
||||
def __init__(self, dim, mult=4, bias=True):
|
||||
def __init__(self, dim, mult=4, bias=True, operations=ops):
|
||||
super().__init__()
|
||||
inner_dim = dim * mult
|
||||
self.w1 = ops.Linear(dim, inner_dim * 2, bias=bias)
|
||||
self.w2 = ops.Linear(inner_dim, dim, bias=bias)
|
||||
self.w1 = operations.Linear(dim, inner_dim * 2, bias=bias)
|
||||
self.w2 = operations.Linear(inner_dim, dim, bias=bias)
|
||||
|
||||
def forward(self, x):
|
||||
gate, x = self.w1(x).chunk(2, dim=-1)
|
||||
|
|
@ -211,15 +211,15 @@ class FeedForward(nn.Module):
|
|||
|
||||
|
||||
class Attention(nn.Module):
|
||||
def __init__(self, heads, dim_head, bias=True, eps=1e-5):
|
||||
def __init__(self, heads, dim_head, bias=True, eps=1e-5, operations=ops):
|
||||
super().__init__()
|
||||
self.dim_head = dim_head
|
||||
self.heads = heads
|
||||
inner_dim = dim_head * heads
|
||||
self.norm_q = ops.RMSNorm(dim_head, eps=eps, elementwise_affine=False)
|
||||
self.norm_k = ops.RMSNorm(dim_head, eps=eps, elementwise_affine=False)
|
||||
self.to_qkv = ops.Linear(inner_dim, inner_dim * 3, bias=bias)
|
||||
self.to_out = ops.Linear(inner_dim, inner_dim, bias=bias)
|
||||
self.to_qkv = operations.Linear(inner_dim, inner_dim * 3, bias=bias)
|
||||
self.to_out = operations.Linear(inner_dim, inner_dim, bias=bias)
|
||||
|
||||
def forward(self, x, rotary_pos_emb=None):
|
||||
batch_size, seq_len, _ = x.shape
|
||||
|
|
@ -242,14 +242,14 @@ class Attention(nn.Module):
|
|||
|
||||
|
||||
class TransformerBlock(nn.Module):
|
||||
def __init__(self, heads, dim_head, bias=True, eps=1e-5):
|
||||
def __init__(self, heads, dim_head, bias=True, eps=1e-5, operations=ops):
|
||||
super().__init__()
|
||||
dim = heads * dim_head
|
||||
self.norm1 = ops.RMSNorm(dim, elementwise_affine=True, eps=eps)
|
||||
self.attn = Attention(heads=heads, dim_head=dim_head, bias=bias, eps=eps)
|
||||
self.attn = Attention(heads=heads, dim_head=dim_head, bias=bias, eps=eps, operations=operations)
|
||||
self.scale1 = nn.Parameter(torch.empty(dim))
|
||||
self.norm2 = ops.RMSNorm(dim, elementwise_affine=True, eps=eps)
|
||||
self.ff = FeedForward(dim=dim, bias=bias)
|
||||
self.ff = FeedForward(dim=dim, bias=bias, operations=operations)
|
||||
self.scale2 = nn.Parameter(torch.empty(dim))
|
||||
|
||||
def forward(self, x, rotary_pos_emb=None):
|
||||
|
|
@ -259,7 +259,7 @@ class TransformerBlock(nn.Module):
|
|||
|
||||
class ViT3DDecoder(nn.Module):
|
||||
def __init__(self, patch_size=16, patch_size_t=4, in_channels=24, out_channels=3, num_layers=36, heads=32, dim_head=64, rope_theta=100.0,
|
||||
rope_dim_ratio=0.75, bias=True, eps=1e-5, num_register_tokens=4):
|
||||
rope_dim_ratio=0.75, bias=True, eps=1e-5, num_register_tokens=4, operations=ops):
|
||||
super().__init__()
|
||||
dim = heads * dim_head
|
||||
self.patch_size = patch_size
|
||||
|
|
@ -274,7 +274,7 @@ class ViT3DDecoder(nn.Module):
|
|||
self.register_buffer("mask_token", torch.empty(1, 1, dim))
|
||||
|
||||
self.transformer_blocks = nn.ModuleList(
|
||||
[TransformerBlock(heads=heads, dim_head=dim_head, bias=bias, eps=eps)
|
||||
[TransformerBlock(heads=heads, dim_head=dim_head, bias=bias, eps=eps, operations=operations)
|
||||
for _ in range(num_layers)]
|
||||
)
|
||||
|
||||
|
|
@ -337,6 +337,7 @@ class MiniMaxH3VideoVAE(nn.Module):
|
|||
tile_size=256,
|
||||
tile_overlap_min=64,
|
||||
tiling=True,
|
||||
operations=ops,
|
||||
):
|
||||
super().__init__()
|
||||
self.vae_ratio = int(math.prod(space_down))
|
||||
|
|
@ -372,6 +373,7 @@ class MiniMaxH3VideoVAE(nn.Module):
|
|||
patch_size_t=self.vae_ratio_t,
|
||||
in_channels=z_channels,
|
||||
out_channels=out_ch,
|
||||
operations=operations,
|
||||
)
|
||||
|
||||
self.register_buffer("latents_mean", torch.tensor(LATENTS_MEAN))
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@ from comfy.ldm.flux.layers import EmbedND
|
|||
from comfy.ldm.flux.math import apply_rope1, rope
|
||||
import comfy.ldm.common_dit
|
||||
import comfy.model_management
|
||||
import comfy.ops
|
||||
import comfy.patcher_extension
|
||||
|
||||
|
||||
|
|
@ -146,7 +147,8 @@ class WanI2VCrossAttention(WanSelfAttention):
|
|||
v = self.v(context)
|
||||
k_img = self.norm_k_img(self.k_img(context_img))
|
||||
v_img = self.v_img(context_img)
|
||||
img_x = optimized_attention(q, k_img, v_img, heads=self.num_heads, transformer_options=transformer_options)
|
||||
# Sageattn can cause Nans here, don't allow it as there is no speed difference anyway as img attention is tiny.
|
||||
img_x = optimized_attention(q, k_img, v_img, heads=self.num_heads, transformer_options=transformer_options, low_precision_attention=False)
|
||||
# compute attention
|
||||
x = optimized_attention(q, k, v, heads=self.num_heads, transformer_options=transformer_options)
|
||||
|
||||
|
|
@ -174,6 +176,13 @@ def repeat_e(e, x):
|
|||
return torch.repeat_interleave(e, repeats + 1, dim=1)[:, :x.size(1)]
|
||||
|
||||
|
||||
class WanFeedForward(nn.Sequential):
|
||||
"""[Linear, GELU(tanh), Linear], with the GELU folded into the down-projection."""
|
||||
|
||||
def forward(self, x):
|
||||
return comfy.ops.linear_input_act(self[2], self[0](x), "gelu_tanh")
|
||||
|
||||
|
||||
class WanAttentionBlock(nn.Module):
|
||||
|
||||
def __init__(self,
|
||||
|
|
@ -207,7 +216,7 @@ class WanAttentionBlock(nn.Module):
|
|||
qk_norm,
|
||||
eps, operation_settings=operation_settings)
|
||||
self.norm2 = operation_settings.get("operations").LayerNorm(dim, eps, elementwise_affine=False, device=operation_settings.get("device"), dtype=operation_settings.get("dtype"))
|
||||
self.ffn = nn.Sequential(
|
||||
self.ffn = WanFeedForward(
|
||||
operation_settings.get("operations").Linear(dim, ffn_dim, device=operation_settings.get("device"), dtype=operation_settings.get("dtype")), nn.GELU(approximate='tanh'),
|
||||
operation_settings.get("operations").Linear(ffn_dim, dim, device=operation_settings.get("device"), dtype=operation_settings.get("dtype")))
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,387 @@
|
|||
# Wan-Animate-2: https://github.com/Wan-Video/Wan-Animate-2
|
||||
"""Wan2.1-I2V-14B weights, driven by a video instead of a motion extractor.
|
||||
|
||||
A pose branch over the pose video's latents runs in lockstep with the generation
|
||||
branch, feeding it K/V per block. The reference image is one extra latent frame at the
|
||||
front of the generation branch, trimmed off by the caller. Upstream calls the pose video
|
||||
the driving video and its branch forward_ref, not to be confused with the reference image.
|
||||
"""
|
||||
|
||||
import torch
|
||||
|
||||
import comfy.ldm.common_dit
|
||||
import comfy.model_management
|
||||
import comfy.quant_ops
|
||||
import comfy.utils
|
||||
from comfy.ldm.flux.math import apply_rope1
|
||||
from comfy.ldm.modules.attention import optimized_attention
|
||||
|
||||
from .model import WanAttentionBlock, WanModel, WanSelfAttention, repeat_e, sinusoidal_embedding_1d
|
||||
|
||||
|
||||
class WanAnimate2SelfAttention(WanSelfAttention):
|
||||
|
||||
def qkv(self, x, freqs):
|
||||
b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
|
||||
q = apply_rope1(self.norm_q(self.q(x)).view(b, s, n, d), freqs)
|
||||
k = apply_rope1(self.norm_k(self.k(x)).view(b, s, n, d), freqs)
|
||||
return q, k, self.v(x).view(b, s, n, d)
|
||||
|
||||
def _attn1_patch(self, x, q, k, transformer_options):
|
||||
for p in transformer_options.get("patches", {}).get("attn1_patch", []):
|
||||
x = p({"x": x, "q": q, "k": k, "transformer_options": transformer_options})
|
||||
return x
|
||||
|
||||
def kv(self, x, freqs):
|
||||
b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
|
||||
return apply_rope1(self.norm_k(self.k(x)).view(b, s, n, d), freqs), self.v(x).view(b, s, n, d)
|
||||
|
||||
def forward_pose(self, x, freqs, transformer_options={}):
|
||||
b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
|
||||
q, k, v = self.qkv(x, freqs)
|
||||
out = optimized_attention(q.reshape(b, s, n * d), k.reshape(b, s, n * d), v.reshape(b, s, n * d), heads=self.num_heads, transformer_options=transformer_options)
|
||||
return self.o(self._attn1_patch(out, q, k, transformer_options)), k, v
|
||||
|
||||
def forward_gen(self, x, freqs, k_pose, v_pose, f_gen, hw, buffers, ref_strength=1.0, transformer_options={}):
|
||||
# frame j attends every gen token plus pose frame j-1 (frame 0 is the reference slot and has none)
|
||||
b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
|
||||
q, k, v = self.qkv(x, freqs)
|
||||
if ref_strength != 1.0:
|
||||
v[:, :hw] *= ref_strength # frame 0 is the reference image's slot
|
||||
|
||||
if k_pose is None: # pose influence windowed out: plain self-attention, no per-frame loop
|
||||
out = optimized_attention(q.reshape(b, s, n * d), k.reshape(b, s, n * d), v.reshape(b, s, n * d), heads=self.num_heads, transformer_options=transformer_options)
|
||||
return self.o(self._attn1_patch(out, q, k, transformer_options))
|
||||
|
||||
# gen half is the same every frame; only the hw-token pose tail is rewritten
|
||||
kbuf, vbuf, out = buffers
|
||||
kbuf[:, :s] = k
|
||||
vbuf[:, :s] = v
|
||||
|
||||
for j in range(f_gen):
|
||||
q_j = q[:, j * hw:(j + 1) * hw].reshape(b, hw, n * d)
|
||||
if j == 0:
|
||||
kk, vv = k, v
|
||||
else:
|
||||
kbuf[:, s:] = k_pose[:, (j - 1) * hw:j * hw]
|
||||
vbuf[:, s:] = v_pose[:, (j - 1) * hw:j * hw]
|
||||
kk, vv = kbuf, vbuf
|
||||
out[:, j * hw:(j + 1) * hw] = optimized_attention(q_j, kk.reshape(b, kk.shape[1], n * d), vv.reshape(b, kk.shape[1], n * d), heads=self.num_heads, transformer_options=transformer_options)
|
||||
return self.o(self._attn1_patch(out, q, k, transformer_options))
|
||||
|
||||
|
||||
class WanAnimate2Block(WanAttentionBlock):
|
||||
|
||||
def __init__(self, cross_attn_type, dim, ffn_dim, num_heads, window_size=(-1, -1), qk_norm=True, cross_attn_norm=False, eps=1e-6, operation_settings={}):
|
||||
super().__init__(cross_attn_type, dim, ffn_dim, num_heads, window_size, qk_norm, cross_attn_norm, eps, operation_settings=operation_settings)
|
||||
self.self_attn = WanAnimate2SelfAttention(dim, num_heads, window_size, qk_norm, eps, operation_settings=operation_settings)
|
||||
|
||||
def _modulation(self, e, x):
|
||||
if e.ndim < 4:
|
||||
return (comfy.model_management.cast_to(self.modulation, dtype=x.dtype, device=x.device) + e).chunk(6, dim=1)
|
||||
return (comfy.model_management.cast_to(self.modulation, dtype=x.dtype, device=x.device).unsqueeze(0) + e).unbind(2)
|
||||
|
||||
def _cross_attn_ffn(self, x, e, context, context_img_len, transformer_options):
|
||||
x = x + self.cross_attn(self.norm3(x), context, context_img_len=context_img_len, transformer_options=transformer_options)
|
||||
for p in transformer_options.get("patches", {}).get("attn2_patch", []):
|
||||
x = p({"x": x, "transformer_options": transformer_options})
|
||||
y = self.ffn(torch.addcmul(repeat_e(e[3], x), self.norm2(x), 1 + repeat_e(e[4], x)))
|
||||
return torch.addcmul(x, y, repeat_e(e[5], x))
|
||||
|
||||
def forward_pose(self, x, e, freqs, context, context_img_len=257, transformer_options={}):
|
||||
e = self._modulation(e, x)
|
||||
x = x.contiguous()
|
||||
y, k, v = self.self_attn.forward_pose(torch.addcmul(repeat_e(e[0], x), self.norm1(x), 1 + repeat_e(e[1], x)), freqs, transformer_options=transformer_options)
|
||||
x = torch.addcmul(x, y, repeat_e(e[2], x))
|
||||
del y
|
||||
return self._cross_attn_ffn(x, e, context, context_img_len, transformer_options), k, v
|
||||
|
||||
def kv_from_input(self, x_pose, e, freqs, transformer_options={}):
|
||||
e = self._modulation(e, x_pose)
|
||||
x_pose = x_pose.contiguous()
|
||||
return self.self_attn.kv(torch.addcmul(repeat_e(e[0], x_pose), self.norm1(x_pose), 1 + repeat_e(e[1], x_pose)), freqs)
|
||||
|
||||
def forward_gen(self, x, e, freqs, context, k_pose, v_pose, f_gen, hw, buffers, ref_strength=1.0, context_img_len=257, transformer_options={}):
|
||||
e = self._modulation(e, x)
|
||||
x = x.contiguous()
|
||||
y = self.self_attn.forward_gen(torch.addcmul(repeat_e(e[0], x), self.norm1(x), 1 + repeat_e(e[1], x)), freqs, k_pose, v_pose, f_gen, hw, buffers, ref_strength=ref_strength, transformer_options=transformer_options)
|
||||
x = torch.addcmul(x, y, repeat_e(e[2], x))
|
||||
del y
|
||||
return self._cross_attn_ffn(x, e, context, context_img_len, transformer_options)
|
||||
|
||||
|
||||
class PoseBranchCache:
|
||||
"""Pose-branch block inputs, reused across the sampling steps of one execution.
|
||||
|
||||
Caching the block input rather than its K/V halves the memory; reprojecting K/V on read
|
||||
costs ~4% of re-running the block. One slot per distinct pose sequence, so under
|
||||
context windows each window keeps its own; least recently used slots are evicted when
|
||||
the store device runs low on memory. Created and freed by WanAnimate2Cache.
|
||||
"""
|
||||
|
||||
CONVROT_GROUPSIZE = 256
|
||||
|
||||
def __init__(self, store_device=None, dtype="default"):
|
||||
self.store_device = torch.device(store_device) if store_device is not None else torch.device("cpu")
|
||||
self.dtype = dtype
|
||||
self.slots = [] # most recently used last
|
||||
self.slot = None
|
||||
self._pending = {}
|
||||
self._staging = {}
|
||||
|
||||
def select(self, pose_latents):
|
||||
# select runs at a forward boundary: an interrupted forward can leave copies in flight that a different slot's forward would then mistake for its own
|
||||
if self._pending:
|
||||
for t, stream in self._pending.values():
|
||||
if stream is not None:
|
||||
stream.synchronize()
|
||||
self._pending = {}
|
||||
# keyed on batch element 0, so a cond batch size change mid-run stays valid
|
||||
k = pose_latents[:1]
|
||||
for s in self.slots:
|
||||
if s["key"].shape == k.shape and torch.equal(s["key"], k.to(s["key"].device)):
|
||||
self.slots.remove(s)
|
||||
self.slots.append(s)
|
||||
self.slot = s
|
||||
return
|
||||
# cache what fits: a filled slot is the size estimate for the next one, and least recently used slots make room when the store device runs low
|
||||
est = max((self._slot_bytes(s) for s in self.slots), default=0) * 1.5
|
||||
while self.slots and comfy.model_management.get_free_memory(self.store_device) < est:
|
||||
self._free_slot(self.slots.pop(0))
|
||||
self.slot = {"key": k.clone().to(self.store_device), "blocks": {}, "params": {}, "shape": None, "pinned": []}
|
||||
self.slots.append(self.slot)
|
||||
|
||||
def _free_slot(self, s):
|
||||
for t, stream in self._pending.values():
|
||||
if stream is not None:
|
||||
stream.synchronize() # an aborted forward can leave a copy in flight, still reading memory we are about to unpin
|
||||
self._pending = {}
|
||||
for t in s["pinned"]:
|
||||
comfy.model_management.unpin_memory(t)
|
||||
|
||||
def free(self):
|
||||
for s in self.slots:
|
||||
self._free_slot(s)
|
||||
self.slots = []
|
||||
self.slot = None
|
||||
self._staging = {}
|
||||
|
||||
def filled(self, num_blocks):
|
||||
return self.slot is not None and len(self.slot["blocks"]) == num_blocks
|
||||
|
||||
def put(self, i, x_pose):
|
||||
t = x_pose[:1]
|
||||
params = None
|
||||
if self.dtype in ("int8", "int4"):
|
||||
# convrot is what lets low-bit survive the ~125x per-channel outliers here, and over a [tokens, dim] view per-row scale means per-token. The kernels want 2D and a power-of-4 group that divides dim.
|
||||
self.slot["shape"] = t.shape
|
||||
g = self.CONVROT_GROUPSIZE
|
||||
while g > 4 and t.shape[-1] % g:
|
||||
g //= 4
|
||||
if self.dtype == "int4":
|
||||
t, params = comfy.quant_ops.TensorCoreConvRotW4A4Layout.quantize(t.reshape(-1, t.shape[-1]), convrot_groupsize=g)
|
||||
else:
|
||||
t, params = comfy.quant_ops.TensorWiseINT8Layout.quantize(t.reshape(-1, t.shape[-1]), is_weight=True, per_channel=True, convrot=True, convrot_groupsize=g)
|
||||
|
||||
t = t.to(self.store_device, copy=True)
|
||||
if comfy.model_management.pin_memory(t):
|
||||
self.slot["pinned"].append(t)
|
||||
self.slot["blocks"][i] = t
|
||||
# the scales follow the blocks off the GPU: per-window slots would otherwise pile them up in VRAM (~200 MB per window at 480p int4)
|
||||
self.slot["params"][i] = params if params is None else params.to_device(self.store_device)
|
||||
|
||||
def prefetch(self, i, device, dtype):
|
||||
# call before the compute this should overlap, so the stream waits only on work already enqueued
|
||||
if i not in self.slot["blocks"] or i in self._pending:
|
||||
return
|
||||
t = self.slot["blocks"][i]
|
||||
cast_dtype = None if self.slot["params"][i] is not None else dtype # int8 entries move in their stored dtype and widen in take()
|
||||
stream = None
|
||||
r = None
|
||||
if t.device != device:
|
||||
stream = comfy.model_management.get_offload_stream(device)
|
||||
cs = comfy.model_management.current_stream(device)
|
||||
if stream is not None and cs is not None:
|
||||
# the handed-out stream last waited on the main stream a full rotation ago, which does not cover the previous consumer's reads of this slot; wait now so the copy cannot overwrite a slot still being read
|
||||
stream.wait_stream(cs)
|
||||
# two persistent staging buffers per tensor shape instead of a fresh allocation per block (~29 GB of churn per pass at 720p); windows of different lengths get their own pair
|
||||
buf_key = (tuple(t.shape), cast_dtype if cast_dtype is not None else t.dtype)
|
||||
if buf_key not in self._staging:
|
||||
self._staging[buf_key] = [torch.empty(t.shape, dtype=buf_key[1], device=device) for _ in range(2)]
|
||||
r = self._staging[buf_key][i % 2]
|
||||
self._pending[i] = (comfy.model_management.cast_to(t, cast_dtype, device, non_blocking=True, stream=stream, r=r), stream)
|
||||
|
||||
def take(self, i, device, dtype, batch_size):
|
||||
if i not in self._pending:
|
||||
self.prefetch(i, device, dtype)
|
||||
t, stream = self._pending.pop(i)
|
||||
comfy.model_management.sync_stream(device, stream)
|
||||
params = self.slot["params"][i]
|
||||
if params is not None:
|
||||
layout = comfy.quant_ops.TensorCoreConvRotW4A4Layout if self.dtype == "int4" else comfy.quant_ops.TensorWiseINT8Layout
|
||||
t = layout.dequantize(t, params.to_device(t.device)).reshape(self.slot["shape"]).to(dtype)
|
||||
return comfy.utils.repeat_to_batch_size(t, batch_size)
|
||||
|
||||
def _slot_bytes(self, s):
|
||||
return sum(t.numel() * t.element_size() for t in s["blocks"].values())
|
||||
|
||||
def memory_bytes(self):
|
||||
return sum(self._slot_bytes(s) for s in self.slots)
|
||||
|
||||
|
||||
class WanAnimate2Model(WanModel):
|
||||
|
||||
def __init__(self,
|
||||
model_type='animate2',
|
||||
patch_size=(1, 2, 2),
|
||||
text_len=512,
|
||||
in_dim=36,
|
||||
dim=5120,
|
||||
ffn_dim=13824,
|
||||
freq_dim=256,
|
||||
text_dim=4096,
|
||||
out_dim=16,
|
||||
num_heads=40,
|
||||
num_layers=40,
|
||||
window_size=(-1, -1),
|
||||
qk_norm=True,
|
||||
cross_attn_norm=True,
|
||||
eps=1e-6,
|
||||
flf_pos_embed_token_number=None,
|
||||
in_dim_ref_conv=None,
|
||||
image_model=None,
|
||||
device=None, dtype=None, operations=None,
|
||||
):
|
||||
# model_type is 'animate2' in unet_config, but the checkpoint is i2v-shaped
|
||||
super().__init__(model_type='i2v', patch_size=patch_size, text_len=text_len, in_dim=in_dim, dim=dim, ffn_dim=ffn_dim, freq_dim=freq_dim,
|
||||
text_dim=text_dim, out_dim=out_dim, num_heads=num_heads, num_layers=num_layers, window_size=window_size, qk_norm=qk_norm,
|
||||
cross_attn_norm=cross_attn_norm, eps=eps, flf_pos_embed_token_number=flf_pos_embed_token_number, in_dim_ref_conv=in_dim_ref_conv,
|
||||
wan_attn_block_class=WanAnimate2Block, image_model=image_model, device=device, dtype=dtype, operations=operations)
|
||||
|
||||
def rope_encode_pose(self, t, h, w, w_patches, device=None, dtype=None):
|
||||
# t_start=1 lines pose frame j up with gen frame j+1, past the reference slot; shift_x parks it in its own strip of rope space.
|
||||
# The caller's rope_options are a user scaling knob and deliberately not forwarded.
|
||||
return super().rope_encode(t, h, w, t_start=1, device=device, dtype=dtype, transformer_options={"rope_options": {"shift_x": float(w_patches)}})
|
||||
|
||||
def _forward(self, x, timestep, context, clip_fea=None, time_dim_concat=None, transformer_options={}, pose_latents=None, clip_fea_pose=None, context_pose=None, **kwargs):
|
||||
bs, c, t, h, w = x.shape
|
||||
x = comfy.ldm.common_dit.pad_to_patch_size(x, self.patch_size)
|
||||
|
||||
# h/w pre-pad: rope_encode's rounding reproduces the post-pad grid
|
||||
freqs = self.rope_encode(t, h, w, device=x.device, dtype=x.dtype, transformer_options=transformer_options)
|
||||
freqs_pose = None
|
||||
if pose_latents is not None: # absent when the node's timestep window excludes this step
|
||||
pose_latents = comfy.ldm.common_dit.pad_to_patch_size(pose_latents.to(x.dtype), self.patch_size)
|
||||
w_patches = (w + (self.patch_size[2] // 2)) // self.patch_size[2]
|
||||
freqs_pose = self.rope_encode_pose(pose_latents.shape[2], h, w, w_patches, device=x.device, dtype=x.dtype)
|
||||
|
||||
return self.forward_orig(x, timestep, context, clip_fea=clip_fea, freqs=freqs, freqs_pose=freqs_pose, pose_latents=pose_latents,
|
||||
clip_fea_pose=clip_fea_pose, context_pose=context_pose, transformer_options=transformer_options, **kwargs)[:, :, :t, :h, :w]
|
||||
|
||||
def forward_orig(self, x, t, context, clip_fea=None, freqs=None, freqs_pose=None, pose_latents=None, clip_fea_pose=None, context_pose=None, pose_strength=1.0, reference_strength=1.0, transformer_options={}, **kwargs):
|
||||
x_input = x[:, :, 1:] # video-only: frame 0 is the reference slot, offset past it below
|
||||
x = self.patch_embedding(x.float()).to(x.dtype)
|
||||
grid_sizes = x.shape[2:]
|
||||
transformer_options["grid_sizes"] = grid_sizes
|
||||
f_gen, gh, gw = grid_sizes
|
||||
hw = gh * gw
|
||||
x = x.flatten(2).transpose(1, 2)
|
||||
|
||||
# the node windows the pose influence via cond timestep ranges: outside the window the cond carries no pose latents, and the branch, its cache traffic and the per-frame attention loop are all skipped
|
||||
apply_pose = pose_latents is not None
|
||||
if apply_pose and pose_latents.shape[2] != f_gen - 1: # before cache.select, which would otherwise keep an empty slot keyed to the rejected latents
|
||||
raise ValueError("pose branch has {} latent frames, expected {} (generation frames minus the reference-image slot)".format(pose_latents.shape[2], f_gen - 1))
|
||||
|
||||
cache = transformer_options.get("animate2_cache", None) if apply_pose else None
|
||||
if cache is not None:
|
||||
cache.select(pose_latents)
|
||||
cached = cache is not None and cache.filled(len(self.blocks))
|
||||
|
||||
x_pose = None
|
||||
if not cached and apply_pose:
|
||||
# 36ch = [latents(16) | mask(4) | latents(16)]; latents twice, and the mask is all ones since every pose frame is known
|
||||
x_pose = self.patch_embedding(torch.cat([pose_latents, torch.ones_like(pose_latents[:, :4]), pose_latents], dim=1).float()).to(x.dtype)
|
||||
x_pose = x_pose.flatten(2).transpose(1, 2)
|
||||
|
||||
# time embeddings
|
||||
e = self.time_embedding(sinusoidal_embedding_1d(self.freq_dim, t.flatten()).to(dtype=x.dtype))
|
||||
e = e.reshape(t.shape[0], -1, e.shape[-1])
|
||||
e0 = self.time_projection(e).unflatten(2, (6, self.dim))
|
||||
|
||||
e0_pose = None
|
||||
if apply_pose:
|
||||
t_pose = torch.ones_like(t.flatten())
|
||||
e_pose = self.time_embedding(sinusoidal_embedding_1d(self.freq_dim, t_pose).to(dtype=x.dtype))
|
||||
e_pose = e_pose.reshape(t.shape[0], -1, e_pose.shape[-1])
|
||||
e0_pose = self.time_projection(e_pose).unflatten(2, (6, self.dim))
|
||||
|
||||
context_gen = self.text_embedding(context)
|
||||
|
||||
context_img_len = None
|
||||
if clip_fea is not None:
|
||||
if self.img_emb is not None:
|
||||
context_gen = torch.cat([self.img_emb(clip_fea), context_gen], dim=1)
|
||||
context_img_len = clip_fea.shape[-2]
|
||||
|
||||
context_img_len_pose = None
|
||||
if not cached and apply_pose:
|
||||
context_pose = self.text_embedding(context if context_pose is None else context_pose)
|
||||
clip_fea_pose = clip_fea if clip_fea_pose is None else clip_fea_pose
|
||||
if clip_fea_pose is not None:
|
||||
if self.img_emb is not None:
|
||||
context_pose = torch.cat([self.img_emb(clip_fea_pose), context_pose], dim=1)
|
||||
context_img_len_pose = clip_fea_pose.shape[-2]
|
||||
|
||||
patches_replace = transformer_options.get("patches_replace", {})
|
||||
patches = transformer_options.get("patches", {})
|
||||
blocks_replace = patches_replace.get("dit", {})
|
||||
transformer_options["total_blocks"] = len(self.blocks)
|
||||
transformer_options["block_type"] = "double"
|
||||
|
||||
if cache is not None and not cached and apply_pose and "context_window" in transformer_options:
|
||||
# pose-only prepass, to avoid inflating dynamic VRAM calibration when using multiple context windows
|
||||
for i, block in enumerate(self.blocks):
|
||||
transformer_options["block_index"] = i
|
||||
cache.put(i, x_pose)
|
||||
x_pose = block.forward_pose(x_pose, e0_pose, freqs_pose, context_pose, context_img_len=context_img_len_pose, transformer_options=transformer_options)[0]
|
||||
x_pose = None
|
||||
cached = True
|
||||
|
||||
buffers = None
|
||||
if apply_pose:
|
||||
# allocated once and reused by every block
|
||||
n, d = self.num_heads, self.dim // self.num_heads
|
||||
buffers = (x.new_empty(x.shape[0], x.shape[1] + hw, n, d), x.new_empty(x.shape[0], x.shape[1] + hw, n, d), x.new_empty(x.shape[0], x.shape[1], self.dim))
|
||||
|
||||
for i, block in enumerate(self.blocks):
|
||||
transformer_options["block_index"] = i
|
||||
|
||||
if not apply_pose:
|
||||
k_pose = v_pose = None
|
||||
elif cached:
|
||||
x_pose_in = cache.take(i, x.device, x.dtype, x.shape[0])
|
||||
cache.prefetch(i + 1, x.device, x.dtype) # queue the next block before the gen compute it should overlap
|
||||
k_pose, v_pose = block.kv_from_input(x_pose_in, e0_pose, freqs_pose, transformer_options=transformer_options)
|
||||
del x_pose_in
|
||||
else:
|
||||
if cache is not None:
|
||||
cache.put(i, x_pose)
|
||||
# runs even under a block replace: its state has to reach block i+1
|
||||
x_pose, k_pose, v_pose = block.forward_pose(x_pose, e0_pose, freqs_pose, context_pose, context_img_len=context_img_len_pose, transformer_options=transformer_options)
|
||||
if v_pose is not None and pose_strength != 1.0:
|
||||
v_pose = v_pose * pose_strength
|
||||
|
||||
if ("double_block", i) in blocks_replace:
|
||||
def block_wrap(args, block=block, k_pose=k_pose, v_pose=v_pose):
|
||||
return {"img": block.forward_gen(args["img"], args["vec"], args["pe"], args["txt"], k_pose, v_pose, f_gen, hw, buffers, ref_strength=reference_strength, context_img_len=context_img_len, transformer_options=args["transformer_options"])}
|
||||
out = blocks_replace[("double_block", i)]({"img": x, "txt": context_gen, "vec": e0, "pe": freqs, "transformer_options": transformer_options}, {"original_block": block_wrap})
|
||||
x = out["img"]
|
||||
else:
|
||||
x = block.forward_gen(x, e0, freqs, context_gen, k_pose, v_pose, f_gen, hw, buffers, ref_strength=reference_strength, context_img_len=context_img_len, transformer_options=transformer_options)
|
||||
|
||||
if "double_block" in patches:
|
||||
for p in patches["double_block"]:
|
||||
out = p({"img": x, "x": x_input, "vec": e, "block_index": i, "img_offset": hw, "transformer_options": transformer_options})
|
||||
x = out["img"]
|
||||
|
||||
return self.unpatchify(self.head(x, e), grid_sizes)
|
||||
|
|
@ -4,7 +4,7 @@ import torch
|
|||
import torch.nn as nn
|
||||
|
||||
from comfy.ldm.flux.layers import EmbedND
|
||||
from .model import WanSelfAttention
|
||||
from .model import WanFeedForward, WanSelfAttention
|
||||
|
||||
|
||||
class Uni3CLayerNormZero(nn.Module):
|
||||
|
|
@ -41,7 +41,7 @@ class Uni3CAttentionBlock(nn.Module):
|
|||
self.norm1 = Uni3CLayerNormZero(time_embed_dim, dim, device=device, dtype=dtype, operations=operations)
|
||||
self.self_attn = WanSelfAttention(dim, num_heads, qk_norm=True, eps=eps, operation_settings=operation_settings)
|
||||
self.norm2 = Uni3CLayerNormZero(time_embed_dim, dim, device=device, dtype=dtype, operations=operations)
|
||||
self.ffn = nn.Sequential(
|
||||
self.ffn = WanFeedForward(
|
||||
operations.Linear(dim, ffn_dim, device=device, dtype=dtype), nn.GELU(approximate='tanh'),
|
||||
operations.Linear(ffn_dim, dim, device=device, dtype=dtype))
|
||||
|
||||
|
|
|
|||
|
|
@ -22,6 +22,7 @@ import torch
|
|||
import logging
|
||||
import comfy.ldm.lightricks.av_model
|
||||
import comfy.ldm.minimax.model
|
||||
import comfy.nested_tensor
|
||||
import comfy.ldm.lightricks.symmetric_patchifier
|
||||
import comfy.context_windows
|
||||
from comfy.ldm.modules.diffusionmodules.openaimodel import UNetModel, Timestep
|
||||
|
|
@ -45,6 +46,7 @@ import comfy.ldm.cosmos.predict2
|
|||
import comfy.ldm.lumina.model
|
||||
import comfy.ldm.wan.model
|
||||
import comfy.ldm.wan.model_animate
|
||||
import comfy.ldm.wan.model_animate2
|
||||
import comfy.ldm.wan.ar_model
|
||||
import comfy.ldm.wan.model_wandancer
|
||||
import comfy.ldm.hunyuan3d.model
|
||||
|
|
@ -100,6 +102,7 @@ class ModelType(Enum):
|
|||
FLOW_COSMOS = 10
|
||||
IMG_TO_IMG_FLOW = 11
|
||||
V_PREDICTION_DDPM = 12
|
||||
FLOW_AV = 13
|
||||
|
||||
|
||||
def model_sampling(model_config, model_type):
|
||||
|
|
@ -136,6 +139,9 @@ def model_sampling(model_config, model_type):
|
|||
c = comfy.model_sampling.IMG_TO_IMG_FLOW
|
||||
elif model_type == ModelType.V_PREDICTION_DDPM:
|
||||
c = comfy.model_sampling.V_PREDICTION_DDPM
|
||||
elif model_type == ModelType.FLOW_AV:
|
||||
c = comfy.model_sampling.CONST
|
||||
s = comfy.model_sampling.ModelSamplingAV
|
||||
|
||||
class ModelSampling(s, c):
|
||||
pass
|
||||
|
|
@ -180,6 +186,7 @@ class BaseModel(torch.nn.Module):
|
|||
|
||||
self.model_type = model_type
|
||||
self.model_sampling = model_sampling(model_config, model_type)
|
||||
self.latent_shapes = None # set by the sampler for models that pack several streams into one latent
|
||||
|
||||
self.adm_channels = unet_config.get("adm_in_channels", None)
|
||||
if self.adm_channels is None:
|
||||
|
|
@ -1807,6 +1814,41 @@ class WAN22_Animate(WAN21):
|
|||
return comfy.context_windows.slice_cond(cond_value, window, x_in, device, temporal_dim=2, temporal_offset=1)
|
||||
return super().resize_cond_for_context_window(cond_key, cond_value, window, x_in, device, retain_index_list=retain_index_list)
|
||||
|
||||
class WAN_Animate2(WAN21):
|
||||
def __init__(self, model_config, model_type=ModelType.FLOW, device=None):
|
||||
super(WAN21, self).__init__(model_config, model_type, device=device, unet_model=comfy.ldm.wan.model_animate2.WanAnimate2Model)
|
||||
self.image_to_video = True
|
||||
|
||||
def extra_conds(self, **kwargs):
|
||||
out = super().extra_conds(**kwargs)
|
||||
|
||||
pose_video_latent = kwargs.get("pose_video_latent", None)
|
||||
if pose_video_latent is not None:
|
||||
out['pose_latents'] = comfy.conds.CONDRegular(self.process_latent_in(pose_video_latent))
|
||||
|
||||
clip_vision_output_pose = kwargs.get("clip_vision_output_pose", None)
|
||||
if clip_vision_output_pose is not None:
|
||||
out['clip_fea_pose'] = comfy.conds.CONDRegular(clip_vision_output_pose.penultimate_hidden_states)
|
||||
|
||||
cross_attn_pose = kwargs.get("cross_attn_pose", None)
|
||||
if cross_attn_pose is not None:
|
||||
out['context_pose'] = comfy.conds.CONDRegular(cross_attn_pose)
|
||||
|
||||
pose_strength = kwargs.get("pose_strength", 1.0)
|
||||
if pose_strength != 1.0:
|
||||
out['pose_strength'] = comfy.conds.CONDConstant(pose_strength)
|
||||
|
||||
reference_strength = kwargs.get("reference_strength", 1.0)
|
||||
if reference_strength != 1.0:
|
||||
out['reference_strength'] = comfy.conds.CONDConstant(reference_strength)
|
||||
|
||||
return out
|
||||
|
||||
def resize_cond_for_context_window(self, cond_key, cond_value, window, x_in, device, retain_index_list=[]):
|
||||
if cond_key == "pose_latents":
|
||||
return comfy.context_windows.slice_cond(cond_value, window, x_in, device, temporal_dim=2, temporal_offset=1)
|
||||
return super().resize_cond_for_context_window(cond_key, cond_value, window, x_in, device, retain_index_list=retain_index_list)
|
||||
|
||||
class WAN22_S2V(WAN21):
|
||||
def __init__(self, model_config, model_type=ModelType.FLOW, device=None):
|
||||
super(WAN21, self).__init__(model_config, model_type, device=device, unet_model=comfy.ldm.wan.model.WanModel_S2V)
|
||||
|
|
@ -2065,9 +2107,33 @@ class Hunyuan3Dv2_1(BaseModel):
|
|||
return out
|
||||
|
||||
class MiniMaxH3(BaseModel):
|
||||
def __init__(self, model_config, model_type=ModelType.FLOW, device=None):
|
||||
def __init__(self, model_config, model_type=ModelType.FLOW_AV, device=None):
|
||||
super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.minimax.model.MiniMaxH3Model)
|
||||
|
||||
def audio_scale(self):
|
||||
"""Scale the sampler carries the audio stream at, 1.0 when not sampling the packed latent."""
|
||||
if self.latent_shapes is None or len(self.latent_shapes) < 2:
|
||||
return 1.0
|
||||
return self.model_sampling.audio_scale
|
||||
|
||||
def _scale_audio_slice(self, latent, scale):
|
||||
# the sampler carries the audio stream scaled onto the video schedule
|
||||
if scale == 1.0:
|
||||
return latent
|
||||
if latent.is_nested: # the x0 output hands back the unpacked view
|
||||
streams = latent.unbind()
|
||||
return comfy.nested_tensor.NestedTensor([streams[0], streams[1] * scale] + list(streams[2:]))
|
||||
n = math.prod(self.latent_shapes[0][1:])
|
||||
latent = latent.clone()
|
||||
latent[..., n:] *= scale
|
||||
return latent
|
||||
|
||||
def process_latent_in(self, latent):
|
||||
return self._scale_audio_slice(super().process_latent_in(latent), self.audio_scale())
|
||||
|
||||
def process_latent_out(self, latent):
|
||||
return super().process_latent_out(self._scale_audio_slice(latent, 1.0 / self.audio_scale()))
|
||||
|
||||
def extra_conds(self, **kwargs):
|
||||
out = super().extra_conds(**kwargs)
|
||||
cross_attn = kwargs.get("cross_attn", None)
|
||||
|
|
@ -2102,6 +2168,8 @@ class MiniMaxH3(BaseModel):
|
|||
if kwargs.get("minimax_audio_cond_noise_aug", None) is not None:
|
||||
payload["audio_cond_noise_aug"] = kwargs["minimax_audio_cond_noise_aug"]
|
||||
payload["seed"] = kwargs.get("seed", 0)
|
||||
# same value process_latent_in/out used, so the model never undoes a scale that was not applied
|
||||
payload["audio_scale"] = self.audio_scale()
|
||||
if cross_attn is not None and latent_shapes is not None and len(latent_shapes) > 1:
|
||||
# packed layout built once per sampling run, h/w rounded up to the DiT's 2x2 patch
|
||||
vs = latent_shapes[0]
|
||||
|
|
@ -2112,9 +2180,6 @@ class MiniMaxH3(BaseModel):
|
|||
out['minimax_payload'] = comfy.conds.CONDConstant(payload)
|
||||
return out
|
||||
|
||||
def scale_latent_inpaint(self, sigma, noise, latent_image, **kwargs):
|
||||
return latent_image
|
||||
|
||||
class TripoSplat(BaseModel):
|
||||
def __init__(self, model_config, model_type=ModelType.FLOW, device=None):
|
||||
super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.triposplat.model.LatentSeqMMFlowModel)
|
||||
|
|
|
|||
|
|
@ -325,6 +325,27 @@ class ModelSamplingDiscreteFlow(torch.nn.Module):
|
|||
return 0.0
|
||||
return time_snr_shift(self.shift, 1.0 - percent)
|
||||
|
||||
class ModelSamplingAV(ModelSamplingDiscreteFlow):
|
||||
"""Flow sampling for packed audio-video latents whose audio stream has its own flow shift.
|
||||
|
||||
Carrying the audio latent scaled onto the video schedule makes the pack an ordinary
|
||||
single-schedule flow latent whose audio target is scaled by audio_scale.
|
||||
"""
|
||||
def __init__(self, model_config=None):
|
||||
super().__init__(model_config)
|
||||
sampling_settings = model_config.sampling_settings if model_config is not None else {}
|
||||
self.audio_shift = sampling_settings.get("audio_shift", None)
|
||||
|
||||
def set_parameters(self, shift=1.0, audio_shift=None, timesteps=1000, multiplier=1000):
|
||||
self.audio_shift = audio_shift
|
||||
super().set_parameters(shift=shift, timesteps=timesteps, multiplier=multiplier)
|
||||
|
||||
@property
|
||||
def audio_scale(self):
|
||||
if self.audio_shift is None:
|
||||
return 1.0
|
||||
return self.shift / self.audio_shift
|
||||
|
||||
class StableCascadeSampling(ModelSamplingDiscrete):
|
||||
def __init__(self, model_config=None):
|
||||
super().__init__()
|
||||
|
|
|
|||
|
|
@ -83,6 +83,9 @@ class NestedTensor:
|
|||
def layout(self):
|
||||
return self.tensors[0].layout
|
||||
|
||||
def __repr__(self):
|
||||
return f"{type(self).__name__}({self.tensors!r})"
|
||||
|
||||
|
||||
def cat_nested(tensors, *args, **kwargs):
|
||||
cated_tensors = []
|
||||
|
|
|
|||
266
comfy/ops.py
266
comfy/ops.py
|
|
@ -73,6 +73,8 @@ try:
|
|||
]
|
||||
|
||||
def scaled_dot_product_attention(q, k, v, *args, **kwargs):
|
||||
if q.nelement() < 1024 * 128: # arbitrary number, for small inputs cudnn attention seems slower
|
||||
return torch.nn.functional.scaled_dot_product_attention(q, k, v, *args, **kwargs)
|
||||
attn_mask = args[0] if len(args) > 0 else kwargs.get("attn_mask")
|
||||
if kwargs.get("enable_gqa", False) and attn_mask is not None and not comfy.model_management.is_nvidia():
|
||||
k, v = repeat_kv_for_gqa(k, v, q.shape[-3], -3)
|
||||
|
|
@ -450,6 +452,26 @@ def uncast_bias_weight(s, weight, bias, offload_stream):
|
|||
device = bias_a.device
|
||||
os.wait_stream(comfy.model_management.current_stream(device))
|
||||
|
||||
class CastBiasWeightContext:
|
||||
# When initialized with no arguments or the first is None, the context
|
||||
# will return the tuple (None, None).
|
||||
def __init__(self, *args, **kwargs):
|
||||
self.slf = args[0] if len(args) else None
|
||||
self.state = (None, None) if self.slf is None else cast_bias_weight(*args, **kwargs)
|
||||
|
||||
def __enter__(self):
|
||||
result = self.state
|
||||
if len(result) < 3 or result[2] is None:
|
||||
# Not offloaded, immediately drop references.
|
||||
self.state = self.slf = None
|
||||
return result[:2]
|
||||
|
||||
def __exit__(self, *_args) -> None:
|
||||
if self.slf is None:
|
||||
return
|
||||
slf, state = self.slf, self.state
|
||||
self.state = self.slf = None
|
||||
uncast_bias_weight(slf, *state)
|
||||
|
||||
class CastWeightBiasOp:
|
||||
comfy_cast_weights = False
|
||||
|
|
@ -536,10 +558,8 @@ class disable_weight_init:
|
|||
return None
|
||||
|
||||
def forward_comfy_cast_weights(self, input):
|
||||
weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True)
|
||||
x = torch.nn.functional.linear(input, weight, bias)
|
||||
uncast_bias_weight(self, weight, bias, offload_stream)
|
||||
return x
|
||||
with CastBiasWeightContext(self, input, offloadable=True) as (weight, bias):
|
||||
return torch.nn.functional.linear(input, weight, bias)
|
||||
|
||||
def forward(self, *args, **kwargs):
|
||||
run_every_op()
|
||||
|
|
@ -553,10 +573,8 @@ class disable_weight_init:
|
|||
return None
|
||||
|
||||
def forward_comfy_cast_weights(self, input):
|
||||
weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True)
|
||||
x = self._conv_forward(input, weight, bias)
|
||||
uncast_bias_weight(self, weight, bias, offload_stream)
|
||||
return x
|
||||
with CastBiasWeightContext(self, input, offloadable=True) as (weight, bias):
|
||||
return self._conv_forward(input, weight, bias)
|
||||
|
||||
def forward(self, *args, **kwargs):
|
||||
run_every_op()
|
||||
|
|
@ -570,10 +588,8 @@ class disable_weight_init:
|
|||
return None
|
||||
|
||||
def forward_comfy_cast_weights(self, input):
|
||||
weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True)
|
||||
x = self._conv_forward(input, weight, bias)
|
||||
uncast_bias_weight(self, weight, bias, offload_stream)
|
||||
return x
|
||||
with CastBiasWeightContext(self, input, offloadable=True) as (weight, bias):
|
||||
return self._conv_forward(input, weight, bias)
|
||||
|
||||
def forward(self, *args, **kwargs):
|
||||
run_every_op()
|
||||
|
|
@ -598,10 +614,8 @@ class disable_weight_init:
|
|||
return super()._conv_forward(input, weight, bias, *args, **kwargs)
|
||||
|
||||
def forward_comfy_cast_weights(self, input, autopad=None):
|
||||
weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True)
|
||||
x = self._conv_forward(input, weight, bias, autopad=autopad)
|
||||
uncast_bias_weight(self, weight, bias, offload_stream)
|
||||
return x
|
||||
with CastBiasWeightContext(self, input, offloadable=True) as (weight, bias):
|
||||
return self._conv_forward(input, weight, bias, autopad=autopad)
|
||||
|
||||
def forward(self, *args, **kwargs):
|
||||
run_every_op()
|
||||
|
|
@ -615,10 +629,8 @@ class disable_weight_init:
|
|||
return None
|
||||
|
||||
def forward_comfy_cast_weights(self, input):
|
||||
weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True)
|
||||
x = torch.nn.functional.group_norm(input, self.num_groups, weight, bias, self.eps)
|
||||
uncast_bias_weight(self, weight, bias, offload_stream)
|
||||
return x
|
||||
with CastBiasWeightContext(self, input, offloadable=True) as (weight, bias):
|
||||
return torch.nn.functional.group_norm(input, self.num_groups, weight, bias, self.eps)
|
||||
|
||||
def forward(self, *args, **kwargs):
|
||||
run_every_op()
|
||||
|
|
@ -632,12 +644,10 @@ class disable_weight_init:
|
|||
return None
|
||||
|
||||
def forward_comfy_cast_weights(self, input):
|
||||
weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True)
|
||||
running_mean = self.running_mean.to(device=input.device, dtype=weight.dtype) if self.running_mean is not None else None
|
||||
running_var = self.running_var.to(device=input.device, dtype=weight.dtype) if self.running_var is not None else None
|
||||
x = torch.nn.functional.batch_norm(input, running_mean, running_var, weight, bias, self.training, self.momentum, self.eps)
|
||||
uncast_bias_weight(self, weight, bias, offload_stream)
|
||||
return x
|
||||
with CastBiasWeightContext(self, input, offloadable=True) as (weight, bias):
|
||||
running_mean = self.running_mean.to(device=input.device, dtype=weight.dtype) if self.running_mean is not None else None
|
||||
running_var = self.running_var.to(device=input.device, dtype=weight.dtype) if self.running_var is not None else None
|
||||
return torch.nn.functional.batch_norm(input, running_mean, running_var, weight, bias, self.training, self.momentum, self.eps)
|
||||
|
||||
def forward(self, *args, **kwargs):
|
||||
run_every_op()
|
||||
|
|
@ -651,15 +661,8 @@ class disable_weight_init:
|
|||
return None
|
||||
|
||||
def forward_comfy_cast_weights(self, input):
|
||||
if self.weight is not None:
|
||||
weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True)
|
||||
else:
|
||||
weight = None
|
||||
bias = None
|
||||
offload_stream = None
|
||||
x = torch.nn.functional.layer_norm(input, self.normalized_shape, weight, bias, self.eps)
|
||||
uncast_bias_weight(self, weight, bias, offload_stream)
|
||||
return x
|
||||
with CastBiasWeightContext(self if self.weight is not None else None, input, offloadable=True) as (weight, bias):
|
||||
return torch.nn.functional.layer_norm(input, self.normalized_shape, weight, bias, self.eps)
|
||||
|
||||
def forward(self, *args, **kwargs):
|
||||
run_every_op()
|
||||
|
|
@ -674,15 +677,8 @@ class disable_weight_init:
|
|||
return None
|
||||
|
||||
def forward_comfy_cast_weights(self, input):
|
||||
if self.weight is not None:
|
||||
weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True)
|
||||
else:
|
||||
weight = None
|
||||
bias = None
|
||||
offload_stream = None
|
||||
x = torch.nn.functional.rms_norm(input, self.normalized_shape, weight, self.eps)
|
||||
uncast_bias_weight(self, weight, bias, offload_stream)
|
||||
return x
|
||||
with CastBiasWeightContext(self if self.weight is not None else None, input, offloadable=True) as (weight, bias):
|
||||
return torch.nn.functional.rms_norm(input, self.normalized_shape, weight, self.eps)
|
||||
|
||||
def forward(self, *args, **kwargs):
|
||||
run_every_op()
|
||||
|
|
@ -701,12 +697,10 @@ class disable_weight_init:
|
|||
input, output_size, self.stride, self.padding, self.kernel_size,
|
||||
num_spatial_dims, self.dilation)
|
||||
|
||||
weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True)
|
||||
x = torch.nn.functional.conv_transpose2d(
|
||||
input, weight, bias, self.stride, self.padding,
|
||||
output_padding, self.groups, self.dilation)
|
||||
uncast_bias_weight(self, weight, bias, offload_stream)
|
||||
return x
|
||||
with CastBiasWeightContext(self, input, offloadable=True) as (weight, bias):
|
||||
return torch.nn.functional.conv_transpose2d(
|
||||
input, weight, bias, self.stride, self.padding,
|
||||
output_padding, self.groups, self.dilation)
|
||||
|
||||
def forward(self, *args, **kwargs):
|
||||
run_every_op()
|
||||
|
|
@ -725,12 +719,10 @@ class disable_weight_init:
|
|||
input, output_size, self.stride, self.padding, self.kernel_size,
|
||||
num_spatial_dims, self.dilation)
|
||||
|
||||
weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True)
|
||||
x = torch.nn.functional.conv_transpose1d(
|
||||
input, weight, bias, self.stride, self.padding,
|
||||
output_padding, self.groups, self.dilation)
|
||||
uncast_bias_weight(self, weight, bias, offload_stream)
|
||||
return x
|
||||
with CastBiasWeightContext(self, input, offloadable=True) as (weight, bias):
|
||||
return torch.nn.functional.conv_transpose1d(
|
||||
input, weight, bias, self.stride, self.padding,
|
||||
output_padding, self.groups, self.dilation)
|
||||
|
||||
def forward(self, *args, **kwargs):
|
||||
run_every_op()
|
||||
|
|
@ -793,10 +785,8 @@ class disable_weight_init:
|
|||
output_dtype = out_dtype
|
||||
if self.weight.dtype == torch.float16 or self.weight.dtype == torch.bfloat16:
|
||||
out_dtype = None
|
||||
weight, bias, offload_stream = cast_bias_weight(self, device=input.device, dtype=out_dtype, offloadable=True)
|
||||
x = torch.nn.functional.embedding(input, weight, self.padding_idx, self.max_norm, self.norm_type, self.scale_grad_by_freq, self.sparse).to(dtype=output_dtype)
|
||||
uncast_bias_weight(self, weight, bias, offload_stream)
|
||||
return x
|
||||
with CastBiasWeightContext(self, device=input.device, dtype=out_dtype, offloadable=True) as (weight, bias):
|
||||
return torch.nn.functional.embedding(input, weight, self.padding_idx, self.max_norm, self.norm_type, self.scale_grad_by_freq, self.sparse).to(dtype=output_dtype)
|
||||
|
||||
|
||||
def forward(self, *args, **kwargs):
|
||||
|
|
@ -872,7 +862,6 @@ def fp8_linear(self, input):
|
|||
if input.ndim != 2:
|
||||
return None
|
||||
lora_compute_dtype=comfy.model_management.lora_compute_dtype(input.device)
|
||||
w, bias, offload_stream = cast_bias_weight(self, input, dtype=dtype, bias_dtype=input_dtype, offloadable=True, compute_dtype=lora_compute_dtype, want_requant=True)
|
||||
scale_weight = torch.ones((), device=input.device, dtype=torch.float32)
|
||||
|
||||
scale_input = torch.ones((), device=input.device, dtype=torch.float32)
|
||||
|
|
@ -881,15 +870,16 @@ def fp8_linear(self, input):
|
|||
layout_params_input = TensorCoreFP8Layout.Params(scale=scale_input, orig_dtype=input_dtype, orig_shape=tuple(input_fp8.shape))
|
||||
quantized_input = QuantizedTensor(input_fp8, "TensorCoreFP8Layout", layout_params_input)
|
||||
|
||||
# Wrap weight in QuantizedTensor - this enables unified dispatch
|
||||
# Call F.linear - __torch_dispatch__ routes to fp8_linear handler in quant_ops.py!
|
||||
layout_params_weight = TensorCoreFP8Layout.Params(scale=scale_weight, orig_dtype=input_dtype, orig_shape=tuple(w.shape))
|
||||
quantized_weight = QuantizedTensor(w, "TensorCoreFP8Layout", layout_params_weight)
|
||||
o = torch.nn.functional.linear(quantized_input, quantized_weight, bias)
|
||||
with CastBiasWeightContext(self, input, dtype=dtype, bias_dtype=input_dtype, offloadable=True, compute_dtype=lora_compute_dtype, want_requant=True) as (w, bias):
|
||||
# Wrap weight in QuantizedTensor - this enables unified dispatch
|
||||
# Call F.linear - __torch_dispatch__ routes to fp8_linear handler in quant_ops.py!
|
||||
w_shape = tuple(w.shape)
|
||||
layout_params_weight = TensorCoreFP8Layout.Params(scale=scale_weight, orig_dtype=input_dtype, orig_shape=w_shape)
|
||||
quantized_weight = QuantizedTensor(w, "TensorCoreFP8Layout", layout_params_weight)
|
||||
o = torch.nn.functional.linear(quantized_input, quantized_weight, bias)
|
||||
|
||||
uncast_bias_weight(self, w, bias, offload_stream)
|
||||
if tensor_3d:
|
||||
o = o.reshape((input_shape[0], input_shape[1], w.shape[0]))
|
||||
o = o.reshape((input_shape[0], input_shape[1], w_shape[0]))
|
||||
|
||||
return o
|
||||
|
||||
|
|
@ -909,10 +899,8 @@ class fp8_ops(manual_cast):
|
|||
except Exception as e:
|
||||
logging.info("Exception during fp8 op: {}".format(e))
|
||||
|
||||
weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True)
|
||||
x = torch.nn.functional.linear(input, weight, bias)
|
||||
uncast_bias_weight(self, weight, bias, offload_stream)
|
||||
return x
|
||||
with CastBiasWeightContext(self, input, offloadable=True) as (weight, bias):
|
||||
return torch.nn.functional.linear(input, weight, bias)
|
||||
|
||||
CUBLAS_IS_AVAILABLE = False
|
||||
try:
|
||||
|
|
@ -928,10 +916,8 @@ if CUBLAS_IS_AVAILABLE:
|
|||
return None
|
||||
|
||||
def forward_comfy_cast_weights(self, input):
|
||||
weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True)
|
||||
x = cublas_half_matmul(input, weight, bias, self._epilogue_str, self.has_bias)
|
||||
uncast_bias_weight(self, weight, bias, offload_stream)
|
||||
return x
|
||||
with CastBiasWeightContext(self, input, offloadable=True) as (weight, bias):
|
||||
return cublas_half_matmul(input, weight, bias, self._epilogue_str, self.has_bias)
|
||||
|
||||
def forward(self, *args, **kwargs):
|
||||
run_every_op()
|
||||
|
|
@ -1209,6 +1195,26 @@ def _load_quantized_module(module, super_load, state_dict, prefix, local_metadat
|
|||
"quant_group_size": 64,
|
||||
"linear_dtype": layer_conf.get("linear_dtype", params_conf.get("linear_dtype", "int4")),
|
||||
}
|
||||
elif module.quant_format == "asym_w4a8_int8":
|
||||
# int4 weight (packed int8 [N,K/2]) + fp8 per-group scale (weight_s_rel),
|
||||
# fp32 per-channel scale (weight_s_channel) + optional Lloyd-Max codebook.
|
||||
scale = pop_scale("weight_s_rel")
|
||||
if scale is None:
|
||||
raise ValueError(f"Missing W4A8 group scale (weight_s_rel) for layer {layer_name}")
|
||||
if scale.dtype == torch.uint8:
|
||||
scale = scale.view(torch.float8_e4m3fn)
|
||||
params_conf = layer_conf.get("params", {})
|
||||
if not isinstance(params_conf, dict):
|
||||
params_conf = {}
|
||||
scales = {
|
||||
"scale": scale,
|
||||
"s_channel": pop_scale("weight_s_channel"),
|
||||
"codebook": pop_scale("weight_codebook"),
|
||||
"group_size": int(layer_conf.get("group_size", params_conf.get("group_size", 16))),
|
||||
"convrot_groupsize": int(
|
||||
layer_conf.get("convrot_groupsize", params_conf.get("convrot_groupsize", 256))
|
||||
),
|
||||
}
|
||||
else:
|
||||
raise ValueError(f"Unsupported quantization format: {module.quant_format}")
|
||||
|
||||
|
|
@ -1260,6 +1266,9 @@ def _quantized_weight_state_dict(module, sd, prefix, extra_quant_conf=None, extr
|
|||
linear_dtype = getattr(params, "linear_dtype", "int4")
|
||||
if linear_dtype != "int4":
|
||||
quant_conf["linear_dtype"] = linear_dtype
|
||||
elif module.quant_format == "asym_w4a8_int8":
|
||||
quant_conf["group_size"] = getattr(params, "group_size", 16)
|
||||
quant_conf["convrot_groupsize"] = getattr(params, "convrot_groupsize", 256)
|
||||
if extra_quant_conf:
|
||||
quant_conf.update(extra_quant_conf)
|
||||
sd[f"{prefix}comfy_quant"] = torch.tensor(list(json.dumps(quant_conf).encode("utf-8")), dtype=torch.uint8)
|
||||
|
|
@ -1319,29 +1328,28 @@ def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_prec
|
|||
want_requant=False,
|
||||
weight_only_quant=False,
|
||||
):
|
||||
if weight_only_quant:
|
||||
weight, bias, offload_stream = cast_bias_weight(
|
||||
self,
|
||||
input=None,
|
||||
dtype=self.weight.dtype,
|
||||
device=input.device,
|
||||
bias_dtype=input.dtype,
|
||||
offloadable=True,
|
||||
compute_dtype=compute_dtype,
|
||||
want_requant=True,
|
||||
)
|
||||
weight = weight.to(dtype=input.dtype)
|
||||
else:
|
||||
weight, bias, offload_stream = cast_bias_weight(
|
||||
if not weight_only_quant:
|
||||
with CastBiasWeightContext(
|
||||
self,
|
||||
input,
|
||||
offloadable=True,
|
||||
compute_dtype=compute_dtype,
|
||||
want_requant=want_requant,
|
||||
)
|
||||
x = self._forward(input, weight, bias)
|
||||
uncast_bias_weight(self, weight, bias, offload_stream)
|
||||
return x
|
||||
) as (weight, bias):
|
||||
return self._forward(input, weight, bias)
|
||||
|
||||
with CastBiasWeightContext(
|
||||
self,
|
||||
input=None,
|
||||
dtype=self.weight.dtype,
|
||||
device=input.device,
|
||||
bias_dtype=input.dtype,
|
||||
offloadable=True,
|
||||
compute_dtype=compute_dtype,
|
||||
want_requant=True,
|
||||
) as (weight, bias):
|
||||
weight = weight.to(dtype=input.dtype)
|
||||
return self._forward(input, weight, bias)
|
||||
|
||||
def forward(self, input, *args, **kwargs):
|
||||
run_every_op()
|
||||
|
|
@ -1366,25 +1374,20 @@ def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_prec
|
|||
|
||||
# Training path: quantized forward with compute_dtype backward via autograd function
|
||||
if (input.requires_grad and _use_quantized and quantize_input):
|
||||
|
||||
weight, bias, offload_stream = cast_bias_weight(
|
||||
with CastBiasWeightContext(
|
||||
self,
|
||||
input,
|
||||
offloadable=True,
|
||||
compute_dtype=compute_dtype,
|
||||
want_requant=True
|
||||
)
|
||||
) as (weight, bias):
|
||||
scale = getattr(self, 'input_scale', None)
|
||||
if scale is not None:
|
||||
scale = comfy.model_management.cast_to_device(scale, input.device, None)
|
||||
|
||||
scale = getattr(self, 'input_scale', None)
|
||||
if scale is not None:
|
||||
scale = comfy.model_management.cast_to_device(scale, input.device, None)
|
||||
|
||||
output = QuantLinearFunc.apply(
|
||||
input, weight, bias, self.layout_type, scale, compute_dtype
|
||||
)
|
||||
|
||||
uncast_bias_weight(self, weight, bias, offload_stream)
|
||||
return output
|
||||
return QuantLinearFunc.apply(
|
||||
input, weight, bias, self.layout_type, scale, compute_dtype
|
||||
)
|
||||
|
||||
# Inference path (unchanged)
|
||||
if _use_quantized and quantize_input:
|
||||
|
|
@ -1495,13 +1498,11 @@ def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_prec
|
|||
"""Cast the whole bank once; expert_linear inside reuses the cast.
|
||||
Not re-entrant — do not nest calls on the same instance.
|
||||
"""
|
||||
weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True)
|
||||
self._resident_bank = (weight, bias)
|
||||
try:
|
||||
yield self
|
||||
finally:
|
||||
self._resident_bank = None
|
||||
uncast_bias_weight(self, weight, bias, offload_stream)
|
||||
with CastBiasWeightContext(self, input, offloadable=True) as self._resident_bank:
|
||||
try:
|
||||
yield self
|
||||
finally:
|
||||
self._resident_bank = None
|
||||
|
||||
def expert_linear(self, input: torch.Tensor, i: int) -> torch.Tensor:
|
||||
"""Linear against expert i's weight (with optional bias)."""
|
||||
|
|
@ -1509,11 +1510,8 @@ def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_prec
|
|||
if resident is not None:
|
||||
weight, bias = resident
|
||||
return self._expert_linear_impl(input, weight, bias, i)
|
||||
weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True)
|
||||
try:
|
||||
with CastBiasWeightContext(self, input, offloadable=True) as (weight, bias):
|
||||
return self._expert_linear_impl(input, weight, bias, i)
|
||||
finally:
|
||||
uncast_bias_weight(self, weight, bias, offload_stream)
|
||||
|
||||
def _expert_linear_impl(self, input, weight, bias, i):
|
||||
if isinstance(weight, QuantizedTensor):
|
||||
|
|
@ -1616,25 +1614,23 @@ def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_prec
|
|||
|
||||
# Optimized path: lookup in fp8/int8, dequantize only the selected rows.
|
||||
if isinstance(weight, QuantizedTensor) and len(self.weight_function) == 0:
|
||||
qdata, _, offload_stream = cast_bias_weight(self, device=input.device, dtype=weight.dtype, offloadable=True)
|
||||
if isinstance(qdata, QuantizedTensor):
|
||||
params = qdata._params
|
||||
scale = params.scale
|
||||
qdata = qdata._qdata
|
||||
else:
|
||||
params = weight._params
|
||||
scale = None
|
||||
with CastBiasWeightContext(self, device=input.device, dtype=weight.dtype, offloadable=True) as (qdata, _bias):
|
||||
if isinstance(qdata, QuantizedTensor):
|
||||
params = qdata._params
|
||||
scale = params.scale
|
||||
qdata = qdata._qdata
|
||||
else:
|
||||
params = weight._params
|
||||
scale = None
|
||||
|
||||
# int8: per-row scale possible ConvRot, so let the layout do the gather
|
||||
if self.quant_format == "int8_tensorwise":
|
||||
x = get_layout_class(self.layout_type).dequantize_embedding(qdata, params, input)
|
||||
uncast_bias_weight(self, qdata, None, offload_stream)
|
||||
return x if out_dtype is None else x.to(dtype=out_dtype)
|
||||
# int8: per-row scale possible ConvRot, so let the layout do the gather
|
||||
if self.quant_format == "int8_tensorwise":
|
||||
x = get_layout_class(self.layout_type).dequantize_embedding(qdata, params, input)
|
||||
return x if out_dtype is None else x.to(dtype=out_dtype)
|
||||
|
||||
x = torch.nn.functional.embedding(
|
||||
input, qdata, self.padding_idx, self.max_norm,
|
||||
self.norm_type, self.scale_grad_by_freq, self.sparse)
|
||||
uncast_bias_weight(self, qdata, None, offload_stream)
|
||||
x = torch.nn.functional.embedding(
|
||||
input, qdata, self.padding_idx, self.max_norm,
|
||||
self.norm_type, self.scale_grad_by_freq, self.sparse)
|
||||
target_dtype = out_dtype if out_dtype is not None else weight._params.orig_dtype
|
||||
x = x.to(dtype=target_dtype)
|
||||
if scale is not None and scale != 1.0:
|
||||
|
|
|
|||
|
|
@ -28,6 +28,7 @@ try:
|
|||
TensorCoreNVFP4Layout as _CKNvfp4Layout,
|
||||
TensorCoreConvRotW4A4Layout as _CKTensorCoreConvRotW4A4Layout,
|
||||
TensorWiseINT8Layout as _CKTensorWiseINT8Layout,
|
||||
AsymW4A8Int8Layout as _CKAsymW4A8Int8Layout,
|
||||
register_layout_op,
|
||||
register_layout_class,
|
||||
get_layout_class,
|
||||
|
|
@ -83,6 +84,9 @@ except ImportError as e:
|
|||
class _CKTensorCoreConvRotW4A4Layout:
|
||||
pass
|
||||
|
||||
class _CKAsymW4A8Int8Layout:
|
||||
pass
|
||||
|
||||
def register_layout_class(name, cls):
|
||||
pass
|
||||
|
||||
|
|
@ -212,7 +216,7 @@ class TensorCoreFP8E5M2Layout(_TensorCoreFP8LayoutBase):
|
|||
TensorCoreFP8Layout = TensorCoreFP8E4M3Layout
|
||||
TensorWiseINT8Layout = _CKTensorWiseINT8Layout
|
||||
TensorCoreConvRotW4A4Layout = _CKTensorCoreConvRotW4A4Layout
|
||||
|
||||
AsymW4A8Int8Layout = _CKAsymW4A8Int8Layout
|
||||
|
||||
# ==============================================================================
|
||||
# Registry
|
||||
|
|
@ -226,6 +230,7 @@ register_layout_class("TensorWiseINT8Layout", _CKTensorWiseINT8Layout)
|
|||
register_layout_class("TensorCoreConvRotW4A4Layout", _CKTensorCoreConvRotW4A4Layout)
|
||||
if _CK_MXFP8_AVAILABLE:
|
||||
register_layout_class("TensorCoreMXFP8Layout", TensorCoreMXFP8Layout)
|
||||
register_layout_class("AsymW4A8Int8Layout", _CKAsymW4A8Int8Layout)
|
||||
|
||||
QUANT_ALGOS = {
|
||||
"float8_e4m3fn": {
|
||||
|
|
@ -268,6 +273,13 @@ QUANT_ALGOS["convrot_w4a4"] = {
|
|||
"quantize_input": False,
|
||||
}
|
||||
|
||||
QUANT_ALGOS["asym_w4a8_int8"] = {
|
||||
"storage_t": torch.int8,
|
||||
"parameters": {"weight_scale"},
|
||||
"comfy_tensor_layout": "AsymW4A8Int8Layout",
|
||||
"quantize_input": False,
|
||||
}
|
||||
|
||||
|
||||
# ==============================================================================
|
||||
# Re-exports for backward compatibility
|
||||
|
|
@ -282,6 +294,7 @@ __all__ = [
|
|||
"TensorCoreNVFP4Layout",
|
||||
"TensorCoreConvRotW4A4Layout",
|
||||
"TensorWiseINT8Layout",
|
||||
"AsymW4A8Int8Layout",
|
||||
"QUANT_ALGOS",
|
||||
"register_layout_op",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -1218,6 +1218,8 @@ class CFGGuider:
|
|||
return sampling_function(self.inner_model, x, timestep, self.conds.get("negative", None), self.conds.get("positive", None), self.cfg, model_options=model_options, seed=seed)
|
||||
|
||||
def inner_sample(self, noise, latent_image, device, sampler, sigmas, denoise_mask, callback, disable_pbar, seed, latent_shapes=None):
|
||||
self.inner_model.latent_shapes = latent_shapes
|
||||
|
||||
if latent_image is not None and torch.count_nonzero(latent_image) > 0: #Don't shift the empty latent image.
|
||||
latent_image = self.inner_model.process_latent_in(latent_image)
|
||||
|
||||
|
|
|
|||
|
|
@ -940,7 +940,11 @@ class VAE:
|
|||
if not comfy.memory_management.aimdo_enabled:
|
||||
self.disable_offload = True
|
||||
elif "decoder.transformer_blocks.0.scale1" in sd and "encoder.down.5.block.0.conv1.weight" in sd: # MiniMax H3 video VAE
|
||||
self.first_stage_model = comfy.ldm.minimax.vae.MiniMaxH3VideoVAE()
|
||||
minimax_ops = comfy.ops.disable_weight_init
|
||||
minimax_quant = comfy.utils.detect_layer_quantization(sd, "")
|
||||
if minimax_quant is not None: # int8+convrot quantized decoder
|
||||
minimax_ops = comfy.ops.mixed_precision_ops(minimax_quant, dtype if dtype is not None else torch.float16)
|
||||
self.first_stage_model = comfy.ldm.minimax.vae.MiniMaxH3VideoVAE(operations=minimax_ops)
|
||||
self.latent_channels = 24
|
||||
self.latent_dim = 3
|
||||
# frames 17k+5 <-> latents 5k+2, 16x spatial
|
||||
|
|
|
|||
|
|
@ -963,6 +963,7 @@ class MiniMaxH3(supported_models_base.BASE):
|
|||
|
||||
sampling_settings = {
|
||||
"shift": 12.0,
|
||||
"audio_shift": 3.0,
|
||||
}
|
||||
|
||||
unet_extra_config = {}
|
||||
|
|
@ -1449,6 +1450,20 @@ class WAN22_Animate(WAN21_T2V):
|
|||
out = model_base.WAN22_Animate(self, device=device)
|
||||
return out
|
||||
|
||||
class WAN_Animate2(WAN21_T2V):
|
||||
unet_config = {
|
||||
"image_model": "wan2.1",
|
||||
"model_type": "animate2",
|
||||
}
|
||||
|
||||
sampling_settings = {
|
||||
"shift": 5.0,
|
||||
}
|
||||
|
||||
def get_model(self, state_dict, prefix="", device=None):
|
||||
out = model_base.WAN_Animate2(self, device=device)
|
||||
return out
|
||||
|
||||
class WAN22_T2V(WAN21_T2V):
|
||||
unet_config = {
|
||||
"image_model": "wan2.1",
|
||||
|
|
@ -2463,6 +2478,7 @@ models = [
|
|||
WAN22_S2V,
|
||||
WAN21_HuMo,
|
||||
WAN22_Animate,
|
||||
WAN_Animate2,
|
||||
WAN21_FlowRVS,
|
||||
WAN21_SCAIL,
|
||||
WAN21_SCAIL2,
|
||||
|
|
|
|||
|
|
@ -868,16 +868,10 @@ class BaseGenerate:
|
|||
else:
|
||||
module = self.model.embed_tokens
|
||||
|
||||
offload_stream = None
|
||||
if module.comfy_cast_weights:
|
||||
weight, _, offload_stream = comfy.ops.cast_bias_weight(module, input, offloadable=True)
|
||||
else:
|
||||
weight = self.model.embed_tokens.weight.to(x)
|
||||
|
||||
x = torch.nn.functional.linear(input, weight, None)
|
||||
|
||||
comfy.ops.uncast_bias_weight(module, weight, None, offload_stream)
|
||||
return x
|
||||
if not module.comfy_cast_weights:
|
||||
return torch.nn.functional.linear(input, self.model.embed_tokens.weight.to(x), None)
|
||||
with comfy.ops.CastBiasWeightContext(module, input, offloadable=True) as (weight, _bias):
|
||||
return torch.nn.functional.linear(input, weight, None)
|
||||
|
||||
def init_kv_cache(self, batch, max_cache_len, device, execution_dtype):
|
||||
model_config = self.model.config
|
||||
|
|
|
|||
|
|
@ -847,6 +847,61 @@ class Load3DAnimation(Load3D):
|
|||
...
|
||||
|
||||
|
||||
@comfytype(io_type="LAYERS")
|
||||
class Layers(ComfyTypeIO):
|
||||
BlendMode = Literal[
|
||||
"normal", "multiply", "screen", "overlay", "darken", "lighten",
|
||||
"color-dodge", "color-burn", "hard-light", "soft-light", "difference",
|
||||
"exclusion", "linear-dodge", "linear-burn", "vivid-light", "pin-light",
|
||||
"linear-light", "hard-mix", "subtract", "divide", "grain-extract",
|
||||
"grain-merge", "hue", "saturation", "color", "luminosity",
|
||||
]
|
||||
|
||||
class LayerItem(TypedDict):
|
||||
image: torch.Tensor
|
||||
type: Literal["raster"]
|
||||
x: NotRequired[int]
|
||||
y: NotRequired[int]
|
||||
mask: NotRequired[torch.Tensor]
|
||||
z_index: int
|
||||
name: NotRequired[str]
|
||||
opacity: NotRequired[float]
|
||||
blend_mode: NotRequired["Layers.BlendMode"]
|
||||
visible: NotRequired[bool]
|
||||
flip_h: NotRequired[bool]
|
||||
flip_v: NotRequired[bool]
|
||||
rotation: NotRequired[float]
|
||||
w: NotRequired[int]
|
||||
h: NotRequired[int]
|
||||
|
||||
class Document(TypedDict):
|
||||
version: int
|
||||
canvas: NotRequired[tuple[int, int]]
|
||||
layers: list["Layers.LayerItem"]
|
||||
|
||||
Type = Document
|
||||
|
||||
|
||||
@comfytype(io_type="COMPOSITOR")
|
||||
class Compositor(ComfyTypeIO):
|
||||
class LayerState(TypedDict):
|
||||
version: NotRequired[int]
|
||||
canvas: dict
|
||||
background: NotRequired[dict]
|
||||
inputs: NotRequired[list[str]]
|
||||
order: NotRequired[list[int]]
|
||||
layers: list[dict]
|
||||
|
||||
Type = LayerState
|
||||
|
||||
class Input(WidgetInput):
|
||||
def __init__(self, id: str, display_name: str=None, optional=False, tooltip: str=None,
|
||||
socketless: bool=True, default: dict=None, advanced: bool=None):
|
||||
super().__init__(id, display_name, optional, tooltip, None, default, socketless, None, None, None, None, advanced)
|
||||
if default is None:
|
||||
self.default = {}
|
||||
|
||||
|
||||
@comfytype(io_type="PHOTOMAKER")
|
||||
class Photomaker(ComfyTypeIO):
|
||||
Type = Any
|
||||
|
|
@ -2403,6 +2458,8 @@ __all__ = [
|
|||
"Load3DModelInfo",
|
||||
"Load3D",
|
||||
"Load3DAnimation",
|
||||
"Compositor",
|
||||
"Layers",
|
||||
"Photomaker",
|
||||
"Point",
|
||||
"FaceAnalysis",
|
||||
|
|
|
|||
|
|
@ -260,6 +260,7 @@ class ImageSaveHelper:
|
|||
class AudioSaveHelper:
|
||||
"""A helper class with static methods to handle audio saving and metadata."""
|
||||
_OPUS_RATES = [8000, 12000, 16000, 24000, 48000]
|
||||
_FORMATS = {"flac", "mp3", "opus"}
|
||||
|
||||
@staticmethod
|
||||
def save_audio(
|
||||
|
|
@ -270,6 +271,9 @@ class AudioSaveHelper:
|
|||
format: str = "flac",
|
||||
quality: str = "128k",
|
||||
) -> list[SavedResult]:
|
||||
if format not in AudioSaveHelper._FORMATS:
|
||||
raise ValueError(f"Unsupported audio format: {format!r}")
|
||||
|
||||
full_output_folder, filename, counter, subfolder, _ = folder_paths.get_save_image_path(
|
||||
filename_prefix, _get_directory_by_folder_type(folder_type)
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
from enum import Enum
|
||||
from typing import Any
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
from pydantic import BaseModel, ConfigDict, Field
|
||||
|
||||
|
||||
class BFLFluxExpandImageRequest(BaseModel):
|
||||
|
|
@ -121,6 +121,8 @@ class BFLFluxProGenerateResponse(BaseModel):
|
|||
class BFLStatus(str, Enum):
|
||||
task_not_found = "Task not found"
|
||||
pending = "Pending"
|
||||
reasoning = "Reasoning"
|
||||
generating = "Generating"
|
||||
request_moderated = "Request Moderated"
|
||||
content_moderated = "Content Moderated"
|
||||
ready = "Ready"
|
||||
|
|
@ -132,3 +134,35 @@ class BFLFluxStatusResponse(BaseModel):
|
|||
status: BFLStatus = Field(...)
|
||||
result: dict[str, Any] | None = Field(None)
|
||||
progress: float | None = Field(None, ge=0.0, le=1.0)
|
||||
|
||||
|
||||
class Flux3VideoRequest(BaseModel):
|
||||
"""Fields shared by every generation mode of /v1/flux-3-video."""
|
||||
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
prompt: str = Field(...)
|
||||
aspect_ratio: str = Field("auto")
|
||||
duration: int | str = Field("auto", description="Whole seconds, or 'auto'.")
|
||||
resolution: str = Field("hd", description="'hd' is the 720p class, 'fhd' the 1080p class.")
|
||||
generate_audio: bool = Field(True)
|
||||
safety_tolerance: int = Field(2, description="0 is the strictest; conditioned modes cap at 2.")
|
||||
|
||||
|
||||
class Flux3TextToVideoRequest(Flux3VideoRequest):
|
||||
mode: str = Field("t2v")
|
||||
|
||||
|
||||
class Flux3ImageToVideoRequest(Flux3VideoRequest):
|
||||
mode: str = Field("i2v")
|
||||
keyframes: list[str] | list[tuple[float, str]] = Field(
|
||||
...,
|
||||
description="Images (URL or base64), or [seconds, image] pairs pinning each to a time.",
|
||||
)
|
||||
|
||||
|
||||
class Flux3VideoContinuationRequest(Flux3VideoRequest):
|
||||
mode: str = Field("v2v")
|
||||
start_video: str = Field(
|
||||
..., description="MP4 (URL or base64); the new clip carries on from its final frames."
|
||||
)
|
||||
|
|
|
|||
|
|
@ -35,6 +35,23 @@ class Seedream4TaskCreationRequest(BaseModel):
|
|||
optimize_prompt_options: Seedream5OptimizePromptOptions | None = None
|
||||
|
||||
|
||||
class Seedream5LayerOptimizePromptOptions(BaseModel):
|
||||
mode: Literal["standard", "fast"] = Field(...)
|
||||
|
||||
|
||||
class Seedream5LayerSeparationRequest(BaseModel):
|
||||
model: str = Field(...)
|
||||
prompt: str | None = Field(None)
|
||||
image: str = Field(..., description="Single image URL")
|
||||
size: str = Field("auto")
|
||||
seed: int = Field(..., ge=0, le=2147483647)
|
||||
response_format: str = Field("url")
|
||||
output_format: str = Field("png")
|
||||
layer_decomposition: bool = Field(True)
|
||||
watermark: bool = Field(False)
|
||||
optimize_prompt_options: Seedream5LayerOptimizePromptOptions | None = Field(None)
|
||||
|
||||
|
||||
class ImageTaskCreationResponse(BaseModel):
|
||||
model: str = Field(...)
|
||||
created: int = Field(..., description="Unix timestamp (in seconds) indicating time when the request was created.")
|
||||
|
|
@ -95,9 +112,10 @@ class Seedance2TaskCreationRequest(BaseModel):
|
|||
generate_audio: bool | None = Field(None)
|
||||
resolution: str | None = Field(None)
|
||||
ratio: str | None = Field(None)
|
||||
duration: int | None = Field(None, ge=4, le=15)
|
||||
duration: int | None = Field(None)
|
||||
seed: int | None = Field(None, ge=0, le=2147483647)
|
||||
watermark: bool | None = Field(None)
|
||||
output_format: str | None = Field(None)
|
||||
|
||||
|
||||
class TaskCreationResponse(BaseModel):
|
||||
|
|
@ -186,6 +204,10 @@ SEEDANCE2_PRICE_PER_1K_TOKENS = {
|
|||
("dreamina-seedance-2-0-mini", True, "480p"): 0.0021,
|
||||
("dreamina-seedance-2-0-mini", False, "720p"): 0.0035,
|
||||
("dreamina-seedance-2-0-mini", True, "720p"): 0.0021,
|
||||
("dreamina-seedance-2-5-260628", False, "480p"): 0.0107,
|
||||
("dreamina-seedance-2-5-260628", True, "480p"): 0.0064,
|
||||
("dreamina-seedance-2-5-260628", False, "720p"): 0.0107,
|
||||
("dreamina-seedance-2-5-260628", True, "720p"): 0.0064,
|
||||
}
|
||||
|
||||
|
||||
|
|
@ -304,8 +326,31 @@ SEEDANCE2_REF_VIDEO_PIXEL_LIMITS = {
|
|||
"480p": {"min": 409_600, "max": 927_408},
|
||||
"720p": {"min": 409_600, "max": 927_408},
|
||||
},
|
||||
"dreamina-seedance-2-5-260628": {
|
||||
"480p": {"min": 409_600, "max": 8_295_044},
|
||||
"720p": {"min": 409_600, "max": 8_295_044},
|
||||
},
|
||||
}
|
||||
|
||||
SEEDANCE2_REFERENCE_LIMITS_DEFAULT = {
|
||||
"max_images": 9,
|
||||
"max_videos": 3,
|
||||
"max_audios": 3,
|
||||
"max_total_seconds": 15.1,
|
||||
}
|
||||
SEEDANCE2_REFERENCE_LIMITS = {
|
||||
"dreamina-seedance-2-5-260628": {
|
||||
"max_images": 30,
|
||||
"max_videos": 10,
|
||||
"max_audios": 10,
|
||||
"max_total_seconds": 30.1,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def seedance2_reference_limits(model_id: str) -> dict:
|
||||
return SEEDANCE2_REFERENCE_LIMITS.get(model_id, SEEDANCE2_REFERENCE_LIMITS_DEFAULT)
|
||||
|
||||
# The time in this dictionary are given for 10 seconds duration.
|
||||
VIDEO_TASKS_EXECUTION_TIME = {
|
||||
"seedance-1-0-lite-t2v-250428": {
|
||||
|
|
|
|||
|
|
@ -20,6 +20,41 @@ class ImageEnhanceRequest(BaseModel):
|
|||
color_preservation: str = Field("true", description="To preserve the original color")
|
||||
|
||||
|
||||
class ImageEnhanceRequestV2(BaseModel):
|
||||
model: str = Field(...)
|
||||
output_format: str = Field("png")
|
||||
source_url: str = Field(...)
|
||||
output_width: Optional[int] = Field(None)
|
||||
output_height: Optional[int] = Field(None)
|
||||
crop_to_fill: Optional[bool] = Field(None, description="Available for Reimagine only")
|
||||
prompt: Optional[str] = Field(None, description="Available for Reimagine and Bloom 2")
|
||||
creativity: Optional[int] = Field(None, description="From 1 to 9; available for Reimagine and Bloom 2")
|
||||
subject_detection: Optional[str] = Field(None, description="Available for Reimagine only")
|
||||
face_enhancement: Optional[bool] = Field(None, description="Available for Reimagine only")
|
||||
face_enhancement_creativity: Optional[float] = Field(None, description="Is ignored if face_enhancement is false")
|
||||
face_enhancement_strength: Optional[float] = Field(None, description="Is ignored if face_enhancement is false")
|
||||
face_preservation: Optional[str] = Field(
|
||||
None, description='String "true" or "false"; available for Reimagine only'
|
||||
)
|
||||
color_preservation: Optional[str] = Field(
|
||||
None, description='String "true" or "false"; available for Reimagine and Bloom 2'
|
||||
)
|
||||
autoprompt: Optional[str] = Field(
|
||||
None, description='String "true" or "false"; auto-generate a prompt, available for Bloom 2 only'
|
||||
)
|
||||
seed: Optional[int] = Field(None, description="Available for Bloom 2 only")
|
||||
enhancement_strength: Optional[str] = Field(
|
||||
None, description="low, medium or high; available for Wonder 3.5 only"
|
||||
)
|
||||
grain: Optional[str] = Field(
|
||||
None, description='String "true" or "false"; available for Bloom 2 and Wonder 3.5'
|
||||
)
|
||||
grain_model: Optional[str] = Field(None, description="silver, gaussian or grey")
|
||||
grain_strength: Optional[float] = Field(None, description="From 0 to 1")
|
||||
grain_size: Optional[float] = Field(None, description="From 1 to 5")
|
||||
grain_density: Optional[float] = Field(None, description="From 0 to 1")
|
||||
|
||||
|
||||
class ImageAsyncTaskResponse(BaseModel):
|
||||
process_id: str = Field(...)
|
||||
|
||||
|
|
|
|||
|
|
@ -1,3 +1,5 @@
|
|||
import math
|
||||
|
||||
import torch
|
||||
from pydantic import BaseModel
|
||||
from typing_extensions import override
|
||||
|
|
@ -14,16 +16,23 @@ from comfy_api_nodes.apis.bfl import (
|
|||
BFLFluxVTORequest,
|
||||
BFLStatus,
|
||||
Flux2ProGenerateRequest,
|
||||
Flux3ImageToVideoRequest,
|
||||
Flux3TextToVideoRequest,
|
||||
Flux3VideoContinuationRequest,
|
||||
Flux3VideoRequest,
|
||||
)
|
||||
from comfy_api_nodes.util import (
|
||||
ApiEndpoint,
|
||||
convert_mask_to_image,
|
||||
download_url_to_image_tensor,
|
||||
download_url_to_video_output,
|
||||
get_number_of_images,
|
||||
poll_op,
|
||||
resize_mask_to_image,
|
||||
sync_op,
|
||||
tensor_to_base64_string,
|
||||
upload_images_to_comfyapi,
|
||||
upload_video_to_comfyapi,
|
||||
validate_aspect_ratio_string,
|
||||
validate_image_dimensions,
|
||||
validate_string,
|
||||
|
|
@ -1007,6 +1016,385 @@ class Flux2ImageNode(IO.ComfyNode):
|
|||
return IO.NodeOutput(await download_url_to_image_tensor(response.result["sample"]))
|
||||
|
||||
|
||||
_FLUX3_ASPECT_RATIOS = ["auto", "21:9", "2:1", "16:9", "4:3", "1:1", "3:4", "9:16"]
|
||||
_FLUX3_MIN_DURATION = 5
|
||||
_FLUX3_MAX_DURATION = 20
|
||||
_FLUX3_DURATIONS = ["auto"] + [str(i) for i in range(_FLUX3_MIN_DURATION, _FLUX3_MAX_DURATION + 1)]
|
||||
_FLUX3_RESOLUTIONS = {"720p": "hd", "1080p": "fhd"}
|
||||
_FLUX3_MAX_IMAGES = 10
|
||||
_FLUX3_MIN_IMAGE_SIDE = 256
|
||||
_FLUX3_MAX_IMAGE_ASPECT = 64
|
||||
|
||||
|
||||
def _flux3_validate_image(image: torch.Tensor) -> None:
|
||||
validate_image_dimensions(image, min_width=_FLUX3_MIN_IMAGE_SIDE, min_height=_FLUX3_MIN_IMAGE_SIDE)
|
||||
height, width = image.shape[-3], image.shape[-2]
|
||||
if max(width, height) > _FLUX3_MAX_IMAGE_ASPECT * min(width, height):
|
||||
raise ValueError(
|
||||
f"Image aspect ratio is too extreme ({width}x{height}); "
|
||||
f"FLUX 3 accepts at most {_FLUX3_MAX_IMAGE_ASPECT}:1."
|
||||
)
|
||||
|
||||
|
||||
def _flux3_collect_images(images: dict | None, field_name: str) -> list[torch.Tensor]:
|
||||
"""Flatten Autogrow slots (each possibly batched) into single images and validate them."""
|
||||
flat: list[torch.Tensor] = []
|
||||
for tensor in (images or {}).values():
|
||||
if tensor is None:
|
||||
continue
|
||||
if tensor.ndim == 4:
|
||||
flat.extend(tensor[i] for i in range(tensor.shape[0]))
|
||||
else:
|
||||
flat.append(tensor)
|
||||
if len(flat) > _FLUX3_MAX_IMAGES:
|
||||
raise ValueError(f"FLUX 3 supports at most {_FLUX3_MAX_IMAGES} {field_name}, got {len(flat)}.")
|
||||
for tensor in flat:
|
||||
_flux3_validate_image(tensor)
|
||||
return flat
|
||||
|
||||
|
||||
def _flux3_parse_times(value: str, image_count: int, duration: int | str) -> list[float]:
|
||||
"""Parse one keyframe time in seconds per image: increasing, inside the clip."""
|
||||
parts = [part.strip() for part in value.split(",") if part.strip()]
|
||||
if len(parts) != image_count:
|
||||
raise ValueError(
|
||||
f"Give one time per keyframe image: got {len(parts)} time(s) for {image_count} image(s)."
|
||||
)
|
||||
try:
|
||||
times = [float(part) for part in parts]
|
||||
except ValueError as exc:
|
||||
raise ValueError(f"Keyframe times must be numbers in seconds, comma-separated; got '{value}'.") from exc
|
||||
if not all(math.isfinite(time) for time in times):
|
||||
raise ValueError(f"Keyframe times must be finite numbers in seconds; got '{value}'.")
|
||||
if any(later <= earlier for earlier, later in zip(times, times[1:])):
|
||||
raise ValueError(f"Keyframe times must increase; got {times}.")
|
||||
if times[0] < 0:
|
||||
raise ValueError(f"Keyframe times cannot be negative; got {times[0]}.")
|
||||
cap = _FLUX3_MAX_DURATION if duration == "auto" else int(duration)
|
||||
if times[-1] > cap:
|
||||
raise ValueError(f"Keyframe time {times[-1]}s is past the end of a {cap}s clip.")
|
||||
return times
|
||||
|
||||
|
||||
class Flux3VideoNodeBase(IO.ComfyNode):
|
||||
"""Shared widgets, request plumbing and polling for the FLUX 3 generation modes."""
|
||||
|
||||
RATE_HD: float
|
||||
RATE_FHD: float
|
||||
|
||||
@classmethod
|
||||
def common_inputs(cls) -> list:
|
||||
return [
|
||||
IO.Combo.Input(
|
||||
"aspect_ratio",
|
||||
options=_FLUX3_ASPECT_RATIOS,
|
||||
default="auto",
|
||||
tooltip="Output aspect ratio. 'auto' picks one from the prompt and inputs.",
|
||||
),
|
||||
IO.Combo.Input(
|
||||
"duration",
|
||||
options=_FLUX3_DURATIONS,
|
||||
default="auto",
|
||||
tooltip="Clip length in seconds. 'auto' fits the length to the content.",
|
||||
),
|
||||
IO.Combo.Input(
|
||||
"resolution",
|
||||
options=list(_FLUX3_RESOLUTIONS),
|
||||
default="720p",
|
||||
tooltip="Output resolution.",
|
||||
),
|
||||
IO.Boolean.Input(
|
||||
"generate_audio",
|
||||
default=True,
|
||||
tooltip="Generate synchronized audio (ambient, speech, effects). "
|
||||
"Off produces a video with no audio track.",
|
||||
),
|
||||
IO.Int.Input(
|
||||
"safety_tolerance",
|
||||
default=2,
|
||||
min=0,
|
||||
max=4,
|
||||
advanced=True,
|
||||
tooltip="Moderation tolerance, 0 is the strictest. Requests that send images or "
|
||||
"video are capped at 2 whatever you set here.",
|
||||
),
|
||||
IO.Int.Input(
|
||||
"seed",
|
||||
default=42,
|
||||
min=0,
|
||||
max=0xFFFFFFFF,
|
||||
control_after_generate=True,
|
||||
tooltip="Seed to determine if node should re-run; FLUX 3 picks its own seed, so "
|
||||
"actual results are nondeterministic regardless of this value.",
|
||||
),
|
||||
]
|
||||
|
||||
@classmethod
|
||||
def common_fields(
|
||||
cls,
|
||||
prompt: str,
|
||||
aspect_ratio: str,
|
||||
duration: str,
|
||||
resolution: str,
|
||||
generate_audio: bool,
|
||||
safety_tolerance: int,
|
||||
) -> dict:
|
||||
validate_string(prompt, field_name="prompt", min_length=1)
|
||||
return {
|
||||
"prompt": prompt,
|
||||
"aspect_ratio": aspect_ratio,
|
||||
"duration": duration if duration == "auto" else int(duration),
|
||||
"resolution": _FLUX3_RESOLUTIONS[resolution],
|
||||
"generate_audio": generate_audio,
|
||||
"safety_tolerance": safety_tolerance,
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def price_badge(cls) -> IO.PriceBadge:
|
||||
return IO.PriceBadge(
|
||||
depends_on=IO.PriceBadgeDepends(widgets=["resolution", "duration"]),
|
||||
expr=f"""
|
||||
(
|
||||
$rate := widgets.resolution = "1080p" ? {cls.RATE_FHD} : {cls.RATE_HD};
|
||||
$type(widgets.duration) = "string" and widgets.duration != "auto"
|
||||
? {{"type":"usd","usd": $rate * $number(widgets.duration)}}
|
||||
: {{"type":"usd","usd": $rate, "format": {{"suffix": "/second"}}}}
|
||||
)
|
||||
""",
|
||||
)
|
||||
|
||||
|
||||
async def _flux3_execute(cls: type[IO.ComfyNode], request: Flux3VideoRequest) -> IO.NodeOutput:
|
||||
initial_response = await sync_op(
|
||||
cls,
|
||||
ApiEndpoint(path="/proxy/bfl/v1/flux-3-video", method="POST"),
|
||||
response_model=BFLFluxProGenerateResponse,
|
||||
data=request,
|
||||
)
|
||||
|
||||
def price_extractor(_r: BaseModel) -> float | None:
|
||||
return None if initial_response.cost is None else initial_response.cost / 100
|
||||
|
||||
response = await poll_op(
|
||||
cls,
|
||||
ApiEndpoint(initial_response.polling_url),
|
||||
response_model=BFLFluxStatusResponse,
|
||||
status_extractor=lambda r: r.status,
|
||||
progress_extractor=lambda r: r.progress,
|
||||
price_extractor=price_extractor,
|
||||
completed_statuses=[BFLStatus.ready],
|
||||
failed_statuses=[
|
||||
BFLStatus.request_moderated,
|
||||
BFLStatus.content_moderated,
|
||||
BFLStatus.error,
|
||||
BFLStatus.task_not_found,
|
||||
],
|
||||
queued_statuses=[BFLStatus.pending],
|
||||
poll_interval=8.0,
|
||||
# a failed task answers the poll with a retryable-class HTTP 5xx (500 and 503 observed);
|
||||
# a small retry budget surfaces real failures quickly
|
||||
max_retries_per_poll=3,
|
||||
)
|
||||
return IO.NodeOutput(await download_url_to_video_output(response.result["sample"]))
|
||||
|
||||
|
||||
class Flux3TextToVideoNode(Flux3VideoNodeBase):
|
||||
RATE_HD = 0.2431
|
||||
RATE_FHD = 0.4147
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls) -> IO.Schema:
|
||||
return IO.Schema(
|
||||
node_id="Flux3TextToVideoNode",
|
||||
display_name="Flux 3 Text to Video",
|
||||
category="partner/video/BFL",
|
||||
description="Generates a video with synchronized audio from a text prompt via FLUX 3.",
|
||||
inputs=[
|
||||
IO.String.Input(
|
||||
"prompt",
|
||||
multiline=True,
|
||||
default="",
|
||||
tooltip="What you want, in plain language; the prompt is interpreted and expanded "
|
||||
"before generation. Describe ambient sound, music and speech separately for layered audio.",
|
||||
),
|
||||
*cls.common_inputs(),
|
||||
],
|
||||
outputs=[IO.Video.Output()],
|
||||
hidden=[
|
||||
IO.Hidden.auth_token_comfy_org,
|
||||
IO.Hidden.api_key_comfy_org,
|
||||
IO.Hidden.unique_id,
|
||||
],
|
||||
is_api_node=True,
|
||||
price_badge=cls.price_badge(),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
async def execute(
|
||||
cls,
|
||||
prompt: str,
|
||||
aspect_ratio: str,
|
||||
duration: str,
|
||||
resolution: str,
|
||||
generate_audio: bool,
|
||||
safety_tolerance: int,
|
||||
seed: int,
|
||||
) -> IO.NodeOutput:
|
||||
request = Flux3TextToVideoRequest(
|
||||
**cls.common_fields(prompt, aspect_ratio, duration, resolution, generate_audio, safety_tolerance)
|
||||
)
|
||||
return await _flux3_execute(cls, request)
|
||||
|
||||
|
||||
class Flux3ImageToVideoNode(Flux3VideoNodeBase):
|
||||
RATE_HD = 0.2431
|
||||
RATE_FHD = 0.4147
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls) -> IO.Schema:
|
||||
return IO.Schema(
|
||||
node_id="Flux3ImageToVideoNode",
|
||||
display_name="Flux 3 Image to Video",
|
||||
category="partner/video/BFL",
|
||||
description="Animates 1 to 10 images with FLUX 3. Each image becomes a frame of the clip: "
|
||||
"one image opens it, two morph from the first to the second, and more are spread across it "
|
||||
"or pinned to times you choose.",
|
||||
inputs=[
|
||||
IO.String.Input(
|
||||
"prompt",
|
||||
multiline=True,
|
||||
default="",
|
||||
tooltip="How the scene should move and sound; the prompt is interpreted and "
|
||||
"expanded before generation.",
|
||||
),
|
||||
IO.Autogrow.Input(
|
||||
"keyframes",
|
||||
template=IO.Autogrow.TemplatePrefix(
|
||||
IO.Image.Input("image", tooltip="Keyframe image."),
|
||||
prefix="image_",
|
||||
min=1,
|
||||
max=_FLUX3_MAX_IMAGES,
|
||||
),
|
||||
tooltip="1 to 10 images, in playback order. Minimum 256x256 pixels each.",
|
||||
),
|
||||
IO.DynamicCombo.Input(
|
||||
"placement",
|
||||
options=[
|
||||
IO.DynamicCombo.Option("spread across the clip", []),
|
||||
IO.DynamicCombo.Option(
|
||||
"at times",
|
||||
[
|
||||
IO.String.Input(
|
||||
"times",
|
||||
default="0",
|
||||
tooltip="One time in seconds per image, comma-separated and "
|
||||
"increasing, e.g. '0, 2.5, 5'.",
|
||||
),
|
||||
],
|
||||
),
|
||||
],
|
||||
tooltip="'spread across the clip' lets FLUX 3 place the images (one opens the clip, "
|
||||
"two become its start and end); 'at times' pins every image to a second you choose.",
|
||||
),
|
||||
*cls.common_inputs(),
|
||||
],
|
||||
outputs=[IO.Video.Output()],
|
||||
hidden=[
|
||||
IO.Hidden.auth_token_comfy_org,
|
||||
IO.Hidden.api_key_comfy_org,
|
||||
IO.Hidden.unique_id,
|
||||
],
|
||||
is_api_node=True,
|
||||
price_badge=cls.price_badge(),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
async def execute(
|
||||
cls,
|
||||
prompt: str,
|
||||
keyframes: IO.Autogrow.Type,
|
||||
placement: dict,
|
||||
aspect_ratio: str,
|
||||
duration: str,
|
||||
resolution: str,
|
||||
generate_audio: bool,
|
||||
safety_tolerance: int,
|
||||
seed: int,
|
||||
) -> IO.NodeOutput:
|
||||
fields = cls.common_fields(prompt, aspect_ratio, duration, resolution, generate_audio, safety_tolerance)
|
||||
images = _flux3_collect_images(keyframes, "keyframes")
|
||||
if not images:
|
||||
raise ValueError("Connect at least one keyframe image.")
|
||||
times = None
|
||||
if placement["placement"] == "at times":
|
||||
times = _flux3_parse_times(placement["times"], len(images), fields["duration"])
|
||||
elif len(images) >= 3 and fields["duration"] == "auto":
|
||||
# spread images land evenly between the first and last, which needs a known length
|
||||
raise ValueError(
|
||||
f"Spreading {len(images)} images across the clip needs an explicit duration: "
|
||||
"set duration, or place the images yourself with 'at times'."
|
||||
)
|
||||
urls = await upload_images_to_comfyapi(
|
||||
cls, images, max_images=_FLUX3_MAX_IMAGES, wait_label="Uploading keyframes"
|
||||
)
|
||||
request = Flux3ImageToVideoRequest(
|
||||
keyframes=list(zip(times, urls)) if times is not None else urls,
|
||||
**fields,
|
||||
)
|
||||
return await _flux3_execute(cls, request)
|
||||
|
||||
|
||||
class Flux3VideoContinuationNode(Flux3VideoNodeBase):
|
||||
RATE_HD = 0.5863
|
||||
RATE_FHD = 0.7579
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls) -> IO.Schema:
|
||||
return IO.Schema(
|
||||
node_id="Flux3VideoContinuationNode",
|
||||
display_name="Flux 3 Video Continuation",
|
||||
category="partner/video/BFL",
|
||||
description="Continues a video with FLUX 3: the new clip carries on from the final frames "
|
||||
"of the one you provide.",
|
||||
inputs=[
|
||||
IO.Video.Input("video", tooltip="The clip to continue."),
|
||||
IO.String.Input(
|
||||
"prompt",
|
||||
multiline=True,
|
||||
default="",
|
||||
tooltip="What the continuation should show; the prompt is interpreted and expanded "
|
||||
"before generation.",
|
||||
),
|
||||
*cls.common_inputs(),
|
||||
],
|
||||
outputs=[IO.Video.Output()],
|
||||
hidden=[
|
||||
IO.Hidden.auth_token_comfy_org,
|
||||
IO.Hidden.api_key_comfy_org,
|
||||
IO.Hidden.unique_id,
|
||||
],
|
||||
is_api_node=True,
|
||||
price_badge=cls.price_badge(),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
async def execute(
|
||||
cls,
|
||||
video: Input.Video,
|
||||
prompt: str,
|
||||
aspect_ratio: str,
|
||||
duration: str,
|
||||
resolution: str,
|
||||
generate_audio: bool,
|
||||
safety_tolerance: int,
|
||||
seed: int,
|
||||
) -> IO.NodeOutput:
|
||||
fields = cls.common_fields(prompt, aspect_ratio, duration, resolution, generate_audio, safety_tolerance)
|
||||
url = await upload_video_to_comfyapi(cls, video, wait_label="Uploading source video")
|
||||
request = Flux3VideoContinuationRequest(start_video=url, **fields)
|
||||
return await _flux3_execute(cls, request)
|
||||
|
||||
|
||||
class BFLExtension(ComfyExtension):
|
||||
@override
|
||||
async def get_node_list(self) -> list[type[IO.ComfyNode]]:
|
||||
|
|
@ -1021,6 +1409,9 @@ class BFLExtension(ComfyExtension):
|
|||
Flux2ProImageNode,
|
||||
Flux2MaxImageNode,
|
||||
Flux2ImageNode,
|
||||
Flux3TextToVideoNode,
|
||||
Flux3ImageToVideoNode,
|
||||
Flux3VideoContinuationNode,
|
||||
]
|
||||
|
||||
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load Diff
|
|
@ -76,8 +76,8 @@ GEMINI_IMAGE_2_PRICE_BADGE = IO.PriceBadge(
|
|||
$m := widgets.model;
|
||||
$r := widgets.resolution;
|
||||
$isFlash := $contains($m, "nano banana 2");
|
||||
$flashPrices := {"1k": 0.0696, "2k": 0.1014, "4k": 0.154};
|
||||
$proPrices := {"1k": 0.134, "2k": 0.134, "4k": 0.24};
|
||||
$flashPrices := {"1k": 0.0835, "2k": 0.1217, "4k": 0.1848};
|
||||
$proPrices := {"1k": 0.1608, "2k": 0.1608, "4k": 0.288};
|
||||
$prices := $isFlash ? $flashPrices : $proPrices;
|
||||
{"type":"usd","usd": $lookup($prices, $r), "format":{"suffix":"/Image","approximate":true}}
|
||||
)
|
||||
|
|
@ -1464,10 +1464,10 @@ class GeminiNanoBanana2V2(IO.ComfyNode):
|
|||
expr="""
|
||||
(
|
||||
$contains(widgets.model, "lite")
|
||||
? {"type":"usd","usd": 0.034, "format":{"suffix":"/Image","approximate":true}}
|
||||
? {"type":"usd","usd": 0.0408, "format":{"suffix":"/Image","approximate":true}}
|
||||
: (
|
||||
$r := $lookup(widgets, "model.resolution");
|
||||
$prices := {"1k": 0.0696, "2k": 0.1014, "4k": 0.154};
|
||||
$prices := {"1k": 0.0835, "2k": 0.1217, "4k": 0.1848};
|
||||
{"type":"usd","usd": $lookup($prices, $r), "format":{"suffix":"/Image","approximate":true}}
|
||||
)
|
||||
)
|
||||
|
|
|
|||
|
|
@ -5,17 +5,13 @@ For source of truth on the allowed permutations of request fields, please refere
|
|||
"""
|
||||
|
||||
import logging
|
||||
import math
|
||||
import re
|
||||
|
||||
import torch
|
||||
from typing_extensions import override
|
||||
|
||||
from comfy_api.latest import IO, ComfyExtension, Input, InputImpl
|
||||
from comfy_api.latest import IO, ComfyExtension, Input
|
||||
from comfy_api_nodes.apis import (
|
||||
KlingCameraControl,
|
||||
KlingCameraConfig,
|
||||
KlingCameraControlType,
|
||||
KlingVideoGenDuration,
|
||||
KlingVideoGenMode,
|
||||
KlingVideoGenAspectRatio,
|
||||
|
|
@ -30,23 +26,12 @@ from comfy_api_nodes.apis import (
|
|||
KlingLipSyncInputObject,
|
||||
KlingLipSyncRequest,
|
||||
KlingLipSyncResponse,
|
||||
KlingVirtualTryOnModelName,
|
||||
KlingVirtualTryOnRequest,
|
||||
KlingVirtualTryOnResponse,
|
||||
KlingVideoResult,
|
||||
KlingImageResult,
|
||||
KlingImageGenerationsRequest,
|
||||
KlingImageGenerationsResponse,
|
||||
KlingImageGenImageReferenceType,
|
||||
KlingImageGenAspectRatio,
|
||||
KlingVideoEffectsRequest,
|
||||
KlingVideoEffectsResponse,
|
||||
KlingDualCharacterEffectsScene,
|
||||
KlingSingleImageEffectsScene,
|
||||
KlingDualCharacterEffectInput,
|
||||
KlingSingleImageEffectInput,
|
||||
KlingCharacterEffectModelName,
|
||||
KlingSingleImageEffectModelName,
|
||||
)
|
||||
from comfy_api_nodes.apis.kling import (
|
||||
ImageToVideoWithAudioRequest,
|
||||
|
|
@ -119,9 +104,6 @@ PATH_TEXT_TO_VIDEO = f"/proxy/kling/{KLING_API_VERSION}/videos/text2video"
|
|||
PATH_IMAGE_TO_VIDEO = f"/proxy/kling/{KLING_API_VERSION}/videos/image2video"
|
||||
PATH_VIDEO_EXTEND = f"/proxy/kling/{KLING_API_VERSION}/videos/video-extend"
|
||||
PATH_LIP_SYNC = f"/proxy/kling/{KLING_API_VERSION}/videos/lip-sync"
|
||||
PATH_VIDEO_EFFECTS = f"/proxy/kling/{KLING_API_VERSION}/videos/effects"
|
||||
PATH_CHARACTER_IMAGE = f"/proxy/kling/{KLING_API_VERSION}/images/generations"
|
||||
PATH_VIRTUAL_TRY_ON = f"/proxy/kling/{KLING_API_VERSION}/images/kolors-virtual-try-on"
|
||||
PATH_IMAGE_GENERATIONS = f"/proxy/kling/{KLING_API_VERSION}/images/generations"
|
||||
|
||||
MAX_PROMPT_LENGTH_T2V = 2500
|
||||
|
|
@ -133,21 +115,11 @@ MAX_PROMPT_LENGTH_LIP_SYNC = 120
|
|||
AVERAGE_DURATION_T2V = 319
|
||||
AVERAGE_DURATION_I2V = 164
|
||||
AVERAGE_DURATION_LIP_SYNC = 455
|
||||
AVERAGE_DURATION_VIRTUAL_TRY_ON = 19
|
||||
AVERAGE_DURATION_IMAGE_GEN = 32
|
||||
AVERAGE_DURATION_VIDEO_EFFECTS = 320
|
||||
AVERAGE_DURATION_VIDEO_EXTEND = 320
|
||||
|
||||
|
||||
MODE_TEXT2VIDEO = {
|
||||
"standard mode / 5s duration / kling-v1-6": ("std", "5", "kling-v1-6"),
|
||||
"standard mode / 10s duration / kling-v1-6": ("std", "10", "kling-v1-6"),
|
||||
"pro mode / 5s duration / kling-v2-master": ("pro", "5", "kling-v2-master"),
|
||||
"pro mode / 10s duration / kling-v2-master": ("pro", "10", "kling-v2-master"),
|
||||
"standard mode / 5s duration / kling-v2-master": ("std", "5", "kling-v2-master"),
|
||||
"standard mode / 10s duration / kling-v2-master": ("std", "10", "kling-v2-master"),
|
||||
"pro mode / 5s duration / kling-v2-1-master": ("pro", "5", "kling-v2-1-master"),
|
||||
"pro mode / 10s duration / kling-v2-1-master": ("pro", "10", "kling-v2-1-master"),
|
||||
"pro mode / 5s duration / kling-v2-5-turbo": ("pro", "5", "kling-v2-5-turbo"),
|
||||
"pro mode / 10s duration / kling-v2-5-turbo": ("pro", "10", "kling-v2-5-turbo"),
|
||||
}
|
||||
|
|
@ -160,12 +132,6 @@ See: [Kling API Docs Capability Map](https://app.klingai.com/global/dev/document
|
|||
|
||||
|
||||
MODE_START_END_FRAME = {
|
||||
"pro mode / 5s duration / kling-v1-5": ("pro", "5", "kling-v1-5"),
|
||||
"pro mode / 10s duration / kling-v1-5": ("pro", "10", "kling-v1-5"),
|
||||
"pro mode / 5s duration / kling-v1-6": ("pro", "5", "kling-v1-6"),
|
||||
"pro mode / 10s duration / kling-v1-6": ("pro", "10", "kling-v1-6"),
|
||||
"pro mode / 5s duration / kling-v2-1": ("pro", "5", "kling-v2-1"),
|
||||
"pro mode / 10s duration / kling-v2-1": ("pro", "10", "kling-v2-1"),
|
||||
"pro mode / 5s duration / kling-v2-5-turbo": ("pro", "5", "kling-v2-5-turbo"),
|
||||
"pro mode / 10s duration / kling-v2-5-turbo": ("pro", "10", "kling-v2-5-turbo"),
|
||||
}
|
||||
|
|
@ -287,11 +253,6 @@ async def finish_omni_video_task(cls: type[IO.ComfyNode], response: TaskStatusRe
|
|||
return IO.NodeOutput(await download_url_to_video_output(final_response.data.task_result.videos[0].url))
|
||||
|
||||
|
||||
def is_valid_camera_control_configs(configs: list[float]) -> bool:
|
||||
"""Verifies that at least one camera control configuration is non-zero."""
|
||||
return any(not math.isclose(value, 0.0) for value in configs)
|
||||
|
||||
|
||||
def is_valid_task_creation_response(response: KlingText2VideoResponse) -> bool:
|
||||
"""Verifies that the initial response contains a task ID."""
|
||||
return bool(response.data.task_id)
|
||||
|
|
@ -307,7 +268,7 @@ def is_valid_video_response(response: KlingText2VideoResponse) -> bool:
|
|||
)
|
||||
|
||||
|
||||
def is_valid_image_response(response: KlingVirtualTryOnResponse) -> bool:
|
||||
def is_valid_image_response(response: KlingImageGenerationsResponse) -> bool:
|
||||
"""Verifies that the response contains a task result with at least one image."""
|
||||
return (
|
||||
response.data is not None
|
||||
|
|
@ -430,7 +391,6 @@ async def execute_text2video(
|
|||
model_mode: str,
|
||||
duration: str,
|
||||
aspect_ratio: str,
|
||||
camera_control: KlingCameraControl | None = None,
|
||||
) -> IO.NodeOutput:
|
||||
validate_prompts(prompt, negative_prompt, MAX_PROMPT_LENGTH_T2V)
|
||||
task_creation_response = await sync_op(
|
||||
|
|
@ -445,7 +405,6 @@ async def execute_text2video(
|
|||
model_name=model_name,
|
||||
cfg_scale=cfg_scale,
|
||||
aspect_ratio=KlingVideoGenAspectRatio(aspect_ratio),
|
||||
camera_control=camera_control,
|
||||
),
|
||||
)
|
||||
|
||||
|
|
@ -475,19 +434,11 @@ async def execute_image2video(
|
|||
model_mode: str,
|
||||
aspect_ratio: str,
|
||||
duration: str,
|
||||
camera_control: KlingCameraControl | None = None,
|
||||
end_frame: torch.Tensor | None = None,
|
||||
) -> IO.NodeOutput:
|
||||
validate_prompts(prompt, negative_prompt, MAX_PROMPT_LENGTH_I2V)
|
||||
validate_input_image(start_frame)
|
||||
|
||||
if camera_control is not None:
|
||||
# Camera control type for image 2 video is always `simple`
|
||||
camera_control.type = KlingCameraControlType.simple
|
||||
|
||||
if model_mode == "std" and model_name == KlingVideoGenModelName.kling_v2_5_turbo.value:
|
||||
model_mode = "pro" # October 5: currently "std" mode is not supported for this model
|
||||
|
||||
task_creation_response = await sync_op(
|
||||
cls,
|
||||
ApiEndpoint(path=PATH_IMAGE_TO_VIDEO, method="POST"),
|
||||
|
|
@ -505,7 +456,6 @@ async def execute_image2video(
|
|||
cfg_scale=cfg_scale,
|
||||
mode=KlingVideoGenMode(model_mode),
|
||||
duration=KlingVideoGenDuration(duration),
|
||||
camera_control=camera_control,
|
||||
),
|
||||
)
|
||||
|
||||
|
|
@ -525,59 +475,6 @@ async def execute_image2video(
|
|||
return IO.NodeOutput(await download_url_to_video_output(str(video.url)), str(video.id), str(video.duration))
|
||||
|
||||
|
||||
async def execute_video_effect(
|
||||
cls: type[IO.ComfyNode],
|
||||
dual_character: bool,
|
||||
effect_scene: KlingDualCharacterEffectsScene | KlingSingleImageEffectsScene,
|
||||
model_name: str,
|
||||
duration: KlingVideoGenDuration,
|
||||
image_1: torch.Tensor,
|
||||
image_2: torch.Tensor | None = None,
|
||||
model_mode: KlingVideoGenMode | None = None,
|
||||
) -> tuple[InputImpl.VideoFromFile, str, str]:
|
||||
if dual_character:
|
||||
request_input_field = KlingDualCharacterEffectInput(
|
||||
model_name=model_name,
|
||||
mode=model_mode,
|
||||
images=[
|
||||
tensor_to_base64_string(image_1),
|
||||
tensor_to_base64_string(image_2),
|
||||
],
|
||||
duration=duration,
|
||||
)
|
||||
else:
|
||||
request_input_field = KlingSingleImageEffectInput(
|
||||
model_name=model_name,
|
||||
image=tensor_to_base64_string(image_1),
|
||||
duration=duration,
|
||||
)
|
||||
|
||||
task_creation_response = await sync_op(
|
||||
cls,
|
||||
endpoint=ApiEndpoint(path=PATH_VIDEO_EFFECTS, method="POST"),
|
||||
response_model=KlingVideoEffectsResponse,
|
||||
data=KlingVideoEffectsRequest(
|
||||
effect_scene=effect_scene,
|
||||
input=request_input_field,
|
||||
),
|
||||
)
|
||||
|
||||
validate_task_creation_response(task_creation_response)
|
||||
task_id = task_creation_response.data.task_id
|
||||
|
||||
final_response = await poll_op(
|
||||
cls,
|
||||
ApiEndpoint(path=f"{PATH_VIDEO_EFFECTS}/{task_id}"),
|
||||
response_model=KlingVideoEffectsResponse,
|
||||
estimated_duration=AVERAGE_DURATION_VIDEO_EFFECTS,
|
||||
status_extractor=lambda r: (r.data.task_status.value if r.data and r.data.task_status else None),
|
||||
)
|
||||
validate_video_result_response(final_response)
|
||||
|
||||
video = get_video_from_response(final_response)
|
||||
return await download_url_to_video_output(str(video.url)), str(video.id), str(video.duration)
|
||||
|
||||
|
||||
async def execute_lipsync(
|
||||
cls: type[IO.ComfyNode],
|
||||
video: Input.Video,
|
||||
|
|
@ -640,125 +537,6 @@ async def execute_lipsync(
|
|||
return IO.NodeOutput(await download_url_to_video_output(str(video.url)), str(video.id), str(video.duration))
|
||||
|
||||
|
||||
class KlingCameraControls(IO.ComfyNode):
|
||||
"""Kling Camera Controls Node"""
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls) -> IO.Schema:
|
||||
return IO.Schema(
|
||||
node_id="KlingCameraControls",
|
||||
display_name="Kling Camera Controls",
|
||||
category="partner/video/Kling",
|
||||
description="Allows specifying configuration options for Kling Camera Controls and motion control effects.",
|
||||
inputs=[
|
||||
IO.Combo.Input("camera_control_type", options=KlingCameraControlType),
|
||||
IO.Float.Input(
|
||||
"horizontal_movement",
|
||||
default=0.0,
|
||||
min=-10.0,
|
||||
max=10.0,
|
||||
step=0.25,
|
||||
display_mode=IO.NumberDisplay.slider,
|
||||
tooltip="Controls camera's movement along horizontal axis (x-axis). Negative indicates left, positive indicates right",
|
||||
),
|
||||
IO.Float.Input(
|
||||
"vertical_movement",
|
||||
default=0.0,
|
||||
min=-10.0,
|
||||
max=10.0,
|
||||
step=0.25,
|
||||
display_mode=IO.NumberDisplay.slider,
|
||||
tooltip="Controls camera's movement along vertical axis (y-axis). Negative indicates downward, positive indicates upward.",
|
||||
),
|
||||
IO.Float.Input(
|
||||
"pan",
|
||||
default=0.5,
|
||||
min=-10.0,
|
||||
max=10.0,
|
||||
step=0.25,
|
||||
display_mode=IO.NumberDisplay.slider,
|
||||
tooltip="Controls camera's rotation in vertical plane (x-axis). Negative indicates downward rotation, positive indicates upward rotation.",
|
||||
),
|
||||
IO.Float.Input(
|
||||
"tilt",
|
||||
default=0.0,
|
||||
min=-10.0,
|
||||
max=10.0,
|
||||
step=0.25,
|
||||
display_mode=IO.NumberDisplay.slider,
|
||||
tooltip="Controls camera's rotation in horizontal plane (y-axis). Negative indicates left rotation, positive indicates right rotation.",
|
||||
),
|
||||
IO.Float.Input(
|
||||
"roll",
|
||||
default=0.0,
|
||||
min=-10.0,
|
||||
max=10.0,
|
||||
step=0.25,
|
||||
display_mode=IO.NumberDisplay.slider,
|
||||
tooltip="Controls camera's rolling amount (z-axis). Negative indicates counterclockwise, positive indicates clockwise.",
|
||||
),
|
||||
IO.Float.Input(
|
||||
"zoom",
|
||||
default=0.0,
|
||||
min=-10.0,
|
||||
max=10.0,
|
||||
step=0.25,
|
||||
display_mode=IO.NumberDisplay.slider,
|
||||
tooltip="Controls change in camera's focal length. Negative indicates narrower field of view, positive indicates wider field of view.",
|
||||
),
|
||||
],
|
||||
outputs=[IO.Custom("CAMERA_CONTROL").Output(display_name="camera_control")],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def validate_inputs(
|
||||
cls,
|
||||
horizontal_movement: float,
|
||||
vertical_movement: float,
|
||||
pan: float,
|
||||
tilt: float,
|
||||
roll: float,
|
||||
zoom: float,
|
||||
) -> bool | str:
|
||||
if not is_valid_camera_control_configs(
|
||||
[
|
||||
horizontal_movement,
|
||||
vertical_movement,
|
||||
pan,
|
||||
tilt,
|
||||
roll,
|
||||
zoom,
|
||||
]
|
||||
):
|
||||
return "Invalid camera control configs: at least one of the values must be non-zero"
|
||||
return True
|
||||
|
||||
@classmethod
|
||||
def execute(
|
||||
cls,
|
||||
camera_control_type: str,
|
||||
horizontal_movement: float,
|
||||
vertical_movement: float,
|
||||
pan: float,
|
||||
tilt: float,
|
||||
roll: float,
|
||||
zoom: float,
|
||||
) -> IO.NodeOutput:
|
||||
return IO.NodeOutput(
|
||||
KlingCameraControl(
|
||||
type=KlingCameraControlType(camera_control_type),
|
||||
config=KlingCameraConfig(
|
||||
horizontal=horizontal_movement,
|
||||
vertical=vertical_movement,
|
||||
pan=pan,
|
||||
roll=roll,
|
||||
tilt=tilt,
|
||||
zoom=zoom,
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
class KlingTextToVideoNode(IO.ComfyNode):
|
||||
"""Kling Text to Video Node"""
|
||||
|
||||
|
|
@ -782,7 +560,7 @@ class KlingTextToVideoNode(IO.ComfyNode):
|
|||
IO.Combo.Input(
|
||||
"mode",
|
||||
options=modes,
|
||||
default=modes[8],
|
||||
default=modes[0],
|
||||
tooltip="The configuration to use for the video generation following the format: mode / duration / model_name.",
|
||||
),
|
||||
],
|
||||
|
|
@ -802,25 +580,7 @@ class KlingTextToVideoNode(IO.ComfyNode):
|
|||
expr="""
|
||||
(
|
||||
$m := widgets.mode;
|
||||
$contains($m,"v2-5-turbo")
|
||||
? ($contains($m,"10") ? {"type":"usd","usd":0.7} : {"type":"usd","usd":0.35})
|
||||
: $contains($m,"v2-1-master")
|
||||
? ($contains($m,"10s") ? {"type":"usd","usd":2.8} : {"type":"usd","usd":1.4})
|
||||
: $contains($m,"v2-master")
|
||||
? ($contains($m,"10s") ? {"type":"usd","usd":2.8} : {"type":"usd","usd":1.4})
|
||||
: $contains($m,"v1-6")
|
||||
? (
|
||||
$contains($m,"pro")
|
||||
? ($contains($m,"10s") ? {"type":"usd","usd":0.98} : {"type":"usd","usd":0.49})
|
||||
: ($contains($m,"10s") ? {"type":"usd","usd":0.56} : {"type":"usd","usd":0.28})
|
||||
)
|
||||
: $contains($m,"v1")
|
||||
? (
|
||||
$contains($m,"pro")
|
||||
? ($contains($m,"10s") ? {"type":"usd","usd":0.98} : {"type":"usd","usd":0.49})
|
||||
: ($contains($m,"10s") ? {"type":"usd","usd":0.28} : {"type":"usd","usd":0.14})
|
||||
)
|
||||
: {"type":"usd","usd":0.14}
|
||||
$contains($m,"10") ? {"type":"usd","usd":0.7} : {"type":"usd","usd":0.35}
|
||||
)
|
||||
""",
|
||||
),
|
||||
|
|
@ -1716,71 +1476,6 @@ class OmniProImageNode(IO.ComfyNode):
|
|||
return IO.NodeOutput(torch.cat(tensors, dim=0))
|
||||
|
||||
|
||||
class KlingCameraControlT2VNode(IO.ComfyNode):
|
||||
"""
|
||||
Kling Text to Video Camera Control Node. This node is a text to video node, but it supports controlling the camera.
|
||||
Duration, mode, and model_name request fields are hard-coded because camera control is only supported in pro mode with the kling-v1-5 model at 5s duration as of 2025-05-02.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls) -> IO.Schema:
|
||||
return IO.Schema(
|
||||
node_id="KlingCameraControlT2VNode",
|
||||
display_name="Kling Text to Video (Camera Control)",
|
||||
category="partner/video/Kling",
|
||||
description="Transform text into cinematic videos with professional camera movements that simulate real-world cinematography. Control virtual camera actions including zoom, rotation, pan, tilt, and first-person view, while maintaining focus on your original text.",
|
||||
inputs=[
|
||||
IO.String.Input("prompt", multiline=True, tooltip="Positive text prompt"),
|
||||
IO.String.Input("negative_prompt", multiline=True, tooltip="Negative text prompt"),
|
||||
IO.Float.Input("cfg_scale", default=0.75, min=0.0, max=1.0),
|
||||
IO.Combo.Input(
|
||||
"aspect_ratio",
|
||||
options=KlingVideoGenAspectRatio,
|
||||
default="16:9",
|
||||
),
|
||||
IO.Custom("CAMERA_CONTROL").Input(
|
||||
"camera_control",
|
||||
tooltip="Can be created using the Kling Camera Controls node. Controls the camera movement and motion during the video generation.",
|
||||
),
|
||||
],
|
||||
outputs=[
|
||||
IO.Video.Output(),
|
||||
IO.String.Output(display_name="video_id"),
|
||||
IO.String.Output(display_name="duration"),
|
||||
],
|
||||
hidden=[
|
||||
IO.Hidden.auth_token_comfy_org,
|
||||
IO.Hidden.api_key_comfy_org,
|
||||
IO.Hidden.unique_id,
|
||||
],
|
||||
is_api_node=True,
|
||||
price_badge=IO.PriceBadge(
|
||||
expr="""{"type":"usd","usd":0.14}""",
|
||||
),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
async def execute(
|
||||
cls,
|
||||
prompt: str,
|
||||
negative_prompt: str,
|
||||
cfg_scale: float,
|
||||
aspect_ratio: str,
|
||||
camera_control: KlingCameraControl | None = None,
|
||||
) -> IO.NodeOutput:
|
||||
return await execute_text2video(
|
||||
cls,
|
||||
model_name=KlingVideoGenModelName.kling_v1,
|
||||
cfg_scale=cfg_scale,
|
||||
model_mode=KlingVideoGenMode.std,
|
||||
aspect_ratio=KlingVideoGenAspectRatio(aspect_ratio),
|
||||
duration=KlingVideoGenDuration.field_5,
|
||||
prompt=prompt,
|
||||
negative_prompt=negative_prompt,
|
||||
camera_control=camera_control,
|
||||
)
|
||||
|
||||
|
||||
class KlingImage2VideoNode(IO.ComfyNode):
|
||||
"""Kling Image to Video Node"""
|
||||
|
||||
|
|
@ -1796,11 +1491,10 @@ class KlingImage2VideoNode(IO.ComfyNode):
|
|||
IO.String.Input("negative_prompt", multiline=True, tooltip="Negative text prompt"),
|
||||
IO.Combo.Input(
|
||||
"model_name",
|
||||
options=KlingVideoGenModelName,
|
||||
default="kling-v2-master",
|
||||
options=["kling-v2-5-turbo"],
|
||||
),
|
||||
IO.Float.Input("cfg_scale", default=0.8, min=0.0, max=1.0),
|
||||
IO.Combo.Input("mode", options=KlingVideoGenMode, default=KlingVideoGenMode.std),
|
||||
IO.Combo.Input("mode", options=["pro"]),
|
||||
IO.Combo.Input(
|
||||
"aspect_ratio",
|
||||
options=KlingVideoGenAspectRatio,
|
||||
|
|
@ -1820,29 +1514,10 @@ class KlingImage2VideoNode(IO.ComfyNode):
|
|||
],
|
||||
is_api_node=True,
|
||||
price_badge=IO.PriceBadge(
|
||||
depends_on=IO.PriceBadgeDepends(widgets=["mode", "model_name", "duration"]),
|
||||
depends_on=IO.PriceBadgeDepends(widgets=["duration"]),
|
||||
expr="""
|
||||
(
|
||||
$mode := widgets.mode;
|
||||
$model := widgets.model_name;
|
||||
$dur := widgets.duration;
|
||||
$contains($model,"v2-5-turbo")
|
||||
? ($contains($dur,"10") ? {"type":"usd","usd":0.7} : {"type":"usd","usd":0.35})
|
||||
: ($contains($model,"v2-1-master") or $contains($model,"v2-master"))
|
||||
? ($contains($dur,"10") ? {"type":"usd","usd":2.8} : {"type":"usd","usd":1.4})
|
||||
: ($contains($model,"v2-1") or $contains($model,"v1-6") or $contains($model,"v1-5"))
|
||||
? (
|
||||
$contains($mode,"pro")
|
||||
? ($contains($dur,"10") ? {"type":"usd","usd":0.98} : {"type":"usd","usd":0.49})
|
||||
: ($contains($dur,"10") ? {"type":"usd","usd":0.56} : {"type":"usd","usd":0.28})
|
||||
)
|
||||
: $contains($model,"v1")
|
||||
? (
|
||||
$contains($mode,"pro")
|
||||
? ($contains($dur,"10") ? {"type":"usd","usd":0.98} : {"type":"usd","usd":0.49})
|
||||
: ($contains($dur,"10") ? {"type":"usd","usd":0.28} : {"type":"usd","usd":0.14})
|
||||
)
|
||||
: {"type":"usd","usd":0.14}
|
||||
$contains(widgets.duration,"10") ? {"type":"usd","usd":0.7} : {"type":"usd","usd":0.35}
|
||||
)
|
||||
""",
|
||||
),
|
||||
|
|
@ -1859,8 +1534,6 @@ class KlingImage2VideoNode(IO.ComfyNode):
|
|||
mode: str,
|
||||
aspect_ratio: str,
|
||||
duration: str,
|
||||
camera_control: KlingCameraControl | None = None,
|
||||
end_frame: torch.Tensor | None = None,
|
||||
) -> IO.NodeOutput:
|
||||
return await execute_image2video(
|
||||
cls,
|
||||
|
|
@ -1872,79 +1545,6 @@ class KlingImage2VideoNode(IO.ComfyNode):
|
|||
aspect_ratio=aspect_ratio,
|
||||
model_mode=mode,
|
||||
duration=duration,
|
||||
camera_control=camera_control,
|
||||
end_frame=end_frame,
|
||||
)
|
||||
|
||||
|
||||
class KlingCameraControlI2VNode(IO.ComfyNode):
|
||||
"""
|
||||
Kling Image to Video Camera Control Node. This node is a image to video node, but it supports controlling the camera.
|
||||
Duration, mode, and model_name request fields are hard-coded because camera control is only supported in pro mode with the kling-v1-5 model at 5s duration as of 2025-05-02.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls) -> IO.Schema:
|
||||
return IO.Schema(
|
||||
node_id="KlingCameraControlI2VNode",
|
||||
display_name="Kling Image to Video (Camera Control)",
|
||||
category="partner/video/Kling",
|
||||
description="Transform still images into cinematic videos with professional camera movements that simulate real-world cinematography. Control virtual camera actions including zoom, rotation, pan, tilt, and first-person view, while maintaining focus on your original image.",
|
||||
inputs=[
|
||||
IO.Image.Input(
|
||||
"start_frame",
|
||||
tooltip="Reference Image - URL or Base64 encoded string, cannot exceed 10MB, resolution not less than 300*300px, aspect ratio between 1:2.5 ~ 2.5:1. Base64 should not include data:image prefix.",
|
||||
),
|
||||
IO.String.Input("prompt", multiline=True, tooltip="Positive text prompt"),
|
||||
IO.String.Input("negative_prompt", multiline=True, tooltip="Negative text prompt"),
|
||||
IO.Float.Input("cfg_scale", default=0.75, min=0.0, max=1.0),
|
||||
IO.Combo.Input(
|
||||
"aspect_ratio",
|
||||
options=KlingVideoGenAspectRatio,
|
||||
default=KlingVideoGenAspectRatio.field_16_9,
|
||||
),
|
||||
IO.Custom("CAMERA_CONTROL").Input(
|
||||
"camera_control",
|
||||
tooltip="Can be created using the Kling Camera Controls node. Controls the camera movement and motion during the video generation.",
|
||||
),
|
||||
],
|
||||
outputs=[
|
||||
IO.Video.Output(),
|
||||
IO.String.Output(display_name="video_id"),
|
||||
IO.String.Output(display_name="duration"),
|
||||
],
|
||||
hidden=[
|
||||
IO.Hidden.auth_token_comfy_org,
|
||||
IO.Hidden.api_key_comfy_org,
|
||||
IO.Hidden.unique_id,
|
||||
],
|
||||
is_api_node=True,
|
||||
price_badge=IO.PriceBadge(
|
||||
expr="""{"type":"usd","usd":0.49}""",
|
||||
),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
async def execute(
|
||||
cls,
|
||||
start_frame: torch.Tensor,
|
||||
prompt: str,
|
||||
negative_prompt: str,
|
||||
cfg_scale: float,
|
||||
aspect_ratio: str,
|
||||
camera_control: KlingCameraControl,
|
||||
) -> IO.NodeOutput:
|
||||
return await execute_image2video(
|
||||
cls,
|
||||
model_name=KlingVideoGenModelName.kling_v1_5,
|
||||
start_frame=start_frame,
|
||||
cfg_scale=cfg_scale,
|
||||
model_mode=KlingVideoGenMode.pro,
|
||||
aspect_ratio=KlingVideoGenAspectRatio(aspect_ratio),
|
||||
duration=KlingVideoGenDuration.field_5,
|
||||
prompt=prompt,
|
||||
negative_prompt=negative_prompt,
|
||||
camera_control=camera_control,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -1977,7 +1577,7 @@ class KlingStartEndFrameNode(IO.ComfyNode):
|
|||
IO.Combo.Input(
|
||||
"mode",
|
||||
options=modes,
|
||||
default=modes[6],
|
||||
default=modes[0],
|
||||
tooltip="The configuration to use for the video generation following the format: mode / duration / model_name.",
|
||||
),
|
||||
],
|
||||
|
|
@ -1997,25 +1597,7 @@ class KlingStartEndFrameNode(IO.ComfyNode):
|
|||
expr="""
|
||||
(
|
||||
$m := widgets.mode;
|
||||
$contains($m,"v2-5-turbo")
|
||||
? ($contains($m,"10") ? {"type":"usd","usd":0.7} : {"type":"usd","usd":0.35})
|
||||
: $contains($m,"v2-1")
|
||||
? ($contains($m,"10s") ? {"type":"usd","usd":0.98} : {"type":"usd","usd":0.49})
|
||||
: $contains($m,"v2-master")
|
||||
? ($contains($m,"10s") ? {"type":"usd","usd":2.8} : {"type":"usd","usd":1.4})
|
||||
: $contains($m,"v1-6")
|
||||
? (
|
||||
$contains($m,"pro")
|
||||
? ($contains($m,"10s") ? {"type":"usd","usd":0.98} : {"type":"usd","usd":0.49})
|
||||
: ($contains($m,"10s") ? {"type":"usd","usd":0.56} : {"type":"usd","usd":0.28})
|
||||
)
|
||||
: $contains($m,"v1")
|
||||
? (
|
||||
$contains($m,"pro")
|
||||
? ($contains($m,"10s") ? {"type":"usd","usd":0.98} : {"type":"usd","usd":0.49})
|
||||
: ($contains($m,"10s") ? {"type":"usd","usd":0.28} : {"type":"usd","usd":0.14})
|
||||
)
|
||||
: {"type":"usd","usd":0.14}
|
||||
$contains($m,"10") ? {"type":"usd","usd":0.7} : {"type":"usd","usd":0.35}
|
||||
)
|
||||
""",
|
||||
),
|
||||
|
|
@ -2126,169 +1708,6 @@ class KlingVideoExtendNode(IO.ComfyNode):
|
|||
return IO.NodeOutput(await download_url_to_video_output(str(video.url)), str(video.id), str(video.duration))
|
||||
|
||||
|
||||
class KlingDualCharacterVideoEffectNode(IO.ComfyNode):
|
||||
"""Kling Dual Character Video Effect Node"""
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls) -> IO.Schema:
|
||||
return IO.Schema(
|
||||
node_id="KlingDualCharacterVideoEffectNode",
|
||||
display_name="Kling Dual Character Video Effects",
|
||||
category="partner/video/Kling",
|
||||
description="Achieve different special effects when generating a video based on the effect_scene. First image will be positioned on left side, second on right side of the composite.",
|
||||
inputs=[
|
||||
IO.Image.Input("image_left", tooltip="Left side image"),
|
||||
IO.Image.Input("image_right", tooltip="Right side image"),
|
||||
IO.Combo.Input(
|
||||
"effect_scene",
|
||||
options=[i.value for i in KlingDualCharacterEffectsScene],
|
||||
),
|
||||
IO.Combo.Input(
|
||||
"model_name",
|
||||
options=[i.value for i in KlingCharacterEffectModelName],
|
||||
default="kling-v1",
|
||||
),
|
||||
IO.Combo.Input(
|
||||
"mode",
|
||||
options=[i.value for i in KlingVideoGenMode],
|
||||
default="std",
|
||||
),
|
||||
IO.Combo.Input(
|
||||
"duration",
|
||||
options=[i.value for i in KlingVideoGenDuration],
|
||||
),
|
||||
],
|
||||
outputs=[
|
||||
IO.Video.Output(),
|
||||
IO.String.Output(display_name="duration"),
|
||||
],
|
||||
hidden=[
|
||||
IO.Hidden.auth_token_comfy_org,
|
||||
IO.Hidden.api_key_comfy_org,
|
||||
IO.Hidden.unique_id,
|
||||
],
|
||||
is_api_node=True,
|
||||
price_badge=IO.PriceBadge(
|
||||
depends_on=IO.PriceBadgeDepends(widgets=["mode", "model_name", "duration"]),
|
||||
expr="""
|
||||
(
|
||||
$mode := widgets.mode;
|
||||
$model := widgets.model_name;
|
||||
$dur := widgets.duration;
|
||||
($contains($model,"v1-6") or $contains($model,"v1-5"))
|
||||
? (
|
||||
$contains($mode,"pro")
|
||||
? ($contains($dur,"10") ? {"type":"usd","usd":0.98} : {"type":"usd","usd":0.49})
|
||||
: ($contains($dur,"10") ? {"type":"usd","usd":0.56} : {"type":"usd","usd":0.28})
|
||||
)
|
||||
: $contains($model,"v1")
|
||||
? (
|
||||
$contains($mode,"pro")
|
||||
? ($contains($dur,"10") ? {"type":"usd","usd":0.98} : {"type":"usd","usd":0.49})
|
||||
: ($contains($dur,"10") ? {"type":"usd","usd":0.28} : {"type":"usd","usd":0.14})
|
||||
)
|
||||
: {"type":"usd","usd":0.14}
|
||||
)
|
||||
""",
|
||||
),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
async def execute(
|
||||
cls,
|
||||
image_left: torch.Tensor,
|
||||
image_right: torch.Tensor,
|
||||
effect_scene: KlingDualCharacterEffectsScene,
|
||||
model_name: KlingCharacterEffectModelName,
|
||||
mode: KlingVideoGenMode,
|
||||
duration: KlingVideoGenDuration,
|
||||
) -> IO.NodeOutput:
|
||||
video, _, duration = await execute_video_effect(
|
||||
cls,
|
||||
dual_character=True,
|
||||
effect_scene=effect_scene,
|
||||
model_name=model_name,
|
||||
model_mode=mode,
|
||||
duration=duration,
|
||||
image_1=image_left,
|
||||
image_2=image_right,
|
||||
)
|
||||
return IO.NodeOutput(video, duration)
|
||||
|
||||
|
||||
class KlingSingleImageVideoEffectNode(IO.ComfyNode):
|
||||
"""Kling Single Image Video Effect Node"""
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls) -> IO.Schema:
|
||||
return IO.Schema(
|
||||
node_id="KlingSingleImageVideoEffectNode",
|
||||
display_name="Kling Video Effects",
|
||||
category="partner/video/Kling",
|
||||
description="Achieve different special effects when generating a video based on the effect_scene.",
|
||||
inputs=[
|
||||
IO.Image.Input(
|
||||
"image",
|
||||
tooltip=" Reference Image. URL or Base64 encoded string (without data:image prefix). File size cannot exceed 10MB, resolution not less than 300*300px, aspect ratio between 1:2.5 ~ 2.5:1",
|
||||
),
|
||||
IO.Combo.Input(
|
||||
"effect_scene",
|
||||
options=[i.value for i in KlingSingleImageEffectsScene],
|
||||
),
|
||||
IO.Combo.Input(
|
||||
"model_name",
|
||||
options=[i.value for i in KlingSingleImageEffectModelName],
|
||||
),
|
||||
IO.Combo.Input(
|
||||
"duration",
|
||||
options=[i.value for i in KlingVideoGenDuration],
|
||||
),
|
||||
],
|
||||
outputs=[
|
||||
IO.Video.Output(),
|
||||
IO.String.Output(display_name="video_id"),
|
||||
IO.String.Output(display_name="duration"),
|
||||
],
|
||||
hidden=[
|
||||
IO.Hidden.auth_token_comfy_org,
|
||||
IO.Hidden.api_key_comfy_org,
|
||||
IO.Hidden.unique_id,
|
||||
],
|
||||
is_api_node=True,
|
||||
price_badge=IO.PriceBadge(
|
||||
depends_on=IO.PriceBadgeDepends(widgets=["effect_scene"]),
|
||||
expr="""
|
||||
(
|
||||
($contains(widgets.effect_scene,"dizzydizzy") or $contains(widgets.effect_scene,"bloombloom"))
|
||||
? {"type":"usd","usd":0.49}
|
||||
: {"type":"usd","usd":0.28}
|
||||
)
|
||||
""",
|
||||
),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
async def execute(
|
||||
cls,
|
||||
image: torch.Tensor,
|
||||
effect_scene: KlingSingleImageEffectsScene,
|
||||
model_name: KlingSingleImageEffectModelName,
|
||||
duration: KlingVideoGenDuration,
|
||||
) -> IO.NodeOutput:
|
||||
return IO.NodeOutput(
|
||||
*(
|
||||
await execute_video_effect(
|
||||
cls,
|
||||
dual_character=False,
|
||||
effect_scene=effect_scene,
|
||||
model_name=model_name,
|
||||
duration=duration,
|
||||
image_1=image,
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
class KlingLipSyncAudioToVideoNode(IO.ComfyNode):
|
||||
"""Kling Lip Sync Audio to Video Node. Syncs mouth movements in a video file to the audio content of an audio file."""
|
||||
|
||||
|
|
@ -2409,73 +1828,6 @@ class KlingLipSyncTextToVideoNode(IO.ComfyNode):
|
|||
)
|
||||
|
||||
|
||||
class KlingVirtualTryOnNode(IO.ComfyNode):
|
||||
"""Kling Virtual Try On Node."""
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls) -> IO.Schema:
|
||||
return IO.Schema(
|
||||
node_id="KlingVirtualTryOnNode",
|
||||
display_name="Kling Virtual Try On",
|
||||
category="partner/image/Kling",
|
||||
description="Kling Virtual Try On Node. Input a human image and a cloth image to try on the cloth on the human. You can merge multiple clothing item pictures into one image with a white background.",
|
||||
inputs=[
|
||||
IO.Image.Input("human_image"),
|
||||
IO.Image.Input("cloth_image"),
|
||||
IO.Combo.Input(
|
||||
"model_name",
|
||||
options=[i.value for i in KlingVirtualTryOnModelName],
|
||||
default="kolors-virtual-try-on-v1",
|
||||
),
|
||||
],
|
||||
outputs=[
|
||||
IO.Image.Output(),
|
||||
],
|
||||
hidden=[
|
||||
IO.Hidden.auth_token_comfy_org,
|
||||
IO.Hidden.api_key_comfy_org,
|
||||
IO.Hidden.unique_id,
|
||||
],
|
||||
is_api_node=True,
|
||||
price_badge=IO.PriceBadge(
|
||||
expr="""{"type":"usd","usd":0.7}""",
|
||||
),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
async def execute(
|
||||
cls,
|
||||
human_image: torch.Tensor,
|
||||
cloth_image: torch.Tensor,
|
||||
model_name: KlingVirtualTryOnModelName,
|
||||
) -> IO.NodeOutput:
|
||||
task_creation_response = await sync_op(
|
||||
cls,
|
||||
ApiEndpoint(path=PATH_VIRTUAL_TRY_ON, method="POST"),
|
||||
response_model=KlingVirtualTryOnResponse,
|
||||
data=KlingVirtualTryOnRequest(
|
||||
human_image=tensor_to_base64_string(human_image),
|
||||
cloth_image=tensor_to_base64_string(cloth_image),
|
||||
model_name=model_name,
|
||||
),
|
||||
)
|
||||
|
||||
validate_task_creation_response(task_creation_response)
|
||||
task_id = task_creation_response.data.task_id
|
||||
|
||||
final_response = await poll_op(
|
||||
cls,
|
||||
ApiEndpoint(path=f"{PATH_VIRTUAL_TRY_ON}/{task_id}"),
|
||||
response_model=KlingVirtualTryOnResponse,
|
||||
estimated_duration=AVERAGE_DURATION_VIRTUAL_TRY_ON,
|
||||
status_extractor=lambda r: (r.data.task_status.value if r.data and r.data.task_status else None),
|
||||
)
|
||||
validate_image_result_response(final_response)
|
||||
|
||||
images = get_images_from_response(final_response)
|
||||
return IO.NodeOutput(await image_result_to_node_output(images))
|
||||
|
||||
|
||||
class KlingImageGenerationNode(IO.ComfyNode):
|
||||
"""Kling Image Generation Node. Generate an image from a text prompt with an optional reference image."""
|
||||
|
||||
|
|
@ -2514,7 +1866,7 @@ class KlingImageGenerationNode(IO.ComfyNode):
|
|||
tooltip="Subject reference similarity",
|
||||
advanced=True,
|
||||
),
|
||||
IO.Combo.Input("model_name", options=["kling-v3", "kling-v2", "kling-v1-5"]),
|
||||
IO.Combo.Input("model_name", options=["kling-v3", "kling-v2"]),
|
||||
IO.Combo.Input(
|
||||
"aspect_ratio",
|
||||
options=[i.value for i in KlingImageGenAspectRatio],
|
||||
|
|
@ -2550,14 +1902,10 @@ class KlingImageGenerationNode(IO.ComfyNode):
|
|||
],
|
||||
is_api_node=True,
|
||||
price_badge=IO.PriceBadge(
|
||||
depends_on=IO.PriceBadgeDepends(widgets=["model_name", "n"], inputs=["image"]),
|
||||
depends_on=IO.PriceBadgeDepends(widgets=["model_name", "n"]),
|
||||
expr="""
|
||||
(
|
||||
$m := widgets.model_name;
|
||||
$base :=
|
||||
$contains($m,"kling-v1-5")
|
||||
? (inputs.image.connected ? 0.028 : 0.014)
|
||||
: $contains($m,"kling-v3") ? 0.028 : 0.014;
|
||||
$base := $contains(widgets.model_name,"kling-v3") ? 0.028 : 0.014;
|
||||
{"type":"usd","usd": $base * widgets.n}
|
||||
)
|
||||
""",
|
||||
|
|
@ -3394,19 +2742,13 @@ class KlingExtension(ComfyExtension):
|
|||
@override
|
||||
async def get_node_list(self) -> list[type[IO.ComfyNode]]:
|
||||
return [
|
||||
KlingCameraControls,
|
||||
KlingTextToVideoNode,
|
||||
KlingImage2VideoNode,
|
||||
KlingCameraControlI2VNode,
|
||||
KlingCameraControlT2VNode,
|
||||
KlingStartEndFrameNode,
|
||||
KlingVideoExtendNode,
|
||||
KlingLipSyncAudioToVideoNode,
|
||||
KlingLipSyncTextToVideoNode,
|
||||
KlingVirtualTryOnNode,
|
||||
KlingImageGenerationNode,
|
||||
KlingSingleImageVideoEffectNode,
|
||||
KlingDualCharacterVideoEffectNode,
|
||||
OmniProTextToVideoNode,
|
||||
OmniProFirstLastFrameNode,
|
||||
OmniProImageToVideoNode,
|
||||
|
|
|
|||
|
|
@ -494,9 +494,9 @@ class OpenAIGPTImage1(IO.ComfyNode):
|
|||
"high": [0.133, 0.22]
|
||||
},
|
||||
"gpt-image-2": {
|
||||
"low": [0.0048, 0.019],
|
||||
"medium": [0.041, 0.168],
|
||||
"high": [0.165, 0.67]
|
||||
"low": [0.0058, 0.0228],
|
||||
"medium": [0.0492, 0.2016],
|
||||
"high": [0.198, 0.804]
|
||||
}
|
||||
};
|
||||
$range := $lookup($lookup($ranges, widgets.model), widgets.quality);
|
||||
|
|
@ -792,9 +792,9 @@ class OpenAIGPTImageNodeV2(IO.ComfyNode):
|
|||
"high": [0.133, 0.22]
|
||||
},
|
||||
"gpt-image-2": {
|
||||
"low": [0.0048, 0.019],
|
||||
"medium": [0.041, 0.168],
|
||||
"high": [0.165, 0.67]
|
||||
"low": [0.0058, 0.0228],
|
||||
"medium": [0.0492, 0.2016],
|
||||
"high": [0.198, 0.804]
|
||||
}
|
||||
};
|
||||
$range := $lookup($lookup($ranges, widgets.model), $lookup(widgets, "model.quality"));
|
||||
|
|
|
|||
|
|
@ -137,6 +137,7 @@ class ReveImageCreateNode(IO.ComfyNode):
|
|||
IO.Hidden.unique_id,
|
||||
],
|
||||
is_api_node=True,
|
||||
is_deprecated=True,
|
||||
price_badge=IO.PriceBadge(
|
||||
depends_on=IO.PriceBadgeDepends(
|
||||
widgets=["upscale", "upscale.upscale_factor"],
|
||||
|
|
@ -228,6 +229,7 @@ class ReveImageEditNode(IO.ComfyNode):
|
|||
IO.Hidden.unique_id,
|
||||
],
|
||||
is_api_node=True,
|
||||
is_deprecated=True,
|
||||
price_badge=IO.PriceBadge(
|
||||
depends_on=IO.PriceBadgeDepends(
|
||||
widgets=["model", "upscale", "upscale.upscale_factor"],
|
||||
|
|
@ -337,6 +339,7 @@ class ReveImageRemixNode(IO.ComfyNode):
|
|||
IO.Hidden.unique_id,
|
||||
],
|
||||
is_api_node=True,
|
||||
is_deprecated=True,
|
||||
price_badge=IO.PriceBadge(
|
||||
depends_on=IO.PriceBadgeDepends(
|
||||
widgets=["model", "upscale", "upscale.upscale_factor"],
|
||||
|
|
|
|||
|
|
@ -12,6 +12,7 @@ from comfy_api_nodes.apis.topaz import (
|
|||
ImageAsyncTaskResponse,
|
||||
ImageDownloadResponse,
|
||||
ImageEnhanceRequest,
|
||||
ImageEnhanceRequestV2,
|
||||
ImageStatusResponse,
|
||||
OutputInformationVideo,
|
||||
Resolution,
|
||||
|
|
@ -51,7 +52,7 @@ class TopazImageEnhance(IO.ComfyNode):
|
|||
def define_schema(cls):
|
||||
return IO.Schema(
|
||||
node_id="TopazImageEnhance",
|
||||
display_name="Topaz Image Enhance",
|
||||
display_name="Topaz Image Enhance (Legacy)",
|
||||
category="partner/image/Topaz",
|
||||
description="Industry-standard upscaling and image enhancement.",
|
||||
inputs=[
|
||||
|
|
@ -162,6 +163,7 @@ class TopazImageEnhance(IO.ComfyNode):
|
|||
IO.Hidden.unique_id,
|
||||
],
|
||||
is_api_node=True,
|
||||
is_deprecated=True,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
|
|
@ -229,6 +231,355 @@ class TopazImageEnhance(IO.ComfyNode):
|
|||
return IO.NodeOutput(await download_url_to_image_tensor(results.download_url))
|
||||
|
||||
|
||||
class TopazImageEnhanceV2(IO.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return IO.Schema(
|
||||
node_id="TopazImageEnhanceV2",
|
||||
display_name="Topaz Image Enhance",
|
||||
category="partner/image/Topaz",
|
||||
description="Industry-standard upscaling and image enhancement.",
|
||||
inputs=[
|
||||
IO.Image.Input("image"),
|
||||
IO.DynamicCombo.Input(
|
||||
"model",
|
||||
options=[
|
||||
IO.DynamicCombo.Option(
|
||||
"Reimagine",
|
||||
[
|
||||
IO.String.Input(
|
||||
"prompt",
|
||||
multiline=True,
|
||||
default="",
|
||||
tooltip="Optional text prompt for creative upscaling guidance.",
|
||||
),
|
||||
IO.Int.Input(
|
||||
"creativity",
|
||||
default=3,
|
||||
min=1,
|
||||
max=9,
|
||||
step=1,
|
||||
display_mode=IO.NumberDisplay.slider,
|
||||
),
|
||||
IO.Combo.Input(
|
||||
"subject_detection",
|
||||
options=["All", "Foreground", "Background"],
|
||||
advanced=True,
|
||||
),
|
||||
IO.Boolean.Input(
|
||||
"face_enhancement",
|
||||
default=True,
|
||||
tooltip="Enhance faces (if present) during processing.",
|
||||
advanced=True,
|
||||
),
|
||||
IO.Float.Input(
|
||||
"face_enhancement_creativity",
|
||||
default=0.0,
|
||||
min=0.0,
|
||||
max=1.0,
|
||||
step=0.01,
|
||||
display_mode=IO.NumberDisplay.number,
|
||||
tooltip="Set the creativity level for face enhancement.",
|
||||
advanced=True,
|
||||
),
|
||||
IO.Float.Input(
|
||||
"face_enhancement_strength",
|
||||
default=1.0,
|
||||
min=0.0,
|
||||
max=1.0,
|
||||
step=0.01,
|
||||
display_mode=IO.NumberDisplay.number,
|
||||
tooltip="Controls how sharp enhanced faces are relative to the background.",
|
||||
advanced=True,
|
||||
),
|
||||
IO.Boolean.Input(
|
||||
"face_preservation",
|
||||
default=True,
|
||||
tooltip="Preserve subjects' facial identity.",
|
||||
advanced=True,
|
||||
),
|
||||
IO.Boolean.Input(
|
||||
"color_preservation",
|
||||
default=True,
|
||||
tooltip="Preserve the original colors.",
|
||||
advanced=True,
|
||||
),
|
||||
IO.Boolean.Input(
|
||||
"crop_to_fill",
|
||||
default=False,
|
||||
tooltip="By default, the image is letterboxed when the output aspect "
|
||||
"ratio differs. Enable to crop the image to fill the output dimensions.",
|
||||
advanced=True,
|
||||
),
|
||||
],
|
||||
),
|
||||
IO.DynamicCombo.Option(
|
||||
"Bloom 2",
|
||||
[
|
||||
IO.String.Input(
|
||||
"prompt",
|
||||
multiline=True,
|
||||
default="",
|
||||
tooltip="Optional text prompt for generation. "
|
||||
"Leave empty to auto-generate a prompt from the input image.",
|
||||
),
|
||||
IO.Int.Input(
|
||||
"creativity",
|
||||
default=3,
|
||||
min=1,
|
||||
max=9,
|
||||
step=1,
|
||||
display_mode=IO.NumberDisplay.slider,
|
||||
tooltip="1 is restrained enhancement, 9 is pronounced reinterpretation "
|
||||
"with newly generated detail.",
|
||||
),
|
||||
IO.Int.Input(
|
||||
"seed",
|
||||
default=2,
|
||||
min=1,
|
||||
max=2000,
|
||||
control_after_generate=True,
|
||||
tooltip="Seed for reproducible generation.",
|
||||
),
|
||||
IO.Boolean.Input(
|
||||
"color_preservation",
|
||||
default=True,
|
||||
tooltip="Preserve the original colors.",
|
||||
advanced=True,
|
||||
),
|
||||
IO.Boolean.Input(
|
||||
"grain",
|
||||
default=False,
|
||||
tooltip="Add grain to the output image.",
|
||||
advanced=True,
|
||||
),
|
||||
IO.Combo.Input(
|
||||
"grain_model",
|
||||
options=["silver", "gaussian", "grey"],
|
||||
tooltip="Is ignored if grain is disabled.",
|
||||
advanced=True,
|
||||
),
|
||||
IO.Float.Input(
|
||||
"grain_strength",
|
||||
default=0.5,
|
||||
min=0.0,
|
||||
max=1.0,
|
||||
step=0.01,
|
||||
display_mode=IO.NumberDisplay.number,
|
||||
tooltip="Strength of the grain effect. Is ignored if grain is disabled.",
|
||||
advanced=True,
|
||||
),
|
||||
IO.Float.Input(
|
||||
"grain_size",
|
||||
default=1.0,
|
||||
min=1.0,
|
||||
max=5.0,
|
||||
step=0.1,
|
||||
display_mode=IO.NumberDisplay.number,
|
||||
tooltip="Size of the grain particles. Is ignored if grain is disabled.",
|
||||
advanced=True,
|
||||
),
|
||||
IO.Float.Input(
|
||||
"grain_density",
|
||||
default=0.5,
|
||||
min=0.0,
|
||||
max=1.0,
|
||||
step=0.01,
|
||||
display_mode=IO.NumberDisplay.number,
|
||||
tooltip="Intensity of the grain effect. Is ignored if grain is disabled.",
|
||||
advanced=True,
|
||||
),
|
||||
],
|
||||
),
|
||||
IO.DynamicCombo.Option(
|
||||
"Wonder 3.5",
|
||||
[
|
||||
IO.Combo.Input(
|
||||
"enhancement_strength",
|
||||
options=["low", "medium", "high"],
|
||||
default="high",
|
||||
tooltip="Enhancement level for varying input conditions.",
|
||||
),
|
||||
IO.Boolean.Input(
|
||||
"grain",
|
||||
default=False,
|
||||
tooltip="Add grain to the output image.",
|
||||
advanced=True,
|
||||
),
|
||||
IO.Combo.Input(
|
||||
"grain_model",
|
||||
options=["silver", "gaussian", "grey"],
|
||||
tooltip="Is ignored if grain is disabled.",
|
||||
advanced=True,
|
||||
),
|
||||
IO.Float.Input(
|
||||
"grain_strength",
|
||||
default=0.5,
|
||||
min=0.0,
|
||||
max=1.0,
|
||||
step=0.01,
|
||||
display_mode=IO.NumberDisplay.number,
|
||||
tooltip="Strength of the grain effect. Is ignored if grain is disabled.",
|
||||
advanced=True,
|
||||
),
|
||||
IO.Float.Input(
|
||||
"grain_size",
|
||||
default=1.0,
|
||||
min=1.0,
|
||||
max=5.0,
|
||||
step=0.1,
|
||||
display_mode=IO.NumberDisplay.number,
|
||||
tooltip="Size of the grain particles. Is ignored if grain is disabled.",
|
||||
advanced=True,
|
||||
),
|
||||
IO.Float.Input(
|
||||
"grain_density",
|
||||
default=0.5,
|
||||
min=0.0,
|
||||
max=1.0,
|
||||
step=0.01,
|
||||
display_mode=IO.NumberDisplay.number,
|
||||
tooltip="Intensity of the grain effect. Is ignored if grain is disabled.",
|
||||
advanced=True,
|
||||
),
|
||||
],
|
||||
),
|
||||
],
|
||||
),
|
||||
IO.Int.Input(
|
||||
"output_width",
|
||||
default=0,
|
||||
min=0,
|
||||
max=32000,
|
||||
step=1,
|
||||
display_mode=IO.NumberDisplay.number,
|
||||
optional=True,
|
||||
tooltip="Zero value means to calculate automatically (usually it will be original size "
|
||||
"or scaled proportionally to output_height if specified). "
|
||||
"Wonder 3.5 supports upscale factors from 1x to 6x only. "
|
||||
"Bloom 2 and Wonder 3.5 preserve the input aspect ratio and treat the "
|
||||
"requested size as a target.",
|
||||
advanced=True,
|
||||
),
|
||||
IO.Int.Input(
|
||||
"output_height",
|
||||
default=0,
|
||||
min=0,
|
||||
max=32000,
|
||||
step=1,
|
||||
display_mode=IO.NumberDisplay.number,
|
||||
optional=True,
|
||||
tooltip="Zero value means to output in the same height as original or scaled "
|
||||
"proportionally to output_width if specified. "
|
||||
"Wonder 3.5 supports upscale factors from 1x to 6x only. "
|
||||
"Bloom 2 and Wonder 3.5 preserve the input aspect ratio and treat the "
|
||||
"requested size as a target.",
|
||||
advanced=True,
|
||||
),
|
||||
],
|
||||
outputs=[
|
||||
IO.Image.Output(),
|
||||
],
|
||||
hidden=[
|
||||
IO.Hidden.auth_token_comfy_org,
|
||||
IO.Hidden.api_key_comfy_org,
|
||||
IO.Hidden.unique_id,
|
||||
],
|
||||
is_api_node=True,
|
||||
price_badge=IO.PriceBadge(
|
||||
depends_on=IO.PriceBadgeDepends(widgets=["model"]),
|
||||
expr="""
|
||||
(
|
||||
$usdPer8Mp := $lookup(
|
||||
{"reimagine": 0.32, "bloom 2": 0.4576, "wonder 3.5": 0.1144},
|
||||
$lookup(widgets, "model")
|
||||
);
|
||||
{"type":"usd","usd": $usdPer8Mp, "format": {"suffix": "/8MP", "approximate": true}}
|
||||
)
|
||||
""",
|
||||
),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
async def execute(
|
||||
cls,
|
||||
image: Input.Image,
|
||||
model: dict,
|
||||
output_width: int = 0,
|
||||
output_height: int = 0,
|
||||
) -> IO.NodeOutput:
|
||||
if get_number_of_images(image) != 1:
|
||||
raise ValueError("Only one input image is supported.")
|
||||
model_choice = model["model"]
|
||||
download_url = await upload_images_to_comfyapi(
|
||||
cls, image, max_images=1, mime_type="image/png", total_pixels=4096 * 4096
|
||||
)
|
||||
request = ImageEnhanceRequestV2(
|
||||
model=model_choice,
|
||||
source_url=download_url[0],
|
||||
output_width=output_width if output_width else None,
|
||||
output_height=output_height if output_height else None,
|
||||
)
|
||||
if model_choice == "Reimagine":
|
||||
request.prompt = model["prompt"]
|
||||
request.creativity = model["creativity"]
|
||||
request.subject_detection = model["subject_detection"]
|
||||
request.face_enhancement = model["face_enhancement"]
|
||||
request.face_enhancement_creativity = model["face_enhancement_creativity"]
|
||||
request.face_enhancement_strength = model["face_enhancement_strength"]
|
||||
request.face_preservation = str(model["face_preservation"]).lower()
|
||||
request.color_preservation = str(model["color_preservation"]).lower()
|
||||
request.crop_to_fill = model["crop_to_fill"]
|
||||
elif model_choice == "Bloom 2":
|
||||
prompt = model["prompt"].strip()
|
||||
if prompt:
|
||||
request.prompt = prompt
|
||||
request.autoprompt = "false"
|
||||
else:
|
||||
request.autoprompt = "true"
|
||||
request.creativity = model["creativity"]
|
||||
request.seed = model["seed"]
|
||||
request.color_preservation = str(model["color_preservation"]).lower()
|
||||
if model["grain"]:
|
||||
request.grain = "true"
|
||||
request.grain_model = model["grain_model"]
|
||||
request.grain_strength = model["grain_strength"]
|
||||
request.grain_size = model["grain_size"]
|
||||
request.grain_density = model["grain_density"]
|
||||
else:
|
||||
request.enhancement_strength = model["enhancement_strength"]
|
||||
if model["grain"]:
|
||||
request.grain = "true"
|
||||
request.grain_model = model["grain_model"]
|
||||
request.grain_strength = model["grain_strength"]
|
||||
request.grain_size = model["grain_size"]
|
||||
request.grain_density = model["grain_density"]
|
||||
initial_response = await sync_op(
|
||||
cls,
|
||||
ApiEndpoint(path="/proxy/topaz/image/v1/enhance-gen/async", method="POST"),
|
||||
response_model=ImageAsyncTaskResponse,
|
||||
data=request,
|
||||
content_type="multipart/form-data",
|
||||
)
|
||||
await poll_op(
|
||||
cls,
|
||||
poll_endpoint=ApiEndpoint(path=f"/proxy/topaz/image/v1/status/{initial_response.process_id}"),
|
||||
response_model=ImageStatusResponse,
|
||||
status_extractor=lambda x: x.status,
|
||||
progress_extractor=lambda x: getattr(x, "progress", 0),
|
||||
price_extractor=lambda x: x.credits * (0.08 if model_choice == "Reimagine" else 0.1144),
|
||||
poll_interval=8.0,
|
||||
estimated_duration=60,
|
||||
)
|
||||
results = await sync_op(
|
||||
cls,
|
||||
ApiEndpoint(path=f"/proxy/topaz/image/v1/download/{initial_response.process_id}"),
|
||||
response_model=ImageDownloadResponse,
|
||||
monitor_progress=False,
|
||||
)
|
||||
return IO.NodeOutput(await download_url_to_image_tensor(results.download_url))
|
||||
|
||||
|
||||
class TopazVideoEnhance(IO.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
|
|
@ -818,6 +1169,7 @@ class TopazExtension(ComfyExtension):
|
|||
async def get_node_list(self) -> list[type[IO.ComfyNode]]:
|
||||
return [
|
||||
TopazImageEnhance,
|
||||
TopazImageEnhanceV2,
|
||||
TopazVideoEnhance,
|
||||
TopazVideoEnhanceV2,
|
||||
]
|
||||
|
|
|
|||
|
|
@ -197,6 +197,7 @@ def normalize_queue_item(item: tuple, status: str) -> dict:
|
|||
'priority': priority,
|
||||
'create_time': create_time,
|
||||
'outputs_count': 0,
|
||||
'previewable_outputs_count': 0,
|
||||
'workflow_id': workflow_id,
|
||||
})
|
||||
|
||||
|
|
@ -215,6 +216,7 @@ def normalize_history_item(prompt_id: str, history_item: dict, include_outputs:
|
|||
|
||||
outputs = history_item.get('outputs', {})
|
||||
outputs_count, preview_output = get_outputs_summary(outputs)
|
||||
previewable_outputs_count = count_previewable_outputs(outputs)
|
||||
|
||||
execution_error = None
|
||||
execution_start_time = None
|
||||
|
|
@ -251,6 +253,7 @@ def normalize_history_item(prompt_id: str, history_item: dict, include_outputs:
|
|||
'execution_end_time': execution_end_time,
|
||||
'execution_error': execution_error,
|
||||
'outputs_count': outputs_count,
|
||||
'previewable_outputs_count': previewable_outputs_count,
|
||||
'preview_output': preview_output,
|
||||
'workflow_id': workflow_id,
|
||||
})
|
||||
|
|
@ -345,6 +348,33 @@ def get_outputs_summary(outputs: dict) -> tuple[int, Optional[dict]]:
|
|||
return count, preview_output or fallback_preview or text_file_fallback or text_fallback
|
||||
|
||||
|
||||
def count_previewable_outputs(outputs: dict) -> int:
|
||||
"""
|
||||
Count only outputs that would actually render in the expanded asset view,
|
||||
i.e. items is_previewable() accepts (image/video/audio/3D/text). Kept
|
||||
separate from get_outputs_summary()'s outputs_count, which counts every
|
||||
output item regardless of media type, so a job with a non-previewable
|
||||
saved file alongside real media (e.g. SaveLatent's .latent output next to
|
||||
a SaveImage output) doesn't inflate the Media Assets badge beyond what
|
||||
the expanded view shows.
|
||||
"""
|
||||
count = 0
|
||||
for node_outputs in outputs.values():
|
||||
if not isinstance(node_outputs, dict):
|
||||
continue
|
||||
for media_type, items in node_outputs.items():
|
||||
if media_type == 'animated' or not isinstance(items, list):
|
||||
continue
|
||||
for item in items:
|
||||
if not isinstance(item, dict):
|
||||
item = normalize_output_item(item)
|
||||
if item is None:
|
||||
continue
|
||||
if is_previewable(media_type, item):
|
||||
count += 1
|
||||
return count
|
||||
|
||||
|
||||
def apply_sorting(jobs: list[dict], sort_by: str, sort_order: str) -> list[dict]:
|
||||
"""Sort jobs list by specified field and order."""
|
||||
reverse = (sort_order == 'desc')
|
||||
|
|
|
|||
|
|
@ -0,0 +1,331 @@
|
|||
import math
|
||||
from typing import NamedTuple, Optional, Union
|
||||
|
||||
import numpy as np
|
||||
|
||||
EPSILON = 1e-6
|
||||
|
||||
LUM_R = 0.2224884
|
||||
LUM_G = 0.71690369
|
||||
LUM_B = 0.06060791
|
||||
|
||||
ArrayLike = Union[np.ndarray, float]
|
||||
|
||||
|
||||
def srgb_to_linear(c: ArrayLike) -> np.ndarray:
|
||||
c = np.asarray(c, dtype=np.float32)
|
||||
high = ((np.maximum(c, 0.0) + 0.055) / 1.055) ** 2.4
|
||||
return np.where(c <= 0.04045, c / 12.92, high).astype(np.float32)
|
||||
|
||||
|
||||
def linear_to_srgb(c: ArrayLike) -> np.ndarray:
|
||||
c = np.asarray(c, dtype=np.float32)
|
||||
high = 1.055 * np.maximum(c, 0.0) ** (1.0 / 2.4) - 0.055
|
||||
return np.where(c <= 0.0031308, 12.92 * c, high).astype(np.float32)
|
||||
|
||||
|
||||
def luminance(rgb: np.ndarray) -> np.ndarray:
|
||||
return rgb[..., 0] * LUM_R + rgb[..., 1] * LUM_G + rgb[..., 2] * LUM_B
|
||||
|
||||
|
||||
def safe_div(a: ArrayLike, b: ArrayLike) -> np.ndarray:
|
||||
a, b = np.broadcast_arrays(
|
||||
np.asarray(a, dtype=np.float32), np.asarray(b, dtype=np.float32)
|
||||
)
|
||||
out = np.zeros(b.shape, dtype=np.float32)
|
||||
np.divide(a, b, out=out, where=np.abs(b) >= EPSILON)
|
||||
return out
|
||||
|
||||
|
||||
CHANNEL_BLEND = {
|
||||
"normal": lambda i, l: l,
|
||||
"multiply": lambda i, l: i * l,
|
||||
"screen": lambda i, l: 1 - (1 - i) * (1 - l),
|
||||
"overlay": lambda i, l: np.where(i < 0.5, 2 * i * l, 1 - 2 * (1 - l) * (1 - i)),
|
||||
"darken": lambda i, l: np.minimum(i, l),
|
||||
"lighten": lambda i, l: np.maximum(i, l),
|
||||
"color-dodge": lambda i, l: np.where(
|
||||
i <= 0,
|
||||
0.0,
|
||||
np.where(1 - l <= EPSILON, 1.0, np.minimum(safe_div(i, 1 - l), 1.0)),
|
||||
),
|
||||
"color-burn": lambda i, l: np.where(
|
||||
i >= 1,
|
||||
1.0,
|
||||
np.where(l <= EPSILON, 0.0, 1 - np.minimum(safe_div(1 - i, l), 1.0)),
|
||||
),
|
||||
"hard-light": lambda i, l: np.where(
|
||||
l > 0.5,
|
||||
np.minimum(1 - (1 - i) * (1 - (l - 0.5) * 2), 1),
|
||||
np.minimum(i * (l * 2), 1),
|
||||
),
|
||||
"soft-light": lambda i, l: (1 - i) * (i * l) + i * (1 - (1 - i) * (1 - l)),
|
||||
"difference": lambda i, l: np.abs(i - l),
|
||||
"exclusion": lambda i, l: 0.5 - 2 * (i - 0.5) * (l - 0.5),
|
||||
"linear-dodge": lambda i, l: i + l,
|
||||
"linear-burn": lambda i, l: i + l - 1,
|
||||
"vivid-light": lambda i, l: np.where(
|
||||
l <= 0.5,
|
||||
np.where(
|
||||
i >= 1,
|
||||
1.0,
|
||||
np.where(
|
||||
2 * l <= EPSILON,
|
||||
0.0,
|
||||
np.maximum(1 - safe_div(1 - i, 2 * l), 0.0),
|
||||
),
|
||||
),
|
||||
np.where(
|
||||
i <= 0,
|
||||
0.0,
|
||||
np.where(
|
||||
2 * (1 - l) <= EPSILON,
|
||||
1.0,
|
||||
np.minimum(safe_div(i, 2 * (1 - l)), 1.0),
|
||||
),
|
||||
),
|
||||
),
|
||||
"pin-light": lambda i, l: np.where(
|
||||
l > 0.5, np.maximum(i, 2 * (l - 0.5)), np.minimum(i, 2 * l)
|
||||
),
|
||||
"linear-light": lambda i, l: i + 2 * l - 1,
|
||||
"hard-mix": lambda i, l: np.where(i + l < 1, 0.0, 1.0),
|
||||
"subtract": lambda i, l: np.maximum(i - l, 0),
|
||||
"divide": lambda i, l: np.clip(i / np.maximum(l, EPSILON), 0, 1),
|
||||
"grain-extract": lambda i, l: i - l + 0.5,
|
||||
"grain-merge": lambda i, l: i + l - 0.5,
|
||||
}
|
||||
|
||||
|
||||
def _blend_hue(i: np.ndarray, l: np.ndarray) -> np.ndarray:
|
||||
src_min = l.min(axis=-1)
|
||||
src_max = l.max(axis=-1)
|
||||
src_delta = src_max - src_min
|
||||
achromatic = src_delta <= EPSILON
|
||||
dest_max = i.max(axis=-1)
|
||||
dest_delta = dest_max - i.min(axis=-1)
|
||||
dest_s = np.where(dest_max != 0, dest_delta / np.where(dest_max != 0, dest_max, 1), 0)
|
||||
ratio = np.where(
|
||||
achromatic, 0, dest_s * dest_max / np.where(achromatic, 1, src_delta)
|
||||
)
|
||||
offset = dest_max - src_max * ratio
|
||||
return np.where(achromatic[..., None], i, l * ratio[..., None] + offset[..., None])
|
||||
|
||||
|
||||
def _blend_saturation(i: np.ndarray, l: np.ndarray) -> np.ndarray:
|
||||
dest_max = i.max(axis=-1)
|
||||
dest_delta = dest_max - i.min(axis=-1)
|
||||
flat = dest_delta <= EPSILON
|
||||
src_max = l.max(axis=-1)
|
||||
src_delta = src_max - l.min(axis=-1)
|
||||
src_s = np.where(src_max != 0, src_delta / np.where(src_max != 0, src_max, 1), 0)
|
||||
ratio = np.where(flat, 0, src_s * dest_max / np.where(flat, 1, dest_delta))
|
||||
offset = (1 - ratio) * dest_max
|
||||
return np.where(
|
||||
flat[..., None],
|
||||
np.broadcast_to(dest_max[..., None], i.shape),
|
||||
i * ratio[..., None] + offset[..., None],
|
||||
)
|
||||
|
||||
|
||||
def _blend_color(i: np.ndarray, l: np.ndarray) -> np.ndarray:
|
||||
dest_l = (i.min(axis=-1) + i.max(axis=-1)) / 2
|
||||
src_l = (l.min(axis=-1) + l.max(axis=-1)) / 2
|
||||
gray = (np.abs(src_l) <= EPSILON) | (np.abs(1 - src_l) <= EPSILON)
|
||||
dest_high = dest_l > 0.5
|
||||
src_high = src_l > 0.5
|
||||
dl = np.minimum(dest_l, 1 - dest_l)
|
||||
sl = np.minimum(src_l, 1 - src_l)
|
||||
ratio = dl / np.where(gray, 1, sl)
|
||||
offset = np.where(dest_high, 1 - 2 * dl, 0) + np.where(src_high, 2 * dl - ratio, 0)
|
||||
return np.where(
|
||||
gray[..., None],
|
||||
np.broadcast_to(dest_l[..., None], i.shape),
|
||||
l * ratio[..., None] + offset[..., None],
|
||||
)
|
||||
|
||||
|
||||
def _blend_luminosity(i: np.ndarray, l: np.ndarray) -> np.ndarray:
|
||||
# Scale the backdrop so it carries the layer's luminance. Where the backdrop
|
||||
# has no luminance to scale there is no hue or saturation to preserve either,
|
||||
# so the result is a neutral grey at the layer's luminance - which is also the
|
||||
# analytic limit of i * lum(l)/lum(i) as a grey backdrop approaches black.
|
||||
# Guarding the numerator here instead (returning black) makes a luminosity
|
||||
# layer disappear over dark backdrops; see tests-unit/comfy_extras_test/
|
||||
# compositor_blend_golden.json.
|
||||
lum_i = luminance(i)
|
||||
lum_l = luminance(l)
|
||||
degenerate = lum_i <= EPSILON
|
||||
ratio = np.where(degenerate, 0.0, lum_l / np.where(degenerate, 1.0, lum_i))
|
||||
return np.where(
|
||||
degenerate[..., None],
|
||||
np.broadcast_to(lum_l[..., None], i.shape),
|
||||
i * ratio[..., None],
|
||||
)
|
||||
|
||||
|
||||
HSL_BLEND = {
|
||||
"hue": _blend_hue,
|
||||
"saturation": _blend_saturation,
|
||||
"color": _blend_color,
|
||||
"luminosity": _blend_luminosity,
|
||||
}
|
||||
|
||||
|
||||
def blend_pixel(blend: str, in_rgb: np.ndarray, layer_rgb: np.ndarray) -> np.ndarray:
|
||||
in_rgb = np.asarray(in_rgb, dtype=np.float32)
|
||||
layer_rgb = np.asarray(layer_rgb, dtype=np.float32)
|
||||
hsl = HSL_BLEND.get(blend)
|
||||
if hsl is not None:
|
||||
return np.asarray(hsl(in_rgb, layer_rgb), dtype=np.float32)
|
||||
fn = CHANNEL_BLEND.get(blend, CHANNEL_BLEND["normal"])
|
||||
return np.asarray(fn(in_rgb, layer_rgb), dtype=np.float32)
|
||||
|
||||
|
||||
def _composite_union(in_c, layer, comp, cov):
|
||||
in_a = in_c[..., 3]
|
||||
layer_a = layer[..., 3] * cov
|
||||
new_a = layer_a + (1 - layer_a) * in_a
|
||||
ratio = np.where(new_a != 0, layer_a / np.where(new_a != 0, new_a, 1), 0)
|
||||
blended = (
|
||||
ratio[..., None]
|
||||
* (in_a[..., None] * (comp - layer[..., :3]) + layer[..., :3] - in_c[..., :3])
|
||||
+ in_c[..., :3]
|
||||
)
|
||||
keep = (layer_a == 0) | (new_a == 0)
|
||||
rgb = np.where(
|
||||
keep[..., None],
|
||||
in_c[..., :3],
|
||||
np.where((in_a == 0)[..., None], layer[..., :3], blended),
|
||||
)
|
||||
return np.concatenate([rgb, new_a[..., None]], axis=-1)
|
||||
|
||||
|
||||
def _composite_clip_to_backdrop(in_c, layer, comp, cov):
|
||||
in_a = in_c[..., 3]
|
||||
layer_a = layer[..., 3] * cov
|
||||
mixed = comp * layer_a[..., None] + in_c[..., :3] * (1 - layer_a[..., None])
|
||||
keep = (in_a == 0) | (layer_a == 0)
|
||||
rgb = np.where(keep[..., None], in_c[..., :3], mixed)
|
||||
return np.concatenate([rgb, in_a[..., None]], axis=-1)
|
||||
|
||||
|
||||
def _composite_clip_to_layer(in_c, layer, comp, cov):
|
||||
in_a = in_c[..., 3]
|
||||
layer_a = layer[..., 3] * cov
|
||||
mixed = comp * in_a[..., None] + layer[..., :3] * (1 - in_a[..., None])
|
||||
rgb = np.where(
|
||||
(layer_a == 0)[..., None],
|
||||
in_c[..., :3],
|
||||
np.where((in_a == 0)[..., None], layer[..., :3], mixed),
|
||||
)
|
||||
return np.concatenate([rgb, layer_a[..., None]], axis=-1)
|
||||
|
||||
|
||||
def _composite_intersection(in_c, layer, comp, cov):
|
||||
new_a = in_c[..., 3] * layer[..., 3] * cov
|
||||
rgb = np.where((new_a == 0)[..., None], in_c[..., :3], comp)
|
||||
return np.concatenate([rgb, new_a[..., None]], axis=-1)
|
||||
|
||||
|
||||
_COMPOSITE = {
|
||||
"union": _composite_union,
|
||||
"clip-to-backdrop": _composite_clip_to_backdrop,
|
||||
"clip-to-layer": _composite_clip_to_layer,
|
||||
"intersection": _composite_intersection,
|
||||
}
|
||||
|
||||
|
||||
def run_composite(mode: str, in_c, layer, comp, cov) -> np.ndarray:
|
||||
fn = _COMPOSITE.get(mode, _composite_union)
|
||||
return fn(in_c, layer, comp, cov)
|
||||
|
||||
|
||||
def _to_space(rgb: np.ndarray, space: str) -> np.ndarray:
|
||||
return rgb if space == "linear" else linear_to_srgb(rgb)
|
||||
|
||||
|
||||
def _from_space(rgb: np.ndarray, space: str) -> np.ndarray:
|
||||
return rgb if space == "linear" else srgb_to_linear(rgb)
|
||||
|
||||
|
||||
class EffectiveMode(NamedTuple):
|
||||
blend: str
|
||||
blend_space: str
|
||||
composite: str
|
||||
|
||||
|
||||
_LAYER_MODES = {
|
||||
"normal": ("linear", "union"),
|
||||
"multiply": ("linear", "clip-to-backdrop"),
|
||||
"screen": ("perceptual", "clip-to-backdrop"),
|
||||
"overlay": ("perceptual", "clip-to-backdrop"),
|
||||
"darken": ("linear", "clip-to-backdrop"),
|
||||
"lighten": ("linear", "clip-to-backdrop"),
|
||||
"color-dodge": ("perceptual", "clip-to-backdrop"),
|
||||
"color-burn": ("perceptual", "clip-to-backdrop"),
|
||||
"hard-light": ("perceptual", "clip-to-backdrop"),
|
||||
"soft-light": ("perceptual", "clip-to-backdrop"),
|
||||
"difference": ("perceptual", "clip-to-backdrop"),
|
||||
"exclusion": ("perceptual", "clip-to-backdrop"),
|
||||
"linear-dodge": ("linear", "clip-to-backdrop"),
|
||||
"linear-burn": ("perceptual", "clip-to-backdrop"),
|
||||
"vivid-light": ("perceptual", "clip-to-backdrop"),
|
||||
"pin-light": ("perceptual", "clip-to-backdrop"),
|
||||
"linear-light": ("perceptual", "clip-to-backdrop"),
|
||||
"hard-mix": ("perceptual", "clip-to-backdrop"),
|
||||
"subtract": ("linear", "clip-to-backdrop"),
|
||||
"divide": ("linear", "clip-to-backdrop"),
|
||||
"grain-extract": ("perceptual", "clip-to-backdrop"),
|
||||
"grain-merge": ("perceptual", "clip-to-backdrop"),
|
||||
"hue": ("perceptual", "clip-to-backdrop"),
|
||||
"saturation": ("perceptual", "clip-to-backdrop"),
|
||||
"color": ("perceptual", "clip-to-backdrop"),
|
||||
"luminosity": ("linear", "clip-to-backdrop"),
|
||||
}
|
||||
|
||||
|
||||
def resolve_mode(blend: str = "normal") -> EffectiveMode:
|
||||
blend_space, composite = _LAYER_MODES.get(blend, _LAYER_MODES["normal"])
|
||||
return EffectiveMode(
|
||||
blend=blend,
|
||||
blend_space=blend_space,
|
||||
composite=composite,
|
||||
)
|
||||
|
||||
|
||||
def blend_composite(
|
||||
mode: EffectiveMode,
|
||||
backdrop: np.ndarray,
|
||||
layer: np.ndarray,
|
||||
opacity: float,
|
||||
mask: Optional[ArrayLike] = None,
|
||||
) -> np.ndarray:
|
||||
backdrop = np.asarray(backdrop, dtype=np.float32)
|
||||
layer = np.asarray(layer, dtype=np.float32)
|
||||
cov = opacity * (1.0 if mask is None else mask)
|
||||
|
||||
in_b = _to_space(backdrop[..., :3], mode.blend_space)
|
||||
layer_b = _to_space(layer[..., :3], mode.blend_space)
|
||||
comp = _from_space(blend_pixel(mode.blend, in_b, layer_b), mode.blend_space)
|
||||
|
||||
return run_composite(mode.composite, backdrop, layer, comp, cov)
|
||||
|
||||
|
||||
def placed_bounds(
|
||||
x: float, y: float, w: float, h: float, rotation: float
|
||||
) -> tuple[int, int, int, int]:
|
||||
cx = x + w / 2
|
||||
cy = y + h / 2
|
||||
cos = math.cos(rotation)
|
||||
sin = math.sin(rotation)
|
||||
hw = w / 2
|
||||
hh = h / 2
|
||||
corners = ((-hw, -hh), (hw, -hh), (hw, hh), (-hw, hh))
|
||||
xs = [cx + dx * cos - dy * sin for dx, dy in corners]
|
||||
ys = [cy + dx * sin + dy * cos for dx, dy in corners]
|
||||
bx = math.floor(min(xs))
|
||||
by = math.floor(min(ys))
|
||||
bw = max(1, math.ceil(max(xs)) - bx)
|
||||
bh = max(1, math.ceil(max(ys)) - by)
|
||||
return bx, by, bw, bh
|
||||
|
|
@ -0,0 +1,855 @@
|
|||
import hashlib
|
||||
import json
|
||||
import math
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
from comfy_api.latest import ComfyExtension, io, UI
|
||||
from comfy_extras.compositor_blend import (
|
||||
_LAYER_MODES,
|
||||
blend_composite,
|
||||
linear_to_srgb,
|
||||
placed_bounds,
|
||||
resolve_mode,
|
||||
srgb_to_linear,
|
||||
)
|
||||
from comfy_extras.color_util import hex_to_rgb
|
||||
from comfy_extras.nodes_bounding_boxes import boxes_from_input
|
||||
from nodes import MAX_RESOLUTION
|
||||
from typing_extensions import override
|
||||
|
||||
|
||||
MAX_LAYERS = 50
|
||||
|
||||
|
||||
def document_items(doc) -> list[dict]:
|
||||
if not isinstance(doc, dict):
|
||||
return []
|
||||
version = doc.get("version")
|
||||
if version is not None and version != 1:
|
||||
raise ValueError(f"LAYERS document version {version!r} is not supported")
|
||||
items = []
|
||||
for item in doc.get("layers") or []:
|
||||
if not isinstance(item, dict):
|
||||
continue
|
||||
item_type = item.get("type", "raster")
|
||||
if item_type != "raster":
|
||||
raise ValueError(f"LAYERS item type {item_type!r} is not supported yet")
|
||||
if not isinstance(item.get("image"), torch.Tensor):
|
||||
continue
|
||||
blend = item.get("blend_mode")
|
||||
if blend is not None and blend not in _LAYER_MODES:
|
||||
raise ValueError(f"LAYERS item blend_mode {blend!r} is not a known blend mode")
|
||||
items.append(item)
|
||||
return sorted(items, key=lambda item: _int(item.get("z_index"), 0))
|
||||
|
||||
|
||||
def document_canvas(doc) -> tuple[int, int] | None:
|
||||
if not isinstance(doc, dict):
|
||||
return None
|
||||
canvas = doc.get("canvas")
|
||||
if not isinstance(canvas, (tuple, list)) or len(canvas) != 2:
|
||||
return None
|
||||
w, h = _int(canvas[0], 0), _int(canvas[1], 0)
|
||||
return (w, h) if w > 0 and h > 0 else None
|
||||
|
||||
|
||||
def _int(value, default: int) -> int:
|
||||
return int(value) if isinstance(value, (int, float)) and not isinstance(value, bool) else default
|
||||
|
||||
|
||||
def _bbox_list(bboxes, canvas_width: int, canvas_height: int) -> list[dict]:
|
||||
if bboxes is None:
|
||||
return []
|
||||
if isinstance(bboxes, str):
|
||||
text = bboxes.strip()
|
||||
if not text:
|
||||
return []
|
||||
try:
|
||||
bboxes = json.loads(text)
|
||||
except (json.JSONDecodeError, ValueError) as exc:
|
||||
raise ValueError(f"bboxes string input is not valid JSON: {exc}") from exc
|
||||
probe = bboxes if isinstance(bboxes, list) else [bboxes]
|
||||
if probe and isinstance(probe[0], list):
|
||||
probe = probe[0]
|
||||
has_elements = any(
|
||||
isinstance(box, dict) and isinstance(box.get("bbox"), (list, tuple))
|
||||
for box in probe
|
||||
)
|
||||
if has_elements and (canvas_width <= 0 or canvas_height <= 0):
|
||||
raise ValueError(
|
||||
"normalized element boxes need canvas_width and canvas_height to resolve to pixels"
|
||||
)
|
||||
return boxes_from_input(bboxes, canvas_width, canvas_height)
|
||||
|
||||
|
||||
def _item_mask_frame(mask, index: int) -> torch.Tensor | None:
|
||||
if not isinstance(mask, torch.Tensor):
|
||||
return None
|
||||
if mask.shape[0] == 1:
|
||||
return mask[:1]
|
||||
if index < mask.shape[0]:
|
||||
return mask[index : index + 1]
|
||||
return None
|
||||
|
||||
|
||||
def expand_item_frames(items: list[dict]) -> list[dict]:
|
||||
frames = []
|
||||
for item in items:
|
||||
image = item["image"]
|
||||
for index in range(image.shape[0]):
|
||||
width = _int(item.get("w"), 0)
|
||||
height = _int(item.get("h"), 0)
|
||||
rotation = item.get("rotation")
|
||||
frames.append({
|
||||
"tensor": image[index : index + 1],
|
||||
"mask": _item_mask_frame(item.get("mask"), index),
|
||||
"name": item.get("name") if isinstance(item.get("name"), str) else None,
|
||||
"x": _int(item.get("x"), 0),
|
||||
"y": _int(item.get("y"), 0),
|
||||
"w": width if width > 0 else int(image.shape[2]),
|
||||
"h": height if height > 0 else int(image.shape[1]),
|
||||
"rotation": float(rotation)
|
||||
if isinstance(rotation, (int, float)) and not isinstance(rotation, bool)
|
||||
else 0.0,
|
||||
"opacity": item.get("opacity", 1.0),
|
||||
"blend": item.get("blend_mode", "normal"),
|
||||
"visible": item.get("visible", True),
|
||||
"flip_h": bool(item.get("flip_h", False)),
|
||||
"flip_v": bool(item.get("flip_v", False)),
|
||||
})
|
||||
if len(frames) > MAX_LAYERS:
|
||||
raise ValueError(
|
||||
f"Compositor supports at most {MAX_LAYERS} layers, got {len(frames)}"
|
||||
)
|
||||
return frames
|
||||
|
||||
|
||||
def frame_alpha(
|
||||
tensor: torch.Tensor, mask: torch.Tensor | None
|
||||
) -> torch.Tensor | None:
|
||||
alpha = tensor[:1, :, :, 3] if tensor.shape[-1] == 4 else None
|
||||
if mask is None:
|
||||
return alpha
|
||||
h, w = tensor.shape[1], tensor.shape[2]
|
||||
m = mask[:1].to(device=tensor.device, dtype=torch.float32)
|
||||
if m.shape[1] != h or m.shape[2] != w:
|
||||
m = torch.nn.functional.interpolate(
|
||||
m.unsqueeze(1), size=(h, w), mode="bilinear"
|
||||
).squeeze(1)
|
||||
inv = torch.clamp(1.0 - m, 0.0, 1.0)
|
||||
return inv if alpha is None else alpha * inv
|
||||
|
||||
|
||||
def layer_preview_tensor(
|
||||
tensor: torch.Tensor, alpha: torch.Tensor | None
|
||||
) -> torch.Tensor:
|
||||
rgb = tensor[:1, :, :, :3]
|
||||
if alpha is None:
|
||||
return rgb
|
||||
return torch.cat([rgb, alpha.unsqueeze(-1)], dim=-1)
|
||||
|
||||
|
||||
def canvas_extent(frames: list[dict]) -> tuple[int, int]:
|
||||
right = 1
|
||||
bottom = 1
|
||||
for frame in frames:
|
||||
bx, by, bw, bh = placed_bounds(
|
||||
frame["x"], frame["y"], frame["w"], frame["h"], frame["rotation"]
|
||||
)
|
||||
right = max(right, bx + bw)
|
||||
bottom = max(bottom, by + bh)
|
||||
return (right, bottom)
|
||||
|
||||
|
||||
def input_fingerprints(
|
||||
frames: list[dict], alphas: list[torch.Tensor | None]
|
||||
) -> list[str]:
|
||||
fingerprints = []
|
||||
for frame, alpha in zip(frames, alphas):
|
||||
tensor = frame["tensor"]
|
||||
rgb = tensor[0, :, :, :3].detach().cpu().numpy()
|
||||
rgb8 = np.clip(np.rint(rgb * 255.0), 0, 255).astype(np.uint8)
|
||||
digest = hashlib.sha256()
|
||||
digest.update(repr(tuple(tensor.shape)).encode())
|
||||
digest.update(rgb8.tobytes())
|
||||
if alpha is not None:
|
||||
alpha8 = np.clip(
|
||||
np.rint(alpha[0].detach().cpu().numpy() * 255.0), 0, 255
|
||||
).astype(np.uint8)
|
||||
digest.update(alpha8.tobytes())
|
||||
digest.update(
|
||||
repr((
|
||||
frame["x"],
|
||||
frame["y"],
|
||||
frame["w"],
|
||||
frame["h"],
|
||||
frame["rotation"],
|
||||
frame["opacity"],
|
||||
frame["blend"],
|
||||
bool(frame["visible"]),
|
||||
frame["flip_h"],
|
||||
frame["flip_v"],
|
||||
)).encode()
|
||||
)
|
||||
fingerprints.append(digest.hexdigest()[:16])
|
||||
return fingerprints
|
||||
|
||||
|
||||
def state_from_items(frames: list[dict], canvas: tuple[int, int]) -> dict:
|
||||
layers = []
|
||||
for frame in frames:
|
||||
layers.append({
|
||||
"name": frame["name"],
|
||||
"visible": bool(frame["visible"]),
|
||||
"opacity": frame["opacity"],
|
||||
"blend": frame["blend"],
|
||||
"flipH": frame["flip_h"],
|
||||
"flipV": frame["flip_v"],
|
||||
"transform": {
|
||||
"x": frame["x"],
|
||||
"y": frame["y"],
|
||||
"w": frame["w"],
|
||||
"h": frame["h"],
|
||||
"rotation": frame["rotation"],
|
||||
},
|
||||
})
|
||||
return {
|
||||
"canvas": canvas,
|
||||
"layers": layers,
|
||||
"inputs": None,
|
||||
"background": {"color": "#ffffff", "opacity": 1.0, "visible": False},
|
||||
}
|
||||
|
||||
|
||||
def layer_ui_entries(frames: list[dict]) -> list:
|
||||
entries = []
|
||||
for frame in frames:
|
||||
entries.append({
|
||||
"x": frame["x"],
|
||||
"y": frame["y"],
|
||||
"width": int(frame["w"]),
|
||||
"height": int(frame["h"]),
|
||||
"rotation": frame["rotation"],
|
||||
"name": frame["name"],
|
||||
"visible": bool(frame["visible"]),
|
||||
"opacity": frame["opacity"] if isinstance(frame["opacity"], (int, float)) else 1.0,
|
||||
"blend": frame["blend"] if isinstance(frame["blend"], str) else "normal",
|
||||
"flipH": frame["flip_h"],
|
||||
"flipV": frame["flip_v"],
|
||||
})
|
||||
return entries
|
||||
|
||||
|
||||
_HEX_DIGITS = set("0123456789abcdef")
|
||||
|
||||
|
||||
def _normalize_hex_color(value) -> str:
|
||||
if isinstance(value, str):
|
||||
text = value.strip().lower()
|
||||
if text.startswith("#"):
|
||||
digits = text[1:]
|
||||
if len(digits) == 3 and set(digits) <= _HEX_DIGITS:
|
||||
digits = "".join(ch * 2 for ch in digits)
|
||||
if len(digits) == 6 and set(digits) <= _HEX_DIGITS:
|
||||
return "#" + digits
|
||||
return "#ffffff"
|
||||
|
||||
|
||||
def _parse_background(entry) -> dict | None:
|
||||
if not isinstance(entry, dict):
|
||||
return None
|
||||
return {
|
||||
"color": _normalize_hex_color(entry.get("color")),
|
||||
"opacity": min(max(_number(entry, "opacity", 1.0), 0.0), 1.0),
|
||||
"visible": bool(entry.get("visible", True)),
|
||||
}
|
||||
|
||||
|
||||
def _parse_order(value, layer_count: int) -> list[int] | None:
|
||||
if not isinstance(value, list) or not value:
|
||||
return None
|
||||
if not all(
|
||||
isinstance(item, int) and not isinstance(item, bool) for item in value
|
||||
):
|
||||
return None
|
||||
if sorted(value) != list(range(layer_count)):
|
||||
return None
|
||||
return value
|
||||
|
||||
|
||||
def layer_state_provided(raw) -> bool:
|
||||
if isinstance(raw, dict):
|
||||
return bool(raw)
|
||||
if isinstance(raw, str):
|
||||
return raw not in ("", "{}")
|
||||
return False
|
||||
|
||||
|
||||
def parse_layer_state(raw) -> dict | None:
|
||||
if isinstance(raw, str):
|
||||
if not raw.strip():
|
||||
return None
|
||||
try:
|
||||
raw = json.loads(raw)
|
||||
except (json.JSONDecodeError, ValueError):
|
||||
return None
|
||||
if not isinstance(raw, dict):
|
||||
return None
|
||||
state = raw
|
||||
version = state.get("version")
|
||||
if version is not None and version != 1:
|
||||
return None
|
||||
canvas = state.get("canvas")
|
||||
layers = state.get("layers")
|
||||
if not isinstance(canvas, dict) or not isinstance(layers, list) or not layers:
|
||||
return None
|
||||
try:
|
||||
w = int(round(float(canvas.get("w"))))
|
||||
h = int(round(float(canvas.get("h"))))
|
||||
except (TypeError, ValueError, OverflowError):
|
||||
return None
|
||||
if w <= 0 or h <= 0:
|
||||
return None
|
||||
inputs = state.get("inputs")
|
||||
if (
|
||||
not isinstance(inputs, list)
|
||||
or len(inputs) != len(layers)
|
||||
or not all(isinstance(entry, str) for entry in inputs)
|
||||
):
|
||||
inputs = None
|
||||
return {
|
||||
"canvas": (w, h),
|
||||
"layers": layers,
|
||||
"inputs": inputs,
|
||||
"background": _parse_background(state.get("background")),
|
||||
"order": _parse_order(state.get("order"), len(layers)),
|
||||
}
|
||||
|
||||
|
||||
def _number(source: dict, key: str, default: float) -> float:
|
||||
value = source.get(key, default)
|
||||
if not isinstance(value, (int, float)) or not math.isfinite(value):
|
||||
return float(default)
|
||||
return float(value)
|
||||
|
||||
|
||||
def _clamped_size(value: float, natural: int) -> float:
|
||||
return float(natural) if value <= 0 else min(value, float(MAX_RESOLUTION))
|
||||
|
||||
|
||||
def _layer_params(entry, natural_w: int, natural_h: int) -> dict:
|
||||
if not isinstance(entry, dict):
|
||||
entry = {}
|
||||
transform = entry.get("transform")
|
||||
if not isinstance(transform, dict):
|
||||
transform = {}
|
||||
blend = entry.get("blend")
|
||||
return {
|
||||
"visible": bool(entry.get("visible", True)),
|
||||
# The layer state is untrusted input: it round-trips through the saved
|
||||
# workflow and can be posted directly to /prompt. An out-of-range opacity
|
||||
# would otherwise reach blend_composite as a raw coverage multiplier and
|
||||
# produce negative or greater-than-white RGB. _parse_background already
|
||||
# clamps the same field.
|
||||
"opacity": min(max(_number(entry, "opacity", 1.0), 0.0), 1.0),
|
||||
"blend": blend if isinstance(blend, str) else "normal",
|
||||
"x": min(max(_number(transform, "x", 0.0), -MAX_RESOLUTION), MAX_RESOLUTION),
|
||||
"y": min(max(_number(transform, "y", 0.0), -MAX_RESOLUTION), MAX_RESOLUTION),
|
||||
"w": _clamped_size(_number(transform, "w", natural_w), natural_w),
|
||||
"h": _clamped_size(_number(transform, "h", natural_h), natural_h),
|
||||
"rotation": _number(transform, "rotation", 0.0),
|
||||
"flip_h": bool(entry.get("flipH", False)),
|
||||
"flip_v": bool(entry.get("flipV", False)),
|
||||
}
|
||||
|
||||
|
||||
def _prepare_layer_bitmap(
|
||||
tensor: torch.Tensor, params: dict, alpha: torch.Tensor | None
|
||||
) -> Image.Image:
|
||||
frame = tensor[0, :, :, :3].detach().cpu().numpy()
|
||||
rgb8 = np.clip(np.rint(frame * 255.0), 0, 255).astype(np.uint8)
|
||||
if alpha is None:
|
||||
img = Image.fromarray(rgb8, "RGB").convert("RGBA")
|
||||
else:
|
||||
alpha8 = np.clip(
|
||||
np.rint(alpha[0].detach().cpu().numpy() * 255.0), 0, 255
|
||||
).astype(np.uint8)
|
||||
img = Image.fromarray(np.dstack([rgb8, alpha8]), "RGBA")
|
||||
if params["flip_h"]:
|
||||
img = img.transpose(Image.Transpose.FLIP_LEFT_RIGHT)
|
||||
if params["flip_v"]:
|
||||
img = img.transpose(Image.Transpose.FLIP_TOP_BOTTOM)
|
||||
target = (max(1, round(params["w"])), max(1, round(params["h"])))
|
||||
if img.size != target:
|
||||
img = img.resize(target, Image.Resampling.LANCZOS)
|
||||
if params["rotation"] != 0:
|
||||
img = img.rotate(
|
||||
-math.degrees(params["rotation"]),
|
||||
expand=True,
|
||||
resample=Image.Resampling.BICUBIC,
|
||||
fillcolor=(0, 0, 0, 0),
|
||||
)
|
||||
return img
|
||||
|
||||
|
||||
def _place_in_bounds(img: Image.Image, bw: int, bh: int) -> np.ndarray:
|
||||
arr = np.asarray(img, dtype=np.float32) / 255.0
|
||||
rgba = np.concatenate([srgb_to_linear(arr[..., :3]), arr[..., 3:4]], axis=-1)
|
||||
aw, ah = img.size
|
||||
buf = np.zeros((bh, bw, 4), dtype=np.float32)
|
||||
ox = (bw - aw) // 2
|
||||
oy = (bh - ah) // 2
|
||||
dx0, dy0 = max(ox, 0), max(oy, 0)
|
||||
dx1, dy1 = min(ox + aw, bw), min(oy + ah, bh)
|
||||
if dx0 < dx1 and dy0 < dy1:
|
||||
buf[dy0:dy1, dx0:dx1] = rgba[dy0 - oy : dy1 - oy, dx0 - ox : dx1 - ox]
|
||||
return buf
|
||||
|
||||
|
||||
def _fill_background(canvas: np.ndarray, background: dict) -> np.ndarray:
|
||||
layer = np.empty(canvas.shape, dtype=np.float32)
|
||||
layer[..., :3] = srgb_to_linear(
|
||||
np.array(hex_to_rgb(background["color"]), dtype=np.float32) / 255.0
|
||||
)
|
||||
layer[..., 3] = 1.0
|
||||
return blend_composite(
|
||||
resolve_mode("normal"), canvas, layer, background["opacity"]
|
||||
)
|
||||
|
||||
|
||||
def composite_from_state(
|
||||
tensors: list[torch.Tensor],
|
||||
state: dict,
|
||||
alphas: list[torch.Tensor | None],
|
||||
) -> torch.Tensor:
|
||||
cw, ch = state["canvas"]
|
||||
if cw > MAX_RESOLUTION or ch > MAX_RESOLUTION:
|
||||
raise ValueError(
|
||||
f"Compositor canvas {cw}x{ch} exceeds the maximum supported size of "
|
||||
f"{MAX_RESOLUTION}x{MAX_RESOLUTION}"
|
||||
)
|
||||
canvas = np.zeros((ch, cw, 4), dtype=np.float32)
|
||||
background = state.get("background")
|
||||
if background is not None and background["visible"] and background["opacity"] > 0:
|
||||
canvas = _fill_background(canvas, background)
|
||||
layers = state["layers"]
|
||||
order = state.get("order") or range(len(tensors))
|
||||
for index in order:
|
||||
if index < 0 or index >= len(tensors):
|
||||
continue
|
||||
tensor = tensors[index]
|
||||
entry = layers[index] if index < len(layers) else None
|
||||
params = _layer_params(entry, tensor.shape[2], tensor.shape[1])
|
||||
if not params["visible"]:
|
||||
continue
|
||||
img = _prepare_layer_bitmap(
|
||||
tensor, params, alphas[index] if index < len(alphas) else None
|
||||
)
|
||||
bx, by, bw, bh = placed_bounds(
|
||||
params["x"], params["y"], params["w"], params["h"], params["rotation"]
|
||||
)
|
||||
buf = _place_in_bounds(img, bw, bh)
|
||||
x0, y0 = max(bx, 0), max(by, 0)
|
||||
x1, y1 = min(bx + bw, cw), min(by + bh, ch)
|
||||
if x0 >= x1 or y0 >= y1:
|
||||
continue
|
||||
region = buf[y0 - by : y1 - by, x0 - bx : x1 - bx]
|
||||
mode = resolve_mode(params["blend"])
|
||||
canvas[y0:y1, x0:x1] = blend_composite(
|
||||
mode, canvas[y0:y1, x0:x1], region, params["opacity"]
|
||||
)
|
||||
rgb = linear_to_srgb(np.clip(canvas[..., :3], 0.0, 1.0))
|
||||
alpha = np.clip(canvas[..., 3:4], 0.0, 1.0)
|
||||
rgba = np.concatenate([rgb, alpha], axis=-1)
|
||||
return torch.from_numpy(rgba.astype(np.float32)).unsqueeze(0)
|
||||
|
||||
|
||||
OPAQUE_EPSILON = 1e-3
|
||||
|
||||
|
||||
def composite_outputs(out: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
if out.shape[-1] != 4:
|
||||
return out, torch.zeros(out.shape[:3], dtype=torch.float32)
|
||||
alpha = out[..., 3]
|
||||
if bool((alpha >= 1.0 - OPAQUE_EPSILON).all()):
|
||||
return out[..., :3], torch.zeros_like(alpha)
|
||||
return out, torch.clamp(1.0 - alpha, 0.0, 1.0)
|
||||
|
||||
|
||||
class ImageCompositor(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id="ImageCompositor",
|
||||
display_name="Create Layered Image",
|
||||
category="image",
|
||||
search_aliases=["compositor", "composite", "layer", "layers", "layer editor", "psd"],
|
||||
is_experimental=True,
|
||||
# both flags on purpose: terminal compositor graphs must execute (the
|
||||
# editor needs a run to open), and cache hits must replay the layer UI
|
||||
is_output_node=True,
|
||||
has_intermediate_output=True,
|
||||
inputs=[
|
||||
io.Layers.Input(
|
||||
"layers",
|
||||
tooltip="Layer stack to composite; build it with Add Layer. Items are stacked by z_index, batch frames inside an item expand to consecutive layers, and item placement, opacity, and blend mode define the initial composition. Without an explicit document canvas the size is a best-effort maximum extent of the placed layers. A saved composition that matches the current inputs takes priority.",
|
||||
),
|
||||
io.Compositor.Input(
|
||||
"compositor",
|
||||
tooltip="Layered composition saved by the compositor editor.",
|
||||
),
|
||||
],
|
||||
outputs=[
|
||||
io.Image.Output(
|
||||
tooltip="Composited image. Carries an alpha channel when the composite has transparent areas (e.g. hidden background), otherwise plain RGB."
|
||||
),
|
||||
io.Mask.Output(
|
||||
tooltip="Transparency of the composite (1 = fully transparent). All zeros when the composite is opaque."
|
||||
),
|
||||
],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, layers: io.Layers.Type, compositor: io.Compositor.Type = None) -> io.NodeOutput:
|
||||
frames = expand_item_frames(document_items(layers))
|
||||
tensors = [frame["tensor"] for frame in frames]
|
||||
alphas = [frame_alpha(frame["tensor"], frame["mask"]) for frame in frames]
|
||||
|
||||
layer_refs = []
|
||||
for tensor, alpha in zip(tensors, alphas):
|
||||
layer_refs.extend(
|
||||
UI.PreviewImage(layer_preview_tensor(tensor, alpha), cls=cls).values
|
||||
)
|
||||
|
||||
fp = input_fingerprints(frames, alphas)
|
||||
raw_state = compositor
|
||||
state = parse_layer_state(raw_state)
|
||||
replay = bool(state is not None and tensors and state["inputs"] == fp)
|
||||
if replay:
|
||||
out = composite_from_state(tensors, state, alphas)
|
||||
elif tensors:
|
||||
canvas = document_canvas(layers) or canvas_extent(frames)
|
||||
out = composite_from_state(
|
||||
tensors, state_from_items(frames, canvas), alphas
|
||||
)
|
||||
else:
|
||||
out = torch.zeros((1, 64, 64, 3), dtype=torch.float32)
|
||||
state_stale = layer_state_provided(raw_state) and not replay
|
||||
out, mask = composite_outputs(out)
|
||||
|
||||
ui_dict = UI.PreviewImage(out, cls=cls).as_dict()
|
||||
ui_dict["compositor_layers"] = layer_refs
|
||||
ui_dict["compositor_inputs"] = fp
|
||||
ui_dict["compositor_bboxes"] = layer_ui_entries(frames)
|
||||
if state_stale:
|
||||
ui_dict["compositor_state_stale"] = [True]
|
||||
return io.NodeOutput(out, mask, ui=ui_dict)
|
||||
|
||||
|
||||
class AddLayer(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id="AddLayer",
|
||||
display_name="Add Layer",
|
||||
category="image",
|
||||
is_experimental=True,
|
||||
inputs=[
|
||||
io.Layers.Input(
|
||||
"layers",
|
||||
optional=True,
|
||||
tooltip="Layer stack to append to. Leave unconnected to start a new stack.",
|
||||
),
|
||||
io.Image.Input(
|
||||
"image",
|
||||
tooltip="Layer content at its native size. A batch expands to consecutive layers.",
|
||||
),
|
||||
io.Mask.Input(
|
||||
"mask",
|
||||
optional=True,
|
||||
tooltip="Transparency mask for this layer. Masked areas (value 1) become transparent, multiplying with any alpha channel the image already carries.",
|
||||
),
|
||||
io.String.Input(
|
||||
"name",
|
||||
optional=True,
|
||||
default="",
|
||||
tooltip="Layer name shown in the compositor editor.",
|
||||
),
|
||||
io.Int.Input(
|
||||
"x",
|
||||
optional=True,
|
||||
default=0,
|
||||
min=-MAX_RESOLUTION,
|
||||
max=MAX_RESOLUTION,
|
||||
tooltip="Initial horizontal placement on the canvas.",
|
||||
),
|
||||
io.Int.Input(
|
||||
"y",
|
||||
optional=True,
|
||||
default=0,
|
||||
min=-MAX_RESOLUTION,
|
||||
max=MAX_RESOLUTION,
|
||||
tooltip="Initial vertical placement on the canvas.",
|
||||
),
|
||||
io.Float.Input(
|
||||
"opacity",
|
||||
optional=True,
|
||||
default=1.0,
|
||||
min=0.0,
|
||||
max=1.0,
|
||||
step=0.01,
|
||||
tooltip="Initial layer opacity.",
|
||||
),
|
||||
io.Combo.Input(
|
||||
"blend_mode",
|
||||
options=list(_LAYER_MODES),
|
||||
default="normal",
|
||||
optional=True,
|
||||
tooltip="Initial blend mode, applied against the layers below. On the bottom layer over the default transparent background, non-normal modes produce transparency.",
|
||||
),
|
||||
io.Float.Input(
|
||||
"rotation",
|
||||
optional=True,
|
||||
default=0.0,
|
||||
min=-360.0,
|
||||
max=360.0,
|
||||
step=1.0,
|
||||
tooltip="Initial rotation in degrees, clockwise.",
|
||||
),
|
||||
io.Int.Input(
|
||||
"width",
|
||||
optional=True,
|
||||
default=0,
|
||||
min=0,
|
||||
max=MAX_RESOLUTION,
|
||||
tooltip="Initial display width. 0 keeps the image's native width.",
|
||||
),
|
||||
io.Int.Input(
|
||||
"height",
|
||||
optional=True,
|
||||
default=0,
|
||||
min=0,
|
||||
max=MAX_RESOLUTION,
|
||||
tooltip="Initial display height. 0 keeps the image's native height.",
|
||||
),
|
||||
io.Int.Input(
|
||||
"z_index",
|
||||
optional=True,
|
||||
default=0,
|
||||
min=-1000,
|
||||
max=1000,
|
||||
tooltip="Stacking override. Layers are stable-sorted by z_index; equal values keep their list order.",
|
||||
),
|
||||
io.Boolean.Input(
|
||||
"flip_h",
|
||||
optional=True,
|
||||
default=False,
|
||||
tooltip="Flip the layer horizontally.",
|
||||
),
|
||||
io.Boolean.Input(
|
||||
"flip_v",
|
||||
optional=True,
|
||||
default=False,
|
||||
tooltip="Flip the layer vertically.",
|
||||
),
|
||||
],
|
||||
outputs=[
|
||||
io.Layers.Output(tooltip="The layer stack with this layer appended."),
|
||||
],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, image: io.Image.Type, layers: io.Layers.Type = None, mask: io.Mask.Type = None, name: str = "", x: int = 0, y: int = 0, opacity: float = 1.0, blend_mode: str = "normal", rotation: float = 0.0, width: int = 0, height: int = 0, z_index: int = 0, flip_h: bool = False, flip_v: bool = False) -> io.NodeOutput:
|
||||
item: dict = {
|
||||
"image": image,
|
||||
"type": "raster",
|
||||
"x": int(x),
|
||||
"y": int(y),
|
||||
"z_index": int(z_index),
|
||||
}
|
||||
if mask is not None:
|
||||
item["mask"] = mask
|
||||
if name:
|
||||
item["name"] = name
|
||||
if opacity != 1.0:
|
||||
item["opacity"] = float(opacity)
|
||||
if blend_mode != "normal":
|
||||
item["blend_mode"] = blend_mode
|
||||
if rotation != 0.0:
|
||||
item["rotation"] = math.radians(rotation)
|
||||
if width > 0:
|
||||
item["w"] = int(width)
|
||||
if height > 0:
|
||||
item["h"] = int(height)
|
||||
if flip_h:
|
||||
item["flip_h"] = True
|
||||
if flip_v:
|
||||
item["flip_v"] = True
|
||||
previous = layers if isinstance(layers, dict) else None
|
||||
document: dict = {
|
||||
"version": 1,
|
||||
"layers": [*(previous.get("layers") or []), item] if previous else [item],
|
||||
}
|
||||
previous_canvas = document_canvas(previous)
|
||||
if previous_canvas:
|
||||
document["canvas"] = previous_canvas
|
||||
return io.NodeOutput(document)
|
||||
|
||||
|
||||
class LayersFromBoundingBoxes(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id="LayersFromBoundingBoxes",
|
||||
display_name="Layers From Bounding Boxes",
|
||||
category="image",
|
||||
is_experimental=True,
|
||||
description=(
|
||||
"Turn an image batch plus its bounding boxes into a layer stack, one layer per frame, "
|
||||
"each placed by its own box. Use this when a node emits layers as a batch - a batch "
|
||||
"carries a single placement for every frame, so the individual positions are otherwise lost."
|
||||
),
|
||||
inputs=[
|
||||
io.Image.Input(
|
||||
"image",
|
||||
tooltip="Image batch; each frame becomes one layer.",
|
||||
),
|
||||
io.MultiType.Input(
|
||||
"bboxes",
|
||||
[io.BoundingBox, io.Array, io.String],
|
||||
tooltip=(
|
||||
"Placement boxes, index-aligned with the image batch. Accepts bounding boxes "
|
||||
"(x, y, width, height), normalized elements (with a 'bbox' - these need "
|
||||
"canvas_width/canvas_height to resolve to pixels), or a JSON string of either. "
|
||||
"Frames without a matching box are placed at the origin. A box's width/height "
|
||||
"scales the layer to fit it. metadata.name (or desc) and metadata.z_index are "
|
||||
"used when present, and metadata.content_rect (frame-relative) crops the frame "
|
||||
"to its real content."
|
||||
),
|
||||
),
|
||||
io.Mask.Input(
|
||||
"mask",
|
||||
optional=True,
|
||||
tooltip=(
|
||||
"Per-frame transparency, index-aligned with the image batch "
|
||||
"(1 = transparent, LoadImage convention)."
|
||||
),
|
||||
),
|
||||
io.Layers.Input(
|
||||
"layers",
|
||||
optional=True,
|
||||
tooltip="Layer stack to append to. Leave unconnected to start a new stack.",
|
||||
),
|
||||
io.Boolean.Input(
|
||||
"crop_to_content",
|
||||
default=True,
|
||||
optional=True,
|
||||
tooltip=(
|
||||
"Crop each frame to metadata.content_rect where present and place the content "
|
||||
"at the box position plus the rect offset. Leave on for batches whose frames "
|
||||
"are padded - it keeps only the real content at its true spot."
|
||||
),
|
||||
),
|
||||
io.Int.Input(
|
||||
"canvas_width",
|
||||
default=0,
|
||||
min=0,
|
||||
max=MAX_RESOLUTION,
|
||||
optional=True,
|
||||
tooltip="Document canvas width. 0 derives it from the placed layers.",
|
||||
),
|
||||
io.Int.Input(
|
||||
"canvas_height",
|
||||
default=0,
|
||||
min=0,
|
||||
max=MAX_RESOLUTION,
|
||||
optional=True,
|
||||
tooltip="Document canvas height. 0 derives it from the placed layers.",
|
||||
),
|
||||
],
|
||||
outputs=[
|
||||
io.Layers.Output(tooltip="The layer stack, ready for Create Layered Image."),
|
||||
],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(
|
||||
cls,
|
||||
image: io.Image.Type,
|
||||
bboxes: io.MultiType.Type,
|
||||
mask: io.Mask.Type = None,
|
||||
layers: io.Layers.Type = None,
|
||||
crop_to_content: bool = True,
|
||||
canvas_width: int = 0,
|
||||
canvas_height: int = 0,
|
||||
) -> io.NodeOutput:
|
||||
boxes = _bbox_list(bboxes, canvas_width, canvas_height)
|
||||
previous = layers if isinstance(layers, dict) else None
|
||||
items: list[dict] = list((previous.get("layers") or []) if previous else [])
|
||||
base_z = max((_int(i.get("z_index"), 0) for i in items), default=-1) + 1
|
||||
|
||||
for index in range(image.shape[0]):
|
||||
box = boxes[index] if index < len(boxes) else {}
|
||||
meta = box.get("metadata") if isinstance(box.get("metadata"), dict) else {}
|
||||
frame = image[index : index + 1]
|
||||
frame_mask = _item_mask_frame(mask, index)
|
||||
|
||||
x, y = _int(box.get("x"), 0), _int(box.get("y"), 0)
|
||||
box_w, box_h = _int(box.get("width"), 0), _int(box.get("height"), 0)
|
||||
cropped = False
|
||||
rect = meta.get("content_rect")
|
||||
if crop_to_content and isinstance(rect, (list, tuple)) and len(rect) == 4:
|
||||
left, top, cw, ch = (_int(v, 0) for v in rect)
|
||||
left = min(max(left, 0), int(frame.shape[2]))
|
||||
top = min(max(top, 0), int(frame.shape[1]))
|
||||
cw = min(max(cw, 0), int(frame.shape[2]) - left)
|
||||
ch = min(max(ch, 0), int(frame.shape[1]) - top)
|
||||
if cw > 0 and ch > 0:
|
||||
frame = frame[:, top : top + ch, left : left + cw]
|
||||
if frame_mask is not None:
|
||||
frame_mask = frame_mask[:, top : top + ch, left : left + cw]
|
||||
x, y = x + left, y + top
|
||||
cropped = True
|
||||
|
||||
item: dict = {
|
||||
"image": frame,
|
||||
"type": "raster",
|
||||
"x": x,
|
||||
"y": y,
|
||||
"z_index": _int(meta.get("z_index"), base_z + index),
|
||||
}
|
||||
if not cropped:
|
||||
if box_w > 0:
|
||||
item["w"] = box_w
|
||||
if box_h > 0:
|
||||
item["h"] = box_h
|
||||
if frame_mask is not None:
|
||||
item["mask"] = frame_mask
|
||||
name = meta.get("name")
|
||||
if not (isinstance(name, str) and name):
|
||||
name = meta.get("desc")
|
||||
if isinstance(name, str) and name:
|
||||
item["name"] = name
|
||||
items.append(item)
|
||||
|
||||
document: dict = {"version": 1, "layers": items}
|
||||
if canvas_width > 0 and canvas_height > 0:
|
||||
document["canvas"] = (canvas_width, canvas_height)
|
||||
else:
|
||||
inherited = document_canvas(previous)
|
||||
if inherited:
|
||||
document["canvas"] = inherited
|
||||
return io.NodeOutput(document)
|
||||
|
||||
|
||||
class CompositorExtension(ComfyExtension):
|
||||
@override
|
||||
async def get_node_list(self) -> list[type[io.ComfyNode]]:
|
||||
return [ImageCompositor, AddLayer, LayersFromBoundingBoxes]
|
||||
|
||||
|
||||
async def comfy_entrypoint() -> CompositorExtension:
|
||||
return CompositorExtension()
|
||||
|
|
@ -591,7 +591,7 @@ class SamplerER_SDE(io.ComfyNode):
|
|||
inputs=[
|
||||
io.Combo.Input("solver_type", options=["ER-SDE", "Reverse-time SDE", "ODE"]),
|
||||
io.Int.Input("max_stage", default=3, min=1, max=3, advanced=True),
|
||||
io.Float.Input("eta", default=1.0, min=0.0, max=100.0, step=0.01, round=False, tooltip="Stochastic strength of reverse-time SDE.\nWhen eta=0, it reduces to deterministic ODE. This setting doesn't apply to ER-SDE solver type.", advanced=True),
|
||||
io.Float.Input("eta", default=1.0, min=0.0, max=10.0, step=0.01, round=False, tooltip="Stochastic strength of SDEs.\nWhen eta=0, they reduce to deterministic ODE.\nLarge eta may cause invalid outputs. If this occurs, try decreasing this value.", advanced=True),
|
||||
io.Float.Input("s_noise", default=1.0, min=0.0, max=100.0, step=0.01, round=False, advanced=True),
|
||||
],
|
||||
outputs=[io.Sampler.Output()]
|
||||
|
|
@ -599,21 +599,35 @@ class SamplerER_SDE(io.ComfyNode):
|
|||
|
||||
@classmethod
|
||||
def execute(cls, solver_type, max_stage, eta, s_noise) -> io.NodeOutput:
|
||||
if solver_type == "ODE" or (solver_type == "Reverse-time SDE" and eta == 0):
|
||||
eta = 0
|
||||
s_noise = 0
|
||||
# Extend existing noise scalers phi(x) with eta-controlled noise scalers:
|
||||
# psi(x) = x**(1-eta) * phi(x)**eta
|
||||
# where eta is constant and directly scales the h^2(t) contribution.
|
||||
|
||||
def reverse_time_sde_noise_scaler(x):
|
||||
def er_sde_noise_scaler(x: torch.Tensor) -> torch.Tensor:
|
||||
return x * ((x ** 0.3).exp() + 10.0) ** eta
|
||||
|
||||
def reverse_time_sde_noise_scaler(x: torch.Tensor) -> torch.Tensor:
|
||||
return x ** (eta + 1)
|
||||
|
||||
if solver_type == "ER-SDE":
|
||||
# Use the default one in sample_er_sde()
|
||||
noise_scaler = None
|
||||
else:
|
||||
noise_scaler = reverse_time_sde_noise_scaler
|
||||
def ode_noise_scaler(x: torch.Tensor) -> torch.Tensor:
|
||||
return x
|
||||
|
||||
solver_scalers = {
|
||||
"ER-SDE": er_sde_noise_scaler,
|
||||
"Reverse-time SDE": reverse_time_sde_noise_scaler,
|
||||
"ODE": ode_noise_scaler,
|
||||
}
|
||||
|
||||
if solver_type == "ODE" or eta == 0:
|
||||
s_noise = 0.0
|
||||
solver_type = "ODE"
|
||||
noise_scaler = solver_scalers[solver_type]
|
||||
|
||||
sampler_name = "er_sde"
|
||||
sampler = comfy.samplers.ksampler(sampler_name, {"s_noise": s_noise, "noise_scaler": noise_scaler, "max_stage": max_stage})
|
||||
sampler = comfy.samplers.ksampler(
|
||||
sampler_name,
|
||||
{"s_noise": s_noise, "noise_scaler": noise_scaler, "max_stage": max_stage},
|
||||
)
|
||||
return io.NodeOutput(sampler)
|
||||
|
||||
get_sampler = execute
|
||||
|
|
|
|||
|
|
@ -195,7 +195,7 @@ class LoadImageDataSetFromFolderNode(io.ComfyNode):
|
|||
|
||||
@classmethod
|
||||
def execute(cls, folder):
|
||||
sub_input_dir = os.path.join(folder_paths.get_input_directory(), folder)
|
||||
sub_input_dir = secure_subfolder_path(folder_paths.get_input_directory(), folder)
|
||||
valid_extensions = [".png", ".jpg", ".jpeg", ".webp"]
|
||||
image_files = [
|
||||
f
|
||||
|
|
@ -241,7 +241,7 @@ class LoadImageTextDataSetFromFolderNode(io.ComfyNode):
|
|||
def execute(cls, folder):
|
||||
logging.info(f"Loading images from folder: {folder}")
|
||||
|
||||
sub_input_dir = os.path.join(folder_paths.get_input_directory(), folder)
|
||||
sub_input_dir = secure_subfolder_path(folder_paths.get_input_directory(), folder)
|
||||
valid_extensions = [".png", ".jpg", ".jpeg", ".webp"]
|
||||
|
||||
image_files = []
|
||||
|
|
@ -310,7 +310,7 @@ class LoadVideoDataSetFromFolderNode(io.ComfyNode):
|
|||
|
||||
@classmethod
|
||||
def execute(cls, folder):
|
||||
sub_input_dir = os.path.join(folder_paths.get_input_directory(), folder)
|
||||
sub_input_dir = secure_subfolder_path(folder_paths.get_input_directory(), folder)
|
||||
video_files = sorted([
|
||||
f for f in os.listdir(sub_input_dir)
|
||||
if any(f.lower().endswith(ext) for ext in VALID_VIDEO_EXTENSIONS)
|
||||
|
|
@ -357,7 +357,7 @@ class LoadVideoTextDataSetFromFolderNode(io.ComfyNode):
|
|||
|
||||
@classmethod
|
||||
def execute(cls, folder):
|
||||
sub_input_dir = os.path.join(folder_paths.get_input_directory(), folder)
|
||||
sub_input_dir = secure_subfolder_path(folder_paths.get_input_directory(), folder)
|
||||
|
||||
video_files = []
|
||||
for item in sorted(os.listdir(sub_input_dir)):
|
||||
|
|
|
|||
|
|
@ -471,6 +471,12 @@ def _mat_to_quat(m):
|
|||
|
||||
|
||||
class SplatToFile3D(IO.ComfyNode):
|
||||
FORMAT_WRITERS = {
|
||||
"ply": _gaussian_ply_bytes,
|
||||
"ksplat": _gaussian_ksplat_bytes,
|
||||
"spz": _gaussian_spz_bytes,
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return IO.Schema(
|
||||
|
|
@ -482,7 +488,7 @@ class SplatToFile3D(IO.ComfyNode):
|
|||
"Supports one item per batch only.",
|
||||
inputs=[
|
||||
IO.Splat.Input("splat"),
|
||||
IO.Combo.Input("format", options=["ply", "ksplat", "spz"], # TODO: add "splat" when we have a writer for it
|
||||
IO.Combo.Input("format", options=list(cls.FORMAT_WRITERS), # TODO: add "splat" when we have a writer for it
|
||||
tooltip="ply: standard 3D Gaussian Splat with full spherical harmonics. "
|
||||
"ksplat: mkkellogg SplatBuffer (level 0, uncompressed), base color only "
|
||||
"spz: Niantic gzip-compressed (~10x smaller), base color only "
|
||||
|
|
@ -493,10 +499,13 @@ class SplatToFile3D(IO.ComfyNode):
|
|||
|
||||
@classmethod
|
||||
def execute(cls, splat, format="ply") -> IO.NodeOutput:
|
||||
writer = cls.FORMAT_WRITERS.get(format)
|
||||
if writer is None:
|
||||
raise ValueError(f"Unsupported splat format: {format!r}")
|
||||
|
||||
if splat.positions.shape[0] > 1:
|
||||
logging.warning("SplatToFile3D supports one item per batch only. Got %d; using first.", splat.positions.shape[0])
|
||||
end = _real_len(splat, 0)
|
||||
writer = {"ksplat": _gaussian_ksplat_bytes, "spz": _gaussian_spz_bytes}.get(format, _gaussian_ply_bytes)
|
||||
data = writer(splat.positions[0, :end], splat.scales[0, :end],
|
||||
splat.rotations[0, :end], splat.opacities[0, :end], splat.sh[0, :end])
|
||||
return IO.NodeOutput(Types.File3D(BytesIO(data), file_format=format))
|
||||
|
|
|
|||
|
|
@ -758,22 +758,60 @@ class LTXVConcatAVLatent(io.ComfyNode):
|
|||
],
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def fit_audio(reference, audio, noise_mask):
|
||||
"""Trim or zero-pad the audio stream to the length of the one it replaces.
|
||||
|
||||
The padded tail is left unmasked so the model generates it, which is what a
|
||||
clip shorter than the video should do.
|
||||
"""
|
||||
dims = [i for i in range(reference.ndim) if reference.shape[i] != audio.shape[i]]
|
||||
if len(dims) == 0:
|
||||
return audio, noise_mask
|
||||
if len(dims) > 1 or dims[0] < 2:
|
||||
raise ValueError("audio latent {} cannot be fitted to {}".format(tuple(audio.shape), tuple(reference.shape)))
|
||||
|
||||
dim, length = dims[0], reference.shape[dims[0]]
|
||||
if noise_mask is not None: # masks carry their own shape until sampling resizes them
|
||||
noise_mask = comfy.utils.reshape_mask(noise_mask, audio.shape)
|
||||
|
||||
if audio.shape[dim] > length:
|
||||
audio = audio.narrow(dim, 0, length)
|
||||
if noise_mask is not None:
|
||||
noise_mask = noise_mask.narrow(dim, 0, length)
|
||||
else:
|
||||
pad = torch.zeros_like(audio.narrow(dim, 0, 1)).repeat(
|
||||
[length - audio.shape[dim] if i == dim else 1 for i in range(audio.ndim)])
|
||||
audio = torch.cat([audio, pad], dim=dim)
|
||||
if noise_mask is not None:
|
||||
noise_mask = torch.cat([noise_mask, torch.ones_like(pad)], dim=dim)
|
||||
return audio, noise_mask
|
||||
|
||||
@classmethod
|
||||
def execute(cls, video_latent, audio_latent) -> io.NodeOutput:
|
||||
output = {}
|
||||
output.update(video_latent)
|
||||
output.update(audio_latent)
|
||||
video_samples = video_latent["samples"]
|
||||
audio_samples = audio_latent["samples"]
|
||||
video_noise_mask = video_latent.get("noise_mask", None)
|
||||
audio_noise_mask = audio_latent.get("noise_mask", None)
|
||||
|
||||
if video_samples.is_nested: # already an AV latent: keep its video and swap the audio stream
|
||||
streams = video_samples.unbind()
|
||||
video_samples = streams[0]
|
||||
if video_noise_mask is not None:
|
||||
video_noise_mask = video_noise_mask.unbind()[0]
|
||||
audio_samples, audio_noise_mask = cls.fit_audio(streams[1], audio_samples, audio_noise_mask)
|
||||
|
||||
if video_noise_mask is not None or audio_noise_mask is not None:
|
||||
if video_noise_mask is None:
|
||||
video_noise_mask = torch.ones_like(video_latent["samples"])
|
||||
video_noise_mask = torch.ones_like(video_samples)
|
||||
if audio_noise_mask is None:
|
||||
audio_noise_mask = torch.ones_like(audio_latent["samples"])
|
||||
audio_noise_mask = torch.ones_like(audio_samples)
|
||||
output["noise_mask"] = comfy.nested_tensor.NestedTensor((video_noise_mask, audio_noise_mask))
|
||||
|
||||
output["samples"] = comfy.nested_tensor.NestedTensor((video_latent["samples"], audio_latent["samples"]))
|
||||
output["samples"] = comfy.nested_tensor.NestedTensor((video_samples, audio_samples))
|
||||
|
||||
return io.NodeOutput(output)
|
||||
|
||||
|
|
|
|||
|
|
@ -283,9 +283,9 @@ class MiniMaxH3ReferenceToVideo(io.ComfyNode):
|
|||
class MiniMaxH3SigmaShift(io.ComfyNode):
|
||||
"""Set the video/audio flow shifts coherently.
|
||||
|
||||
The video shift drives the sampler's sigma schedule; both values are also
|
||||
handed to the DiT, which inverts the video schedule to the shared base grid
|
||||
and derives the audio schedule from it.
|
||||
The video shift drives the sampler's sigma schedule (ModelSamplingAV); both
|
||||
values are also handed to the DiT, which inverts the video schedule to the
|
||||
shared base grid and derives the audio schedule from it.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
|
|
@ -293,7 +293,8 @@ class MiniMaxH3SigmaShift(io.ComfyNode):
|
|||
return io.Schema(
|
||||
node_id="MiniMaxH3SigmaShift",
|
||||
description="Set the video/audio flow shifts.",
|
||||
display_name="MiniMax H3 Sigma Shift",
|
||||
display_name="ModelSamplingMiniMaxH3",
|
||||
search_aliases=["sigma shift", "minimax shift"],
|
||||
category="model/patch/minimax",
|
||||
inputs=[
|
||||
io.Model.Input("model"),
|
||||
|
|
@ -307,12 +308,12 @@ class MiniMaxH3SigmaShift(io.ComfyNode):
|
|||
def execute(cls, model, shift_video, shift_audio) -> io.NodeOutput:
|
||||
m = model.clone()
|
||||
|
||||
class ModelSamplingAdvanced(comfy.model_sampling.ModelSamplingDiscreteFlow, comfy.model_sampling.CONST):
|
||||
class ModelSamplingAdvanced(comfy.model_sampling.ModelSamplingAV, comfy.model_sampling.CONST):
|
||||
pass
|
||||
|
||||
original = m.get_model_object("model_sampling")
|
||||
model_sampling = ModelSamplingAdvanced(model.model.model_config)
|
||||
model_sampling.set_parameters(shift=shift_video)
|
||||
model_sampling.set_parameters(shift=shift_video, audio_shift=shift_audio)
|
||||
if hasattr(original, "noise_scale"):
|
||||
model_sampling.set_noise_scale(original.noise_scale)
|
||||
m.add_object_patch("model_sampling", model_sampling)
|
||||
|
|
@ -329,7 +330,7 @@ class MiniMaxH3Extension(ComfyExtension):
|
|||
EmptyMiniMaxH3LatentAV,
|
||||
MiniMaxH3ImageToVideo,
|
||||
MiniMaxH3ReferenceToVideo,
|
||||
MiniMaxH3SigmaShift
|
||||
MiniMaxH3SigmaShift,
|
||||
]
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -72,7 +72,7 @@ class ImageUpscaleWithModel(io.ComfyNode):
|
|||
|
||||
memory_required = (512 * 512 * 3) * image.element_size() * max(upscale_model.scale, 1.0) * 384.0 #The 384.0 is an estimate of how much some of these models take, TODO: make it more accurate
|
||||
memory_required += image.nelement() * image.element_size()
|
||||
model_management.load_models_gpu([upscale_model.patcher], memory_required=memory_required)
|
||||
model_management.load_models_gpu([upscale_model.patcher], memory_required=memory_required, force_full_load=True)
|
||||
|
||||
in_img = image.movedim(-1,-3).to(device)
|
||||
|
||||
|
|
|
|||
|
|
@ -6,6 +6,8 @@ import comfy.model_management
|
|||
import comfy.utils
|
||||
import comfy.latent_formats
|
||||
import comfy.clip_vision
|
||||
import comfy.ldm.wan.model_animate2
|
||||
import comfy.patcher_extension
|
||||
import json
|
||||
import numpy as np
|
||||
from typing import Tuple, TypedDict
|
||||
|
|
@ -1248,6 +1250,167 @@ class WanAnimateToVideo(io.ComfyNode):
|
|||
out_latent["samples"] = latent
|
||||
return io.NodeOutput(positive, negative, out_latent, trim_latent, max(0, ref_motion_latent_length * 4 - 3), video_frame_offset + length)
|
||||
|
||||
class WanAnimate2ToVideo(io.ComfyNode):
|
||||
|
||||
CONTINUE_MOTION_FRAMES = 1
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id="WanAnimate2ToVideo",
|
||||
category="model/conditioning/wan/animate",
|
||||
description="Animate a character in a reference image using a video, effectively replicating the facial animation, body motion and hands gesture from the video.",
|
||||
inputs=[
|
||||
io.Conditioning.Input("positive"),
|
||||
io.Conditioning.Input("negative"),
|
||||
io.Vae.Input("vae"),
|
||||
io.Int.Input("width", default=832, min=16, max=nodes.MAX_RESOLUTION, step=16, tooltip="Output video width in pixels."),
|
||||
io.Int.Input("height", default=480, min=16, max=nodes.MAX_RESOLUTION, step=16, tooltip="Output video height in pixels."),
|
||||
io.Int.Input("length", default=81, min=1, max=nodes.MAX_RESOLUTION, step=4, tooltip="Number of frames to generate."),
|
||||
io.Int.Input("batch_size", default=1, min=1, max=4096, tooltip="Number of videos to generate simultaneously."),
|
||||
io.Image.Input("reference_image", optional=True, tooltip="The character to animate."),
|
||||
io.Image.Input("pose_video", optional=True, tooltip="The video whose motion is transferred to the reference character."),
|
||||
io.ClipVisionOutput.Input("clip_vision_output", optional=True, tooltip="CLIP vision of the reference image."),
|
||||
io.Conditioning.Input("positive_pose", optional=True, tooltip="Prompt for the pose-video branch, describing the motion rather than the character. Defaults to positive. Used for both the cond and uncond passes."),
|
||||
io.ClipVisionOutput.Input("clip_vision_output_pose", optional=True, tooltip="CLIP vision of the pose video's first frame. Defaults to clip_vision_output."),
|
||||
io.Image.Input("continue_motion", optional=True, tooltip="Previous motion sequence to continue from for temporal consistency."),
|
||||
io.Int.Input("video_frame_offset", default=0, min=0, max=nodes.MAX_RESOLUTION, step=1, tooltip="Frames to seek into the pose video. Connect to the video_frame_offset output of the previous node when extending."),
|
||||
io.Float.Input("pose_strength", default=1.0, min=0.0, max=10.0, step=0.01, tooltip="Scales the pose video's influence on the motion. 1.0 is the trained behavior; below weakens adherence, above amplifies. 0.0 mutes it but does not fully remove it."),
|
||||
io.Float.Input("pose_start_percent", default=0.0, min=0.0, max=1.0, step=0.01, tooltip="Sampling percent at which the pose influence starts. Outside the window the pose branch is skipped entirely, which also speeds those steps up."),
|
||||
io.Float.Input("pose_end_percent", default=1.0, min=0.0, max=1.0, step=0.01, tooltip="Sampling percent at which the pose influence ends. Motion is mostly established early, so e.g. 0.7 can loosen fine detail while keeping the choreography."),
|
||||
io.Float.Input("reference_image_strength", default=1.0, min=0.0, max=10.0, step=0.01, tooltip="Scales how strongly generated frames attend to the reference image's latent frame. Below 1.0 loosens identity/appearance adherence (e.g. to let the prompt restyle), above tightens it against drift."),
|
||||
],
|
||||
outputs=[
|
||||
io.Conditioning.Output(display_name="positive"),
|
||||
io.Conditioning.Output(display_name="negative"),
|
||||
io.Latent.Output(display_name="latent"),
|
||||
io.Int.Output(display_name="trim_latent", tooltip="Number of latent frames that should be trimmed before decoding."),
|
||||
io.Int.Output(display_name="trim_image", tooltip="Number of overlapping image frames when extending a video."),
|
||||
io.Int.Output(display_name="video_frame_offset", tooltip="Frames to seek into the pose video."),
|
||||
],
|
||||
is_experimental=True,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, positive, negative, vae, width, height, length, batch_size, video_frame_offset, reference_image=None, pose_video=None, clip_vision_output=None, positive_pose=None, clip_vision_output_pose=None, continue_motion=None, pose_strength=1.0, pose_start_percent=0.0, pose_end_percent=1.0, reference_image_strength=1.0) -> io.NodeOutput:
|
||||
if pose_start_percent > pose_end_percent:
|
||||
raise ValueError("pose_start_percent ({}) must not be greater than pose_end_percent ({}).".format(pose_start_percent, pose_end_percent))
|
||||
latent_length = ((length - 1) // 4) + 1
|
||||
latent_width = width // 8
|
||||
latent_height = height // 8
|
||||
|
||||
if reference_image is None:
|
||||
reference_image = torch.zeros((1, height, width, 3))
|
||||
|
||||
ref_image = comfy.utils.common_upscale(reference_image[:1].movedim(-1, 1), width, height, "area", "center").movedim(1, -1)
|
||||
ref_latent = vae.encode(ref_image[:, :, :, :3])
|
||||
trim_latent = ref_latent.shape[2]
|
||||
|
||||
ref_motion_latent_length = 0
|
||||
if continue_motion is None:
|
||||
image = torch.ones((length, height, width, 3)) * 0.5
|
||||
else:
|
||||
continue_motion = continue_motion[-cls.CONTINUE_MOTION_FRAMES:]
|
||||
video_frame_offset = max(0, video_frame_offset - continue_motion.shape[0])
|
||||
continue_motion = comfy.utils.common_upscale(continue_motion[-length:].movedim(-1, 1), width, height, "area", "center").movedim(1, -1)
|
||||
# 0.5 is mid-grey, matching upstream's zeros in [-1, 1] pixel space
|
||||
image = torch.ones((length, height, width, continue_motion.shape[-1]), device=continue_motion.device, dtype=continue_motion.dtype) * 0.5
|
||||
image[:continue_motion.shape[0]] = continue_motion
|
||||
ref_motion_latent_length += ((continue_motion.shape[0] - 1) // 4) + 1
|
||||
|
||||
concat_latent_image = torch.cat((ref_latent, vae.encode(image[:, :, :, :3])), dim=2)
|
||||
|
||||
# 1-channel, 0 == known; concat_cond inverts and repeats it to the 4 mask channels
|
||||
mask = torch.ones((1, 1, latent_length + trim_latent, latent_height, latent_width), device=concat_latent_image.device, dtype=concat_latent_image.dtype)
|
||||
mask[:, :, :trim_latent + ref_motion_latent_length] = 0.0
|
||||
|
||||
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})
|
||||
|
||||
if clip_vision_output is not None:
|
||||
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})
|
||||
|
||||
# not windowed with the pose values: the reference frame is part of the latent on every step
|
||||
if reference_image_strength != 1.0:
|
||||
positive = node_helpers.conditioning_set_values(positive, {"reference_strength": reference_image_strength})
|
||||
negative = node_helpers.conditioning_set_values(negative, {"reference_strength": reference_image_strength})
|
||||
|
||||
# set on the negative too: upstream runs the pose branch once, outside the CFG loop, so it never sees the negative prompt
|
||||
pose_values = {}
|
||||
if pose_video is not None:
|
||||
if pose_video.shape[0] <= video_frame_offset:
|
||||
raise ValueError("pose_video has {} frames but video_frame_offset is {} -- nothing left to read.".format(pose_video.shape[0], video_frame_offset))
|
||||
pose_video = pose_video[video_frame_offset:]
|
||||
pose_video = comfy.utils.common_upscale(pose_video[:length].movedim(-1, 1), width, height, "area", "center").movedim(1, -1)
|
||||
if pose_video.shape[0] < length: # hold the last frame, as upstream pads its clips
|
||||
pose_video = torch.cat((pose_video,) + (pose_video[-1:],) * (length - pose_video.shape[0]), dim=0)
|
||||
pose_values["pose_video_latent"] = vae.encode(pose_video[:, :, :, :3])
|
||||
|
||||
pose_clip = clip_vision_output_pose if clip_vision_output_pose is not None else clip_vision_output
|
||||
if pose_clip is not None:
|
||||
pose_values["clip_vision_output_pose"] = pose_clip
|
||||
|
||||
pose_cond = positive_pose if positive_pose is not None else positive
|
||||
if len(pose_cond) > 0:
|
||||
pose_values["cross_attn_pose"] = pose_cond[0][0]
|
||||
|
||||
if pose_strength != 1.0:
|
||||
pose_values["pose_strength"] = pose_strength
|
||||
|
||||
if pose_start_percent > 0.0 or pose_end_percent < 1.0:
|
||||
# windowed via cond timestep ranges: the pose values ride a cond limited to the window, and complement conds without them cover the rest, where the model runs without the pose branch at all
|
||||
def windowed(cond):
|
||||
parts = node_helpers.conditioning_set_values(cond, {**pose_values, "start_percent": pose_start_percent, "end_percent": pose_end_percent})
|
||||
if pose_start_percent > 0.0:
|
||||
parts = parts + node_helpers.conditioning_set_values(cond, {"start_percent": 0.0, "end_percent": pose_start_percent})
|
||||
if pose_end_percent < 1.0:
|
||||
parts = parts + node_helpers.conditioning_set_values(cond, {"start_percent": pose_end_percent, "end_percent": 1.0})
|
||||
return parts
|
||||
positive = windowed(positive)
|
||||
negative = windowed(negative)
|
||||
else:
|
||||
positive = node_helpers.conditioning_set_values(positive, pose_values)
|
||||
negative = node_helpers.conditioning_set_values(negative, pose_values)
|
||||
|
||||
latent = torch.zeros([batch_size, 16, latent_length + trim_latent, latent_height, latent_width], device=comfy.model_management.intermediate_device())
|
||||
out_latent = {}
|
||||
out_latent["samples"] = latent
|
||||
return io.NodeOutput(positive, negative, out_latent, trim_latent, max(0, ref_motion_latent_length * 4 - 3), video_frame_offset + length)
|
||||
|
||||
|
||||
class WanAnimate2Cache(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id="WanAnimate2Cache",
|
||||
category="model/conditioning/wan/animate",
|
||||
description=(
|
||||
"Caches the pose-video's per-block activations so it runs once instead of on every sampling step. Roughly halves generation time "
|
||||
"Tradeoff is ~12.5 GB of system RAM at 480x832/81 frames in bf16 (scales with resolution and length). "
|
||||
"With context windows each window is cached separately, so RAM scales with the window count; use the static_standard schedule, as uniform schedules shift the windows every step and nothing ever recurs to hit the cache."
|
||||
),
|
||||
inputs=[
|
||||
io.Model.Input("model"),
|
||||
io.Combo.Input("device", options=["cpu", "gpu"], default="cpu",
|
||||
tooltip="Where to keep the cache. cpu (RAM) is the safe choice, the cache will not fit in VRAM alongside the model at typical sizes. gpu (VRAM) can be faster if it fits."),
|
||||
io.Combo.Input("dtype", options=["default", "int8", "int4"], default="default",
|
||||
tooltip="Storage precision. default stores the activations in the model's compute dtype. int8 halves the cache, int4 quarters it, convrot is used to retain accuracy."),
|
||||
],
|
||||
outputs=[io.Model.Output()],
|
||||
is_experimental=True,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, model, device, dtype="default") -> io.NodeOutput:
|
||||
store = comfy.model_management.get_torch_device() if device == "gpu" else torch.device("cpu")
|
||||
cache = comfy.ldm.wan.model_animate2.PoseBranchCache(store_device=store, dtype=dtype)
|
||||
m = model.clone()
|
||||
m.model_options["transformer_options"]["animate2_cache"] = cache
|
||||
m.add_callback(comfy.patcher_extension.CallbacksMP.ON_CLEANUP, lambda patcher: cache.free())
|
||||
return io.NodeOutput(m)
|
||||
|
||||
|
||||
class Wan22ImageToVideoLatent(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
|
|
@ -1475,6 +1638,8 @@ class WanExtension(ComfyExtension):
|
|||
WanSoundImageToVideoExtend,
|
||||
WanHuMoImageToVideo,
|
||||
WanAnimateToVideo,
|
||||
WanAnimate2ToVideo,
|
||||
WanAnimate2Cache,
|
||||
Wan22ImageToVideoLatent,
|
||||
WanInfiniteTalkToVideo,
|
||||
]
|
||||
|
|
|
|||
|
|
@ -1,3 +1,3 @@
|
|||
# This file is automatically generated by the build process when version is
|
||||
# updated in pyproject.toml.
|
||||
__version__ = "0.30.0"
|
||||
__version__ = "0.31.0"
|
||||
|
|
|
|||
20
nodes.py
20
nodes.py
|
|
@ -633,15 +633,18 @@ class DiffusersLoader:
|
|||
SEARCH_ALIASES = ["load diffusers model"]
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
def _model_paths(cls):
|
||||
paths = []
|
||||
for search_path in folder_paths.get_folder_paths("diffusers"):
|
||||
if os.path.exists(search_path):
|
||||
for root, subdir, files in os.walk(search_path, followlinks=True):
|
||||
if "model_index.json" in files:
|
||||
paths.append(os.path.relpath(root, start=search_path))
|
||||
return paths
|
||||
|
||||
return {"required": {"model_path": (paths,), }}
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {"required": {"model_path": (cls._model_paths(),), }}
|
||||
RETURN_TYPES = ("MODEL", "CLIP", "VAE")
|
||||
FUNCTION = "load_checkpoint"
|
||||
DEPRECATED = True
|
||||
|
|
@ -649,14 +652,20 @@ class DiffusersLoader:
|
|||
CATEGORY = "model/loaders"
|
||||
|
||||
def load_checkpoint(self, model_path, output_vae=True, output_clip=True):
|
||||
if model_path not in self._model_paths():
|
||||
raise ValueError(f"Invalid diffusers model path: {model_path!r}")
|
||||
|
||||
resolved_model_path = None
|
||||
for search_path in folder_paths.get_folder_paths("diffusers"):
|
||||
if os.path.exists(search_path):
|
||||
path = os.path.join(search_path, model_path)
|
||||
if os.path.exists(path):
|
||||
model_path = path
|
||||
if os.path.isfile(os.path.join(path, "model_index.json")):
|
||||
resolved_model_path = path
|
||||
break
|
||||
if resolved_model_path is None:
|
||||
raise FileNotFoundError(f"Diffusers model {model_path!r} not found.")
|
||||
|
||||
return comfy.diffusers_load.load_diffusers(model_path, output_vae=output_vae, output_clip=output_clip, embedding_directory=folder_paths.get_folder_paths("embeddings"))
|
||||
return comfy.diffusers_load.load_diffusers(resolved_model_path, output_vae=output_vae, output_clip=output_clip, embedding_directory=folder_paths.get_folder_paths("embeddings"))
|
||||
|
||||
|
||||
class unCLIPCheckpointLoader:
|
||||
|
|
@ -2492,6 +2501,7 @@ async def init_builtin_extra_nodes():
|
|||
"nodes_math.py",
|
||||
"nodes_number_convert.py",
|
||||
"nodes_painter.py",
|
||||
"nodes_compositor.py",
|
||||
"nodes_curve.py",
|
||||
"nodes_bg_removal.py",
|
||||
"nodes_rtdetr.py",
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
[project]
|
||||
name = "ComfyUI"
|
||||
version = "0.30.0"
|
||||
version = "0.31.0"
|
||||
readme = "README.md"
|
||||
license = { file = "LICENSE" }
|
||||
requires-python = ">=3.10"
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
comfyui-frontend-package==1.47.12
|
||||
comfyui-workflow-templates==0.11.27
|
||||
comfyui-frontend-package==1.48.7
|
||||
comfyui-workflow-templates==0.11.37
|
||||
comfyui-embedded-docs==0.5.9
|
||||
torch
|
||||
torchsde
|
||||
|
|
@ -22,8 +22,8 @@ alembic
|
|||
SQLAlchemy>=2.0.0
|
||||
filelock
|
||||
av>=16.0.0
|
||||
comfy-kitchen==0.2.26
|
||||
comfy-aimdo==0.4.11
|
||||
comfy-kitchen==0.2.28
|
||||
comfy-aimdo==0.4.13
|
||||
requests
|
||||
simpleeval>=1.0.0
|
||||
blake3
|
||||
|
|
|
|||
|
|
@ -0,0 +1,111 @@
|
|||
"""Regenerate ``compositor_blend_golden.json``.
|
||||
|
||||
The golden file is the *shared contract* for layer blending. Every
|
||||
implementation of these 26 modes must reproduce it within ``tolerance``:
|
||||
|
||||
* ``comfy_extras/compositor_blend.py`` - numpy, server-side compositing
|
||||
* ``layerBlend.frag`` - GLSL, the live preview in the layer editor
|
||||
* any future CPU reference in the frontend
|
||||
|
||||
Run from the repository root::
|
||||
|
||||
python tests-unit/comfy_extras_test/compositor_blend_fixture_gen.py
|
||||
|
||||
and review the diff. A change to this file is a change to user-visible
|
||||
blending behaviour in every implementation, so it should never be
|
||||
regenerated just to make a test pass.
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
|
||||
import numpy as np
|
||||
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", ".."))
|
||||
|
||||
from comfy_extras.compositor_blend import CHANNEL_BLEND, HSL_BLEND, blend_pixel # noqa: E402
|
||||
|
||||
GOLDEN_PATH = os.path.join(os.path.dirname(__file__), "compositor_blend_golden.json")
|
||||
|
||||
# Scalar grid for the per-channel modes: both endpoints, the midpoint, values
|
||||
# just inside each endpoint, and values inside the 1e-6 epsilon guards.
|
||||
SCALARS = [0.0, 1e-7, 0.001, 0.25, 0.5, 0.75, 0.999, 1.0 - 1e-7, 1.0]
|
||||
|
||||
# Colour pairs for the HSL modes, which read all three channels at once.
|
||||
COLORS = [
|
||||
[0.0, 0.0, 0.0],
|
||||
[1.0, 1.0, 1.0],
|
||||
[0.5, 0.5, 0.5],
|
||||
[1.0, 0.0, 0.0],
|
||||
[0.0, 0.0, 1.0],
|
||||
[0.2, 0.4, 0.6],
|
||||
[0.9, 0.1, 0.35],
|
||||
[1e-7, 1e-7, 1e-7],
|
||||
[1e-7, 0.0, 0.0],
|
||||
[0.05, 0.05, 0.05],
|
||||
]
|
||||
|
||||
|
||||
def _round(value) -> float:
|
||||
return round(float(value), 7)
|
||||
|
||||
|
||||
def build() -> dict:
|
||||
channel = {}
|
||||
for mode in CHANNEL_BLEND:
|
||||
rows = []
|
||||
for i in SCALARS:
|
||||
for l in SCALARS:
|
||||
out = blend_pixel(mode, np.float32([i] * 3), np.float32([l] * 3))
|
||||
rows.append([_round(i), _round(l), _round(np.asarray(out).reshape(3)[0])])
|
||||
channel[mode] = rows
|
||||
hsl = {}
|
||||
for mode in HSL_BLEND:
|
||||
rows = []
|
||||
for i in COLORS:
|
||||
for l in COLORS:
|
||||
out = blend_pixel(mode, np.float32(i), np.float32(l))
|
||||
rows.append([
|
||||
[_round(v) for v in i],
|
||||
[_round(v) for v in l],
|
||||
[_round(v) for v in np.asarray(out).reshape(3)],
|
||||
])
|
||||
hsl[mode] = rows
|
||||
return {
|
||||
"_comment": (
|
||||
"Golden blend values shared by comfy_extras/compositor_blend.py and "
|
||||
"layerBlend.frag. Inputs are unpremultiplied colours already in the "
|
||||
"blend space; outputs are unclamped (the compositor clamps once, at "
|
||||
"the end). 'channel' rows are [i, l, out] applied per channel; 'hsl' "
|
||||
"rows are [rgb_backdrop, rgb_layer, rgb_out]. Regenerate with "
|
||||
"tests-unit/comfy_extras_test/compositor_blend_fixture_gen.py."
|
||||
),
|
||||
"tolerance": 1e-4,
|
||||
"channel": channel,
|
||||
"hsl": hsl,
|
||||
}
|
||||
|
||||
|
||||
def dumps(data: dict) -> str:
|
||||
"""One row per line, so a behaviour change shows up as a readable diff."""
|
||||
lines = ["{", f' "_comment": {json.dumps(data["_comment"])},', f' "tolerance": {data["tolerance"]},']
|
||||
for section in ("channel", "hsl"):
|
||||
lines.append(f' "{section}": {{')
|
||||
modes = sorted(data[section])
|
||||
for m_index, mode in enumerate(modes):
|
||||
lines.append(f' "{mode}": [')
|
||||
rows = data[section][mode]
|
||||
for r_index, row in enumerate(rows):
|
||||
comma = "" if r_index == len(rows) - 1 else ","
|
||||
lines.append(f" {json.dumps(row)}{comma}")
|
||||
lines.append(" ]" + ("" if m_index == len(modes) - 1 else ","))
|
||||
lines.append(" }" + ("," if section == "channel" else ""))
|
||||
lines.append("}")
|
||||
return "\n".join(lines) + "\n"
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
with open(GOLDEN_PATH, "w") as handle:
|
||||
handle.write(dumps(build()))
|
||||
sys.stdout.write(f"wrote {GOLDEN_PATH}\n")
|
||||
File diff suppressed because it is too large
Load Diff
|
|
@ -0,0 +1,157 @@
|
|||
"""Blend-mode parity tests for the compositor.
|
||||
|
||||
The compositor blends in three places: this numpy module (server-side), the
|
||||
``layerBlend.frag`` GLSL shader (the live preview the user actually sees), and
|
||||
anything the frontend adds later. They have diverged before, silently, and the
|
||||
divergences only show up as "the render does not look like the preview".
|
||||
|
||||
``compositor_blend_golden.json`` is the shared contract. This file pins the
|
||||
numpy implementation to it and additionally spells out, by hand, the boundary
|
||||
rules that the epsilon guards exist to enforce - so a future refactor of
|
||||
``safe_div`` cannot quietly re-introduce the old behaviour by regenerating the
|
||||
fixture.
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from comfy_extras.compositor_blend import (
|
||||
CHANNEL_BLEND,
|
||||
HSL_BLEND,
|
||||
EffectiveMode,
|
||||
blend_composite,
|
||||
blend_pixel,
|
||||
resolve_mode,
|
||||
)
|
||||
|
||||
GOLDEN_PATH = os.path.join(os.path.dirname(__file__), "compositor_blend_golden.json")
|
||||
|
||||
with open(GOLDEN_PATH) as _handle:
|
||||
GOLDEN = json.load(_handle)
|
||||
|
||||
TOLERANCE = GOLDEN["tolerance"]
|
||||
|
||||
|
||||
def _blend(mode: str, i, l) -> np.ndarray:
|
||||
return np.asarray(
|
||||
blend_pixel(mode, np.float32(i), np.float32(l)), dtype=np.float64
|
||||
).reshape(3)
|
||||
|
||||
|
||||
def test_golden_covers_every_mode():
|
||||
"""A new blend mode must arrive with golden values, not silently."""
|
||||
assert set(GOLDEN["channel"]) == set(CHANNEL_BLEND)
|
||||
assert set(GOLDEN["hsl"]) == set(HSL_BLEND)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mode", sorted(CHANNEL_BLEND))
|
||||
def test_channel_modes_match_golden(mode):
|
||||
for i, l, expected in GOLDEN["channel"][mode]:
|
||||
actual = _blend(mode, [i] * 3, [l] * 3)
|
||||
assert actual == pytest.approx([expected] * 3, abs=TOLERANCE), (
|
||||
f"{mode}(i={i}, l={l}) -> {actual.tolist()}, golden {expected}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mode", sorted(HSL_BLEND))
|
||||
def test_hsl_modes_match_golden(mode):
|
||||
for i, l, expected in GOLDEN["hsl"][mode]:
|
||||
actual = _blend(mode, i, l)
|
||||
assert actual == pytest.approx(expected, abs=TOLERANCE), (
|
||||
f"{mode}(i={i}, l={l}) -> {actual.tolist()}, golden {expected}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mode", sorted(set(CHANNEL_BLEND) | set(HSL_BLEND)))
|
||||
def test_no_mode_produces_nan_or_inf(mode):
|
||||
edges = [0.0, 1e-7, 1e-6, 0.5, 1.0 - 1e-7, 1.0]
|
||||
for i in edges:
|
||||
for l in edges:
|
||||
out = _blend(mode, [i, 0.0, 1.0], [l, 1.0, 0.0])
|
||||
assert np.all(np.isfinite(out)), f"{mode}(i={i}, l={l}) -> {out.tolist()}"
|
||||
|
||||
|
||||
class TestBoundaryRules:
|
||||
"""The rules the epsilon guards encode, written out independently of the fixture."""
|
||||
|
||||
def test_color_dodge_full_layer_is_white_not_black(self):
|
||||
# Guarding the denominator returns 0 here, which reads as "the dodge
|
||||
# layer turned the image black" - the exact inversion CodeRabbit flagged.
|
||||
assert _blend("color-dodge", [0.5] * 3, [1.0] * 3) == pytest.approx([1.0] * 3)
|
||||
|
||||
def test_color_dodge_black_backdrop_stays_black(self):
|
||||
assert _blend("color-dodge", [0.0] * 3, [1.0] * 3) == pytest.approx([0.0] * 3)
|
||||
|
||||
def test_color_dodge_is_clamped(self):
|
||||
assert _blend("color-dodge", [0.6] * 3, [0.9] * 3) == pytest.approx([1.0] * 3)
|
||||
|
||||
def test_color_burn_empty_layer_is_black_not_white(self):
|
||||
assert _blend("color-burn", [0.5] * 3, [0.0] * 3) == pytest.approx([0.0] * 3)
|
||||
|
||||
def test_color_burn_white_backdrop_stays_white(self):
|
||||
assert _blend("color-burn", [1.0] * 3, [0.0] * 3) == pytest.approx([1.0] * 3)
|
||||
|
||||
def test_vivid_light_boundaries(self):
|
||||
assert _blend("vivid-light", [0.5] * 3, [0.0] * 3) == pytest.approx([0.0] * 3)
|
||||
assert _blend("vivid-light", [0.5] * 3, [1.0] * 3) == pytest.approx([1.0] * 3)
|
||||
assert _blend("vivid-light", [1.0] * 3, [0.0] * 3) == pytest.approx([1.0] * 3)
|
||||
assert _blend("vivid-light", [0.0] * 3, [1.0] * 3) == pytest.approx([0.0] * 3)
|
||||
|
||||
def test_divide_by_zero_is_clamped_to_one(self):
|
||||
assert _blend("divide", [0.5] * 3, [0.0] * 3) == pytest.approx([1.0] * 3)
|
||||
|
||||
def test_luminosity_over_black_takes_the_layer_luminance(self):
|
||||
# A luminosity layer over a black backdrop must not vanish. There is no
|
||||
# hue or saturation in the backdrop to preserve, so the result is a
|
||||
# neutral grey at the layer's luminance.
|
||||
assert _blend("luminosity", [0.0] * 3, [1.0] * 3) == pytest.approx([1.0] * 3)
|
||||
assert _blend("luminosity", [0.0] * 3, [0.5] * 3) == pytest.approx([0.5] * 3)
|
||||
|
||||
def test_luminosity_is_continuous_approaching_black(self):
|
||||
near = _blend("luminosity", [1e-7] * 3, [1.0] * 3)
|
||||
at = _blend("luminosity", [0.0] * 3, [1.0] * 3)
|
||||
assert near == pytest.approx(at, abs=TOLERANCE)
|
||||
|
||||
def test_luminosity_preserves_backdrop_chroma(self):
|
||||
out = _blend("luminosity", [0.4, 0.2, 0.1], [0.5] * 3)
|
||||
assert out[0] > out[1] > out[2]
|
||||
|
||||
|
||||
class TestCompositeAndModeTable:
|
||||
def test_unknown_blend_mode_falls_back_to_normal(self):
|
||||
unknown = resolve_mode("not-a-mode")
|
||||
assert (unknown.blend_space, unknown.composite) == (
|
||||
resolve_mode("normal").blend_space,
|
||||
resolve_mode("normal").composite,
|
||||
)
|
||||
assert _blend("not-a-mode", [0.1, 0.2, 0.3], [0.4, 0.5, 0.6]) == pytest.approx(
|
||||
_blend("normal", [0.1, 0.2, 0.3], [0.4, 0.5, 0.6])
|
||||
)
|
||||
|
||||
def test_every_blend_mode_has_a_composite_entry(self):
|
||||
for mode in set(CHANNEL_BLEND) | set(HSL_BLEND):
|
||||
resolved = resolve_mode(mode)
|
||||
assert isinstance(resolved, EffectiveMode)
|
||||
assert resolved.blend == mode
|
||||
assert resolved.blend_space in ("linear", "perceptual")
|
||||
assert resolved.composite in (
|
||||
"union",
|
||||
"clip-to-backdrop",
|
||||
"clip-to-layer",
|
||||
"intersection",
|
||||
)
|
||||
|
||||
def test_normal_over_transparent_backdrop_keeps_the_layer(self):
|
||||
backdrop = np.zeros((1, 1, 4), dtype=np.float32)
|
||||
layer = np.float32([[[0.25, 0.5, 0.75, 1.0]]])
|
||||
out = blend_composite(resolve_mode("normal"), backdrop, layer, 1.0)
|
||||
assert out[0, 0].tolist() == pytest.approx([0.25, 0.5, 0.75, 1.0])
|
||||
|
||||
def test_zero_opacity_is_a_no_op(self):
|
||||
backdrop = np.float32([[[0.1, 0.2, 0.3, 1.0]]])
|
||||
layer = np.float32([[[1.0, 1.0, 1.0, 1.0]]])
|
||||
out = blend_composite(resolve_mode("multiply"), backdrop, layer, 0.0)
|
||||
assert out[0, 0].tolist() == pytest.approx([0.1, 0.2, 0.3, 1.0])
|
||||
|
|
@ -0,0 +1,72 @@
|
|||
"""Regression tests for ImageCompositor's handling of untrusted layer state.
|
||||
|
||||
The compositor's `compositor` widget value is persisted into the saved workflow
|
||||
and is accepted verbatim on `POST /prompt`, so every field in it is untrusted
|
||||
input, not an internal invariant.
|
||||
"""
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from comfy_extras.nodes_compositor import (
|
||||
_layer_params,
|
||||
composite_from_state,
|
||||
expand_item_frames,
|
||||
state_from_items,
|
||||
)
|
||||
|
||||
|
||||
def _solid(color, w=4, h=4) -> torch.Tensor:
|
||||
frame = np.zeros((h, w, len(color)), dtype=np.float32)
|
||||
frame[:] = color
|
||||
return torch.from_numpy(frame).unsqueeze(0)
|
||||
|
||||
|
||||
class TestLayerOpacity:
|
||||
@pytest.mark.parametrize(
|
||||
("raw", "expected"),
|
||||
[(-0.5, 0.0), (0.0, 0.0), (0.25, 0.25), (1.0, 1.0), (3.0, 1.0)],
|
||||
)
|
||||
def test_opacity_is_clamped(self, raw, expected):
|
||||
assert _layer_params({"opacity": raw}, 4, 4)["opacity"] == expected
|
||||
|
||||
def test_opacity_defaults_to_opaque(self):
|
||||
assert _layer_params({}, 4, 4)["opacity"] == 1.0
|
||||
|
||||
def test_out_of_range_opacity_does_not_leak_into_the_next_layer(self):
|
||||
# The canvas is only clamped once, after every layer has been composited,
|
||||
# so an out-of-range coverage multiplier on one layer changes the *blend*
|
||||
# of the layer above it. White at opacity 3.0 over black leaves the canvas
|
||||
# at 3.0; the multiply above it then reads 3.0 as its backdrop and the
|
||||
# result is visibly lighter than the same stack at opacity 1.0.
|
||||
def run(opacity):
|
||||
state = {
|
||||
"canvas": (2, 2),
|
||||
"layers": [{"opacity": opacity}, {"opacity": 1.0, "blend": "multiply"}],
|
||||
"inputs": None,
|
||||
"background": {"color": "#000000", "opacity": 1.0, "visible": True},
|
||||
"order": None,
|
||||
}
|
||||
tensors = [_solid([1.0, 1.0, 1.0], 2, 2), _solid([0.5, 0.5, 0.5], 2, 2)]
|
||||
return composite_from_state(tensors, state, [None, None])[0, 0, 0, :3]
|
||||
|
||||
assert run(3.0).tolist() == pytest.approx(run(1.0).tolist(), abs=1e-6)
|
||||
|
||||
|
||||
class TestGraphOnlyBackground:
|
||||
def test_default_layout_background_is_hidden(self):
|
||||
# A visible white background here would make every graph-only run emit a
|
||||
# white matte instead of transparency.
|
||||
frames = expand_item_frames([{"image": _solid([1.0, 0.0, 0.0])}])
|
||||
state = state_from_items(frames, (4, 4))
|
||||
assert state["background"]["visible"] is False
|
||||
|
||||
def test_uncovered_canvas_stays_transparent(self):
|
||||
tensors = [_solid([1.0, 0.0, 0.0], w=2, h=2)]
|
||||
frames = expand_item_frames([{"image": tensors[0]}])
|
||||
state = state_from_items(frames, (4, 4))
|
||||
out = composite_from_state(tensors, state, [None])[0]
|
||||
assert out.shape[-1] == 4
|
||||
assert float(out[0, 0, 3]) == pytest.approx(1.0)
|
||||
assert float(out[3, 3, 3]) == pytest.approx(0.0)
|
||||
|
|
@ -10,6 +10,7 @@ from comfy_execution.jobs import (
|
|||
normalize_output_item,
|
||||
normalize_outputs,
|
||||
get_outputs_summary,
|
||||
count_previewable_outputs,
|
||||
apply_sorting,
|
||||
has_3d_extension,
|
||||
validate_job_id,
|
||||
|
|
@ -361,6 +362,79 @@ class TestGetOutputsSummary:
|
|||
assert preview['mediaType'] == 'files'
|
||||
|
||||
|
||||
class TestCountPreviewableOutputs:
|
||||
"""Unit tests for count_previewable_outputs()
|
||||
|
||||
Kept separate from get_outputs_summary()'s outputs_count: the Media Assets
|
||||
badge should reflect only what the expanded asset view actually renders
|
||||
(previewable outputs), while outputs_count keeps counting every output
|
||||
item for other consumers.
|
||||
"""
|
||||
|
||||
def test_empty_outputs(self):
|
||||
assert count_previewable_outputs({}) == 0
|
||||
|
||||
def test_previewable_outputs_all_counted(self):
|
||||
"""When every output is previewable, the two counts should match."""
|
||||
outputs = {
|
||||
'node1': {'images': [{'filename': 'a.png', 'type': 'output'}]},
|
||||
'node2': {'images': [{'filename': 'b.png', 'type': 'output'}]},
|
||||
}
|
||||
outputs_count, _ = get_outputs_summary(outputs)
|
||||
assert count_previewable_outputs(outputs) == outputs_count == 2
|
||||
|
||||
def test_save_latent_counted_but_not_previewable(self):
|
||||
"""SaveLatent (nodes.py) emits a real saved file under the 'latents'
|
||||
media type: {'latents': [{'filename': '..._00001_.latent',
|
||||
'subfolder': '', 'type': 'output'}]}. It has no previewable media
|
||||
type, format, or extension, so it inflates outputs_count without
|
||||
ever rendering in the expanded asset view."""
|
||||
outputs = {
|
||||
'node1': {
|
||||
'images': [{'filename': 'ComfyUI_00001_.png', 'subfolder': '', 'type': 'output'}]
|
||||
},
|
||||
'node2': {
|
||||
'latents': [{'filename': 'ComfyUI_00001_.latent', 'subfolder': '', 'type': 'output'}]
|
||||
},
|
||||
}
|
||||
outputs_count, _ = get_outputs_summary(outputs)
|
||||
assert outputs_count == 2
|
||||
assert count_previewable_outputs(outputs) == 1
|
||||
|
||||
def test_save_text_file_output_is_previewable_by_extension(self):
|
||||
"""SaveText (comfy_extras/nodes_text.py) emits its saved file under a
|
||||
'files' media type via ui.SavedResult: {'files': [{'filename':
|
||||
'..._00001.txt', 'subfolder': ..., 'type': 'output'}]}. The .txt
|
||||
extension makes it previewable even though 'files' itself isn't a
|
||||
previewable media type."""
|
||||
outputs = {
|
||||
'node1': {
|
||||
'files': [{'filename': 'ComfyUI_00001.txt', 'subfolder': '', 'type': 'output'}]
|
||||
}
|
||||
}
|
||||
assert count_previewable_outputs(outputs) == 1
|
||||
|
||||
def test_preview_any_text_tuple_not_counted(self):
|
||||
"""PreviewAny (comfy_extras/nodes_preview_any.py) emits only
|
||||
{'text': (value,)} with no saved file. Since the value is a tuple,
|
||||
not a list, it is excluded from both outputs_count and
|
||||
previewable_outputs_count — matching get_outputs_summary()."""
|
||||
outputs = {
|
||||
'node1': {'text': ('some previewed value',)}
|
||||
}
|
||||
outputs_count, _ = get_outputs_summary(outputs)
|
||||
assert outputs_count == 0
|
||||
assert count_previewable_outputs(outputs) == 0
|
||||
|
||||
def test_string_3d_filename_previewable(self):
|
||||
"""String 3D filenames (e.g. Preview3D) normalize into a previewable
|
||||
item just like they do for outputs_count."""
|
||||
outputs = {
|
||||
'node1': {'result': ['preview3d_abc123.glb', None]}
|
||||
}
|
||||
assert count_previewable_outputs(outputs) == 1
|
||||
|
||||
|
||||
class TestHas3DExtension:
|
||||
"""Unit tests for has_3d_extension()"""
|
||||
|
||||
|
|
@ -447,6 +521,7 @@ class TestNormalizeQueueItem:
|
|||
assert 'execution_error' not in job
|
||||
assert 'preview_output' not in job
|
||||
assert job['outputs_count'] == 0
|
||||
assert job['previewable_outputs_count'] == 0
|
||||
assert job['workflow_id'] == 'workflow-abc'
|
||||
|
||||
|
||||
|
|
@ -635,6 +710,54 @@ class TestNormalizeHistoryItem:
|
|||
{'filename': 'photo.png', 'type': 'output', 'subfolder': ''},
|
||||
]
|
||||
|
||||
def test_previewable_outputs_count_excludes_non_previewable_outputs(self):
|
||||
"""Regression test for the Media Assets badge overcount: a job with an
|
||||
image (SaveImage) and a SaveLatent output should report previewable_
|
||||
outputs_count == 1 while outputs_count == 2, so the frontend badge
|
||||
(once switched to previewable_outputs_count) matches what the
|
||||
expanded asset view actually renders."""
|
||||
history_item = {
|
||||
'prompt': (
|
||||
5,
|
||||
'prompt-mixed',
|
||||
{'nodes': {}},
|
||||
{'create_time': 1234567890},
|
||||
['node1', 'node2'],
|
||||
),
|
||||
'status': {'status_str': 'success', 'completed': True, 'messages': []},
|
||||
'outputs': {
|
||||
'node1': {
|
||||
'images': [{'filename': 'ComfyUI_00001_.png', 'subfolder': '', 'type': 'output'}]
|
||||
},
|
||||
'node2': {
|
||||
'latents': [{'filename': 'ComfyUI_00001_.latent', 'subfolder': '', 'type': 'output'}]
|
||||
},
|
||||
},
|
||||
}
|
||||
job = normalize_history_item('prompt-mixed', history_item)
|
||||
|
||||
assert job['outputs_count'] == 2
|
||||
assert job['previewable_outputs_count'] == 1
|
||||
|
||||
def test_previewable_outputs_count_zero_pruned_by_prune_dict(self):
|
||||
"""A job with no outputs at all should still report both counts as 0,
|
||||
not omit the field (prune_dict only strips None, not 0)."""
|
||||
history_item = {
|
||||
'prompt': (
|
||||
5,
|
||||
'prompt-empty',
|
||||
{'nodes': {}},
|
||||
{'create_time': 1234567890},
|
||||
['node1'],
|
||||
),
|
||||
'status': {'status_str': 'success', 'completed': True, 'messages': []},
|
||||
'outputs': {},
|
||||
}
|
||||
job = normalize_history_item('prompt-empty', history_item)
|
||||
|
||||
assert job['outputs_count'] == 0
|
||||
assert job['previewable_outputs_count'] == 0
|
||||
|
||||
|
||||
class TestNormalizeOutputItem:
|
||||
"""Unit tests for normalize_output_item()"""
|
||||
|
|
|
|||
Loading…
Reference in New Issue