Delete comfy

comfy/checkpoint_pickle.py

@@ -1,13 +0,0 @@
-import pickle
-
-load = pickle.load
-
-class Empty:
-    pass
-
-class Unpickler(pickle.Unpickler):
-    def find_class(self, module, name):
-        #TODO: safe unpickle
-        if module.startswith("pytorch_lightning"):
-            return Empty
-        return super().find_class(module, name)

comfy/cldm/cldm.py

@@ -1,433 +0,0 @@
-#taken from: https://github.com/lllyasviel/ControlNet
-#and modified
-
-import torch
-import torch.nn as nn
-
-from ..ldm.modules.diffusionmodules.util import (
-    timestep_embedding,
-)
-
-from ..ldm.modules.attention import SpatialTransformer
-from ..ldm.modules.diffusionmodules.openaimodel import UNetModel, TimestepEmbedSequential, ResBlock, Downsample
-from ..ldm.util import exists
-from .control_types import UNION_CONTROLNET_TYPES
-from collections import OrderedDict
-import comfy.ops
-from comfy.ldm.modules.attention import optimized_attention
-
-class OptimizedAttention(nn.Module):
-    def __init__(self, c, nhead, dropout=0.0, dtype=None, device=None, operations=None):
-        super().__init__()
-        self.heads = nhead
-        self.c = c
-
-        self.in_proj = operations.Linear(c, c * 3, bias=True, dtype=dtype, device=device)
-        self.out_proj = operations.Linear(c, c, bias=True, dtype=dtype, device=device)
-
-    def forward(self, x):
-        x = self.in_proj(x)
-        q, k, v = x.split(self.c, dim=2)
-        out = optimized_attention(q, k, v, self.heads)
-        return self.out_proj(out)
-
-class QuickGELU(nn.Module):
-    def forward(self, x: torch.Tensor):
-        return x * torch.sigmoid(1.702 * x)
-
-class ResBlockUnionControlnet(nn.Module):
-    def __init__(self, dim, nhead, dtype=None, device=None, operations=None):
-        super().__init__()
-        self.attn = OptimizedAttention(dim, nhead, dtype=dtype, device=device, operations=operations)
-        self.ln_1 = operations.LayerNorm(dim, dtype=dtype, device=device)
-        self.mlp = nn.Sequential(
-            OrderedDict([("c_fc", operations.Linear(dim, dim * 4, dtype=dtype, device=device)), ("gelu", QuickGELU()),
-                         ("c_proj", operations.Linear(dim * 4, dim, dtype=dtype, device=device))]))
-        self.ln_2 = operations.LayerNorm(dim, dtype=dtype, device=device)
-
-    def attention(self, x: torch.Tensor):
-        return self.attn(x)
-
-    def forward(self, x: torch.Tensor):
-        x = x + self.attention(self.ln_1(x))
-        x = x + self.mlp(self.ln_2(x))
-        return x
-
-class ControlledUnetModel(UNetModel):
-    #implemented in the ldm unet
-    pass
-
-class ControlNet(nn.Module):
-    def __init__(
-        self,
-        image_size,
-        in_channels,
-        model_channels,
-        hint_channels,
-        num_res_blocks,
-        dropout=0,
-        channel_mult=(1, 2, 4, 8),
-        conv_resample=True,
-        dims=2,
-        num_classes=None,
-        use_checkpoint=False,
-        dtype=torch.float32,
-        num_heads=-1,
-        num_head_channels=-1,
-        num_heads_upsample=-1,
-        use_scale_shift_norm=False,
-        resblock_updown=False,
-        use_new_attention_order=False,
-        use_spatial_transformer=False,    # custom transformer support
-        transformer_depth=1,              # custom transformer support
-        context_dim=None,                 # custom transformer support
-        n_embed=None,                     # custom support for prediction of discrete ids into codebook of first stage vq model
-        legacy=True,
-        disable_self_attentions=None,
-        num_attention_blocks=None,
-        disable_middle_self_attn=False,
-        use_linear_in_transformer=False,
-        adm_in_channels=None,
-        transformer_depth_middle=None,
-        transformer_depth_output=None,
-        attn_precision=None,
-        union_controlnet_num_control_type=None,
-        device=None,
-        operations=comfy.ops.disable_weight_init,
-        **kwargs,
-    ):
-        super().__init__()
-        assert use_spatial_transformer == True, "use_spatial_transformer has to be true"
-        if use_spatial_transformer:
-            assert context_dim is not None, 'Fool!! You forgot to include the dimension of your cross-attention conditioning...'
-
-        if context_dim is not None:
-            assert use_spatial_transformer, 'Fool!! You forgot to use the spatial transformer for your cross-attention conditioning...'
-            # from omegaconf.listconfig import ListConfig
-            # if type(context_dim) == ListConfig:
-            #     context_dim = list(context_dim)
-
-        if num_heads_upsample == -1:
-            num_heads_upsample = num_heads
-
-        if num_heads == -1:
-            assert num_head_channels != -1, 'Either num_heads or num_head_channels has to be set'
-
-        if num_head_channels == -1:
-            assert num_heads != -1, 'Either num_heads or num_head_channels has to be set'
-
-        self.dims = dims
-        self.image_size = image_size
-        self.in_channels = in_channels
-        self.model_channels = model_channels
-
-        if isinstance(num_res_blocks, int):
-            self.num_res_blocks = len(channel_mult) * [num_res_blocks]
-        else:
-            if len(num_res_blocks) != len(channel_mult):
-                raise ValueError("provide num_res_blocks either as an int (globally constant) or "
-                                 "as a list/tuple (per-level) with the same length as channel_mult")
-            self.num_res_blocks = num_res_blocks
-
-        if disable_self_attentions is not None:
-            # should be a list of booleans, indicating whether to disable self-attention in TransformerBlocks or not
-            assert len(disable_self_attentions) == len(channel_mult)
-        if num_attention_blocks is not None:
-            assert len(num_attention_blocks) == len(self.num_res_blocks)
-            assert all(map(lambda i: self.num_res_blocks[i] >= num_attention_blocks[i], range(len(num_attention_blocks))))
-
-        transformer_depth = transformer_depth[:]
-
-        self.dropout = dropout
-        self.channel_mult = channel_mult
-        self.conv_resample = conv_resample
-        self.num_classes = num_classes
-        self.use_checkpoint = use_checkpoint
-        self.dtype = dtype
-        self.num_heads = num_heads
-        self.num_head_channels = num_head_channels
-        self.num_heads_upsample = num_heads_upsample
-        self.predict_codebook_ids = n_embed is not None
-
-        time_embed_dim = model_channels * 4
-        self.time_embed = nn.Sequential(
-            operations.Linear(model_channels, time_embed_dim, dtype=self.dtype, device=device),
-            nn.SiLU(),
-            operations.Linear(time_embed_dim, time_embed_dim, dtype=self.dtype, device=device),
-        )
-
-        if self.num_classes is not None:
-            if isinstance(self.num_classes, int):
-                self.label_emb = nn.Embedding(num_classes, time_embed_dim)
-            elif self.num_classes == "continuous":
-                self.label_emb = nn.Linear(1, time_embed_dim)
-            elif self.num_classes == "sequential":
-                assert adm_in_channels is not None
-                self.label_emb = nn.Sequential(
-                    nn.Sequential(
-                        operations.Linear(adm_in_channels, time_embed_dim, dtype=self.dtype, device=device),
-                        nn.SiLU(),
-                        operations.Linear(time_embed_dim, time_embed_dim, dtype=self.dtype, device=device),
-                    )
-                )
-            else:
-                raise ValueError()
-
-        self.input_blocks = nn.ModuleList(
-            [
-                TimestepEmbedSequential(
-                    operations.conv_nd(dims, in_channels, model_channels, 3, padding=1, dtype=self.dtype, device=device)
-                )
-            ]
-        )
-        self.zero_convs = nn.ModuleList([self.make_zero_conv(model_channels, operations=operations, dtype=self.dtype, device=device)])
-
-        self.input_hint_block = TimestepEmbedSequential(
-                    operations.conv_nd(dims, hint_channels, 16, 3, padding=1, dtype=self.dtype, device=device),
-                    nn.SiLU(),
-                    operations.conv_nd(dims, 16, 16, 3, padding=1, dtype=self.dtype, device=device),
-                    nn.SiLU(),
-                    operations.conv_nd(dims, 16, 32, 3, padding=1, stride=2, dtype=self.dtype, device=device),
-                    nn.SiLU(),
-                    operations.conv_nd(dims, 32, 32, 3, padding=1, dtype=self.dtype, device=device),
-                    nn.SiLU(),
-                    operations.conv_nd(dims, 32, 96, 3, padding=1, stride=2, dtype=self.dtype, device=device),
-                    nn.SiLU(),
-                    operations.conv_nd(dims, 96, 96, 3, padding=1, dtype=self.dtype, device=device),
-                    nn.SiLU(),
-                    operations.conv_nd(dims, 96, 256, 3, padding=1, stride=2, dtype=self.dtype, device=device),
-                    nn.SiLU(),
-                    operations.conv_nd(dims, 256, model_channels, 3, padding=1, dtype=self.dtype, device=device)
-        )
-
-        self._feature_size = model_channels
-        input_block_chans = [model_channels]
-        ch = model_channels
-        ds = 1
-        for level, mult in enumerate(channel_mult):
-            for nr in range(self.num_res_blocks[level]):
-                layers = [
-                    ResBlock(
-                        ch,
-                        time_embed_dim,
-                        dropout,
-                        out_channels=mult * model_channels,
-                        dims=dims,
-                        use_checkpoint=use_checkpoint,
-                        use_scale_shift_norm=use_scale_shift_norm,
-                        dtype=self.dtype,
-                        device=device,
-                        operations=operations,
-                    )
-                ]
-                ch = mult * model_channels
-                num_transformers = transformer_depth.pop(0)
-                if num_transformers > 0:
-                    if num_head_channels == -1:
-                        dim_head = ch // num_heads
-                    else:
-                        num_heads = ch // num_head_channels
-                        dim_head = num_head_channels
-                    if legacy:
-                        #num_heads = 1
-                        dim_head = ch // num_heads if use_spatial_transformer else num_head_channels
-                    if exists(disable_self_attentions):
-                        disabled_sa = disable_self_attentions[level]
-                    else:
-                        disabled_sa = False
-
-                    if not exists(num_attention_blocks) or nr < num_attention_blocks[level]:
-                        layers.append(
-                            SpatialTransformer(
-                                ch, num_heads, dim_head, depth=num_transformers, context_dim=context_dim,
-                                disable_self_attn=disabled_sa, use_linear=use_linear_in_transformer,
-                                use_checkpoint=use_checkpoint, attn_precision=attn_precision, dtype=self.dtype, device=device, operations=operations
-                            )
-                        )
-                self.input_blocks.append(TimestepEmbedSequential(*layers))
-                self.zero_convs.append(self.make_zero_conv(ch, operations=operations, dtype=self.dtype, device=device))
-                self._feature_size += ch
-                input_block_chans.append(ch)
-            if level != len(channel_mult) - 1:
-                out_ch = ch
-                self.input_blocks.append(
-                    TimestepEmbedSequential(
-                        ResBlock(
-                            ch,
-                            time_embed_dim,
-                            dropout,
-                            out_channels=out_ch,
-                            dims=dims,
-                            use_checkpoint=use_checkpoint,
-                            use_scale_shift_norm=use_scale_shift_norm,
-                            down=True,
-                            dtype=self.dtype,
-                            device=device,
-                            operations=operations
-                        )
-                        if resblock_updown
-                        else Downsample(
-                            ch, conv_resample, dims=dims, out_channels=out_ch, dtype=self.dtype, device=device, operations=operations
-                        )
-                    )
-                )
-                ch = out_ch
-                input_block_chans.append(ch)
-                self.zero_convs.append(self.make_zero_conv(ch, operations=operations, dtype=self.dtype, device=device))
-                ds *= 2
-                self._feature_size += ch
-
-        if num_head_channels == -1:
-            dim_head = ch // num_heads
-        else:
-            num_heads = ch // num_head_channels
-            dim_head = num_head_channels
-        if legacy:
-            #num_heads = 1
-            dim_head = ch // num_heads if use_spatial_transformer else num_head_channels
-        mid_block = [
-            ResBlock(
-                ch,
-                time_embed_dim,
-                dropout,
-                dims=dims,
-                use_checkpoint=use_checkpoint,
-                use_scale_shift_norm=use_scale_shift_norm,
-                dtype=self.dtype,
-                device=device,
-                operations=operations
-            )]
-        if transformer_depth_middle >= 0:
-            mid_block += [SpatialTransformer(  # always uses a self-attn
-                            ch, num_heads, dim_head, depth=transformer_depth_middle, context_dim=context_dim,
-                            disable_self_attn=disable_middle_self_attn, use_linear=use_linear_in_transformer,
-                            use_checkpoint=use_checkpoint, attn_precision=attn_precision, dtype=self.dtype, device=device, operations=operations
-                        ),
-            ResBlock(
-                ch,
-                time_embed_dim,
-                dropout,
-                dims=dims,
-                use_checkpoint=use_checkpoint,
-                use_scale_shift_norm=use_scale_shift_norm,
-                dtype=self.dtype,
-                device=device,
-                operations=operations
-            )]
-        self.middle_block = TimestepEmbedSequential(*mid_block)
-        self.middle_block_out = self.make_zero_conv(ch, operations=operations, dtype=self.dtype, device=device)
-        self._feature_size += ch
-
-        if union_controlnet_num_control_type is not None:
-            self.num_control_type = union_controlnet_num_control_type
-            num_trans_channel = 320
-            num_trans_head = 8
-            num_trans_layer = 1
-            num_proj_channel = 320
-            # task_scale_factor = num_trans_channel ** 0.5
-            self.task_embedding = nn.Parameter(torch.empty(self.num_control_type, num_trans_channel, dtype=self.dtype, device=device))
-
-            self.transformer_layes = nn.Sequential(*[ResBlockUnionControlnet(num_trans_channel, num_trans_head, dtype=self.dtype, device=device, operations=operations) for _ in range(num_trans_layer)])
-            self.spatial_ch_projs = operations.Linear(num_trans_channel, num_proj_channel, dtype=self.dtype, device=device)
-            #-----------------------------------------------------------------------------------------------------
-
-            control_add_embed_dim = 256
-            class ControlAddEmbedding(nn.Module):
-                def __init__(self, in_dim, out_dim, num_control_type, dtype=None, device=None, operations=None):
-                    super().__init__()
-                    self.num_control_type = num_control_type
-                    self.in_dim = in_dim
-                    self.linear_1 = operations.Linear(in_dim * num_control_type, out_dim, dtype=dtype, device=device)
-                    self.linear_2 = operations.Linear(out_dim, out_dim, dtype=dtype, device=device)
-                def forward(self, control_type, dtype, device):
-                    c_type = torch.zeros((self.num_control_type,), device=device)
-                    c_type[control_type] = 1.0
-                    c_type = timestep_embedding(c_type.flatten(), self.in_dim, repeat_only=False).to(dtype).reshape((-1, self.num_control_type * self.in_dim))
-                    return self.linear_2(torch.nn.functional.silu(self.linear_1(c_type)))
-
-            self.control_add_embedding = ControlAddEmbedding(control_add_embed_dim, time_embed_dim, self.num_control_type, dtype=self.dtype, device=device, operations=operations)
-        else:
-            self.task_embedding = None
-            self.control_add_embedding = None
-
-    def union_controlnet_merge(self, hint, control_type, emb, context):
-        # Equivalent to: https://github.com/xinsir6/ControlNetPlus/tree/main
-        inputs = []
-        condition_list = []
-
-        for idx in range(min(1, len(control_type))):
-            controlnet_cond = self.input_hint_block(hint[idx], emb, context)
-            feat_seq = torch.mean(controlnet_cond, dim=(2, 3))
-            if idx < len(control_type):
-                feat_seq += self.task_embedding[control_type[idx]].to(dtype=feat_seq.dtype, device=feat_seq.device)
-
-            inputs.append(feat_seq.unsqueeze(1))
-            condition_list.append(controlnet_cond)
-
-        x = torch.cat(inputs, dim=1)
-        x = self.transformer_layes(x)
-        controlnet_cond_fuser = None
-        for idx in range(len(control_type)):
-            alpha = self.spatial_ch_projs(x[:, idx])
-            alpha = alpha.unsqueeze(-1).unsqueeze(-1)
-            o = condition_list[idx] + alpha
-            if controlnet_cond_fuser is None:
-                controlnet_cond_fuser = o
-            else:
-                controlnet_cond_fuser += o
-        return controlnet_cond_fuser
-
-    def make_zero_conv(self, channels, operations=None, dtype=None, device=None):
-        return TimestepEmbedSequential(operations.conv_nd(self.dims, channels, channels, 1, padding=0, dtype=dtype, device=device))
-
-    def forward(self, x, hint, timesteps, context, y=None, **kwargs):
-        t_emb = timestep_embedding(timesteps, self.model_channels, repeat_only=False).to(x.dtype)
-        emb = self.time_embed(t_emb)
-
-        guided_hint = None
-        if self.control_add_embedding is not None: #Union Controlnet
-            control_type = kwargs.get("control_type", [])
-
-            if any([c >= self.num_control_type for c in control_type]):
-                max_type = max(control_type)
-                max_type_name = {
-                    v: k for k, v in UNION_CONTROLNET_TYPES.items()
-                }[max_type]
-                raise ValueError(
-                    f"Control type {max_type_name}({max_type}) is out of range for the number of control types" +
-                    f"({self.num_control_type}) supported.\n" +
-                    "Please consider using the ProMax ControlNet Union model.\n" +
-                    "https://huggingface.co/xinsir/controlnet-union-sdxl-1.0/tree/main"
-                )
-
-            emb += self.control_add_embedding(control_type, emb.dtype, emb.device)
-            if len(control_type) > 0:
-                if len(hint.shape) < 5:
-                    hint = hint.unsqueeze(dim=0)
-                guided_hint = self.union_controlnet_merge(hint, control_type, emb, context)
-
-        if guided_hint is None:
-            guided_hint = self.input_hint_block(hint, emb, context)
-
-        out_output = []
-        out_middle = []
-
-        if self.num_classes is not None:
-            assert y.shape[0] == x.shape[0]
-            emb = emb + self.label_emb(y)
-
-        h = x
-        for module, zero_conv in zip(self.input_blocks, self.zero_convs):
-            if guided_hint is not None:
-                h = module(h, emb, context)
-                h += guided_hint
-                guided_hint = None
-            else:
-                h = module(h, emb, context)
-            out_output.append(zero_conv(h, emb, context))
-
-        h = self.middle_block(h, emb, context)
-        out_middle.append(self.middle_block_out(h, emb, context))
-
-        return {"middle": out_middle, "output": out_output}
-

comfy/cldm/control_types.py

@@ -1,10 +0,0 @@
-UNION_CONTROLNET_TYPES = {
-    "openpose": 0,
-    "depth": 1,
-    "hed/pidi/scribble/ted": 2,
-    "canny/lineart/anime_lineart/mlsd": 3,
-    "normal": 4,
-    "segment": 5,
-    "tile": 6,
-    "repaint": 7,
-}

comfy/cldm/dit_embedder.py

@@ -1,120 +0,0 @@
-import math
-from typing import List, Optional, Tuple
-
-import torch
-import torch.nn as nn
-from torch import Tensor
-
-from comfy.ldm.modules.diffusionmodules.mmdit import DismantledBlock, PatchEmbed, VectorEmbedder, TimestepEmbedder, get_2d_sincos_pos_embed_torch
-
-
-class ControlNetEmbedder(nn.Module):
-
-    def __init__(
-        self,
-        img_size: int,
-        patch_size: int,
-        in_chans: int,
-        attention_head_dim: int,
-        num_attention_heads: int,
-        adm_in_channels: int,
-        num_layers: int,
-        main_model_double: int,
-        double_y_emb: bool,
-        device: torch.device,
-        dtype: torch.dtype,
-        pos_embed_max_size: Optional[int] = None,
-        operations = None,
-    ):
-        super().__init__()
-        self.main_model_double = main_model_double
-        self.dtype = dtype
-        self.hidden_size = num_attention_heads * attention_head_dim
-        self.patch_size = patch_size
-        self.x_embedder = PatchEmbed(
-            img_size=img_size,
-            patch_size=patch_size,
-            in_chans=in_chans,
-            embed_dim=self.hidden_size,
-            strict_img_size=pos_embed_max_size is None,
-            device=device,
-            dtype=dtype,
-            operations=operations,
-        )
-
-        self.t_embedder = TimestepEmbedder(self.hidden_size, dtype=dtype, device=device, operations=operations)
-
-        self.double_y_emb = double_y_emb
-        if self.double_y_emb:
-            self.orig_y_embedder = VectorEmbedder(
-                adm_in_channels, self.hidden_size, dtype, device, operations=operations
-            )
-            self.y_embedder = VectorEmbedder(
-                self.hidden_size, self.hidden_size, dtype, device, operations=operations
-            )
-        else:
-            self.y_embedder = VectorEmbedder(
-                adm_in_channels, self.hidden_size, dtype, device, operations=operations
-            )
-
-        self.transformer_blocks = nn.ModuleList(
-            DismantledBlock(
-                hidden_size=self.hidden_size, num_heads=num_attention_heads, qkv_bias=True,
-                dtype=dtype, device=device, operations=operations
-            )
-            for _ in range(num_layers)
-        )
-
-        # self.use_y_embedder = pooled_projection_dim != self.time_text_embed.text_embedder.linear_1.in_features
-        # TODO double check this logic when 8b
-        self.use_y_embedder = True
-
-        self.controlnet_blocks = nn.ModuleList([])
-        for _ in range(len(self.transformer_blocks)):
-            controlnet_block = operations.Linear(self.hidden_size, self.hidden_size, dtype=dtype, device=device)
-            self.controlnet_blocks.append(controlnet_block)
-
-        self.pos_embed_input = PatchEmbed(
-            img_size=img_size,
-            patch_size=patch_size,
-            in_chans=in_chans,
-            embed_dim=self.hidden_size,
-            strict_img_size=False,
-            device=device,
-            dtype=dtype,
-            operations=operations,
-        )
-
-    def forward(
-        self,
-        x: torch.Tensor,
-        timesteps: torch.Tensor,
-        y: Optional[torch.Tensor] = None,
-        context: Optional[torch.Tensor] = None,
-        hint = None,
-    ) -> Tuple[Tensor, List[Tensor]]:
-        x_shape = list(x.shape)
-        x = self.x_embedder(x)
-        if not self.double_y_emb:
-            h = (x_shape[-2] + 1) // self.patch_size
-            w = (x_shape[-1] + 1) // self.patch_size
-            x += get_2d_sincos_pos_embed_torch(self.hidden_size, w, h, device=x.device)
-        c = self.t_embedder(timesteps, dtype=x.dtype)
-        if y is not None and self.y_embedder is not None:
-            if self.double_y_emb:
-                y = self.orig_y_embedder(y)
-            y = self.y_embedder(y)
-            c = c + y
-
-        x = x + self.pos_embed_input(hint)
-
-        block_out = ()
-
-        repeat = math.ceil(self.main_model_double / len(self.transformer_blocks))
-        for i in range(len(self.transformer_blocks)):
-            out = self.transformer_blocks[i](x, c)
-            if not self.double_y_emb:
-                x = out
-            block_out += (self.controlnet_blocks[i](out),) * repeat
-
-        return {"output": block_out}

comfy/cldm/mmdit.py

@@ -1,81 +0,0 @@
-import torch
-from typing import Optional
-import comfy.ldm.modules.diffusionmodules.mmdit
-
-class ControlNet(comfy.ldm.modules.diffusionmodules.mmdit.MMDiT):
-    def __init__(
-        self,
-        num_blocks = None,
-        control_latent_channels = None,
-        dtype = None,
-        device = None,
-        operations = None,
-        **kwargs,
-    ):
-        super().__init__(dtype=dtype, device=device, operations=operations, final_layer=False, num_blocks=num_blocks, **kwargs)
-        # controlnet_blocks
-        self.controlnet_blocks = torch.nn.ModuleList([])
-        for _ in range(len(self.joint_blocks)):
-            self.controlnet_blocks.append(operations.Linear(self.hidden_size, self.hidden_size, device=device, dtype=dtype))
-
-        if control_latent_channels is None:
-            control_latent_channels = self.in_channels
-
-        self.pos_embed_input = comfy.ldm.modules.diffusionmodules.mmdit.PatchEmbed(
-            None,
-            self.patch_size,
-            control_latent_channels,
-            self.hidden_size,
-            bias=True,
-            strict_img_size=False,
-            dtype=dtype,
-            device=device,
-            operations=operations
-        )
-
-    def forward(
-        self,
-        x: torch.Tensor,
-        timesteps: torch.Tensor,
-        y: Optional[torch.Tensor] = None,
-        context: Optional[torch.Tensor] = None,
-        hint = None,
-    ) -> torch.Tensor:
-
-        #weird sd3 controlnet specific stuff
-        y = torch.zeros_like(y)
-
-        if self.context_processor is not None:
-            context = self.context_processor(context)
-
-        hw = x.shape[-2:]
-        x = self.x_embedder(x) + self.cropped_pos_embed(hw, device=x.device).to(dtype=x.dtype, device=x.device)
-        x += self.pos_embed_input(hint)
-
-        c = self.t_embedder(timesteps, dtype=x.dtype)
-        if y is not None and self.y_embedder is not None:
-            y = self.y_embedder(y)
-            c = c + y
-
-        if context is not None:
-            context = self.context_embedder(context)
-
-        output = []
-
-        blocks = len(self.joint_blocks)
-        for i in range(blocks):
-            context, x = self.joint_blocks[i](
-                context,
-                x,
-                c=c,
-                use_checkpoint=self.use_checkpoint,
-            )
-
-            out = self.controlnet_blocks[i](x)
-            count = self.depth // blocks
-            if i == blocks - 1:
-                count -= 1
-            for j in range(count):
-                output.append(out)
-
-        return {"output": output}

comfy/cli_args.py

@@ -1,236 +0,0 @@
-import argparse
-import enum
-import os
-import comfy.options
-
-
-class EnumAction(argparse.Action):
-    """
-    Argparse action for handling Enums
-    """
-    def __init__(self, **kwargs):
-        # Pop off the type value
-        enum_type = kwargs.pop("type", None)
-
-        # Ensure an Enum subclass is provided
-        if enum_type is None:
-            raise ValueError("type must be assigned an Enum when using EnumAction")
-        if not issubclass(enum_type, enum.Enum):
-            raise TypeError("type must be an Enum when using EnumAction")
-
-        # Generate choices from the Enum
-        choices = tuple(e.value for e in enum_type)
-        kwargs.setdefault("choices", choices)
-        kwargs.setdefault("metavar", f"[{','.join(list(choices))}]")
-
-        super(EnumAction, self).__init__(**kwargs)
-
-        self._enum = enum_type
-
-    def __call__(self, parser, namespace, values, option_string=None):
-        # Convert value back into an Enum
-        value = self._enum(values)
-        setattr(namespace, self.dest, value)
-
-
-parser = argparse.ArgumentParser()
-
-parser.add_argument("--listen", type=str, default="127.0.0.1", metavar="IP", nargs="?", const="0.0.0.0,::", help="Specify the IP address to listen on (default: 127.0.0.1). You can give a list of ip addresses by separating them with a comma like: 127.2.2.2,127.3.3.3 If --listen is provided without an argument, it defaults to 0.0.0.0,:: (listens on all ipv4 and ipv6)")
-parser.add_argument("--port", type=int, default=8188, help="Set the listen port.")
-parser.add_argument("--tls-keyfile", type=str, help="Path to TLS (SSL) key file. Enables TLS, makes app accessible at https://... requires --tls-certfile to function")
-parser.add_argument("--tls-certfile", type=str, help="Path to TLS (SSL) certificate file. Enables TLS, makes app accessible at https://... requires --tls-keyfile to function")
-parser.add_argument("--enable-cors-header", type=str, default=None, metavar="ORIGIN", nargs="?", const="*", help="Enable CORS (Cross-Origin Resource Sharing) with optional origin or allow all with default '*'.")
-parser.add_argument("--max-upload-size", type=float, default=100, help="Set the maximum upload size in MB.")
-
-parser.add_argument("--base-directory", type=str, default=None, help="Set the ComfyUI base directory for models, custom_nodes, input, output, temp, and user directories.")
-parser.add_argument("--extra-model-paths-config", type=str, default=None, metavar="PATH", nargs='+', action='append', help="Load one or more extra_model_paths.yaml files.")
-parser.add_argument("--output-directory", type=str, default=None, help="Set the ComfyUI output directory. Overrides --base-directory.")
-parser.add_argument("--temp-directory", type=str, default=None, help="Set the ComfyUI temp directory (default is in the ComfyUI directory). Overrides --base-directory.")
-parser.add_argument("--input-directory", type=str, default=None, help="Set the ComfyUI input directory. Overrides --base-directory.")
-parser.add_argument("--auto-launch", action="store_true", help="Automatically launch ComfyUI in the default browser.")
-parser.add_argument("--disable-auto-launch", action="store_true", help="Disable auto launching the browser.")
-parser.add_argument("--cuda-device", type=int, default=None, metavar="DEVICE_ID", help="Set the id of the cuda device this instance will use.")
-cm_group = parser.add_mutually_exclusive_group()
-cm_group.add_argument("--cuda-malloc", action="store_true", help="Enable cudaMallocAsync (enabled by default for torch 2.0 and up).")
-cm_group.add_argument("--disable-cuda-malloc", action="store_true", help="Disable cudaMallocAsync.")
-
-
-fp_group = parser.add_mutually_exclusive_group()
-fp_group.add_argument("--force-fp32", action="store_true", help="Force fp32 (If this makes your GPU work better please report it).")
-fp_group.add_argument("--force-fp16", action="store_true", help="Force fp16.")
-
-fpunet_group = parser.add_mutually_exclusive_group()
-fpunet_group.add_argument("--fp32-unet", action="store_true", help="Run the diffusion model in fp32.")
-fpunet_group.add_argument("--fp64-unet", action="store_true", help="Run the diffusion model in fp64.")
-fpunet_group.add_argument("--bf16-unet", action="store_true", help="Run the diffusion model in bf16.")
-fpunet_group.add_argument("--fp16-unet", action="store_true", help="Run the diffusion model in fp16")
-fpunet_group.add_argument("--fp8_e4m3fn-unet", action="store_true", help="Store unet weights in fp8_e4m3fn.")
-fpunet_group.add_argument("--fp8_e5m2-unet", action="store_true", help="Store unet weights in fp8_e5m2.")
-fpunet_group.add_argument("--fp8_e8m0fnu-unet", action="store_true", help="Store unet weights in fp8_e8m0fnu.")
-
-fpvae_group = parser.add_mutually_exclusive_group()
-fpvae_group.add_argument("--fp16-vae", action="store_true", help="Run the VAE in fp16, might cause black images.")
-fpvae_group.add_argument("--fp32-vae", action="store_true", help="Run the VAE in full precision fp32.")
-fpvae_group.add_argument("--bf16-vae", action="store_true", help="Run the VAE in bf16.")
-
-parser.add_argument("--cpu-vae", action="store_true", help="Run the VAE on the CPU.")
-
-fpte_group = parser.add_mutually_exclusive_group()
-fpte_group.add_argument("--fp8_e4m3fn-text-enc", action="store_true", help="Store text encoder weights in fp8 (e4m3fn variant).")
-fpte_group.add_argument("--fp8_e5m2-text-enc", action="store_true", help="Store text encoder weights in fp8 (e5m2 variant).")
-fpte_group.add_argument("--fp16-text-enc", action="store_true", help="Store text encoder weights in fp16.")
-fpte_group.add_argument("--fp32-text-enc", action="store_true", help="Store text encoder weights in fp32.")
-fpte_group.add_argument("--bf16-text-enc", action="store_true", help="Store text encoder weights in bf16.")
-
-parser.add_argument("--force-channels-last", action="store_true", help="Force channels last format when inferencing the models.")
-
-parser.add_argument("--directml", type=int, nargs="?", metavar="DIRECTML_DEVICE", const=-1, help="Use torch-directml.")
-
-parser.add_argument("--oneapi-device-selector", type=str, default=None, metavar="SELECTOR_STRING", help="Sets the oneAPI device(s) this instance will use.")
-parser.add_argument("--disable-ipex-optimize", action="store_true", help="Disables ipex.optimize default when loading models with Intel's Extension for Pytorch.")
-parser.add_argument("--supports-fp8-compute", action="store_true", help="ComfyUI will act like if the device supports fp8 compute.")
-
-class LatentPreviewMethod(enum.Enum):
-    NoPreviews = "none"
-    Auto = "auto"
-    Latent2RGB = "latent2rgb"
-    TAESD = "taesd"
-
-parser.add_argument("--preview-method", type=LatentPreviewMethod, default=LatentPreviewMethod.NoPreviews, help="Default preview method for sampler nodes.", action=EnumAction)
-
-parser.add_argument("--preview-size", type=int, default=512, help="Sets the maximum preview size for sampler nodes.")
-
-cache_group = parser.add_mutually_exclusive_group()
-cache_group.add_argument("--cache-classic", action="store_true", help="Use the old style (aggressive) caching.")
-cache_group.add_argument("--cache-lru", type=int, default=0, help="Use LRU caching with a maximum of N node results cached. May use more RAM/VRAM.")
-cache_group.add_argument("--cache-none", action="store_true", help="Reduced RAM/VRAM usage at the expense of executing every node for each run.")
-
-attn_group = parser.add_mutually_exclusive_group()
-attn_group.add_argument("--use-split-cross-attention", action="store_true", help="Use the split cross attention optimization. Ignored when xformers is used.")
-attn_group.add_argument("--use-quad-cross-attention", action="store_true", help="Use the sub-quadratic cross attention optimization . Ignored when xformers is used.")
-attn_group.add_argument("--use-pytorch-cross-attention", action="store_true", help="Use the new pytorch 2.0 cross attention function.")
-attn_group.add_argument("--use-sage-attention", action="store_true", help="Use sage attention.")
-attn_group.add_argument("--use-flash-attention", action="store_true", help="Use FlashAttention.")
-
-parser.add_argument("--disable-xformers", action="store_true", help="Disable xformers.")
-
-upcast = parser.add_mutually_exclusive_group()
-upcast.add_argument("--force-upcast-attention", action="store_true", help="Force enable attention upcasting, please report if it fixes black images.")
-upcast.add_argument("--dont-upcast-attention", action="store_true", help="Disable all upcasting of attention. Should be unnecessary except for debugging.")
-
-
-vram_group = parser.add_mutually_exclusive_group()
-vram_group.add_argument("--gpu-only", action="store_true", help="Store and run everything (text encoders/CLIP models, etc... on the GPU).")
-vram_group.add_argument("--highvram", action="store_true", help="By default models will be unloaded to CPU memory after being used. This option keeps them in GPU memory.")
-vram_group.add_argument("--normalvram", action="store_true", help="Used to force normal vram use if lowvram gets automatically enabled.")
-vram_group.add_argument("--lowvram", action="store_true", help="Split the unet in parts to use less vram.")
-vram_group.add_argument("--novram", action="store_true", help="When lowvram isn't enough.")
-vram_group.add_argument("--cpu", action="store_true", help="To use the CPU for everything (slow).")
-
-parser.add_argument("--reserve-vram", type=float, default=None, help="Set the amount of vram in GB you want to reserve for use by your OS/other software. By default some amount is reserved depending on your OS.")
-
-parser.add_argument("--async-offload", action="store_true", help="Use async weight offloading.")
-
-parser.add_argument("--default-hashing-function", type=str, choices=['md5', 'sha1', 'sha256', 'sha512'], default='sha256', help="Allows you to choose the hash function to use for duplicate filename / contents comparison. Default is sha256.")
-
-parser.add_argument("--disable-smart-memory", action="store_true", help="Force ComfyUI to agressively offload to regular ram instead of keeping models in vram when it can.")
-parser.add_argument("--deterministic", action="store_true", help="Make pytorch use slower deterministic algorithms when it can. Note that this might not make images deterministic in all cases.")
-
-class PerformanceFeature(enum.Enum):
-    Fp16Accumulation = "fp16_accumulation"
-    Fp8MatrixMultiplication = "fp8_matrix_mult"
-    CublasOps = "cublas_ops"
-
-parser.add_argument("--fast", nargs="*", type=PerformanceFeature, help="Enable some untested and potentially quality deteriorating optimizations. --fast with no arguments enables everything. You can pass a list specific optimizations if you only want to enable specific ones. Current valid optimizations: fp16_accumulation fp8_matrix_mult cublas_ops")
-
-parser.add_argument("--mmap-torch-files", action="store_true", help="Use mmap when loading ckpt/pt files.")
-parser.add_argument("--disable-mmap", action="store_true", help="Don't use mmap when loading safetensors.")
-
-parser.add_argument("--dont-print-server", action="store_true", help="Don't print server output.")
-parser.add_argument("--quick-test-for-ci", action="store_true", help="Quick test for CI.")
-parser.add_argument("--windows-standalone-build", action="store_true", help="Windows standalone build: Enable convenient things that most people using the standalone windows build will probably enjoy (like auto opening the page on startup).")
-
-parser.add_argument("--disable-metadata", action="store_true", help="Disable saving prompt metadata in files.")
-parser.add_argument("--disable-all-custom-nodes", action="store_true", help="Disable loading all custom nodes.")
-parser.add_argument("--whitelist-custom-nodes", type=str, nargs='+', default=[], help="Specify custom node folders to load even when --disable-all-custom-nodes is enabled.")
-parser.add_argument("--disable-api-nodes", action="store_true", help="Disable loading all api nodes.")
-
-parser.add_argument("--multi-user", action="store_true", help="Enables per-user storage.")
-
-parser.add_argument("--verbose", default='INFO', const='DEBUG', nargs="?", choices=['DEBUG', 'INFO', 'WARNING', 'ERROR', 'CRITICAL'], help='Set the logging level')
-parser.add_argument("--log-stdout", action="store_true", help="Send normal process output to stdout instead of stderr (default).")
-
-# The default built-in provider hosted under web/
-DEFAULT_VERSION_STRING = "comfyanonymous/ComfyUI@latest"
-
-parser.add_argument(
-    "--front-end-version",
-    type=str,
-    default=DEFAULT_VERSION_STRING,
-    help="""
-    Specifies the version of the frontend to be used. This command needs internet connectivity to query and
-    download available frontend implementations from GitHub releases.
-
-    The version string should be in the format of:
-    [repoOwner]/[repoName]@[version]
-    where version is one of: "latest" or a valid version number (e.g. "1.0.0")
-    """,
-)
-
-def is_valid_directory(path: str) -> str:
-    """Validate if the given path is a directory, and check permissions."""
-    if not os.path.exists(path):
-        raise argparse.ArgumentTypeError(f"The path '{path}' does not exist.")
-    if not os.path.isdir(path):
-        raise argparse.ArgumentTypeError(f"'{path}' is not a directory.")
-    if not os.access(path, os.R_OK):
-        raise argparse.ArgumentTypeError(f"You do not have read permissions for '{path}'.")
-    return path
-
-parser.add_argument(
-    "--front-end-root",
-    type=is_valid_directory,
-    default=None,
-    help="The local filesystem path to the directory where the frontend is located. Overrides --front-end-version.",
-)
-
-parser.add_argument("--user-directory", type=is_valid_directory, default=None, help="Set the ComfyUI user directory with an absolute path. Overrides --base-directory.")
-
-parser.add_argument("--enable-compress-response-body", action="store_true", help="Enable compressing response body.")
-
-parser.add_argument(
-    "--comfy-api-base",
-    type=str,
-    default="https://api.comfy.org",
-    help="Set the base URL for the ComfyUI API.  (default: https://api.comfy.org)",
-)
-
-database_default_path = os.path.abspath(
-    os.path.join(os.path.dirname(__file__), "..", "user", "comfyui.db")
-)
-parser.add_argument("--database-url", type=str, default=f"sqlite:///{database_default_path}", help="Specify the database URL, e.g. for an in-memory database you can use 'sqlite:///:memory:'.")
-
-if comfy.options.args_parsing:
-    args = parser.parse_args()
-else:
-    args = parser.parse_args([])
-
-if args.windows_standalone_build:
-    args.auto_launch = True
-
-if args.disable_auto_launch:
-    args.auto_launch = False
-
-if args.force_fp16:
-    args.fp16_unet = True
-
-
-# '--fast' is not provided, use an empty set
-if args.fast is None:
-    args.fast = set()
-# '--fast' is provided with an empty list, enable all optimizations
-elif args.fast == []:
-    args.fast = set(PerformanceFeature)
-# '--fast' is provided with a list of performance features, use that list
-else:
-    args.fast = set(args.fast)

comfy/clip_config_bigg.json

@@ -1,23 +0,0 @@
-{
-  "architectures": [
-    "CLIPTextModel"
-  ],
-  "attention_dropout": 0.0,
-  "bos_token_id": 0,
-  "dropout": 0.0,
-  "eos_token_id": 49407,
-  "hidden_act": "gelu",
-  "hidden_size": 1280,
-  "initializer_factor": 1.0,
-  "initializer_range": 0.02,
-  "intermediate_size": 5120,
-  "layer_norm_eps": 1e-05,
-  "max_position_embeddings": 77,
-  "model_type": "clip_text_model",
-  "num_attention_heads": 20,
-  "num_hidden_layers": 32,
-  "pad_token_id": 1,
-  "projection_dim": 1280,
-  "torch_dtype": "float32",
-  "vocab_size": 49408
-}

comfy/clip_model.py

@@ -1,244 +0,0 @@
-import torch
-from comfy.ldm.modules.attention import optimized_attention_for_device
-import comfy.ops
-
-class CLIPAttention(torch.nn.Module):
-    def __init__(self, embed_dim, heads, dtype, device, operations):
-        super().__init__()
-
-        self.heads = heads
-        self.q_proj = operations.Linear(embed_dim, embed_dim, bias=True, dtype=dtype, device=device)
-        self.k_proj = operations.Linear(embed_dim, embed_dim, bias=True, dtype=dtype, device=device)
-        self.v_proj = operations.Linear(embed_dim, embed_dim, bias=True, dtype=dtype, device=device)
-
-        self.out_proj = operations.Linear(embed_dim, embed_dim, bias=True, dtype=dtype, device=device)
-
-    def forward(self, x, mask=None, optimized_attention=None):
-        q = self.q_proj(x)
-        k = self.k_proj(x)
-        v = self.v_proj(x)
-
-        out = optimized_attention(q, k, v, self.heads, mask)
-        return self.out_proj(out)
-
-ACTIVATIONS = {"quick_gelu": lambda a: a * torch.sigmoid(1.702 * a),
-               "gelu": torch.nn.functional.gelu,
-               "gelu_pytorch_tanh": lambda a: torch.nn.functional.gelu(a, approximate="tanh"),
-}
-
-class CLIPMLP(torch.nn.Module):
-    def __init__(self, embed_dim, intermediate_size, activation, dtype, device, operations):
-        super().__init__()
-        self.fc1 = operations.Linear(embed_dim, intermediate_size, bias=True, dtype=dtype, device=device)
-        self.activation = ACTIVATIONS[activation]
-        self.fc2 = operations.Linear(intermediate_size, embed_dim, bias=True, dtype=dtype, device=device)
-
-    def forward(self, x):
-        x = self.fc1(x)
-        x = self.activation(x)
-        x = self.fc2(x)
-        return x
-
-class CLIPLayer(torch.nn.Module):
-    def __init__(self, embed_dim, heads, intermediate_size, intermediate_activation, dtype, device, operations):
-        super().__init__()
-        self.layer_norm1 = operations.LayerNorm(embed_dim, dtype=dtype, device=device)
-        self.self_attn = CLIPAttention(embed_dim, heads, dtype, device, operations)
-        self.layer_norm2 = operations.LayerNorm(embed_dim, dtype=dtype, device=device)
-        self.mlp = CLIPMLP(embed_dim, intermediate_size, intermediate_activation, dtype, device, operations)
-
-    def forward(self, x, mask=None, optimized_attention=None):
-        x += self.self_attn(self.layer_norm1(x), mask, optimized_attention)
-        x += self.mlp(self.layer_norm2(x))
-        return x
-
-
-class CLIPEncoder(torch.nn.Module):
-    def __init__(self, num_layers, embed_dim, heads, intermediate_size, intermediate_activation, dtype, device, operations):
-        super().__init__()
-        self.layers = torch.nn.ModuleList([CLIPLayer(embed_dim, heads, intermediate_size, intermediate_activation, dtype, device, operations) for i in range(num_layers)])
-
-    def forward(self, x, mask=None, intermediate_output=None):
-        optimized_attention = optimized_attention_for_device(x.device, mask=mask is not None, small_input=True)
-
-        if intermediate_output is not None:
-            if intermediate_output < 0:
-                intermediate_output = len(self.layers) + intermediate_output
-
-        intermediate = None
-        for i, l in enumerate(self.layers):
-            x = l(x, mask, optimized_attention)
-            if i == intermediate_output:
-                intermediate = x.clone()
-        return x, intermediate
-
-class CLIPEmbeddings(torch.nn.Module):
-    def __init__(self, embed_dim, vocab_size=49408, num_positions=77, dtype=None, device=None, operations=None):
-        super().__init__()
-        self.token_embedding = operations.Embedding(vocab_size, embed_dim, dtype=dtype, device=device)
-        self.position_embedding = operations.Embedding(num_positions, embed_dim, dtype=dtype, device=device)
-
-    def forward(self, input_tokens, dtype=torch.float32):
-        return self.token_embedding(input_tokens, out_dtype=dtype) + comfy.ops.cast_to(self.position_embedding.weight, dtype=dtype, device=input_tokens.device)
-
-
-class CLIPTextModel_(torch.nn.Module):
-    def __init__(self, config_dict, dtype, device, operations):
-        num_layers = config_dict["num_hidden_layers"]
-        embed_dim = config_dict["hidden_size"]
-        heads = config_dict["num_attention_heads"]
-        intermediate_size = config_dict["intermediate_size"]
-        intermediate_activation = config_dict["hidden_act"]
-        num_positions = config_dict["max_position_embeddings"]
-        self.eos_token_id = config_dict["eos_token_id"]
-
-        super().__init__()
-        self.embeddings = CLIPEmbeddings(embed_dim, num_positions=num_positions, dtype=dtype, device=device, operations=operations)
-        self.encoder = CLIPEncoder(num_layers, embed_dim, heads, intermediate_size, intermediate_activation, dtype, device, operations)
-        self.final_layer_norm = operations.LayerNorm(embed_dim, dtype=dtype, device=device)
-
-    def forward(self, input_tokens=None, attention_mask=None, embeds=None, num_tokens=None, intermediate_output=None, final_layer_norm_intermediate=True, dtype=torch.float32):
-        if embeds is not None:
-            x = embeds + comfy.ops.cast_to(self.embeddings.position_embedding.weight, dtype=dtype, device=embeds.device)
-        else:
-            x = self.embeddings(input_tokens, dtype=dtype)
-
-        mask = None
-        if attention_mask is not None:
-            mask = 1.0 - attention_mask.to(x.dtype).reshape((attention_mask.shape[0], 1, -1, attention_mask.shape[-1])).expand(attention_mask.shape[0], 1, attention_mask.shape[-1], attention_mask.shape[-1])
-            mask = mask.masked_fill(mask.to(torch.bool), -torch.finfo(x.dtype).max)
-
-        causal_mask = torch.full((x.shape[1], x.shape[1]), -torch.finfo(x.dtype).max, dtype=x.dtype, device=x.device).triu_(1)
-
-        if mask is not None:
-            mask += causal_mask
-        else:
-            mask = causal_mask
-
-        x, i = self.encoder(x, mask=mask, intermediate_output=intermediate_output)
-        x = self.final_layer_norm(x)
-        if i is not None and final_layer_norm_intermediate:
-            i = self.final_layer_norm(i)
-
-        if num_tokens is not None:
-            pooled_output = x[list(range(x.shape[0])), list(map(lambda a: a - 1, num_tokens))]
-        else:
-            pooled_output = x[torch.arange(x.shape[0], device=x.device), (torch.round(input_tokens).to(dtype=torch.int, device=x.device) == self.eos_token_id).int().argmax(dim=-1),]
-        return x, i, pooled_output
-
-class CLIPTextModel(torch.nn.Module):
-    def __init__(self, config_dict, dtype, device, operations):
-        super().__init__()
-        self.num_layers = config_dict["num_hidden_layers"]
-        self.text_model = CLIPTextModel_(config_dict, dtype, device, operations)
-        embed_dim = config_dict["hidden_size"]
-        self.text_projection = operations.Linear(embed_dim, embed_dim, bias=False, dtype=dtype, device=device)
-        self.dtype = dtype
-
-    def get_input_embeddings(self):
-        return self.text_model.embeddings.token_embedding
-
-    def set_input_embeddings(self, embeddings):
-        self.text_model.embeddings.token_embedding = embeddings
-
-    def forward(self, *args, **kwargs):
-        x = self.text_model(*args, **kwargs)
-        out = self.text_projection(x[2])
-        return (x[0], x[1], out, x[2])
-
-
-class CLIPVisionEmbeddings(torch.nn.Module):
-    def __init__(self, embed_dim, num_channels=3, patch_size=14, image_size=224, model_type="", dtype=None, device=None, operations=None):
-        super().__init__()
-
-        num_patches = (image_size // patch_size) ** 2
-        if model_type == "siglip_vision_model":
-            self.class_embedding = None
-            patch_bias = True
-        else:
-            num_patches = num_patches + 1
-            self.class_embedding = torch.nn.Parameter(torch.empty(embed_dim, dtype=dtype, device=device))
-            patch_bias = False
-
-        self.patch_embedding = operations.Conv2d(
-            in_channels=num_channels,
-            out_channels=embed_dim,
-            kernel_size=patch_size,
-            stride=patch_size,
-            bias=patch_bias,
-            dtype=dtype,
-            device=device
-        )
-
-        self.position_embedding = operations.Embedding(num_patches, embed_dim, dtype=dtype, device=device)
-
-    def forward(self, pixel_values):
-        embeds = self.patch_embedding(pixel_values).flatten(2).transpose(1, 2)
-        if self.class_embedding is not None:
-            embeds = torch.cat([comfy.ops.cast_to_input(self.class_embedding, embeds).expand(pixel_values.shape[0], 1, -1), embeds], dim=1)
-        return embeds + comfy.ops.cast_to_input(self.position_embedding.weight, embeds)
-
-
-class CLIPVision(torch.nn.Module):
-    def __init__(self, config_dict, dtype, device, operations):
-        super().__init__()
-        num_layers = config_dict["num_hidden_layers"]
-        embed_dim = config_dict["hidden_size"]
-        heads = config_dict["num_attention_heads"]
-        intermediate_size = config_dict["intermediate_size"]
-        intermediate_activation = config_dict["hidden_act"]
-        model_type = config_dict["model_type"]
-
-        self.embeddings = CLIPVisionEmbeddings(embed_dim, config_dict["num_channels"], config_dict["patch_size"], config_dict["image_size"], model_type=model_type, dtype=dtype, device=device, operations=operations)
-        if model_type == "siglip_vision_model":
-            self.pre_layrnorm = lambda a: a
-            self.output_layernorm = True
-        else:
-            self.pre_layrnorm = operations.LayerNorm(embed_dim)
-            self.output_layernorm = False
-        self.encoder = CLIPEncoder(num_layers, embed_dim, heads, intermediate_size, intermediate_activation, dtype, device, operations)
-        self.post_layernorm = operations.LayerNorm(embed_dim)
-
-    def forward(self, pixel_values, attention_mask=None, intermediate_output=None):
-        x = self.embeddings(pixel_values)
-        x = self.pre_layrnorm(x)
-        #TODO: attention_mask?
-        x, i = self.encoder(x, mask=None, intermediate_output=intermediate_output)
-        if self.output_layernorm:
-            x = self.post_layernorm(x)
-            pooled_output = x
-        else:
-            pooled_output = self.post_layernorm(x[:, 0, :])
-        return x, i, pooled_output
-
-class LlavaProjector(torch.nn.Module):
-    def __init__(self, in_dim, out_dim, dtype, device, operations):
-        super().__init__()
-        self.linear_1 = operations.Linear(in_dim, out_dim, bias=True, device=device, dtype=dtype)
-        self.linear_2 = operations.Linear(out_dim, out_dim, bias=True, device=device, dtype=dtype)
-
-    def forward(self, x):
-        return self.linear_2(torch.nn.functional.gelu(self.linear_1(x[:, 1:])))
-
-class CLIPVisionModelProjection(torch.nn.Module):
-    def __init__(self, config_dict, dtype, device, operations):
-        super().__init__()
-        self.vision_model = CLIPVision(config_dict, dtype, device, operations)
-        if "projection_dim" in config_dict:
-            self.visual_projection = operations.Linear(config_dict["hidden_size"], config_dict["projection_dim"], bias=False)
-        else:
-            self.visual_projection = lambda a: a
-
-        if "llava3" == config_dict.get("projector_type", None):
-            self.multi_modal_projector = LlavaProjector(config_dict["hidden_size"], 4096, dtype, device, operations)
-        else:
-            self.multi_modal_projector = None
-
-    def forward(self, *args, **kwargs):
-        x = self.vision_model(*args, **kwargs)
-        out = self.visual_projection(x[2])
-        projected = None
-        if self.multi_modal_projector is not None:
-            projected = self.multi_modal_projector(x[1])
-
-        return (x[0], x[1], out, projected)

comfy/clip_vision.py

@@ -1,148 +0,0 @@
-from .utils import load_torch_file, transformers_convert, state_dict_prefix_replace
-import os
-import torch
-import json
-import logging
-
-import comfy.ops
-import comfy.model_patcher
-import comfy.model_management
-import comfy.utils
-import comfy.clip_model
-import comfy.image_encoders.dino2
-
-class Output:
-    def __getitem__(self, key):
-        return getattr(self, key)
-    def __setitem__(self, key, item):
-        setattr(self, key, item)
-
-def clip_preprocess(image, size=224, mean=[0.48145466, 0.4578275, 0.40821073], std=[0.26862954, 0.26130258, 0.27577711], crop=True):
-    image = image[:, :, :, :3] if image.shape[3] > 3 else image
-    mean = torch.tensor(mean, device=image.device, dtype=image.dtype)
-    std = torch.tensor(std, device=image.device, dtype=image.dtype)
-    image = image.movedim(-1, 1)
-    if not (image.shape[2] == size and image.shape[3] == size):
-        if crop:
-            scale = (size / min(image.shape[2], image.shape[3]))
-            scale_size = (round(scale * image.shape[2]), round(scale * image.shape[3]))
-        else:
-            scale_size = (size, size)
-
-        image = torch.nn.functional.interpolate(image, size=scale_size, mode="bicubic", antialias=True)
-        h = (image.shape[2] - size)//2
-        w = (image.shape[3] - size)//2
-        image = image[:,:,h:h+size,w:w+size]
-    image = torch.clip((255. * image), 0, 255).round() / 255.0
-    return (image - mean.view([3,1,1])) / std.view([3,1,1])
-
-IMAGE_ENCODERS = {
-    "clip_vision_model": comfy.clip_model.CLIPVisionModelProjection,
-    "siglip_vision_model": comfy.clip_model.CLIPVisionModelProjection,
-    "dinov2": comfy.image_encoders.dino2.Dinov2Model,
-}
-
-class ClipVisionModel():
-    def __init__(self, json_config):
-        with open(json_config) as f:
-            config = json.load(f)
-
-        self.image_size = config.get("image_size", 224)
-        self.image_mean = config.get("image_mean", [0.48145466, 0.4578275, 0.40821073])
-        self.image_std = config.get("image_std", [0.26862954, 0.26130258, 0.27577711])
-        model_class = IMAGE_ENCODERS.get(config.get("model_type", "clip_vision_model"))
-        self.load_device = comfy.model_management.text_encoder_device()
-        offload_device = comfy.model_management.text_encoder_offload_device()
-        self.dtype = comfy.model_management.text_encoder_dtype(self.load_device)
-        self.model = model_class(config, self.dtype, offload_device, comfy.ops.manual_cast)
-        self.model.eval()
-
-        self.patcher = comfy.model_patcher.ModelPatcher(self.model, load_device=self.load_device, offload_device=offload_device)
-
-    def load_sd(self, sd):
-        return self.model.load_state_dict(sd, strict=False)
-
-    def get_sd(self):
-        return self.model.state_dict()
-
-    def encode_image(self, image, crop=True):
-        comfy.model_management.load_model_gpu(self.patcher)
-        pixel_values = clip_preprocess(image.to(self.load_device), size=self.image_size, mean=self.image_mean, std=self.image_std, crop=crop).float()
-        out = self.model(pixel_values=pixel_values, intermediate_output=-2)
-
-        outputs = Output()
-        outputs["last_hidden_state"] = out[0].to(comfy.model_management.intermediate_device())
-        outputs["image_embeds"] = out[2].to(comfy.model_management.intermediate_device())
-        outputs["penultimate_hidden_states"] = out[1].to(comfy.model_management.intermediate_device())
-        outputs["mm_projected"] = out[3]
-        return outputs
-
-def convert_to_transformers(sd, prefix):
-    sd_k = sd.keys()
-    if "{}transformer.resblocks.0.attn.in_proj_weight".format(prefix) in sd_k:
-        keys_to_replace = {
-            "{}class_embedding".format(prefix): "vision_model.embeddings.class_embedding",
-            "{}conv1.weight".format(prefix): "vision_model.embeddings.patch_embedding.weight",
-            "{}positional_embedding".format(prefix): "vision_model.embeddings.position_embedding.weight",
-            "{}ln_post.bias".format(prefix): "vision_model.post_layernorm.bias",
-            "{}ln_post.weight".format(prefix): "vision_model.post_layernorm.weight",
-            "{}ln_pre.bias".format(prefix): "vision_model.pre_layrnorm.bias",
-            "{}ln_pre.weight".format(prefix): "vision_model.pre_layrnorm.weight",
-        }
-
-        for x in keys_to_replace:
-            if x in sd_k:
-                sd[keys_to_replace[x]] = sd.pop(x)
-
-        if "{}proj".format(prefix) in sd_k:
-            sd['visual_projection.weight'] = sd.pop("{}proj".format(prefix)).transpose(0, 1)
-
-        sd = transformers_convert(sd, prefix, "vision_model.", 48)
-    else:
-        replace_prefix = {prefix: ""}
-        sd = state_dict_prefix_replace(sd, replace_prefix)
-    return sd
-
-def load_clipvision_from_sd(sd, prefix="", convert_keys=False):
-    if convert_keys:
-        sd = convert_to_transformers(sd, prefix)
-    if "vision_model.encoder.layers.47.layer_norm1.weight" in sd:
-        json_config = os.path.join(os.path.dirname(os.path.realpath(__file__)), "clip_vision_config_g.json")
-    elif "vision_model.encoder.layers.30.layer_norm1.weight" in sd:
-        json_config = os.path.join(os.path.dirname(os.path.realpath(__file__)), "clip_vision_config_h.json")
-    elif "vision_model.encoder.layers.22.layer_norm1.weight" in sd:
-        embed_shape = sd["vision_model.embeddings.position_embedding.weight"].shape[0]
-        if sd["vision_model.encoder.layers.0.layer_norm1.weight"].shape[0] == 1152:
-            if embed_shape == 729:
-                json_config = os.path.join(os.path.dirname(os.path.realpath(__file__)), "clip_vision_siglip_384.json")
-            elif embed_shape == 1024:
-                json_config = os.path.join(os.path.dirname(os.path.realpath(__file__)), "clip_vision_siglip_512.json")
-        elif embed_shape == 577:
-            if "multi_modal_projector.linear_1.bias" in sd:
-                json_config = os.path.join(os.path.dirname(os.path.realpath(__file__)), "clip_vision_config_vitl_336_llava.json")
-            else:
-                json_config = os.path.join(os.path.dirname(os.path.realpath(__file__)), "clip_vision_config_vitl_336.json")
-        else:
-            json_config = os.path.join(os.path.dirname(os.path.realpath(__file__)), "clip_vision_config_vitl.json")
-    elif "embeddings.patch_embeddings.projection.weight" in sd:
-        json_config = os.path.join(os.path.join(os.path.dirname(os.path.realpath(__file__)), "image_encoders"), "dino2_giant.json")
-    else:
-        return None
-
-    clip = ClipVisionModel(json_config)
-    m, u = clip.load_sd(sd)
-    if len(m) > 0:
-        logging.warning("missing clip vision: {}".format(m))
-    u = set(u)
-    keys = list(sd.keys())
-    for k in keys:
-        if k not in u:
-            sd.pop(k)
-    return clip
-
-def load(ckpt_path):
-    sd = load_torch_file(ckpt_path)
-    if "visual.transformer.resblocks.0.attn.in_proj_weight" in sd:
-        return load_clipvision_from_sd(sd, prefix="visual.", convert_keys=True)
-    else:
-        return load_clipvision_from_sd(sd)

comfy/clip_vision_config_g.json

@@ -1,18 +0,0 @@
-{
-  "attention_dropout": 0.0,
-  "dropout": 0.0,
-  "hidden_act": "gelu",
-  "hidden_size": 1664,
-  "image_size": 224,
-  "initializer_factor": 1.0,
-  "initializer_range": 0.02,
-  "intermediate_size": 8192,
-  "layer_norm_eps": 1e-05,
-  "model_type": "clip_vision_model",
-  "num_attention_heads": 16,
-  "num_channels": 3,
-  "num_hidden_layers": 48,
-  "patch_size": 14,
-  "projection_dim": 1280,
-  "torch_dtype": "float32"
-}

comfy/clip_vision_config_h.json

@@ -1,18 +0,0 @@
-{
-  "attention_dropout": 0.0,
-  "dropout": 0.0,
-  "hidden_act": "gelu",
-  "hidden_size": 1280,
-  "image_size": 224,
-  "initializer_factor": 1.0,
-  "initializer_range": 0.02,
-  "intermediate_size": 5120,
-  "layer_norm_eps": 1e-05,
-  "model_type": "clip_vision_model",
-  "num_attention_heads": 16,
-  "num_channels": 3,
-  "num_hidden_layers": 32,
-  "patch_size": 14,
-  "projection_dim": 1024,
-  "torch_dtype": "float32"
-}

comfy/clip_vision_config_vitl.json

@@ -1,18 +0,0 @@
-{
-  "attention_dropout": 0.0,
-  "dropout": 0.0,
-  "hidden_act": "quick_gelu",
-  "hidden_size": 1024,
-  "image_size": 224,
-  "initializer_factor": 1.0,
-  "initializer_range": 0.02,
-  "intermediate_size": 4096,
-  "layer_norm_eps": 1e-05,
-  "model_type": "clip_vision_model",
-  "num_attention_heads": 16,
-  "num_channels": 3,
-  "num_hidden_layers": 24,
-  "patch_size": 14,
-  "projection_dim": 768,
-  "torch_dtype": "float32"
-}

comfy/clip_vision_config_vitl_336.json

@@ -1,18 +0,0 @@
-{
-  "attention_dropout": 0.0,
-  "dropout": 0.0,
-  "hidden_act": "quick_gelu",
-  "hidden_size": 1024,
-  "image_size": 336,
-  "initializer_factor": 1.0,
-  "initializer_range": 0.02,
-  "intermediate_size": 4096,
-  "layer_norm_eps": 1e-5,
-  "model_type": "clip_vision_model",
-  "num_attention_heads": 16,
-  "num_channels": 3,
-  "num_hidden_layers": 24,
-  "patch_size": 14,
-  "projection_dim": 768,
-  "torch_dtype": "float32"
-}

comfy/clip_vision_config_vitl_336_llava.json

@@ -1,19 +0,0 @@
-{
-  "attention_dropout": 0.0,
-  "dropout": 0.0,
-  "hidden_act": "quick_gelu",
-  "hidden_size": 1024,
-  "image_size": 336,
-  "initializer_factor": 1.0,
-  "initializer_range": 0.02,
-  "intermediate_size": 4096,
-  "layer_norm_eps": 1e-5,
-  "model_type": "clip_vision_model",
-  "num_attention_heads": 16,
-  "num_channels": 3,
-  "num_hidden_layers": 24,
-  "patch_size": 14,
-  "projection_dim": 768,
-  "projector_type": "llava3",
-  "torch_dtype": "float32"
-}

comfy/clip_vision_siglip_384.json

@@ -1,13 +0,0 @@
-{
-  "num_channels": 3,
-  "hidden_act": "gelu_pytorch_tanh",
-  "hidden_size": 1152,
-  "image_size": 384,
-  "intermediate_size": 4304,
-  "model_type": "siglip_vision_model",
-  "num_attention_heads": 16,
-  "num_hidden_layers": 27,
-  "patch_size": 14,
-  "image_mean": [0.5, 0.5, 0.5],
-  "image_std": [0.5, 0.5, 0.5]
-}

comfy/clip_vision_siglip_512.json

@@ -1,13 +0,0 @@
-{
-  "num_channels": 3,
-  "hidden_act": "gelu_pytorch_tanh",
-  "hidden_size": 1152,
-  "image_size": 512,
-  "intermediate_size": 4304,
-  "model_type": "siglip_vision_model",
-  "num_attention_heads": 16,
-  "num_hidden_layers": 27,
-  "patch_size": 16,
-  "image_mean": [0.5, 0.5, 0.5],
-  "image_std": [0.5, 0.5, 0.5]
-}

comfy/comfy_types/README.md

@@ -1,43 +0,0 @@
-# Comfy Typing
-## Type hinting for ComfyUI Node development
-
-This module provides type hinting and concrete convenience types for node developers.
-If cloned to the custom_nodes directory of ComfyUI, types can be imported using:
-
-```python
-from comfy.comfy_types import IO, ComfyNodeABC, CheckLazyMixin
-
-class ExampleNode(ComfyNodeABC):
-    @classmethod
-    def INPUT_TYPES(s) -> InputTypeDict:
-        return {"required": {}}
-```
-
-Full example is in [examples/example_nodes.py](examples/example_nodes.py).
-
-# Types
-A few primary types are documented below.  More complete information is available via the docstrings on each type.
-
-## `IO`
-
-A string enum of built-in and a few custom data types.  Includes the following special types and their requisite plumbing:
-
-- `ANY`: `"*"`
-- `NUMBER`: `"FLOAT,INT"`
-- `PRIMITIVE`: `"STRING,FLOAT,INT,BOOLEAN"`
-
-## `ComfyNodeABC`
-
-An abstract base class for nodes, offering type-hinting / autocomplete, and somewhat-alright docstrings.
-
-### Type hinting for `INPUT_TYPES`
-
-![INPUT_TYPES auto-completion in Visual Studio Code](examples/input_types.png)
-
-### `INPUT_TYPES` return dict
-
-![INPUT_TYPES return value type hinting in Visual Studio Code](examples/required_hint.png)
-
-### Options for individual inputs
-
-![INPUT_TYPES return value option auto-completion in Visual Studio Code](examples/input_options.png)

comfy/comfy_types/__init__.py

@@ -1,46 +0,0 @@
-import torch
-from typing import Callable, Protocol, TypedDict, Optional, List
-from .node_typing import IO, InputTypeDict, ComfyNodeABC, CheckLazyMixin, FileLocator
-
-
-class UnetApplyFunction(Protocol):
-    """Function signature protocol on comfy.model_base.BaseModel.apply_model"""
-
-    def __call__(self, x: torch.Tensor, t: torch.Tensor, **kwargs) -> torch.Tensor:
-        pass
-
-
-class UnetApplyConds(TypedDict):
-    """Optional conditions for unet apply function."""
-
-    c_concat: Optional[torch.Tensor]
-    c_crossattn: Optional[torch.Tensor]
-    control: Optional[torch.Tensor]
-    transformer_options: Optional[dict]
-
-
-class UnetParams(TypedDict):
-    # Tensor of shape [B, C, H, W]
-    input: torch.Tensor
-    # Tensor of shape [B]
-    timestep: torch.Tensor
-    c: UnetApplyConds
-    # List of [0, 1], [0], [1], ...
-    # 0 means conditional, 1 means conditional unconditional
-    cond_or_uncond: List[int]
-
-
-UnetWrapperFunction = Callable[[UnetApplyFunction, UnetParams], torch.Tensor]
-
-
-__all__ = [
-    "UnetWrapperFunction",
-    UnetApplyConds.__name__,
-    UnetParams.__name__,
-    UnetApplyFunction.__name__,
-    IO.__name__,
-    InputTypeDict.__name__,
-    ComfyNodeABC.__name__,
-    CheckLazyMixin.__name__,
-    FileLocator.__name__,
-]

comfy/comfy_types/examples/example_nodes.py

@@ -1,28 +0,0 @@
-from comfy.comfy_types import IO, ComfyNodeABC, InputTypeDict
-from inspect import cleandoc
-
-
-class ExampleNode(ComfyNodeABC):
-    """An example node that just adds 1 to an input integer.
-
-    * Requires a modern IDE to provide any benefit (detail: an IDE configured with analysis paths etc).
-    * This node is intended as an example for developers only.
-    """
-
-    DESCRIPTION = cleandoc(__doc__)
-    CATEGORY = "examples"
-
-    @classmethod
-    def INPUT_TYPES(s) -> InputTypeDict:
-        return {
-            "required": {
-                "input_int": (IO.INT, {"defaultInput": True}),
-            }
-        }
-
-    RETURN_TYPES = (IO.INT,)
-    RETURN_NAMES = ("input_plus_one",)
-    FUNCTION = "execute"
-
-    def execute(self, input_int: int):
-        return (input_int + 1,)

comfy/comfy_types/node_typing.py

@@ -1,350 +0,0 @@
-"""Comfy-specific type hinting"""
-
-from __future__ import annotations
-from typing import Literal, TypedDict, Optional
-from typing_extensions import NotRequired
-from abc import ABC, abstractmethod
-from enum import Enum
-
-
-class StrEnum(str, Enum):
-    """Base class for string enums. Python's StrEnum is not available until 3.11."""
-
-    def __str__(self) -> str:
-        return self.value
-
-
-class IO(StrEnum):
-    """Node input/output data types.
-
-    Includes functionality for ``"*"`` (`ANY`) and ``"MULTI,TYPES"``.
-    """
-
-    STRING = "STRING"
-    IMAGE = "IMAGE"
-    MASK = "MASK"
-    LATENT = "LATENT"
-    BOOLEAN = "BOOLEAN"
-    INT = "INT"
-    FLOAT = "FLOAT"
-    COMBO = "COMBO"
-    CONDITIONING = "CONDITIONING"
-    SAMPLER = "SAMPLER"
-    SIGMAS = "SIGMAS"
-    GUIDER = "GUIDER"
-    NOISE = "NOISE"
-    CLIP = "CLIP"
-    CONTROL_NET = "CONTROL_NET"
-    VAE = "VAE"
-    MODEL = "MODEL"
-    LORA_MODEL = "LORA_MODEL"
-    LOSS_MAP = "LOSS_MAP"
-    CLIP_VISION = "CLIP_VISION"
-    CLIP_VISION_OUTPUT = "CLIP_VISION_OUTPUT"
-    STYLE_MODEL = "STYLE_MODEL"
-    GLIGEN = "GLIGEN"
-    UPSCALE_MODEL = "UPSCALE_MODEL"
-    AUDIO = "AUDIO"
-    WEBCAM = "WEBCAM"
-    POINT = "POINT"
-    FACE_ANALYSIS = "FACE_ANALYSIS"
-    BBOX = "BBOX"
-    SEGS = "SEGS"
-    VIDEO = "VIDEO"
-
-    ANY = "*"
-    """Always matches any type, but at a price.
-
-    Causes some functionality issues (e.g. reroutes, link types), and should be avoided whenever possible.
-    """
-    NUMBER = "FLOAT,INT"
-    """A float or an int - could be either"""
-    PRIMITIVE = "STRING,FLOAT,INT,BOOLEAN"
-    """Could be any of: string, float, int, or bool"""
-
-    def __ne__(self, value: object) -> bool:
-        if self == "*" or value == "*":
-            return False
-        if not isinstance(value, str):
-            return True
-        a = frozenset(self.split(","))
-        b = frozenset(value.split(","))
-        return not (b.issubset(a) or a.issubset(b))
-
-
-class RemoteInputOptions(TypedDict):
-    route: str
-    """The route to the remote source."""
-    refresh_button: bool
-    """Specifies whether to show a refresh button in the UI below the widget."""
-    control_after_refresh: Literal["first", "last"]
-    """Specifies the control after the refresh button is clicked. If "first", the first item will be automatically selected, and so on."""
-    timeout: int
-    """The maximum amount of time to wait for a response from the remote source in milliseconds."""
-    max_retries: int
-    """The maximum number of retries before aborting the request."""
-    refresh: int
-    """The TTL of the remote input's value in milliseconds. Specifies the interval at which the remote input's value is refreshed."""
-
-
-class MultiSelectOptions(TypedDict):
-    placeholder: NotRequired[str]
-    """The placeholder text to display in the multi-select widget when no items are selected."""
-    chip: NotRequired[bool]
-    """Specifies whether to use chips instead of comma separated values for the multi-select widget."""
-
-
-class InputTypeOptions(TypedDict):
-    """Provides type hinting for the return type of the INPUT_TYPES node function.
-
-    Due to IDE limitations with unions, for now all options are available for all types (e.g. `label_on` is hinted even when the type is not `IO.BOOLEAN`).
-
-    Comfy Docs: https://docs.comfy.org/custom-nodes/backend/datatypes
-    """
-
-    default: NotRequired[bool | str | float | int | list | tuple]
-    """The default value of the widget"""
-    defaultInput: NotRequired[bool]
-    """@deprecated in v1.16 frontend. v1.16 frontend allows input socket and widget to co-exist.
-    - defaultInput on required inputs should be dropped.
-    - defaultInput on optional inputs should be replaced with forceInput.
-    Ref: https://github.com/Comfy-Org/ComfyUI_frontend/pull/3364
-    """
-    forceInput: NotRequired[bool]
-    """Forces the input to be an input slot rather than a widget even a widget is available for the input type."""
-    lazy: NotRequired[bool]
-    """Declares that this input uses lazy evaluation"""
-    rawLink: NotRequired[bool]
-    """When a link exists, rather than receiving the evaluated value, you will receive the link (i.e. `["nodeId", <outputIndex>]`). Designed for node expansion."""
-    tooltip: NotRequired[str]
-    """Tooltip for the input (or widget), shown on pointer hover"""
-    socketless: NotRequired[bool]
-    """All inputs (including widgets) have an input socket to connect links. When ``true``, if there is a widget for this input, no socket will be created.
-    Available from frontend v1.17.5
-    Ref: https://github.com/Comfy-Org/ComfyUI_frontend/pull/3548
-    """
-    widgetType: NotRequired[str]
-    """Specifies a type to be used for widget initialization if different from the input type.
-    Available from frontend v1.18.0
-    https://github.com/Comfy-Org/ComfyUI_frontend/pull/3550"""
-    # class InputTypeNumber(InputTypeOptions):
-    # default: float | int
-    min: NotRequired[float]
-    """The minimum value of a number (``FLOAT`` | ``INT``)"""
-    max: NotRequired[float]
-    """The maximum value of a number (``FLOAT`` | ``INT``)"""
-    step: NotRequired[float]
-    """The amount to increment or decrement a widget by when stepping up/down (``FLOAT`` | ``INT``)"""
-    round: NotRequired[float]
-    """Floats are rounded by this value (``FLOAT``)"""
-    # class InputTypeBoolean(InputTypeOptions):
-    # default: bool
-    label_on: NotRequired[str]
-    """The label to use in the UI when the bool is True (``BOOLEAN``)"""
-    label_off: NotRequired[str]
-    """The label to use in the UI when the bool is False (``BOOLEAN``)"""
-    # class InputTypeString(InputTypeOptions):
-    # default: str
-    multiline: NotRequired[bool]
-    """Use a multiline text box (``STRING``)"""
-    placeholder: NotRequired[str]
-    """Placeholder text to display in the UI when empty (``STRING``)"""
-    # Deprecated:
-    # defaultVal: str
-    dynamicPrompts: NotRequired[bool]
-    """Causes the front-end to evaluate dynamic prompts (``STRING``)"""
-    # class InputTypeCombo(InputTypeOptions):
-    image_upload: NotRequired[bool]
-    """Specifies whether the input should have an image upload button and image preview attached to it. Requires that the input's name is `image`."""
-    image_folder: NotRequired[Literal["input", "output", "temp"]]
-    """Specifies which folder to get preview images from if the input has the ``image_upload`` flag.
-    """
-    remote: NotRequired[RemoteInputOptions]
-    """Specifies the configuration for a remote input.
-    Available after ComfyUI frontend v1.9.7
-    https://github.com/Comfy-Org/ComfyUI_frontend/pull/2422"""
-    control_after_generate: NotRequired[bool]
-    """Specifies whether a control widget should be added to the input, adding options to automatically change the value after each prompt is queued. Currently only used for INT and COMBO types."""
-    options: NotRequired[list[str | int | float]]
-    """COMBO type only. Specifies the selectable options for the combo widget.
-    Prefer:
-    ["COMBO", {"options": ["Option 1", "Option 2", "Option 3"]}]
-    Over:
-    [["Option 1", "Option 2", "Option 3"]]
-    """
-    multi_select: NotRequired[MultiSelectOptions]
-    """COMBO type only. Specifies the configuration for a multi-select widget.
-    Available after ComfyUI frontend v1.13.4
-    https://github.com/Comfy-Org/ComfyUI_frontend/pull/2987"""
-
-
-class HiddenInputTypeDict(TypedDict):
-    """Provides type hinting for the hidden entry of node INPUT_TYPES."""
-
-    node_id: NotRequired[Literal["UNIQUE_ID"]]
-    """UNIQUE_ID is the unique identifier of the node, and matches the id property of the node on the client side. It is commonly used in client-server communications (see messages)."""
-    unique_id: NotRequired[Literal["UNIQUE_ID"]]
-    """UNIQUE_ID is the unique identifier of the node, and matches the id property of the node on the client side. It is commonly used in client-server communications (see messages)."""
-    prompt: NotRequired[Literal["PROMPT"]]
-    """PROMPT is the complete prompt sent by the client to the server. See the prompt object for a full description."""
-    extra_pnginfo: NotRequired[Literal["EXTRA_PNGINFO"]]
-    """EXTRA_PNGINFO is a dictionary that will be copied into the metadata of any .png files saved. Custom nodes can store additional information in this dictionary for saving (or as a way to communicate with a downstream node)."""
-    dynprompt: NotRequired[Literal["DYNPROMPT"]]
-    """DYNPROMPT is an instance of comfy_execution.graph.DynamicPrompt. It differs from PROMPT in that it may mutate during the course of execution in response to Node Expansion."""
-
-
-class InputTypeDict(TypedDict):
-    """Provides type hinting for node INPUT_TYPES.
-
-    Comfy Docs: https://docs.comfy.org/custom-nodes/backend/more_on_inputs
-    """
-
-    required: NotRequired[dict[str, tuple[IO, InputTypeOptions]]]
-    """Describes all inputs that must be connected for the node to execute."""
-    optional: NotRequired[dict[str, tuple[IO, InputTypeOptions]]]
-    """Describes inputs which do not need to be connected."""
-    hidden: NotRequired[HiddenInputTypeDict]
-    """Offers advanced functionality and server-client communication.
-
-    Comfy Docs: https://docs.comfy.org/custom-nodes/backend/more_on_inputs#hidden-inputs
-    """
-
-
-class ComfyNodeABC(ABC):
-    """Abstract base class for Comfy nodes.  Includes the names and expected types of attributes.
-
-    Comfy Docs: https://docs.comfy.org/custom-nodes/backend/server_overview
-    """
-
-    DESCRIPTION: str
-    """Node description, shown as a tooltip when hovering over the node.
-
-    Usage::
-
-        # Explicitly define the description
-        DESCRIPTION = "Example description here."
-
-        # Use the docstring of the node class.
-        DESCRIPTION = cleandoc(__doc__)
-    """
-    CATEGORY: str
-    """The category of the node, as per the "Add Node" menu.
-
-    Comfy Docs: https://docs.comfy.org/custom-nodes/backend/server_overview#category
-    """
-    EXPERIMENTAL: bool
-    """Flags a node as experimental, informing users that it may change or not work as expected."""
-    DEPRECATED: bool
-    """Flags a node as deprecated, indicating to users that they should find alternatives to this node."""
-    API_NODE: Optional[bool]
-    """Flags a node as an API node. See: https://docs.comfy.org/tutorials/api-nodes/overview."""
-
-    @classmethod
-    @abstractmethod
-    def INPUT_TYPES(s) -> InputTypeDict:
-        """Defines node inputs.
-
-        * Must include the ``required`` key, which describes all inputs that must be connected for the node to execute.
-        * The ``optional`` key can be added to describe inputs which do not need to be connected.
-        * The ``hidden`` key offers some advanced functionality.  More info at: https://docs.comfy.org/custom-nodes/backend/more_on_inputs#hidden-inputs
-
-        Comfy Docs: https://docs.comfy.org/custom-nodes/backend/server_overview#input-types
-        """
-        return {"required": {}}
-
-    OUTPUT_NODE: bool
-    """Flags this node as an output node, causing any inputs it requires to be executed.
-
-    If a node is not connected to any output nodes, that node will not be executed.  Usage::
-
-        OUTPUT_NODE = True
-
-    From the docs:
-
-    By default, a node is not considered an output. Set ``OUTPUT_NODE = True`` to specify that it is.
-
-    Comfy Docs: https://docs.comfy.org/custom-nodes/backend/server_overview#output-node
-    """
-    INPUT_IS_LIST: bool
-    """A flag indicating if this node implements the additional code necessary to deal with OUTPUT_IS_LIST nodes.
-
-    All inputs of ``type`` will become ``list[type]``, regardless of how many items are passed in.  This also affects ``check_lazy_status``.
-
-    From the docs:
-
-    A node can also override the default input behaviour and receive the whole list in a single call. This is done by setting a class attribute `INPUT_IS_LIST` to ``True``.
-
-    Comfy Docs: https://docs.comfy.org/custom-nodes/backend/lists#list-processing
-    """
-    OUTPUT_IS_LIST: tuple[bool, ...]
-    """A tuple indicating which node outputs are lists, but will be connected to nodes that expect individual items.
-
-    Connected nodes that do not implement `INPUT_IS_LIST` will be executed once for every item in the list.
-
-    A ``tuple[bool]``, where the items match those in `RETURN_TYPES`::
-
-        RETURN_TYPES = (IO.INT, IO.INT, IO.STRING)
-        OUTPUT_IS_LIST = (True, True, False) # The string output will be handled normally
-
-    From the docs:
-
-    In order to tell Comfy that the list being returned should not be wrapped, but treated as a series of data for sequential processing,
-    the node should provide a class attribute `OUTPUT_IS_LIST`, which is a ``tuple[bool]``, of the same length as `RETURN_TYPES`,
-    specifying which outputs which should be so treated.
-
-    Comfy Docs: https://docs.comfy.org/custom-nodes/backend/lists#list-processing
-    """
-
-    RETURN_TYPES: tuple[IO, ...]
-    """A tuple representing the outputs of this node.
-
-    Usage::
-
-        RETURN_TYPES = (IO.INT, "INT", "CUSTOM_TYPE")
-
-    Comfy Docs: https://docs.comfy.org/custom-nodes/backend/server_overview#return-types
-    """
-    RETURN_NAMES: tuple[str, ...]
-    """The output slot names for each item in `RETURN_TYPES`, e.g. ``RETURN_NAMES = ("count", "filter_string")``
-
-    Comfy Docs: https://docs.comfy.org/custom-nodes/backend/server_overview#return-names
-    """
-    OUTPUT_TOOLTIPS: tuple[str, ...]
-    """A tuple of strings to use as tooltips for node outputs, one for each item in `RETURN_TYPES`."""
-    FUNCTION: str
-    """The name of the function to execute as a literal string, e.g. `FUNCTION = "execute"`
-
-    Comfy Docs: https://docs.comfy.org/custom-nodes/backend/server_overview#function
-    """
-
-
-class CheckLazyMixin:
-    """Provides a basic check_lazy_status implementation and type hinting for nodes that use lazy inputs."""
-
-    def check_lazy_status(self, **kwargs) -> list[str]:
-        """Returns a list of input names that should be evaluated.
-
-        This basic mixin impl. requires all inputs.
-
-        :kwargs: All node inputs will be included here.  If the input is ``None``, it should be assumed that it has not yet been evaluated.  \
-            When using ``INPUT_IS_LIST = True``, unevaluated will instead be ``(None,)``.
-
-        Params should match the nodes execution ``FUNCTION`` (self, and all inputs by name).
-        Will be executed repeatedly until it returns an empty list, or all requested items were already evaluated (and sent as params).
-
-        Comfy Docs: https://docs.comfy.org/custom-nodes/backend/lazy_evaluation#defining-check-lazy-status
-        """
-
-        need = [name for name in kwargs if kwargs[name] is None]
-        return need
-
-
-class FileLocator(TypedDict):
-    """Provides type hinting for the file location"""
-
-    filename: str
-    """The filename of the file."""
-    subfolder: str
-    """The subfolder of the file."""
-    type: Literal["input", "output", "temp"]
-    """The root folder of the file."""

comfy/conds.py

@@ -1,130 +0,0 @@
-import torch
-import math
-import comfy.utils
-
-
-class CONDRegular:
-    def __init__(self, cond):
-        self.cond = cond
-
-    def _copy_with(self, cond):
-        return self.__class__(cond)
-
-    def process_cond(self, batch_size, device, **kwargs):
-        return self._copy_with(comfy.utils.repeat_to_batch_size(self.cond, batch_size).to(device))
-
-    def can_concat(self, other):
-        if self.cond.shape != other.cond.shape:
-            return False
-        return True
-
-    def concat(self, others):
-        conds = [self.cond]
-        for x in others:
-            conds.append(x.cond)
-        return torch.cat(conds)
-
-    def size(self):
-        return list(self.cond.size())
-
-
-class CONDNoiseShape(CONDRegular):
-    def process_cond(self, batch_size, device, area, **kwargs):
-        data = self.cond
-        if area is not None:
-            dims = len(area) // 2
-            for i in range(dims):
-                data = data.narrow(i + 2, area[i + dims], area[i])
-
-        return self._copy_with(comfy.utils.repeat_to_batch_size(data, batch_size).to(device))
-
-
-class CONDCrossAttn(CONDRegular):
-    def can_concat(self, other):
-        s1 = self.cond.shape
-        s2 = other.cond.shape
-        if s1 != s2:
-            if s1[0] != s2[0] or s1[2] != s2[2]: #these 2 cases should not happen
-                return False
-
-            mult_min = math.lcm(s1[1], s2[1])
-            diff = mult_min // min(s1[1], s2[1])
-            if diff > 4: #arbitrary limit on the padding because it's probably going to impact performance negatively if it's too much
-                return False
-        return True
-
-    def concat(self, others):
-        conds = [self.cond]
-        crossattn_max_len = self.cond.shape[1]
-        for x in others:
-            c = x.cond
-            crossattn_max_len = math.lcm(crossattn_max_len, c.shape[1])
-            conds.append(c)
-
-        out = []
-        for c in conds:
-            if c.shape[1] < crossattn_max_len:
-                c = c.repeat(1, crossattn_max_len // c.shape[1], 1) #padding with repeat doesn't change result
-            out.append(c)
-        return torch.cat(out)
-
-
-class CONDConstant(CONDRegular):
-    def __init__(self, cond):
-        self.cond = cond
-
-    def process_cond(self, batch_size, device, **kwargs):
-        return self._copy_with(self.cond)
-
-    def can_concat(self, other):
-        if self.cond != other.cond:
-            return False
-        return True
-
-    def concat(self, others):
-        return self.cond
-
-    def size(self):
-        return [1]
-
-
-class CONDList(CONDRegular):
-    def __init__(self, cond):
-        self.cond = cond
-
-    def process_cond(self, batch_size, device, **kwargs):
-        out = []
-        for c in self.cond:
-            out.append(comfy.utils.repeat_to_batch_size(c, batch_size).to(device))
-
-        return self._copy_with(out)
-
-    def can_concat(self, other):
-        if len(self.cond) != len(other.cond):
-            return False
-        for i in range(len(self.cond)):
-            if self.cond[i].shape != other.cond[i].shape:
-                return False
-
-        return True
-
-    def concat(self, others):
-        out = []
-        for i in range(len(self.cond)):
-            o = [self.cond[i]]
-            for x in others:
-                o.append(x.cond[i])
-            out.append(torch.cat(o))
-
-        return out
-
-    def size(self):  # hackish implementation to make the mem estimation work
-        o = 0
-        c = 1
-        for c in self.cond:
-            size = c.size()
-            o += math.prod(size)
-            if len(size) > 1:
-                c = size[1]
-
-        return [1, c, o // c]

comfy/controlnet.py

@@ -1,858 +0,0 @@
-"""
-    This file is part of ComfyUI.
-    Copyright (C) 2024 Comfy
-
-    This program is free software: you can redistribute it and/or modify
-    it under the terms of the GNU General Public License as published by
-    the Free Software Foundation, either version 3 of the License, or
-    (at your option) any later version.
-
-    This program is distributed in the hope that it will be useful,
-    but WITHOUT ANY WARRANTY; without even the implied warranty of
-    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
-    GNU General Public License for more details.
-
-    You should have received a copy of the GNU General Public License
-    along with this program.  If not, see <https://www.gnu.org/licenses/>.
-"""
-
-
-import torch
-from enum import Enum
-import math
-import os
-import logging
-import comfy.utils
-import comfy.model_management
-import comfy.model_detection
-import comfy.model_patcher
-import comfy.ops
-import comfy.latent_formats
-
-import comfy.cldm.cldm
-import comfy.t2i_adapter.adapter
-import comfy.ldm.cascade.controlnet
-import comfy.cldm.mmdit
-import comfy.ldm.hydit.controlnet
-import comfy.ldm.flux.controlnet
-import comfy.cldm.dit_embedder
-from typing import TYPE_CHECKING
-if TYPE_CHECKING:
-    from comfy.hooks import HookGroup
-
-
-def broadcast_image_to(tensor, target_batch_size, batched_number):
-    current_batch_size = tensor.shape[0]
-    #print(current_batch_size, target_batch_size)
-    if current_batch_size == 1:
-        return tensor
-
-    per_batch = target_batch_size // batched_number
-    tensor = tensor[:per_batch]
-
-    if per_batch > tensor.shape[0]:
-        tensor = torch.cat([tensor] * (per_batch // tensor.shape[0]) + [tensor[:(per_batch % tensor.shape[0])]], dim=0)
-
-    current_batch_size = tensor.shape[0]
-    if current_batch_size == target_batch_size:
-        return tensor
-    else:
-        return torch.cat([tensor] * batched_number, dim=0)
-
-class StrengthType(Enum):
-    CONSTANT = 1
-    LINEAR_UP = 2
-
-class ControlBase:
-    def __init__(self):
-        self.cond_hint_original = None
-        self.cond_hint = None
-        self.strength = 1.0
-        self.timestep_percent_range = (0.0, 1.0)
-        self.latent_format = None
-        self.vae = None
-        self.global_average_pooling = False
-        self.timestep_range = None
-        self.compression_ratio = 8
-        self.upscale_algorithm = 'nearest-exact'
-        self.extra_args = {}
-        self.previous_controlnet = None
-        self.extra_conds = []
-        self.strength_type = StrengthType.CONSTANT
-        self.concat_mask = False
-        self.extra_concat_orig = []
-        self.extra_concat = None
-        self.extra_hooks: HookGroup = None
-        self.preprocess_image = lambda a: a
-
-    def set_cond_hint(self, cond_hint, strength=1.0, timestep_percent_range=(0.0, 1.0), vae=None, extra_concat=[]):
-        self.cond_hint_original = cond_hint
-        self.strength = strength
-        self.timestep_percent_range = timestep_percent_range
-        if self.latent_format is not None:
-            if vae is None:
-                logging.warning("WARNING: no VAE provided to the controlnet apply node when this controlnet requires one.")
-            self.vae = vae
-        self.extra_concat_orig = extra_concat.copy()
-        if self.concat_mask and len(self.extra_concat_orig) == 0:
-            self.extra_concat_orig.append(torch.tensor([[[[1.0]]]]))
-        return self
-
-    def pre_run(self, model, percent_to_timestep_function):
-        self.timestep_range = (percent_to_timestep_function(self.timestep_percent_range[0]), percent_to_timestep_function(self.timestep_percent_range[1]))
-        if self.previous_controlnet is not None:
-            self.previous_controlnet.pre_run(model, percent_to_timestep_function)
-
-    def set_previous_controlnet(self, controlnet):
-        self.previous_controlnet = controlnet
-        return self
-
-    def cleanup(self):
-        if self.previous_controlnet is not None:
-            self.previous_controlnet.cleanup()
-
-        self.cond_hint = None
-        self.extra_concat = None
-        self.timestep_range = None
-
-    def get_models(self):
-        out = []
-        if self.previous_controlnet is not None:
-            out += self.previous_controlnet.get_models()
-        return out
-
-    def get_extra_hooks(self):
-        out = []
-        if self.extra_hooks is not None:
-            out.append(self.extra_hooks)
-        if self.previous_controlnet is not None:
-            out += self.previous_controlnet.get_extra_hooks()
-        return out
-
-    def copy_to(self, c):
-        c.cond_hint_original = self.cond_hint_original
-        c.strength = self.strength
-        c.timestep_percent_range = self.timestep_percent_range
-        c.global_average_pooling = self.global_average_pooling
-        c.compression_ratio = self.compression_ratio
-        c.upscale_algorithm = self.upscale_algorithm
-        c.latent_format = self.latent_format
-        c.extra_args = self.extra_args.copy()
-        c.vae = self.vae
-        c.extra_conds = self.extra_conds.copy()
-        c.strength_type = self.strength_type
-        c.concat_mask = self.concat_mask
-        c.extra_concat_orig = self.extra_concat_orig.copy()
-        c.extra_hooks = self.extra_hooks.clone() if self.extra_hooks else None
-        c.preprocess_image = self.preprocess_image
-
-    def inference_memory_requirements(self, dtype):
-        if self.previous_controlnet is not None:
-            return self.previous_controlnet.inference_memory_requirements(dtype)
-        return 0
-
-    def control_merge(self, control, control_prev, output_dtype):
-        out = {'input':[], 'middle':[], 'output': []}
-
-        for key in control:
-            control_output = control[key]
-            applied_to = set()
-            for i in range(len(control_output)):
-                x = control_output[i]
-                if x is not None:
-                    if self.global_average_pooling:
-                        x = torch.mean(x, dim=(2, 3), keepdim=True).repeat(1, 1, x.shape[2], x.shape[3])
-
-                    if x not in applied_to: #memory saving strategy, allow shared tensors and only apply strength to shared tensors once
-                        applied_to.add(x)
-                        if self.strength_type == StrengthType.CONSTANT:
-                            x *= self.strength
-                        elif self.strength_type == StrengthType.LINEAR_UP:
-                            x *= (self.strength ** float(len(control_output) - i))
-
-                    if output_dtype is not None and x.dtype != output_dtype:
-                        x = x.to(output_dtype)
-
-                out[key].append(x)
-
-        if control_prev is not None:
-            for x in ['input', 'middle', 'output']:
-                o = out[x]
-                for i in range(len(control_prev[x])):
-                    prev_val = control_prev[x][i]
-                    if i >= len(o):
-                        o.append(prev_val)
-                    elif prev_val is not None:
-                        if o[i] is None:
-                            o[i] = prev_val
-                        else:
-                            if o[i].shape[0] < prev_val.shape[0]:
-                                o[i] = prev_val + o[i]
-                            else:
-                                o[i] = prev_val + o[i] #TODO: change back to inplace add if shared tensors stop being an issue
-        return out
-
-    def set_extra_arg(self, argument, value=None):
-        self.extra_args[argument] = value
-
-
-class ControlNet(ControlBase):
-    def __init__(self, control_model=None, global_average_pooling=False, compression_ratio=8, latent_format=None, load_device=None, manual_cast_dtype=None, extra_conds=["y"], strength_type=StrengthType.CONSTANT, concat_mask=False, preprocess_image=lambda a: a):
-        super().__init__()
-        self.control_model = control_model
-        self.load_device = load_device
-        if control_model is not None:
-            self.control_model_wrapped = comfy.model_patcher.ModelPatcher(self.control_model, load_device=load_device, offload_device=comfy.model_management.unet_offload_device())
-
-        self.compression_ratio = compression_ratio
-        self.global_average_pooling = global_average_pooling
-        self.model_sampling_current = None
-        self.manual_cast_dtype = manual_cast_dtype
-        self.latent_format = latent_format
-        self.extra_conds += extra_conds
-        self.strength_type = strength_type
-        self.concat_mask = concat_mask
-        self.preprocess_image = preprocess_image
-
-    def get_control(self, x_noisy, t, cond, batched_number, transformer_options):
-        control_prev = None
-        if self.previous_controlnet is not None:
-            control_prev = self.previous_controlnet.get_control(x_noisy, t, cond, batched_number, transformer_options)
-
-        if self.timestep_range is not None:
-            if t[0] > self.timestep_range[0] or t[0] < self.timestep_range[1]:
-                if control_prev is not None:
-                    return control_prev
-                else:
-                    return None
-
-        dtype = self.control_model.dtype
-        if self.manual_cast_dtype is not None:
-            dtype = self.manual_cast_dtype
-
-        if self.cond_hint is None or x_noisy.shape[2] * self.compression_ratio != self.cond_hint.shape[2] or x_noisy.shape[3] * self.compression_ratio != self.cond_hint.shape[3]:
-            if self.cond_hint is not None:
-                del self.cond_hint
-            self.cond_hint = None
-            compression_ratio = self.compression_ratio
-            if self.vae is not None:
-                compression_ratio *= self.vae.downscale_ratio
-            else:
-                if self.latent_format is not None:
-                    raise ValueError("This Controlnet needs a VAE but none was provided, please use a ControlNetApply node with a VAE input and connect it.")
-            self.cond_hint = comfy.utils.common_upscale(self.cond_hint_original, x_noisy.shape[3] * compression_ratio, x_noisy.shape[2] * compression_ratio, self.upscale_algorithm, "center")
-            self.cond_hint = self.preprocess_image(self.cond_hint)
-            if self.vae is not None:
-                loaded_models = comfy.model_management.loaded_models(only_currently_used=True)
-                self.cond_hint = self.vae.encode(self.cond_hint.movedim(1, -1))
-                comfy.model_management.load_models_gpu(loaded_models)
-            if self.latent_format is not None:
-                self.cond_hint = self.latent_format.process_in(self.cond_hint)
-            if len(self.extra_concat_orig) > 0:
-                to_concat = []
-                for c in self.extra_concat_orig:
-                    c = c.to(self.cond_hint.device)
-                    c = comfy.utils.common_upscale(c, self.cond_hint.shape[3], self.cond_hint.shape[2], self.upscale_algorithm, "center")
-                    to_concat.append(comfy.utils.repeat_to_batch_size(c, self.cond_hint.shape[0]))
-                self.cond_hint = torch.cat([self.cond_hint] + to_concat, dim=1)
-
-            self.cond_hint = self.cond_hint.to(device=x_noisy.device, dtype=dtype)
-        if x_noisy.shape[0] != self.cond_hint.shape[0]:
-            self.cond_hint = broadcast_image_to(self.cond_hint, x_noisy.shape[0], batched_number)
-
-        context = cond.get('crossattn_controlnet', cond['c_crossattn'])
-        extra = self.extra_args.copy()
-        for c in self.extra_conds:
-            temp = cond.get(c, None)
-            if temp is not None:
-                extra[c] = temp.to(dtype)
-
-        timestep = self.model_sampling_current.timestep(t)
-        x_noisy = self.model_sampling_current.calculate_input(t, x_noisy)
-
-        control = self.control_model(x=x_noisy.to(dtype), hint=self.cond_hint, timesteps=timestep.to(dtype), context=context.to(dtype), **extra)
-        return self.control_merge(control, control_prev, output_dtype=None)
-
-    def copy(self):
-        c = ControlNet(None, global_average_pooling=self.global_average_pooling, load_device=self.load_device, manual_cast_dtype=self.manual_cast_dtype)
-        c.control_model = self.control_model
-        c.control_model_wrapped = self.control_model_wrapped
-        self.copy_to(c)
-        return c
-
-    def get_models(self):
-        out = super().get_models()
-        out.append(self.control_model_wrapped)
-        return out
-
-    def pre_run(self, model, percent_to_timestep_function):
-        super().pre_run(model, percent_to_timestep_function)
-        self.model_sampling_current = model.model_sampling
-
-    def cleanup(self):
-        self.model_sampling_current = None
-        super().cleanup()
-
-class ControlLoraOps:
-    class Linear(torch.nn.Module, comfy.ops.CastWeightBiasOp):
-        def __init__(self, in_features: int, out_features: int, bias: bool = True,
-                    device=None, dtype=None) -> None:
-            super().__init__()
-            self.in_features = in_features
-            self.out_features = out_features
-            self.weight = None
-            self.up = None
-            self.down = None
-            self.bias = None
-
-        def forward(self, input):
-            weight, bias = comfy.ops.cast_bias_weight(self, input)
-            if self.up is not None:
-                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)
-            else:
-                return torch.nn.functional.linear(input, weight, bias)
-
-    class Conv2d(torch.nn.Module, comfy.ops.CastWeightBiasOp):
-        def __init__(
-            self,
-            in_channels,
-            out_channels,
-            kernel_size,
-            stride=1,
-            padding=0,
-            dilation=1,
-            groups=1,
-            bias=True,
-            padding_mode='zeros',
-            device=None,
-            dtype=None
-        ):
-            super().__init__()
-            self.in_channels = in_channels
-            self.out_channels = out_channels
-            self.kernel_size = kernel_size
-            self.stride = stride
-            self.padding = padding
-            self.dilation = dilation
-            self.transposed = False
-            self.output_padding = 0
-            self.groups = groups
-            self.padding_mode = padding_mode
-
-            self.weight = None
-            self.bias = None
-            self.up = None
-            self.down = None
-
-
-        def forward(self, input):
-            weight, bias = comfy.ops.cast_bias_weight(self, input)
-            if self.up is not None:
-                return torch.nn.functional.conv2d(input, weight + (torch.mm(self.up.flatten(start_dim=1), self.down.flatten(start_dim=1))).reshape(self.weight.shape).type(input.dtype), bias, self.stride, self.padding, self.dilation, self.groups)
-            else:
-                return torch.nn.functional.conv2d(input, weight, bias, self.stride, self.padding, self.dilation, self.groups)
-
-
-class ControlLora(ControlNet):
-    def __init__(self, control_weights, global_average_pooling=False, model_options={}): #TODO? model_options
-        ControlBase.__init__(self)
-        self.control_weights = control_weights
-        self.global_average_pooling = global_average_pooling
-        self.extra_conds += ["y"]
-
-    def pre_run(self, model, percent_to_timestep_function):
-        super().pre_run(model, percent_to_timestep_function)
-        controlnet_config = model.model_config.unet_config.copy()
-        controlnet_config.pop("out_channels")
-        controlnet_config["hint_channels"] = self.control_weights["input_hint_block.0.weight"].shape[1]
-        self.manual_cast_dtype = model.manual_cast_dtype
-        dtype = model.get_dtype()
-        if self.manual_cast_dtype is None:
-            class control_lora_ops(ControlLoraOps, comfy.ops.disable_weight_init):
-                pass
-        else:
-            class control_lora_ops(ControlLoraOps, comfy.ops.manual_cast):
-                pass
-            dtype = self.manual_cast_dtype
-
-        controlnet_config["operations"] = control_lora_ops
-        controlnet_config["dtype"] = dtype
-        self.control_model = comfy.cldm.cldm.ControlNet(**controlnet_config)
-        self.control_model.to(comfy.model_management.get_torch_device())
-        diffusion_model = model.diffusion_model
-        sd = diffusion_model.state_dict()
-
-        for k in sd:
-            weight = sd[k]
-            try:
-                comfy.utils.set_attr_param(self.control_model, k, weight)
-            except:
-                pass
-
-        for k in self.control_weights:
-            if (k not in {"lora_controlnet"}):
-                if (k.endswith(".up") or k.endswith(".down") or k.endswith(".weight") or k.endswith(".bias")) and ("__" not in k):
-                    comfy.utils.set_attr_param(self.control_model, k, self.control_weights[k].to(dtype).to(comfy.model_management.get_torch_device()))
-
-    def copy(self):
-        c = ControlLora(self.control_weights, global_average_pooling=self.global_average_pooling)
-        self.copy_to(c)
-        return c
-
-    def cleanup(self):
-        del self.control_model
-        self.control_model = None
-        super().cleanup()
-
-    def get_models(self):
-        out = ControlBase.get_models(self)
-        return out
-
-    def inference_memory_requirements(self, dtype):
-        return comfy.utils.calculate_parameters(self.control_weights) * comfy.model_management.dtype_size(dtype) + ControlBase.inference_memory_requirements(self, dtype)
-
-def controlnet_config(sd, model_options={}):
-    model_config = comfy.model_detection.model_config_from_unet(sd, "", True)
-
-    unet_dtype = model_options.get("dtype", None)
-    if unet_dtype is None:
-        weight_dtype = comfy.utils.weight_dtype(sd)
-
-        supported_inference_dtypes = list(model_config.supported_inference_dtypes)
-        unet_dtype = comfy.model_management.unet_dtype(model_params=-1, supported_dtypes=supported_inference_dtypes, weight_dtype=weight_dtype)
-
-    load_device = comfy.model_management.get_torch_device()
-    manual_cast_dtype = comfy.model_management.unet_manual_cast(unet_dtype, load_device)
-
-    operations = model_options.get("custom_operations", None)
-    if operations is None:
-        operations = comfy.ops.pick_operations(unet_dtype, manual_cast_dtype, disable_fast_fp8=True)
-
-    offload_device = comfy.model_management.unet_offload_device()
-    return model_config, operations, load_device, unet_dtype, manual_cast_dtype, offload_device
-
-def controlnet_load_state_dict(control_model, sd):
-    missing, unexpected = control_model.load_state_dict(sd, strict=False)
-
-    if len(missing) > 0:
-        logging.warning("missing controlnet keys: {}".format(missing))
-
-    if len(unexpected) > 0:
-        logging.debug("unexpected controlnet keys: {}".format(unexpected))
-    return control_model
-
-
-def load_controlnet_mmdit(sd, model_options={}):
-    new_sd = comfy.model_detection.convert_diffusers_mmdit(sd, "")
-    model_config, operations, load_device, unet_dtype, manual_cast_dtype, offload_device = controlnet_config(new_sd, model_options=model_options)
-    num_blocks = comfy.model_detection.count_blocks(new_sd, 'joint_blocks.{}.')
-    for k in sd:
-        new_sd[k] = sd[k]
-
-    concat_mask = False
-    control_latent_channels = new_sd.get("pos_embed_input.proj.weight").shape[1]
-    if control_latent_channels == 17: #inpaint controlnet
-        concat_mask = True
-
-    control_model = comfy.cldm.mmdit.ControlNet(num_blocks=num_blocks, control_latent_channels=control_latent_channels, operations=operations, device=offload_device, dtype=unet_dtype, **model_config.unet_config)
-    control_model = controlnet_load_state_dict(control_model, new_sd)
-
-    latent_format = comfy.latent_formats.SD3()
-    latent_format.shift_factor = 0 #SD3 controlnet weirdness
-    control = ControlNet(control_model, compression_ratio=1, latent_format=latent_format, concat_mask=concat_mask, load_device=load_device, manual_cast_dtype=manual_cast_dtype)
-    return control
-
-
-class ControlNetSD35(ControlNet):
-    def pre_run(self, model, percent_to_timestep_function):
-        if self.control_model.double_y_emb:
-            missing, unexpected = self.control_model.orig_y_embedder.load_state_dict(model.diffusion_model.y_embedder.state_dict(), strict=False)
-        else:
-            missing, unexpected = self.control_model.x_embedder.load_state_dict(model.diffusion_model.x_embedder.state_dict(), strict=False)
-        super().pre_run(model, percent_to_timestep_function)
-
-    def copy(self):
-        c = ControlNetSD35(None, global_average_pooling=self.global_average_pooling, load_device=self.load_device, manual_cast_dtype=self.manual_cast_dtype)
-        c.control_model = self.control_model
-        c.control_model_wrapped = self.control_model_wrapped
-        self.copy_to(c)
-        return c
-
-def load_controlnet_sd35(sd, model_options={}):
-    control_type = -1
-    if "control_type" in sd:
-        control_type = round(sd.pop("control_type").item())
-
-    # blur_cnet = control_type == 0
-    canny_cnet = control_type == 1
-    depth_cnet = control_type == 2
-
-    new_sd = {}
-    for k in comfy.utils.MMDIT_MAP_BASIC:
-        if k[1] in sd:
-            new_sd[k[0]] = sd.pop(k[1])
-    for k in sd:
-        new_sd[k] = sd[k]
-    sd = new_sd
-
-    y_emb_shape = sd["y_embedder.mlp.0.weight"].shape
-    depth = y_emb_shape[0] // 64
-    hidden_size = 64 * depth
-    num_heads = depth
-    head_dim = hidden_size // num_heads
-    num_blocks = comfy.model_detection.count_blocks(new_sd, 'transformer_blocks.{}.')
-
-    load_device = comfy.model_management.get_torch_device()
-    offload_device = comfy.model_management.unet_offload_device()
-    unet_dtype = comfy.model_management.unet_dtype(model_params=-1)
-
-    manual_cast_dtype = comfy.model_management.unet_manual_cast(unet_dtype, load_device)
-
-    operations = model_options.get("custom_operations", None)
-    if operations is None:
-        operations = comfy.ops.pick_operations(unet_dtype, manual_cast_dtype, disable_fast_fp8=True)
-
-    control_model = comfy.cldm.dit_embedder.ControlNetEmbedder(img_size=None,
-                                                               patch_size=2,
-                                                               in_chans=16,
-                                                               num_layers=num_blocks,
-                                                               main_model_double=depth,
-                                                               double_y_emb=y_emb_shape[0] == y_emb_shape[1],
-                                                               attention_head_dim=head_dim,
-                                                               num_attention_heads=num_heads,
-                                                               adm_in_channels=2048,
-                                                               device=offload_device,
-                                                               dtype=unet_dtype,
-                                                               operations=operations)
-
-    control_model = controlnet_load_state_dict(control_model, sd)
-
-    latent_format = comfy.latent_formats.SD3()
-    preprocess_image = lambda a: a
-    if canny_cnet:
-        preprocess_image = lambda a: (a * 255 * 0.5 + 0.5)
-    elif depth_cnet:
-        preprocess_image = lambda a: 1.0 - a
-
-    control = ControlNetSD35(control_model, compression_ratio=1, latent_format=latent_format, load_device=load_device, manual_cast_dtype=manual_cast_dtype, preprocess_image=preprocess_image)
-    return control
-
-
-
-def load_controlnet_hunyuandit(controlnet_data, model_options={}):
-    model_config, operations, load_device, unet_dtype, manual_cast_dtype, offload_device = controlnet_config(controlnet_data, model_options=model_options)
-
-    control_model = comfy.ldm.hydit.controlnet.HunYuanControlNet(operations=operations, device=offload_device, dtype=unet_dtype)
-    control_model = controlnet_load_state_dict(control_model, controlnet_data)
-
-    latent_format = comfy.latent_formats.SDXL()
-    extra_conds = ['text_embedding_mask', 'encoder_hidden_states_t5', 'text_embedding_mask_t5', 'image_meta_size', 'style', 'cos_cis_img', 'sin_cis_img']
-    control = ControlNet(control_model, compression_ratio=1, latent_format=latent_format, load_device=load_device, manual_cast_dtype=manual_cast_dtype, extra_conds=extra_conds, strength_type=StrengthType.CONSTANT)
-    return control
-
-def load_controlnet_flux_xlabs_mistoline(sd, mistoline=False, model_options={}):
-    model_config, operations, load_device, unet_dtype, manual_cast_dtype, offload_device = controlnet_config(sd, model_options=model_options)
-    control_model = comfy.ldm.flux.controlnet.ControlNetFlux(mistoline=mistoline, operations=operations, device=offload_device, dtype=unet_dtype, **model_config.unet_config)
-    control_model = controlnet_load_state_dict(control_model, sd)
-    extra_conds = ['y', 'guidance']
-    control = ControlNet(control_model, load_device=load_device, manual_cast_dtype=manual_cast_dtype, extra_conds=extra_conds)
-    return control
-
-def load_controlnet_flux_instantx(sd, model_options={}):
-    new_sd = comfy.model_detection.convert_diffusers_mmdit(sd, "")
-    model_config, operations, load_device, unet_dtype, manual_cast_dtype, offload_device = controlnet_config(new_sd, model_options=model_options)
-    for k in sd:
-        new_sd[k] = sd[k]
-
-    num_union_modes = 0
-    union_cnet = "controlnet_mode_embedder.weight"
-    if union_cnet in new_sd:
-        num_union_modes = new_sd[union_cnet].shape[0]
-
-    control_latent_channels = new_sd.get("pos_embed_input.weight").shape[1] // 4
-    concat_mask = False
-    if control_latent_channels == 17:
-        concat_mask = True
-
-    control_model = comfy.ldm.flux.controlnet.ControlNetFlux(latent_input=True, num_union_modes=num_union_modes, control_latent_channels=control_latent_channels, operations=operations, device=offload_device, dtype=unet_dtype, **model_config.unet_config)
-    control_model = controlnet_load_state_dict(control_model, new_sd)
-
-    latent_format = comfy.latent_formats.Flux()
-    extra_conds = ['y', 'guidance']
-    control = ControlNet(control_model, compression_ratio=1, latent_format=latent_format, concat_mask=concat_mask, load_device=load_device, manual_cast_dtype=manual_cast_dtype, extra_conds=extra_conds)
-    return control
-
-def convert_mistoline(sd):
-    return comfy.utils.state_dict_prefix_replace(sd, {"single_controlnet_blocks.": "controlnet_single_blocks."})
-
-
-def load_controlnet_state_dict(state_dict, model=None, model_options={}):
-    controlnet_data = state_dict
-    if 'after_proj_list.18.bias' in controlnet_data.keys(): #Hunyuan DiT
-        return load_controlnet_hunyuandit(controlnet_data, model_options=model_options)
-
-    if "lora_controlnet" in controlnet_data:
-        return ControlLora(controlnet_data, model_options=model_options)
-
-    controlnet_config = None
-    supported_inference_dtypes = None
-
-    if "controlnet_cond_embedding.conv_in.weight" in controlnet_data: #diffusers format
-        controlnet_config = comfy.model_detection.unet_config_from_diffusers_unet(controlnet_data)
-        diffusers_keys = comfy.utils.unet_to_diffusers(controlnet_config)
-        diffusers_keys["controlnet_mid_block.weight"] = "middle_block_out.0.weight"
-        diffusers_keys["controlnet_mid_block.bias"] = "middle_block_out.0.bias"
-
-        count = 0
-        loop = True
-        while loop:
-            suffix = [".weight", ".bias"]
-            for s in suffix:
-                k_in = "controlnet_down_blocks.{}{}".format(count, s)
-                k_out = "zero_convs.{}.0{}".format(count, s)
-                if k_in not in controlnet_data:
-                    loop = False
-                    break
-                diffusers_keys[k_in] = k_out
-            count += 1
-
-        count = 0
-        loop = True
-        while loop:
-            suffix = [".weight", ".bias"]
-            for s in suffix:
-                if count == 0:
-                    k_in = "controlnet_cond_embedding.conv_in{}".format(s)
-                else:
-                    k_in = "controlnet_cond_embedding.blocks.{}{}".format(count - 1, s)
-                k_out = "input_hint_block.{}{}".format(count * 2, s)
-                if k_in not in controlnet_data:
-                    k_in = "controlnet_cond_embedding.conv_out{}".format(s)
-                    loop = False
-                diffusers_keys[k_in] = k_out
-            count += 1
-
-        new_sd = {}
-        for k in diffusers_keys:
-            if k in controlnet_data:
-                new_sd[diffusers_keys[k]] = controlnet_data.pop(k)
-
-        if "control_add_embedding.linear_1.bias" in controlnet_data: #Union Controlnet
-            controlnet_config["union_controlnet_num_control_type"] = controlnet_data["task_embedding"].shape[0]
-            for k in list(controlnet_data.keys()):
-                new_k = k.replace('.attn.in_proj_', '.attn.in_proj.')
-                new_sd[new_k] = controlnet_data.pop(k)
-
-        leftover_keys = controlnet_data.keys()
-        if len(leftover_keys) > 0:
-            logging.warning("leftover keys: {}".format(leftover_keys))
-        controlnet_data = new_sd
-    elif "controlnet_blocks.0.weight" in controlnet_data:
-        if "double_blocks.0.img_attn.norm.key_norm.scale" in controlnet_data:
-            return load_controlnet_flux_xlabs_mistoline(controlnet_data, model_options=model_options)
-        elif "pos_embed_input.proj.weight" in controlnet_data:
-            if "transformer_blocks.0.adaLN_modulation.1.bias" in controlnet_data:
-                return load_controlnet_sd35(controlnet_data, model_options=model_options) #Stability sd3.5 format
-            else:
-                return load_controlnet_mmdit(controlnet_data, model_options=model_options) #SD3 diffusers controlnet
-        elif "controlnet_x_embedder.weight" in controlnet_data:
-            return load_controlnet_flux_instantx(controlnet_data, model_options=model_options)
-    elif "controlnet_blocks.0.linear.weight" in controlnet_data: #mistoline flux
-        return load_controlnet_flux_xlabs_mistoline(convert_mistoline(controlnet_data), mistoline=True, model_options=model_options)
-
-    pth_key = 'control_model.zero_convs.0.0.weight'
-    pth = False
-    key = 'zero_convs.0.0.weight'
-    if pth_key in controlnet_data:
-        pth = True
-        key = pth_key
-        prefix = "control_model."
-    elif key in controlnet_data:
-        prefix = ""
-    else:
-        net = load_t2i_adapter(controlnet_data, model_options=model_options)
-        if net is None:
-            logging.error("error could not detect control model type.")
-        return net
-
-    if controlnet_config is None:
-        model_config = comfy.model_detection.model_config_from_unet(controlnet_data, prefix, True)
-        supported_inference_dtypes = list(model_config.supported_inference_dtypes)
-        controlnet_config = model_config.unet_config
-
-    unet_dtype = model_options.get("dtype", None)
-    if unet_dtype is None:
-        weight_dtype = comfy.utils.weight_dtype(controlnet_data)
-
-        if supported_inference_dtypes is None:
-            supported_inference_dtypes = [comfy.model_management.unet_dtype()]
-
-        unet_dtype = comfy.model_management.unet_dtype(model_params=-1, supported_dtypes=supported_inference_dtypes, weight_dtype=weight_dtype)
-
-    load_device = comfy.model_management.get_torch_device()
-
-    manual_cast_dtype = comfy.model_management.unet_manual_cast(unet_dtype, load_device)
-    operations = model_options.get("custom_operations", None)
-    if operations is None:
-        operations = comfy.ops.pick_operations(unet_dtype, manual_cast_dtype)
-
-    controlnet_config["operations"] = operations
-    controlnet_config["dtype"] = unet_dtype
-    controlnet_config["device"] = comfy.model_management.unet_offload_device()
-    controlnet_config.pop("out_channels")
-    controlnet_config["hint_channels"] = controlnet_data["{}input_hint_block.0.weight".format(prefix)].shape[1]
-    control_model = comfy.cldm.cldm.ControlNet(**controlnet_config)
-
-    if pth:
-        if 'difference' in controlnet_data:
-            if model is not None:
-                comfy.model_management.load_models_gpu([model])
-                model_sd = model.model_state_dict()
-                for x in controlnet_data:
-                    c_m = "control_model."
-                    if x.startswith(c_m):
-                        sd_key = "diffusion_model.{}".format(x[len(c_m):])
-                        if sd_key in model_sd:
-                            cd = controlnet_data[x]
-                            cd += model_sd[sd_key].type(cd.dtype).to(cd.device)
-            else:
-                logging.warning("WARNING: Loaded a diff controlnet without a model. It will very likely not work.")
-
-        class WeightsLoader(torch.nn.Module):
-            pass
-        w = WeightsLoader()
-        w.control_model = control_model
-        missing, unexpected = w.load_state_dict(controlnet_data, strict=False)
-    else:
-        missing, unexpected = control_model.load_state_dict(controlnet_data, strict=False)
-
-    if len(missing) > 0:
-        logging.warning("missing controlnet keys: {}".format(missing))
-
-    if len(unexpected) > 0:
-        logging.debug("unexpected controlnet keys: {}".format(unexpected))
-
-    global_average_pooling = model_options.get("global_average_pooling", False)
-    control = ControlNet(control_model, global_average_pooling=global_average_pooling, load_device=load_device, manual_cast_dtype=manual_cast_dtype)
-    return control
-
-def load_controlnet(ckpt_path, model=None, model_options={}):
-    model_options = model_options.copy()
-    if "global_average_pooling" not in model_options:
-        filename = os.path.splitext(ckpt_path)[0]
-        if filename.endswith("_shuffle") or filename.endswith("_shuffle_fp16"): #TODO: smarter way of enabling global_average_pooling
-            model_options["global_average_pooling"] = True
-
-    cnet = load_controlnet_state_dict(comfy.utils.load_torch_file(ckpt_path, safe_load=True), model=model, model_options=model_options)
-    if cnet is None:
-        logging.error("error checkpoint does not contain controlnet or t2i adapter data {}".format(ckpt_path))
-    return cnet
-
-class T2IAdapter(ControlBase):
-    def __init__(self, t2i_model, channels_in, compression_ratio, upscale_algorithm, device=None):
-        super().__init__()
-        self.t2i_model = t2i_model
-        self.channels_in = channels_in
-        self.control_input = None
-        self.compression_ratio = compression_ratio
-        self.upscale_algorithm = upscale_algorithm
-        if device is None:
-            device = comfy.model_management.get_torch_device()
-        self.device = device
-
-    def scale_image_to(self, width, height):
-        unshuffle_amount = self.t2i_model.unshuffle_amount
-        width = math.ceil(width / unshuffle_amount) * unshuffle_amount
-        height = math.ceil(height / unshuffle_amount) * unshuffle_amount
-        return width, height
-
-    def get_control(self, x_noisy, t, cond, batched_number, transformer_options):
-        control_prev = None
-        if self.previous_controlnet is not None:
-            control_prev = self.previous_controlnet.get_control(x_noisy, t, cond, batched_number, transformer_options)
-
-        if self.timestep_range is not None:
-            if t[0] > self.timestep_range[0] or t[0] < self.timestep_range[1]:
-                if control_prev is not None:
-                    return control_prev
-                else:
-                    return None
-
-        if self.cond_hint is None or x_noisy.shape[2] * self.compression_ratio != self.cond_hint.shape[2] or x_noisy.shape[3] * self.compression_ratio != self.cond_hint.shape[3]:
-            if self.cond_hint is not None:
-                del self.cond_hint
-            self.control_input = None
-            self.cond_hint = None
-            width, height = self.scale_image_to(x_noisy.shape[3] * self.compression_ratio, x_noisy.shape[2] * self.compression_ratio)
-            self.cond_hint = comfy.utils.common_upscale(self.cond_hint_original, width, height, self.upscale_algorithm, "center").float().to(self.device)
-            if self.channels_in == 1 and self.cond_hint.shape[1] > 1:
-                self.cond_hint = torch.mean(self.cond_hint, 1, keepdim=True)
-        if x_noisy.shape[0] != self.cond_hint.shape[0]:
-            self.cond_hint = broadcast_image_to(self.cond_hint, x_noisy.shape[0], batched_number)
-        if self.control_input is None:
-            self.t2i_model.to(x_noisy.dtype)
-            self.t2i_model.to(self.device)
-            self.control_input = self.t2i_model(self.cond_hint.to(x_noisy.dtype))
-            self.t2i_model.cpu()
-
-        control_input = {}
-        for k in self.control_input:
-            control_input[k] = list(map(lambda a: None if a is None else a.clone(), self.control_input[k]))
-
-        return self.control_merge(control_input, control_prev, x_noisy.dtype)
-
-    def copy(self):
-        c = T2IAdapter(self.t2i_model, self.channels_in, self.compression_ratio, self.upscale_algorithm)
-        self.copy_to(c)
-        return c
-
-def load_t2i_adapter(t2i_data, model_options={}): #TODO: model_options
-    compression_ratio = 8
-    upscale_algorithm = 'nearest-exact'
-
-    if 'adapter' in t2i_data:
-        t2i_data = t2i_data['adapter']
-    if 'adapter.body.0.resnets.0.block1.weight' in t2i_data: #diffusers format
-        prefix_replace = {}
-        for i in range(4):
-            for j in range(2):
-                prefix_replace["adapter.body.{}.resnets.{}.".format(i, j)] = "body.{}.".format(i * 2 + j)
-            prefix_replace["adapter.body.{}.".format(i, )] = "body.{}.".format(i * 2)
-        prefix_replace["adapter."] = ""
-        t2i_data = comfy.utils.state_dict_prefix_replace(t2i_data, prefix_replace)
-    keys = t2i_data.keys()
-
-    if "body.0.in_conv.weight" in keys:
-        cin = t2i_data['body.0.in_conv.weight'].shape[1]
-        model_ad = comfy.t2i_adapter.adapter.Adapter_light(cin=cin, channels=[320, 640, 1280, 1280], nums_rb=4)
-    elif 'conv_in.weight' in keys:
-        cin = t2i_data['conv_in.weight'].shape[1]
-        channel = t2i_data['conv_in.weight'].shape[0]
-        ksize = t2i_data['body.0.block2.weight'].shape[2]
-        use_conv = False
-        down_opts = list(filter(lambda a: a.endswith("down_opt.op.weight"), keys))
-        if len(down_opts) > 0:
-            use_conv = True
-        xl = False
-        if cin == 256 or cin == 768:
-            xl = True
-        model_ad = comfy.t2i_adapter.adapter.Adapter(cin=cin, channels=[channel, channel*2, channel*4, channel*4][:4], nums_rb=2, ksize=ksize, sk=True, use_conv=use_conv, xl=xl)
-    elif "backbone.0.0.weight" in keys:
-        model_ad = comfy.ldm.cascade.controlnet.ControlNet(c_in=t2i_data['backbone.0.0.weight'].shape[1], proj_blocks=[0, 4, 8, 12, 51, 55, 59, 63])
-        compression_ratio = 32
-        upscale_algorithm = 'bilinear'
-    elif "backbone.10.blocks.0.weight" in keys:
-        model_ad = comfy.ldm.cascade.controlnet.ControlNet(c_in=t2i_data['backbone.0.weight'].shape[1], bottleneck_mode="large", proj_blocks=[0, 4, 8, 12, 51, 55, 59, 63])
-        compression_ratio = 1
-        upscale_algorithm = 'nearest-exact'
-    else:
-        return None
-
-    missing, unexpected = model_ad.load_state_dict(t2i_data)
-    if len(missing) > 0:
-        logging.warning("t2i missing {}".format(missing))
-
-    if len(unexpected) > 0:
-        logging.debug("t2i unexpected {}".format(unexpected))
-
-    return T2IAdapter(model_ad, model_ad.input_channels, compression_ratio, upscale_algorithm)

comfy/diffusers_convert.py

@@ -1,189 +0,0 @@
-import re
-import torch
-import logging
-
-# conversion code from https://github.com/huggingface/diffusers/blob/main/scripts/convert_diffusers_to_original_stable_diffusion.py
-
-# ================#
-# VAE Conversion #
-# ================#
-
-vae_conversion_map = [
-    # (stable-diffusion, HF Diffusers)
-    ("nin_shortcut", "conv_shortcut"),
-    ("norm_out", "conv_norm_out"),
-    ("mid.attn_1.", "mid_block.attentions.0."),
-]
-
-for i in range(4):
-    # down_blocks have two resnets
-    for j in range(2):
-        hf_down_prefix = f"encoder.down_blocks.{i}.resnets.{j}."
-        sd_down_prefix = f"encoder.down.{i}.block.{j}."
-        vae_conversion_map.append((sd_down_prefix, hf_down_prefix))
-
-    if i < 3:
-        hf_downsample_prefix = f"down_blocks.{i}.downsamplers.0."
-        sd_downsample_prefix = f"down.{i}.downsample."
-        vae_conversion_map.append((sd_downsample_prefix, hf_downsample_prefix))
-
-        hf_upsample_prefix = f"up_blocks.{i}.upsamplers.0."
-        sd_upsample_prefix = f"up.{3 - i}.upsample."
-        vae_conversion_map.append((sd_upsample_prefix, hf_upsample_prefix))
-
-    # up_blocks have three resnets
-    # also, up blocks in hf are numbered in reverse from sd
-    for j in range(3):
-        hf_up_prefix = f"decoder.up_blocks.{i}.resnets.{j}."
-        sd_up_prefix = f"decoder.up.{3 - i}.block.{j}."
-        vae_conversion_map.append((sd_up_prefix, hf_up_prefix))
-
-# this part accounts for mid blocks in both the encoder and the decoder
-for i in range(2):
-    hf_mid_res_prefix = f"mid_block.resnets.{i}."
-    sd_mid_res_prefix = f"mid.block_{i + 1}."
-    vae_conversion_map.append((sd_mid_res_prefix, hf_mid_res_prefix))
-
-vae_conversion_map_attn = [
-    # (stable-diffusion, HF Diffusers)
-    ("norm.", "group_norm."),
-    ("q.", "query."),
-    ("k.", "key."),
-    ("v.", "value."),
-    ("q.", "to_q."),
-    ("k.", "to_k."),
-    ("v.", "to_v."),
-    ("proj_out.", "to_out.0."),
-    ("proj_out.", "proj_attn."),
-]
-
-
-def reshape_weight_for_sd(w, conv3d=False):
-    # convert HF linear weights to SD conv2d weights
-    if conv3d:
-        return w.reshape(*w.shape, 1, 1, 1)
-    else:
-        return w.reshape(*w.shape, 1, 1)
-
-
-def convert_vae_state_dict(vae_state_dict):
-    mapping = {k: k for k in vae_state_dict.keys()}
-    conv3d = False
-    for k, v in mapping.items():
-        for sd_part, hf_part in vae_conversion_map:
-            v = v.replace(hf_part, sd_part)
-        if v.endswith(".conv.weight"):
-            if not conv3d and vae_state_dict[k].ndim == 5:
-                conv3d = True
-        mapping[k] = v
-    for k, v in mapping.items():
-        if "attentions" in k:
-            for sd_part, hf_part in vae_conversion_map_attn:
-                v = v.replace(hf_part, sd_part)
-            mapping[k] = v
-    new_state_dict = {v: vae_state_dict[k] for k, v in mapping.items()}
-    weights_to_convert = ["q", "k", "v", "proj_out"]
-    for k, v in new_state_dict.items():
-        for weight_name in weights_to_convert:
-            if f"mid.attn_1.{weight_name}.weight" in k:
-                logging.debug(f"Reshaping {k} for SD format")
-                new_state_dict[k] = reshape_weight_for_sd(v, conv3d=conv3d)
-    return new_state_dict
-
-
-# =========================#
-# Text Encoder Conversion #
-# =========================#
-
-
-textenc_conversion_lst = [
-    # (stable-diffusion, HF Diffusers)
-    ("resblocks.", "text_model.encoder.layers."),
-    ("ln_1", "layer_norm1"),
-    ("ln_2", "layer_norm2"),
-    (".c_fc.", ".fc1."),
-    (".c_proj.", ".fc2."),
-    (".attn", ".self_attn"),
-    ("ln_final.", "transformer.text_model.final_layer_norm."),
-    ("token_embedding.weight", "transformer.text_model.embeddings.token_embedding.weight"),
-    ("positional_embedding", "transformer.text_model.embeddings.position_embedding.weight"),
-]
-protected = {re.escape(x[1]): x[0] for x in textenc_conversion_lst}
-textenc_pattern = re.compile("|".join(protected.keys()))
-
-# Ordering is from https://github.com/pytorch/pytorch/blob/master/test/cpp/api/modules.cpp
-code2idx = {"q": 0, "k": 1, "v": 2}
-
-
-# This function exists because at the time of writing torch.cat can't do fp8 with cuda
-def cat_tensors(tensors):
-    x = 0
-    for t in tensors:
-        x += t.shape[0]
-
-    shape = [x] + list(tensors[0].shape)[1:]
-    out = torch.empty(shape, device=tensors[0].device, dtype=tensors[0].dtype)
-
-    x = 0
-    for t in tensors:
-        out[x:x + t.shape[0]] = t
-        x += t.shape[0]
-
-    return out
-
-
-def convert_text_enc_state_dict_v20(text_enc_dict, prefix=""):
-    new_state_dict = {}
-    capture_qkv_weight = {}
-    capture_qkv_bias = {}
-    for k, v in text_enc_dict.items():
-        if not k.startswith(prefix):
-            continue
-        if (
-                k.endswith(".self_attn.q_proj.weight")
-                or k.endswith(".self_attn.k_proj.weight")
-                or k.endswith(".self_attn.v_proj.weight")
-        ):
-            k_pre = k[: -len(".q_proj.weight")]
-            k_code = k[-len("q_proj.weight")]
-            if k_pre not in capture_qkv_weight:
-                capture_qkv_weight[k_pre] = [None, None, None]
-            capture_qkv_weight[k_pre][code2idx[k_code]] = v
-            continue
-
-        if (
-                k.endswith(".self_attn.q_proj.bias")
-                or k.endswith(".self_attn.k_proj.bias")
-                or k.endswith(".self_attn.v_proj.bias")
-        ):
-            k_pre = k[: -len(".q_proj.bias")]
-            k_code = k[-len("q_proj.bias")]
-            if k_pre not in capture_qkv_bias:
-                capture_qkv_bias[k_pre] = [None, None, None]
-            capture_qkv_bias[k_pre][code2idx[k_code]] = v
-            continue
-
-        text_proj = "transformer.text_projection.weight"
-        if k.endswith(text_proj):
-            new_state_dict[k.replace(text_proj, "text_projection")] = v.transpose(0, 1).contiguous()
-        else:
-            relabelled_key = textenc_pattern.sub(lambda m: protected[re.escape(m.group(0))], k)
-            new_state_dict[relabelled_key] = v
-
-    for k_pre, tensors in capture_qkv_weight.items():
-        if None in tensors:
-            raise Exception("CORRUPTED MODEL: one of the q-k-v values for the text encoder was missing")
-        relabelled_key = textenc_pattern.sub(lambda m: protected[re.escape(m.group(0))], k_pre)
-        new_state_dict[relabelled_key + ".in_proj_weight"] = cat_tensors(tensors)
-
-    for k_pre, tensors in capture_qkv_bias.items():
-        if None in tensors:
-            raise Exception("CORRUPTED MODEL: one of the q-k-v values for the text encoder was missing")
-        relabelled_key = textenc_pattern.sub(lambda m: protected[re.escape(m.group(0))], k_pre)
-        new_state_dict[relabelled_key + ".in_proj_bias"] = cat_tensors(tensors)
-
-    return new_state_dict
-
-
-def convert_text_enc_state_dict(text_enc_dict):
-    return text_enc_dict

comfy/diffusers_load.py

@@ -1,36 +0,0 @@
-import os
-
-import comfy.sd
-
-def first_file(path, filenames):
-    for f in filenames:
-        p = os.path.join(path, f)
-        if os.path.exists(p):
-            return p
-    return None
-
-def load_diffusers(model_path, output_vae=True, output_clip=True, embedding_directory=None):
-    diffusion_model_names = ["diffusion_pytorch_model.fp16.safetensors", "diffusion_pytorch_model.safetensors", "diffusion_pytorch_model.fp16.bin", "diffusion_pytorch_model.bin"]
-    unet_path = first_file(os.path.join(model_path, "unet"), diffusion_model_names)
-    vae_path = first_file(os.path.join(model_path, "vae"), diffusion_model_names)
-
-    text_encoder_model_names = ["model.fp16.safetensors", "model.safetensors", "pytorch_model.fp16.bin", "pytorch_model.bin"]
-    text_encoder1_path = first_file(os.path.join(model_path, "text_encoder"), text_encoder_model_names)
-    text_encoder2_path = first_file(os.path.join(model_path, "text_encoder_2"), text_encoder_model_names)
-
-    text_encoder_paths = [text_encoder1_path]
-    if text_encoder2_path is not None:
-        text_encoder_paths.append(text_encoder2_path)
-
-    unet = comfy.sd.load_diffusion_model(unet_path)
-
-    clip = None
-    if output_clip:
-        clip = comfy.sd.load_clip(text_encoder_paths, embedding_directory=embedding_directory)
-
-    vae = None
-    if output_vae:
-        sd = comfy.utils.load_torch_file(vae_path)
-        vae = comfy.sd.VAE(sd=sd)
-
-    return (unet, clip, vae)

comfy/extra_samplers/uni_pc.py

@@ -1,873 +0,0 @@
-#code taken from: https://github.com/wl-zhao/UniPC and modified
-
-import torch
-import math
-import logging
-
-from tqdm.auto import trange
-
-
-class NoiseScheduleVP:
-    def __init__(
-            self,
-            schedule='discrete',
-            betas=None,
-            alphas_cumprod=None,
-            continuous_beta_0=0.1,
-            continuous_beta_1=20.,
-        ):
-        r"""Create a wrapper class for the forward SDE (VP type).
-
-        ***
-        Update: We support discrete-time diffusion models by implementing a picewise linear interpolation for log_alpha_t.
-                We recommend to use schedule='discrete' for the discrete-time diffusion models, especially for high-resolution images.
-        ***
-
-        The forward SDE ensures that the condition distribution q_{t|0}(x_t | x_0) = N ( alpha_t * x_0, sigma_t^2 * I ).
-        We further define lambda_t = log(alpha_t) - log(sigma_t), which is the half-logSNR (described in the DPM-Solver paper).
-        Therefore, we implement the functions for computing alpha_t, sigma_t and lambda_t. For t in [0, T], we have:
-
-            log_alpha_t = self.marginal_log_mean_coeff(t)
-            sigma_t = self.marginal_std(t)
-            lambda_t = self.marginal_lambda(t)
-
-        Moreover, as lambda(t) is an invertible function, we also support its inverse function:
-
-            t = self.inverse_lambda(lambda_t)
-
-        ===============================================================
-
-        We support both discrete-time DPMs (trained on n = 0, 1, ..., N-1) and continuous-time DPMs (trained on t in [t_0, T]).
-
-        1. For discrete-time DPMs:
-
-            For discrete-time DPMs trained on n = 0, 1, ..., N-1, we convert the discrete steps to continuous time steps by:
-                t_i = (i + 1) / N
-            e.g. for N = 1000, we have t_0 = 1e-3 and T = t_{N-1} = 1.
-            We solve the corresponding diffusion ODE from time T = 1 to time t_0 = 1e-3.
-
-            Args:
-                betas: A `torch.Tensor`. The beta array for the discrete-time DPM. (See the original DDPM paper for details)
-                alphas_cumprod: A `torch.Tensor`. The cumprod alphas for the discrete-time DPM. (See the original DDPM paper for details)
-
-            Note that we always have alphas_cumprod = cumprod(betas). Therefore, we only need to set one of `betas` and `alphas_cumprod`.
-
-            **Important**:  Please pay special attention for the args for `alphas_cumprod`:
-                The `alphas_cumprod` is the \hat{alpha_n} arrays in the notations of DDPM. Specifically, DDPMs assume that
-                    q_{t_n | 0}(x_{t_n} | x_0) = N ( \sqrt{\hat{alpha_n}} * x_0, (1 - \hat{alpha_n}) * I ).
-                Therefore, the notation \hat{alpha_n} is different from the notation alpha_t in DPM-Solver. In fact, we have
-                    alpha_{t_n} = \sqrt{\hat{alpha_n}},
-                and
-                    log(alpha_{t_n}) = 0.5 * log(\hat{alpha_n}).
-
-
-        2. For continuous-time DPMs:
-
-            We support two types of VPSDEs: linear (DDPM) and cosine (improved-DDPM). The hyperparameters for the noise
-            schedule are the default settings in DDPM and improved-DDPM:
-
-            Args:
-                beta_min: A `float` number. The smallest beta for the linear schedule.
-                beta_max: A `float` number. The largest beta for the linear schedule.
-                cosine_s: A `float` number. The hyperparameter in the cosine schedule.
-                cosine_beta_max: A `float` number. The hyperparameter in the cosine schedule.
-                T: A `float` number. The ending time of the forward process.
-
-        ===============================================================
-
-        Args:
-            schedule: A `str`. The noise schedule of the forward SDE. 'discrete' for discrete-time DPMs,
-                    'linear' or 'cosine' for continuous-time DPMs.
-        Returns:
-            A wrapper object of the forward SDE (VP type).
-
-        ===============================================================
-
-        Example:
-
-        # For discrete-time DPMs, given betas (the beta array for n = 0, 1, ..., N - 1):
-        >>> ns = NoiseScheduleVP('discrete', betas=betas)
-
-        # For discrete-time DPMs, given alphas_cumprod (the \hat{alpha_n} array for n = 0, 1, ..., N - 1):
-        >>> ns = NoiseScheduleVP('discrete', alphas_cumprod=alphas_cumprod)
-
-        # For continuous-time DPMs (VPSDE), linear schedule:
-        >>> ns = NoiseScheduleVP('linear', continuous_beta_0=0.1, continuous_beta_1=20.)
-
-        """
-
-        if schedule not in ['discrete', 'linear', 'cosine']:
-            raise ValueError("Unsupported noise schedule {}. The schedule needs to be 'discrete' or 'linear' or 'cosine'".format(schedule))
-
-        self.schedule = schedule
-        if schedule == 'discrete':
-            if betas is not None:
-                log_alphas = 0.5 * torch.log(1 - betas).cumsum(dim=0)
-            else:
-                assert alphas_cumprod is not None
-                log_alphas = 0.5 * torch.log(alphas_cumprod)
-            self.total_N = len(log_alphas)
-            self.T = 1.
-            self.t_array = torch.linspace(0., 1., self.total_N + 1)[1:].reshape((1, -1))
-            self.log_alpha_array = log_alphas.reshape((1, -1,))
-        else:
-            self.total_N = 1000
-            self.beta_0 = continuous_beta_0
-            self.beta_1 = continuous_beta_1
-            self.cosine_s = 0.008
-            self.cosine_beta_max = 999.
-            self.cosine_t_max = math.atan(self.cosine_beta_max * (1. + self.cosine_s) / math.pi) * 2. * (1. + self.cosine_s) / math.pi - self.cosine_s
-            self.cosine_log_alpha_0 = math.log(math.cos(self.cosine_s / (1. + self.cosine_s) * math.pi / 2.))
-            self.schedule = schedule
-            if schedule == 'cosine':
-                # For the cosine schedule, T = 1 will have numerical issues. So we manually set the ending time T.
-                # Note that T = 0.9946 may be not the optimal setting. However, we find it works well.
-                self.T = 0.9946
-            else:
-                self.T = 1.
-
-    def marginal_log_mean_coeff(self, t):
-        """
-        Compute log(alpha_t) of a given continuous-time label t in [0, T].
-        """
-        if self.schedule == 'discrete':
-            return interpolate_fn(t.reshape((-1, 1)), self.t_array.to(t.device), self.log_alpha_array.to(t.device)).reshape((-1))
-        elif self.schedule == 'linear':
-            return -0.25 * t ** 2 * (self.beta_1 - self.beta_0) - 0.5 * t * self.beta_0
-        elif self.schedule == 'cosine':
-            log_alpha_fn = lambda s: torch.log(torch.cos((s + self.cosine_s) / (1. + self.cosine_s) * math.pi / 2.))
-            log_alpha_t =  log_alpha_fn(t) - self.cosine_log_alpha_0
-            return log_alpha_t
-
-    def marginal_alpha(self, t):
-        """
-        Compute alpha_t of a given continuous-time label t in [0, T].
-        """
-        return torch.exp(self.marginal_log_mean_coeff(t))
-
-    def marginal_std(self, t):
-        """
-        Compute sigma_t of a given continuous-time label t in [0, T].
-        """
-        return torch.sqrt(1. - torch.exp(2. * self.marginal_log_mean_coeff(t)))
-
-    def marginal_lambda(self, t):
-        """
-        Compute lambda_t = log(alpha_t) - log(sigma_t) of a given continuous-time label t in [0, T].
-        """
-        log_mean_coeff = self.marginal_log_mean_coeff(t)
-        log_std = 0.5 * torch.log(1. - torch.exp(2. * log_mean_coeff))
-        return log_mean_coeff - log_std
-
-    def inverse_lambda(self, lamb):
-        """
-        Compute the continuous-time label t in [0, T] of a given half-logSNR lambda_t.
-        """
-        if self.schedule == 'linear':
-            tmp = 2. * (self.beta_1 - self.beta_0) * torch.logaddexp(-2. * lamb, torch.zeros((1,)).to(lamb))
-            Delta = self.beta_0**2 + tmp
-            return tmp / (torch.sqrt(Delta) + self.beta_0) / (self.beta_1 - self.beta_0)
-        elif self.schedule == 'discrete':
-            log_alpha = -0.5 * torch.logaddexp(torch.zeros((1,)).to(lamb.device), -2. * lamb)
-            t = interpolate_fn(log_alpha.reshape((-1, 1)), torch.flip(self.log_alpha_array.to(lamb.device), [1]), torch.flip(self.t_array.to(lamb.device), [1]))
-            return t.reshape((-1,))
-        else:
-            log_alpha = -0.5 * torch.logaddexp(-2. * lamb, torch.zeros((1,)).to(lamb))
-            t_fn = lambda log_alpha_t: torch.arccos(torch.exp(log_alpha_t + self.cosine_log_alpha_0)) * 2. * (1. + self.cosine_s) / math.pi - self.cosine_s
-            t = t_fn(log_alpha)
-            return t
-
-
-def model_wrapper(
-    model,
-    noise_schedule,
-    model_type="noise",
-    model_kwargs={},
-    guidance_type="uncond",
-    condition=None,
-    unconditional_condition=None,
-    guidance_scale=1.,
-    classifier_fn=None,
-    classifier_kwargs={},
-):
-    """Create a wrapper function for the noise prediction model.
-
-    DPM-Solver needs to solve the continuous-time diffusion ODEs. For DPMs trained on discrete-time labels, we need to
-    firstly wrap the model function to a noise prediction model that accepts the continuous time as the input.
-
-    We support four types of the diffusion model by setting `model_type`:
-
-        1. "noise": noise prediction model. (Trained by predicting noise).
-
-        2. "x_start": data prediction model. (Trained by predicting the data x_0 at time 0).
-
-        3. "v": velocity prediction model. (Trained by predicting the velocity).
-            The "v" prediction is derivation detailed in Appendix D of [1], and is used in Imagen-Video [2].
-
-            [1] Salimans, Tim, and Jonathan Ho. "Progressive distillation for fast sampling of diffusion models."
-                arXiv preprint arXiv:2202.00512 (2022).
-            [2] Ho, Jonathan, et al. "Imagen Video: High Definition Video Generation with Diffusion Models."
-                arXiv preprint arXiv:2210.02303 (2022).
-
-        4. "score": marginal score function. (Trained by denoising score matching).
-            Note that the score function and the noise prediction model follows a simple relationship:
-            ```
-                noise(x_t, t) = -sigma_t * score(x_t, t)
-            ```
-
-    We support three types of guided sampling by DPMs by setting `guidance_type`:
-        1. "uncond": unconditional sampling by DPMs.
-            The input `model` has the following format:
-            ``
-                model(x, t_input, **model_kwargs) -> noise | x_start | v | score
-            ``
-
-        2. "classifier": classifier guidance sampling [3] by DPMs and another classifier.
-            The input `model` has the following format:
-            ``
-                model(x, t_input, **model_kwargs) -> noise | x_start | v | score
-            ``
-
-            The input `classifier_fn` has the following format:
-            ``
-                classifier_fn(x, t_input, cond, **classifier_kwargs) -> logits(x, t_input, cond)
-            ``
-
-            [3] P. Dhariwal and A. Q. Nichol, "Diffusion models beat GANs on image synthesis,"
-                in Advances in Neural Information Processing Systems, vol. 34, 2021, pp. 8780-8794.
-
-        3. "classifier-free": classifier-free guidance sampling by conditional DPMs.
-            The input `model` has the following format:
-            ``
-                model(x, t_input, cond, **model_kwargs) -> noise | x_start | v | score
-            ``
-            And if cond == `unconditional_condition`, the model output is the unconditional DPM output.
-
-            [4] Ho, Jonathan, and Tim Salimans. "Classifier-free diffusion guidance."
-                arXiv preprint arXiv:2207.12598 (2022).
-
-
-    The `t_input` is the time label of the model, which may be discrete-time labels (i.e. 0 to 999)
-    or continuous-time labels (i.e. epsilon to T).
-
-    We wrap the model function to accept only `x` and `t_continuous` as inputs, and outputs the predicted noise:
-    ``
-        def model_fn(x, t_continuous) -> noise:
-            t_input = get_model_input_time(t_continuous)
-            return noise_pred(model, x, t_input, **model_kwargs)
-    ``
-    where `t_continuous` is the continuous time labels (i.e. epsilon to T). And we use `model_fn` for DPM-Solver.
-
-    ===============================================================
-
-    Args:
-        model: A diffusion model with the corresponding format described above.
-        noise_schedule: A noise schedule object, such as NoiseScheduleVP.
-        model_type: A `str`. The parameterization type of the diffusion model.
-                    "noise" or "x_start" or "v" or "score".
-        model_kwargs: A `dict`. A dict for the other inputs of the model function.
-        guidance_type: A `str`. The type of the guidance for sampling.
-                    "uncond" or "classifier" or "classifier-free".
-        condition: A pytorch tensor. The condition for the guided sampling.
-                    Only used for "classifier" or "classifier-free" guidance type.
-        unconditional_condition: A pytorch tensor. The condition for the unconditional sampling.
-                    Only used for "classifier-free" guidance type.
-        guidance_scale: A `float`. The scale for the guided sampling.
-        classifier_fn: A classifier function. Only used for the classifier guidance.
-        classifier_kwargs: A `dict`. A dict for the other inputs of the classifier function.
-    Returns:
-        A noise prediction model that accepts the noised data and the continuous time as the inputs.
-    """
-
-    def get_model_input_time(t_continuous):
-        """
-        Convert the continuous-time `t_continuous` (in [epsilon, T]) to the model input time.
-        For discrete-time DPMs, we convert `t_continuous` in [1 / N, 1] to `t_input` in [0, 1000 * (N - 1) / N].
-        For continuous-time DPMs, we just use `t_continuous`.
-        """
-        if noise_schedule.schedule == 'discrete':
-            return (t_continuous - 1. / noise_schedule.total_N) * 1000.
-        else:
-            return t_continuous
-
-    def noise_pred_fn(x, t_continuous, cond=None):
-        if t_continuous.reshape((-1,)).shape[0] == 1:
-            t_continuous = t_continuous.expand((x.shape[0]))
-        t_input = get_model_input_time(t_continuous)
-        output = model(x, t_input, **model_kwargs)
-        if model_type == "noise":
-            return output
-        elif model_type == "x_start":
-            alpha_t, sigma_t = noise_schedule.marginal_alpha(t_continuous), noise_schedule.marginal_std(t_continuous)
-            dims = x.dim()
-            return (x - expand_dims(alpha_t, dims) * output) / expand_dims(sigma_t, dims)
-        elif model_type == "v":
-            alpha_t, sigma_t = noise_schedule.marginal_alpha(t_continuous), noise_schedule.marginal_std(t_continuous)
-            dims = x.dim()
-            return expand_dims(alpha_t, dims) * output + expand_dims(sigma_t, dims) * x
-        elif model_type == "score":
-            sigma_t = noise_schedule.marginal_std(t_continuous)
-            dims = x.dim()
-            return -expand_dims(sigma_t, dims) * output
-
-    def cond_grad_fn(x, t_input):
-        """
-        Compute the gradient of the classifier, i.e. nabla_{x} log p_t(cond | x_t).
-        """
-        with torch.enable_grad():
-            x_in = x.detach().requires_grad_(True)
-            log_prob = classifier_fn(x_in, t_input, condition, **classifier_kwargs)
-            return torch.autograd.grad(log_prob.sum(), x_in)[0]
-
-    def model_fn(x, t_continuous):
-        """
-        The noise predicition model function that is used for DPM-Solver.
-        """
-        if t_continuous.reshape((-1,)).shape[0] == 1:
-            t_continuous = t_continuous.expand((x.shape[0]))
-        if guidance_type == "uncond":
-            return noise_pred_fn(x, t_continuous)
-        elif guidance_type == "classifier":
-            assert classifier_fn is not None
-            t_input = get_model_input_time(t_continuous)
-            cond_grad = cond_grad_fn(x, t_input)
-            sigma_t = noise_schedule.marginal_std(t_continuous)
-            noise = noise_pred_fn(x, t_continuous)
-            return noise - guidance_scale * expand_dims(sigma_t, dims=cond_grad.dim()) * cond_grad
-        elif guidance_type == "classifier-free":
-            if guidance_scale == 1. or unconditional_condition is None:
-                return noise_pred_fn(x, t_continuous, cond=condition)
-            else:
-                x_in = torch.cat([x] * 2)
-                t_in = torch.cat([t_continuous] * 2)
-                c_in = torch.cat([unconditional_condition, condition])
-                noise_uncond, noise = noise_pred_fn(x_in, t_in, cond=c_in).chunk(2)
-                return noise_uncond + guidance_scale * (noise - noise_uncond)
-
-    assert model_type in ["noise", "x_start", "v"]
-    assert guidance_type in ["uncond", "classifier", "classifier-free"]
-    return model_fn
-
-
-class UniPC:
-    def __init__(
-        self,
-        model_fn,
-        noise_schedule,
-        predict_x0=True,
-        thresholding=False,
-        max_val=1.,
-        variant='bh1',
-    ):
-        """Construct a UniPC.
-
-        We support both data_prediction and noise_prediction.
-        """
-        self.model = model_fn
-        self.noise_schedule = noise_schedule
-        self.variant = variant
-        self.predict_x0 = predict_x0
-        self.thresholding = thresholding
-        self.max_val = max_val
-
-    def dynamic_thresholding_fn(self, x0, t=None):
-        """
-        The dynamic thresholding method.
-        """
-        dims = x0.dim()
-        p = self.dynamic_thresholding_ratio
-        s = torch.quantile(torch.abs(x0).reshape((x0.shape[0], -1)), p, dim=1)
-        s = expand_dims(torch.maximum(s, self.thresholding_max_val * torch.ones_like(s).to(s.device)), dims)
-        x0 = torch.clamp(x0, -s, s) / s
-        return x0
-
-    def noise_prediction_fn(self, x, t):
-        """
-        Return the noise prediction model.
-        """
-        return self.model(x, t)
-
-    def data_prediction_fn(self, x, t):
-        """
-        Return the data prediction model (with thresholding).
-        """
-        noise = self.noise_prediction_fn(x, t)
-        dims = x.dim()
-        alpha_t, sigma_t = self.noise_schedule.marginal_alpha(t), self.noise_schedule.marginal_std(t)
-        x0 = (x - expand_dims(sigma_t, dims) * noise) / expand_dims(alpha_t, dims)
-        if self.thresholding:
-            p = 0.995   # A hyperparameter in the paper of "Imagen" [1].
-            s = torch.quantile(torch.abs(x0).reshape((x0.shape[0], -1)), p, dim=1)
-            s = expand_dims(torch.maximum(s, self.max_val * torch.ones_like(s).to(s.device)), dims)
-            x0 = torch.clamp(x0, -s, s) / s
-        return x0
-
-    def model_fn(self, x, t):
-        """
-        Convert the model to the noise prediction model or the data prediction model.
-        """
-        if self.predict_x0:
-            return self.data_prediction_fn(x, t)
-        else:
-            return self.noise_prediction_fn(x, t)
-
-    def get_time_steps(self, skip_type, t_T, t_0, N, device):
-        """Compute the intermediate time steps for sampling.
-        """
-        if skip_type == 'logSNR':
-            lambda_T = self.noise_schedule.marginal_lambda(torch.tensor(t_T).to(device))
-            lambda_0 = self.noise_schedule.marginal_lambda(torch.tensor(t_0).to(device))
-            logSNR_steps = torch.linspace(lambda_T.cpu().item(), lambda_0.cpu().item(), N + 1).to(device)
-            return self.noise_schedule.inverse_lambda(logSNR_steps)
-        elif skip_type == 'time_uniform':
-            return torch.linspace(t_T, t_0, N + 1).to(device)
-        elif skip_type == 'time_quadratic':
-            t_order = 2
-            t = torch.linspace(t_T**(1. / t_order), t_0**(1. / t_order), N + 1).pow(t_order).to(device)
-            return t
-        else:
-            raise ValueError("Unsupported skip_type {}, need to be 'logSNR' or 'time_uniform' or 'time_quadratic'".format(skip_type))
-
-    def get_orders_and_timesteps_for_singlestep_solver(self, steps, order, skip_type, t_T, t_0, device):
-        """
-        Get the order of each step for sampling by the singlestep DPM-Solver.
-        """
-        if order == 3:
-            K = steps // 3 + 1
-            if steps % 3 == 0:
-                orders = [3,] * (K - 2) + [2, 1]
-            elif steps % 3 == 1:
-                orders = [3,] * (K - 1) + [1]
-            else:
-                orders = [3,] * (K - 1) + [2]
-        elif order == 2:
-            if steps % 2 == 0:
-                K = steps // 2
-                orders = [2,] * K
-            else:
-                K = steps // 2 + 1
-                orders = [2,] * (K - 1) + [1]
-        elif order == 1:
-            K = steps
-            orders = [1,] * steps
-        else:
-            raise ValueError("'order' must be '1' or '2' or '3'.")
-        if skip_type == 'logSNR':
-            # To reproduce the results in DPM-Solver paper
-            timesteps_outer = self.get_time_steps(skip_type, t_T, t_0, K, device)
-        else:
-            timesteps_outer = self.get_time_steps(skip_type, t_T, t_0, steps, device)[torch.cumsum(torch.tensor([0,] + orders), 0).to(device)]
-        return timesteps_outer, orders
-
-    def denoise_to_zero_fn(self, x, s):
-        """
-        Denoise at the final step, which is equivalent to solve the ODE from lambda_s to infty by first-order discretization.
-        """
-        return self.data_prediction_fn(x, s)
-
-    def multistep_uni_pc_update(self, x, model_prev_list, t_prev_list, t, order, **kwargs):
-        if len(t.shape) == 0:
-            t = t.view(-1)
-        if 'bh' in self.variant:
-            return self.multistep_uni_pc_bh_update(x, model_prev_list, t_prev_list, t, order, **kwargs)
-        else:
-            assert self.variant == 'vary_coeff'
-            return self.multistep_uni_pc_vary_update(x, model_prev_list, t_prev_list, t, order, **kwargs)
-
-    def multistep_uni_pc_vary_update(self, x, model_prev_list, t_prev_list, t, order, use_corrector=True):
-        logging.info(f'using unified predictor-corrector with order {order} (solver type: vary coeff)')
-        ns = self.noise_schedule
-        assert order <= len(model_prev_list)
-
-        # first compute rks
-        t_prev_0 = t_prev_list[-1]
-        lambda_prev_0 = ns.marginal_lambda(t_prev_0)
-        lambda_t = ns.marginal_lambda(t)
-        model_prev_0 = model_prev_list[-1]
-        sigma_prev_0, sigma_t = ns.marginal_std(t_prev_0), ns.marginal_std(t)
-        log_alpha_t = ns.marginal_log_mean_coeff(t)
-        alpha_t = torch.exp(log_alpha_t)
-
-        h = lambda_t - lambda_prev_0
-
-        rks = []
-        D1s = []
-        for i in range(1, order):
-            t_prev_i = t_prev_list[-(i + 1)]
-            model_prev_i = model_prev_list[-(i + 1)]
-            lambda_prev_i = ns.marginal_lambda(t_prev_i)
-            rk = (lambda_prev_i - lambda_prev_0) / h
-            rks.append(rk)
-            D1s.append((model_prev_i - model_prev_0) / rk)
-
-        rks.append(1.)
-        rks = torch.tensor(rks, device=x.device)
-
-        K = len(rks)
-        # build C matrix
-        C = []
-
-        col = torch.ones_like(rks)
-        for k in range(1, K + 1):
-            C.append(col)
-            col = col * rks / (k + 1)
-        C = torch.stack(C, dim=1)
-
-        if len(D1s) > 0:
-            D1s = torch.stack(D1s, dim=1) # (B, K)
-            C_inv_p = torch.linalg.inv(C[:-1, :-1])
-            A_p = C_inv_p
-
-        if use_corrector:
-            C_inv = torch.linalg.inv(C)
-            A_c = C_inv
-
-        hh = -h if self.predict_x0 else h
-        h_phi_1 = torch.expm1(hh)
-        h_phi_ks = []
-        factorial_k = 1
-        h_phi_k = h_phi_1
-        for k in range(1, K + 2):
-            h_phi_ks.append(h_phi_k)
-            h_phi_k = h_phi_k / hh - 1 / factorial_k
-            factorial_k *= (k + 1)
-
-        model_t = None
-        if self.predict_x0:
-            x_t_ = (
-                sigma_t / sigma_prev_0 * x
-                - alpha_t * h_phi_1 * model_prev_0
-            )
-            # now predictor
-            x_t = x_t_
-            if len(D1s) > 0:
-                # compute the residuals for predictor
-                for k in range(K - 1):
-                    x_t = x_t - alpha_t * h_phi_ks[k + 1] * torch.einsum('bkchw,k->bchw', D1s, A_p[k])
-            # now corrector
-            if use_corrector:
-                model_t = self.model_fn(x_t, t)
-                D1_t = (model_t - model_prev_0)
-                x_t = x_t_
-                k = 0
-                for k in range(K - 1):
-                    x_t = x_t - alpha_t * h_phi_ks[k + 1] * torch.einsum('bkchw,k->bchw', D1s, A_c[k][:-1])
-                x_t = x_t - alpha_t * h_phi_ks[K] * (D1_t * A_c[k][-1])
-        else:
-            log_alpha_prev_0, log_alpha_t = ns.marginal_log_mean_coeff(t_prev_0), ns.marginal_log_mean_coeff(t)
-            x_t_ = (
-                (torch.exp(log_alpha_t - log_alpha_prev_0)) * x
-                - (sigma_t * h_phi_1) * model_prev_0
-            )
-            # now predictor
-            x_t = x_t_
-            if len(D1s) > 0:
-                # compute the residuals for predictor
-                for k in range(K - 1):
-                    x_t = x_t - sigma_t * h_phi_ks[k + 1] * torch.einsum('bkchw,k->bchw', D1s, A_p[k])
-            # now corrector
-            if use_corrector:
-                model_t = self.model_fn(x_t, t)
-                D1_t = (model_t - model_prev_0)
-                x_t = x_t_
-                k = 0
-                for k in range(K - 1):
-                    x_t = x_t - sigma_t * h_phi_ks[k + 1] * torch.einsum('bkchw,k->bchw', D1s, A_c[k][:-1])
-                x_t = x_t - sigma_t * h_phi_ks[K] * (D1_t * A_c[k][-1])
-        return x_t, model_t
-
-    def multistep_uni_pc_bh_update(self, x, model_prev_list, t_prev_list, t, order, x_t=None, use_corrector=True):
-        # print(f'using unified predictor-corrector with order {order} (solver type: B(h))')
-        ns = self.noise_schedule
-        assert order <= len(model_prev_list)
-        dims = x.dim()
-
-        # first compute rks
-        t_prev_0 = t_prev_list[-1]
-        lambda_prev_0 = ns.marginal_lambda(t_prev_0)
-        lambda_t = ns.marginal_lambda(t)
-        model_prev_0 = model_prev_list[-1]
-        sigma_prev_0, sigma_t = ns.marginal_std(t_prev_0), ns.marginal_std(t)
-        log_alpha_prev_0, log_alpha_t = ns.marginal_log_mean_coeff(t_prev_0), ns.marginal_log_mean_coeff(t)
-        alpha_t = torch.exp(log_alpha_t)
-
-        h = lambda_t - lambda_prev_0
-
-        rks = []
-        D1s = []
-        for i in range(1, order):
-            t_prev_i = t_prev_list[-(i + 1)]
-            model_prev_i = model_prev_list[-(i + 1)]
-            lambda_prev_i = ns.marginal_lambda(t_prev_i)
-            rk = ((lambda_prev_i - lambda_prev_0) / h)[0]
-            rks.append(rk)
-            D1s.append((model_prev_i - model_prev_0) / rk)
-
-        rks.append(1.)
-        rks = torch.tensor(rks, device=x.device)
-
-        R = []
-        b = []
-
-        hh = -h[0] if self.predict_x0 else h[0]
-        h_phi_1 = torch.expm1(hh) # h\phi_1(h) = e^h - 1
-        h_phi_k = h_phi_1 / hh - 1
-
-        factorial_i = 1
-
-        if self.variant == 'bh1':
-            B_h = hh
-        elif self.variant == 'bh2':
-            B_h = torch.expm1(hh)
-        else:
-            raise NotImplementedError()
-
-        for i in range(1, order + 1):
-            R.append(torch.pow(rks, i - 1))
-            b.append(h_phi_k * factorial_i / B_h)
-            factorial_i *= (i + 1)
-            h_phi_k = h_phi_k / hh - 1 / factorial_i
-
-        R = torch.stack(R)
-        b = torch.tensor(b, device=x.device)
-
-        # now predictor
-        use_predictor = len(D1s) > 0 and x_t is None
-        if len(D1s) > 0:
-            D1s = torch.stack(D1s, dim=1) # (B, K)
-            if x_t is None:
-                # for order 2, we use a simplified version
-                if order == 2:
-                    rhos_p = torch.tensor([0.5], device=b.device)
-                else:
-                    rhos_p = torch.linalg.solve(R[:-1, :-1], b[:-1])
-        else:
-            D1s = None
-
-        if use_corrector:
-            # print('using corrector')
-            # for order 1, we use a simplified version
-            if order == 1:
-                rhos_c = torch.tensor([0.5], device=b.device)
-            else:
-                rhos_c = torch.linalg.solve(R, b)
-
-        model_t = None
-        if self.predict_x0:
-            x_t_ = (
-                expand_dims(sigma_t / sigma_prev_0, dims) * x
-                - expand_dims(alpha_t * h_phi_1, dims)* model_prev_0
-            )
-
-            if x_t is None:
-                if use_predictor:
-                    pred_res = torch.tensordot(D1s, rhos_p, dims=([1], [0]))  # torch.einsum('k,bkchw->bchw', rhos_p, D1s)
-                else:
-                    pred_res = 0
-                x_t = x_t_ - expand_dims(alpha_t * B_h, dims) * pred_res
-
-            if use_corrector:
-                model_t = self.model_fn(x_t, t)
-                if D1s is not None:
-                    corr_res = torch.tensordot(D1s, rhos_c[:-1], dims=([1], [0]))  # torch.einsum('k,bkchw->bchw', rhos_c[:-1], D1s)
-                else:
-                    corr_res = 0
-                D1_t = (model_t - model_prev_0)
-                x_t = x_t_ - expand_dims(alpha_t * B_h, dims) * (corr_res + rhos_c[-1] * D1_t)
-        else:
-            x_t_ = (
-                expand_dims(torch.exp(log_alpha_t - log_alpha_prev_0), dims) * x
-                - expand_dims(sigma_t * h_phi_1, dims) * model_prev_0
-            )
-            if x_t is None:
-                if use_predictor:
-                    pred_res = torch.einsum('k,bkchw->bchw', rhos_p, D1s)
-                else:
-                    pred_res = 0
-                x_t = x_t_ - expand_dims(sigma_t * B_h, dims) * pred_res
-
-            if use_corrector:
-                model_t = self.model_fn(x_t, t)
-                if D1s is not None:
-                    corr_res = torch.einsum('k,bkchw->bchw', rhos_c[:-1], D1s)
-                else:
-                    corr_res = 0
-                D1_t = (model_t - model_prev_0)
-                x_t = x_t_ - expand_dims(sigma_t * B_h, dims) * (corr_res + rhos_c[-1] * D1_t)
-        return x_t, model_t
-
-
-    def sample(self, x, timesteps, t_start=None, t_end=None, order=3, skip_type='time_uniform',
-        method='singlestep', lower_order_final=True, denoise_to_zero=False, solver_type='dpm_solver',
-        atol=0.0078, rtol=0.05, corrector=False, callback=None, disable_pbar=False
-    ):
-        # t_0 = 1. / self.noise_schedule.total_N if t_end is None else t_end
-        # t_T = self.noise_schedule.T if t_start is None else t_start
-        steps = len(timesteps) - 1
-        if method == 'multistep':
-            assert steps >= order
-            # timesteps = self.get_time_steps(skip_type=skip_type, t_T=t_T, t_0=t_0, N=steps, device=device)
-            assert timesteps.shape[0] - 1 == steps
-            # with torch.no_grad():
-            for step_index in trange(steps, disable=disable_pbar):
-                if step_index == 0:
-                    vec_t = timesteps[0].expand((x.shape[0]))
-                    model_prev_list = [self.model_fn(x, vec_t)]
-                    t_prev_list = [vec_t]
-                elif step_index < order:
-                    init_order = step_index
-                # Init the first `order` values by lower order multistep DPM-Solver.
-                # for init_order in range(1, order):
-                    vec_t = timesteps[init_order].expand(x.shape[0])
-                    x, model_x = self.multistep_uni_pc_update(x, model_prev_list, t_prev_list, vec_t, init_order, use_corrector=True)
-                    if model_x is None:
-                        model_x = self.model_fn(x, vec_t)
-                    model_prev_list.append(model_x)
-                    t_prev_list.append(vec_t)
-                else:
-                    extra_final_step = 0
-                    if step_index == (steps - 1):
-                        extra_final_step = 1
-                    for step in range(step_index, step_index + 1 + extra_final_step):
-                        vec_t = timesteps[step].expand(x.shape[0])
-                        if lower_order_final:
-                            step_order = min(order, steps + 1 - step)
-                        else:
-                            step_order = order
-                        # print('this step order:', step_order)
-                        if step == steps:
-                            # print('do not run corrector at the last step')
-                            use_corrector = False
-                        else:
-                            use_corrector = True
-                        x, model_x =  self.multistep_uni_pc_update(x, model_prev_list, t_prev_list, vec_t, step_order, use_corrector=use_corrector)
-                        for i in range(order - 1):
-                            t_prev_list[i] = t_prev_list[i + 1]
-                            model_prev_list[i] = model_prev_list[i + 1]
-                        t_prev_list[-1] = vec_t
-                        # We do not need to evaluate the final model value.
-                        if step < steps:
-                            if model_x is None:
-                                model_x = self.model_fn(x, vec_t)
-                            model_prev_list[-1] = model_x
-                if callback is not None:
-                    callback({'x': x, 'i': step_index, 'denoised': model_prev_list[-1]})
-        else:
-            raise NotImplementedError()
-        # if denoise_to_zero:
-        #     x = self.denoise_to_zero_fn(x, torch.ones((x.shape[0],)).to(device) * t_0)
-        return x
-
-
-#############################################################
-# other utility functions
-#############################################################
-
-def interpolate_fn(x, xp, yp):
-    """
-    A piecewise linear function y = f(x), using xp and yp as keypoints.
-    We implement f(x) in a differentiable way (i.e. applicable for autograd).
-    The function f(x) is well-defined for all x-axis. (For x beyond the bounds of xp, we use the outmost points of xp to define the linear function.)
-
-    Args:
-        x: PyTorch tensor with shape [N, C], where N is the batch size, C is the number of channels (we use C = 1 for DPM-Solver).
-        xp: PyTorch tensor with shape [C, K], where K is the number of keypoints.
-        yp: PyTorch tensor with shape [C, K].
-    Returns:
-        The function values f(x), with shape [N, C].
-    """
-    N, K = x.shape[0], xp.shape[1]
-    all_x = torch.cat([x.unsqueeze(2), xp.unsqueeze(0).repeat((N, 1, 1))], dim=2)
-    sorted_all_x, x_indices = torch.sort(all_x, dim=2)
-    x_idx = torch.argmin(x_indices, dim=2)
-    cand_start_idx = x_idx - 1
-    start_idx = torch.where(
-        torch.eq(x_idx, 0),
-        torch.tensor(1, device=x.device),
-        torch.where(
-            torch.eq(x_idx, K), torch.tensor(K - 2, device=x.device), cand_start_idx,
-        ),
-    )
-    end_idx = torch.where(torch.eq(start_idx, cand_start_idx), start_idx + 2, start_idx + 1)
-    start_x = torch.gather(sorted_all_x, dim=2, index=start_idx.unsqueeze(2)).squeeze(2)
-    end_x = torch.gather(sorted_all_x, dim=2, index=end_idx.unsqueeze(2)).squeeze(2)
-    start_idx2 = torch.where(
-        torch.eq(x_idx, 0),
-        torch.tensor(0, device=x.device),
-        torch.where(
-            torch.eq(x_idx, K), torch.tensor(K - 2, device=x.device), cand_start_idx,
-        ),
-    )
-    y_positions_expanded = yp.unsqueeze(0).expand(N, -1, -1)
-    start_y = torch.gather(y_positions_expanded, dim=2, index=start_idx2.unsqueeze(2)).squeeze(2)
-    end_y = torch.gather(y_positions_expanded, dim=2, index=(start_idx2 + 1).unsqueeze(2)).squeeze(2)
-    cand = start_y + (x - start_x) * (end_y - start_y) / (end_x - start_x)
-    return cand
-
-
-def expand_dims(v, dims):
-    """
-    Expand the tensor `v` to the dim `dims`.
-
-    Args:
-        `v`: a PyTorch tensor with shape [N].
-        `dim`: a `int`.
-    Returns:
-        a PyTorch tensor with shape [N, 1, 1, ..., 1] and the total dimension is `dims`.
-    """
-    return v[(...,) + (None,)*(dims - 1)]
-
-
-class SigmaConvert:
-    schedule = ""
-    def marginal_log_mean_coeff(self, sigma):
-        return 0.5 * torch.log(1 / ((sigma * sigma) + 1))
-
-    def marginal_alpha(self, t):
-        return torch.exp(self.marginal_log_mean_coeff(t))
-
-    def marginal_std(self, t):
-        return torch.sqrt(1. - torch.exp(2. * self.marginal_log_mean_coeff(t)))
-
-    def marginal_lambda(self, t):
-        """
-        Compute lambda_t = log(alpha_t) - log(sigma_t) of a given continuous-time label t in [0, T].
-        """
-        log_mean_coeff = self.marginal_log_mean_coeff(t)
-        log_std = 0.5 * torch.log(1. - torch.exp(2. * log_mean_coeff))
-        return log_mean_coeff - log_std
-
-def predict_eps_sigma(model, input, sigma_in, **kwargs):
-    sigma = sigma_in.view(sigma_in.shape[:1] + (1,) * (input.ndim - 1))
-    input = input * ((sigma ** 2 + 1.0) ** 0.5)
-    return  (input - model(input, sigma_in, **kwargs)) / sigma
-
-
-def sample_unipc(model, noise, sigmas, extra_args=None, callback=None, disable=False, variant='bh1'):
-        timesteps = sigmas.clone()
-        if sigmas[-1] == 0:
-            timesteps = sigmas[:]
-            timesteps[-1] = 0.001
-        else:
-            timesteps = sigmas.clone()
-        ns = SigmaConvert()
-
-        noise = noise / torch.sqrt(1.0 + timesteps[0] ** 2.0)
-        model_type = "noise"
-
-        model_fn = model_wrapper(
-            lambda input, sigma, **kwargs: predict_eps_sigma(model, input, sigma, **kwargs),
-            ns,
-            model_type=model_type,
-            guidance_type="uncond",
-            model_kwargs=extra_args,
-        )
-
-        order = min(3, len(timesteps) - 2)
-        uni_pc = UniPC(model_fn, ns, predict_x0=True, thresholding=False, variant=variant)
-        x = uni_pc.sample(noise, timesteps=timesteps, skip_type="time_uniform", method="multistep", order=order, lower_order_final=True, callback=callback, disable_pbar=disable)
-        x /= ns.marginal_alpha(timesteps[-1])
-        return x
-
-def sample_unipc_bh2(model, noise, sigmas, extra_args=None, callback=None, disable=False):
-    return sample_unipc(model, noise, sigmas, extra_args, callback, disable, variant='bh2')

comfy/float.py

@@ -1,67 +0,0 @@
-import torch
-
-def calc_mantissa(abs_x, exponent, normal_mask, MANTISSA_BITS, EXPONENT_BIAS, generator=None):
-    mantissa_scaled = torch.where(
-        normal_mask,
-        (abs_x / (2.0 ** (exponent - EXPONENT_BIAS)) - 1.0) * (2**MANTISSA_BITS),
-        (abs_x / (2.0 ** (-EXPONENT_BIAS + 1 - MANTISSA_BITS)))
-    )
-
-    mantissa_scaled += torch.rand(mantissa_scaled.size(), dtype=mantissa_scaled.dtype, layout=mantissa_scaled.layout, device=mantissa_scaled.device, generator=generator)
-    return mantissa_scaled.floor() / (2**MANTISSA_BITS)
-
-#Not 100% sure about this
-def manual_stochastic_round_to_float8(x, dtype, generator=None):
-    if dtype == torch.float8_e4m3fn:
-        EXPONENT_BITS, MANTISSA_BITS, EXPONENT_BIAS = 4, 3, 7
-    elif dtype == torch.float8_e5m2:
-        EXPONENT_BITS, MANTISSA_BITS, EXPONENT_BIAS = 5, 2, 15
-    else:
-        raise ValueError("Unsupported dtype")
-
-    x = x.half()
-    sign = torch.sign(x)
-    abs_x = x.abs()
-    sign = torch.where(abs_x == 0, 0, sign)
-
-    # Combine exponent calculation and clamping
-    exponent = torch.clamp(
-        torch.floor(torch.log2(abs_x)) + EXPONENT_BIAS,
-        0, 2**EXPONENT_BITS - 1
-    )
-
-    # Combine mantissa calculation and rounding
-    normal_mask = ~(exponent == 0)
-
-    abs_x[:] = calc_mantissa(abs_x, exponent, normal_mask, MANTISSA_BITS, EXPONENT_BIAS, generator=generator)
-
-    sign *= torch.where(
-        normal_mask,
-        (2.0 ** (exponent - EXPONENT_BIAS)) * (1.0 + abs_x),
-        (2.0 ** (-EXPONENT_BIAS + 1)) * abs_x
-    )
-
-    inf = torch.finfo(dtype)
-    torch.clamp(sign, min=inf.min, max=inf.max, out=sign)
-    return sign
-
-
-
-def stochastic_rounding(value, dtype, seed=0):
-    if dtype == torch.float32:
-        return value.to(dtype=torch.float32)
-    if dtype == torch.float16:
-        return value.to(dtype=torch.float16)
-    if dtype == torch.bfloat16:
-        return value.to(dtype=torch.bfloat16)
-    if dtype == torch.float8_e4m3fn or dtype == torch.float8_e5m2:
-        generator = torch.Generator(device=value.device)
-        generator.manual_seed(seed)
-        output = torch.empty_like(value, dtype=dtype)
-        num_slices = max(1, (value.numel() / (4096 * 4096)))
-        slice_size = max(1, round(value.shape[0] / num_slices))
-        for i in range(0, value.shape[0], slice_size):
-            output[i:i+slice_size].copy_(manual_stochastic_round_to_float8(value[i:i+slice_size], dtype, generator=generator))
-        return output
-
-    return value.to(dtype=dtype)

comfy/gligen.py

@@ -1,299 +0,0 @@
-import math
-import torch
-from torch import nn
-from .ldm.modules.attention import CrossAttention, FeedForward
-import comfy.ops
-ops = comfy.ops.manual_cast
-
-
-class GatedCrossAttentionDense(nn.Module):
-    def __init__(self, query_dim, context_dim, n_heads, d_head):
-        super().__init__()
-
-        self.attn = CrossAttention(
-            query_dim=query_dim,
-            context_dim=context_dim,
-            heads=n_heads,
-            dim_head=d_head,
-            operations=ops)
-        self.ff = FeedForward(query_dim, glu=True)
-
-        self.norm1 = ops.LayerNorm(query_dim)
-        self.norm2 = ops.LayerNorm(query_dim)
-
-        self.register_parameter('alpha_attn', nn.Parameter(torch.tensor(0.)))
-        self.register_parameter('alpha_dense', nn.Parameter(torch.tensor(0.)))
-
-        # this can be useful: we can externally change magnitude of tanh(alpha)
-        # for example, when it is set to 0, then the entire model is same as
-        # original one
-        self.scale = 1
-
-    def forward(self, x, objs):
-
-        x = x + self.scale * \
-            torch.tanh(self.alpha_attn) * self.attn(self.norm1(x), objs, objs)
-        x = x + self.scale * \
-            torch.tanh(self.alpha_dense) * self.ff(self.norm2(x))
-
-        return x
-
-
-class GatedSelfAttentionDense(nn.Module):
-    def __init__(self, query_dim, context_dim, n_heads, d_head):
-        super().__init__()
-
-        # we need a linear projection since we need cat visual feature and obj
-        # feature
-        self.linear = ops.Linear(context_dim, query_dim)
-
-        self.attn = CrossAttention(
-            query_dim=query_dim,
-            context_dim=query_dim,
-            heads=n_heads,
-            dim_head=d_head,
-            operations=ops)
-        self.ff = FeedForward(query_dim, glu=True)
-
-        self.norm1 = ops.LayerNorm(query_dim)
-        self.norm2 = ops.LayerNorm(query_dim)
-
-        self.register_parameter('alpha_attn', nn.Parameter(torch.tensor(0.)))
-        self.register_parameter('alpha_dense', nn.Parameter(torch.tensor(0.)))
-
-        # this can be useful: we can externally change magnitude of tanh(alpha)
-        # for example, when it is set to 0, then the entire model is same as
-        # original one
-        self.scale = 1
-
-    def forward(self, x, objs):
-
-        N_visual = x.shape[1]
-        objs = self.linear(objs)
-
-        x = x + self.scale * torch.tanh(self.alpha_attn) * self.attn(
-            self.norm1(torch.cat([x, objs], dim=1)))[:, 0:N_visual, :]
-        x = x + self.scale * \
-            torch.tanh(self.alpha_dense) * self.ff(self.norm2(x))
-
-        return x
-
-
-class GatedSelfAttentionDense2(nn.Module):
-    def __init__(self, query_dim, context_dim, n_heads, d_head):
-        super().__init__()
-
-        # we need a linear projection since we need cat visual feature and obj
-        # feature
-        self.linear = ops.Linear(context_dim, query_dim)
-
-        self.attn = CrossAttention(
-            query_dim=query_dim, context_dim=query_dim, dim_head=d_head, operations=ops)
-        self.ff = FeedForward(query_dim, glu=True)
-
-        self.norm1 = ops.LayerNorm(query_dim)
-        self.norm2 = ops.LayerNorm(query_dim)
-
-        self.register_parameter('alpha_attn', nn.Parameter(torch.tensor(0.)))
-        self.register_parameter('alpha_dense', nn.Parameter(torch.tensor(0.)))
-
-        # this can be useful: we can externally change magnitude of tanh(alpha)
-        # for example, when it is set to 0, then the entire model is same as
-        # original one
-        self.scale = 1
-
-    def forward(self, x, objs):
-
-        B, N_visual, _ = x.shape
-        B, N_ground, _ = objs.shape
-
-        objs = self.linear(objs)
-
-        # sanity check
-        size_v = math.sqrt(N_visual)
-        size_g = math.sqrt(N_ground)
-        assert int(size_v) == size_v, "Visual tokens must be square rootable"
-        assert int(size_g) == size_g, "Grounding tokens must be square rootable"
-        size_v = int(size_v)
-        size_g = int(size_g)
-
-        # select grounding token and resize it to visual token size as residual
-        out = self.attn(self.norm1(torch.cat([x, objs], dim=1)))[
-            :, N_visual:, :]
-        out = out.permute(0, 2, 1).reshape(B, -1, size_g, size_g)
-        out = torch.nn.functional.interpolate(
-            out, (size_v, size_v), mode='bicubic')
-        residual = out.reshape(B, -1, N_visual).permute(0, 2, 1)
-
-        # add residual to visual feature
-        x = x + self.scale * torch.tanh(self.alpha_attn) * residual
-        x = x + self.scale * \
-            torch.tanh(self.alpha_dense) * self.ff(self.norm2(x))
-
-        return x
-
-
-class FourierEmbedder():
-    def __init__(self, num_freqs=64, temperature=100):
-
-        self.num_freqs = num_freqs
-        self.temperature = temperature
-        self.freq_bands = temperature ** (torch.arange(num_freqs) / num_freqs)
-
-    @torch.no_grad()
-    def __call__(self, x, cat_dim=-1):
-        "x: arbitrary shape of tensor. dim: cat dim"
-        out = []
-        for freq in self.freq_bands:
-            out.append(torch.sin(freq * x))
-            out.append(torch.cos(freq * x))
-        return torch.cat(out, cat_dim)
-
-
-class PositionNet(nn.Module):
-    def __init__(self, in_dim, out_dim, fourier_freqs=8):
-        super().__init__()
-        self.in_dim = in_dim
-        self.out_dim = out_dim
-
-        self.fourier_embedder = FourierEmbedder(num_freqs=fourier_freqs)
-        self.position_dim = fourier_freqs * 2 * 4  # 2 is sin&cos, 4 is xyxy
-
-        self.linears = nn.Sequential(
-            ops.Linear(self.in_dim + self.position_dim, 512),
-            nn.SiLU(),
-            ops.Linear(512, 512),
-            nn.SiLU(),
-            ops.Linear(512, out_dim),
-        )
-
-        self.null_positive_feature = torch.nn.Parameter(
-            torch.zeros([self.in_dim]))
-        self.null_position_feature = torch.nn.Parameter(
-            torch.zeros([self.position_dim]))
-
-    def forward(self, boxes, masks, positive_embeddings):
-        B, N, _ = boxes.shape
-        masks = masks.unsqueeze(-1)
-        positive_embeddings = positive_embeddings
-
-        # embedding position (it may includes padding as placeholder)
-        xyxy_embedding = self.fourier_embedder(boxes)  # B*N*4 --> B*N*C
-
-        # learnable null embedding
-        positive_null = self.null_positive_feature.to(device=boxes.device, dtype=boxes.dtype).view(1, 1, -1)
-        xyxy_null = self.null_position_feature.to(device=boxes.device, dtype=boxes.dtype).view(1, 1, -1)
-
-        # replace padding with learnable null embedding
-        positive_embeddings = positive_embeddings * \
-            masks + (1 - masks) * positive_null
-        xyxy_embedding = xyxy_embedding * masks + (1 - masks) * xyxy_null
-
-        objs = self.linears(
-            torch.cat([positive_embeddings, xyxy_embedding], dim=-1))
-        assert objs.shape == torch.Size([B, N, self.out_dim])
-        return objs
-
-
-class Gligen(nn.Module):
-    def __init__(self, modules, position_net, key_dim):
-        super().__init__()
-        self.module_list = nn.ModuleList(modules)
-        self.position_net = position_net
-        self.key_dim = key_dim
-        self.max_objs = 30
-        self.current_device = torch.device("cpu")
-
-    def _set_position(self, boxes, masks, positive_embeddings):
-        objs = self.position_net(boxes, masks, positive_embeddings)
-        def func(x, extra_options):
-            key = extra_options["transformer_index"]
-            module = self.module_list[key]
-            return module(x, objs.to(device=x.device, dtype=x.dtype))
-        return func
-
-    def set_position(self, latent_image_shape, position_params, device):
-        batch, c, h, w = latent_image_shape
-        masks = torch.zeros([self.max_objs], device="cpu")
-        boxes = []
-        positive_embeddings = []
-        for p in position_params:
-            x1 = (p[4]) / w
-            y1 = (p[3]) / h
-            x2 = (p[4] + p[2]) / w
-            y2 = (p[3] + p[1]) / h
-            masks[len(boxes)] = 1.0
-            boxes += [torch.tensor((x1, y1, x2, y2)).unsqueeze(0)]
-            positive_embeddings += [p[0]]
-        append_boxes = []
-        append_conds = []
-        if len(boxes) < self.max_objs:
-            append_boxes = [torch.zeros(
-                [self.max_objs - len(boxes), 4], device="cpu")]
-            append_conds = [torch.zeros(
-                [self.max_objs - len(boxes), self.key_dim], device="cpu")]
-
-        box_out = torch.cat(
-            boxes + append_boxes).unsqueeze(0).repeat(batch, 1, 1)
-        masks = masks.unsqueeze(0).repeat(batch, 1)
-        conds = torch.cat(positive_embeddings +
-                          append_conds).unsqueeze(0).repeat(batch, 1, 1)
-        return self._set_position(
-            box_out.to(device),
-            masks.to(device),
-            conds.to(device))
-
-    def set_empty(self, latent_image_shape, device):
-        batch, c, h, w = latent_image_shape
-        masks = torch.zeros([self.max_objs], device="cpu").repeat(batch, 1)
-        box_out = torch.zeros([self.max_objs, 4],
-                              device="cpu").repeat(batch, 1, 1)
-        conds = torch.zeros([self.max_objs, self.key_dim],
-                            device="cpu").repeat(batch, 1, 1)
-        return self._set_position(
-            box_out.to(device),
-            masks.to(device),
-            conds.to(device))
-
-
-def load_gligen(sd):
-    sd_k = sd.keys()
-    output_list = []
-    key_dim = 768
-    for a in ["input_blocks", "middle_block", "output_blocks"]:
-        for b in range(20):
-            k_temp = filter(lambda k: "{}.{}.".format(a, b)
-                            in k and ".fuser." in k, sd_k)
-            k_temp = map(lambda k: (k, k.split(".fuser.")[-1]), k_temp)
-
-            n_sd = {}
-            for k in k_temp:
-                n_sd[k[1]] = sd[k[0]]
-            if len(n_sd) > 0:
-                query_dim = n_sd["linear.weight"].shape[0]
-                key_dim = n_sd["linear.weight"].shape[1]
-
-                if key_dim == 768:  # SD1.x
-                    n_heads = 8
-                    d_head = query_dim // n_heads
-                else:
-                    d_head = 64
-                    n_heads = query_dim // d_head
-
-                gated = GatedSelfAttentionDense(
-                    query_dim, key_dim, n_heads, d_head)
-                gated.load_state_dict(n_sd, strict=False)
-                output_list.append(gated)
-
-    if "position_net.null_positive_feature" in sd_k:
-        in_dim = sd["position_net.null_positive_feature"].shape[0]
-        out_dim = sd["position_net.linears.4.weight"].shape[0]
-
-        class WeightsLoader(torch.nn.Module):
-            pass
-        w = WeightsLoader()
-        w.position_net = PositionNet(in_dim, out_dim)
-        w.load_state_dict(sd, strict=False)
-
-    gligen = Gligen(output_list, w.position_net, key_dim)
-    return gligen

comfy/hooks.py

@@ -1,785 +0,0 @@
-from __future__ import annotations
-from typing import TYPE_CHECKING, Callable
-import enum
-import math
-import torch
-import numpy as np
-import itertools
-import logging
-
-if TYPE_CHECKING:
-    from comfy.model_patcher import ModelPatcher, PatcherInjection
-    from comfy.model_base import BaseModel
-    from comfy.sd import CLIP
-import comfy.lora
-import comfy.model_management
-import comfy.patcher_extension
-from node_helpers import conditioning_set_values
-
-# #######################################################################################################
-# Hooks explanation
-# -------------------
-# The purpose of hooks is to allow conds to influence sampling without the need for ComfyUI core code to
-# make explicit special cases like it does for ControlNet and GLIGEN.
-#
-# This is necessary for nodes/features that are intended for use with masked or scheduled conds, or those
-# that should run special code when a 'marked' cond is used in sampling.
-# #######################################################################################################
-
-class EnumHookMode(enum.Enum):
-    '''
-    Priority of hook memory optimization vs. speed, mostly related to WeightHooks.
-
-    MinVram: No caching will occur for any operations related to hooks.
-    MaxSpeed: Excess VRAM (and RAM, once VRAM is sufficiently depleted) will be used to cache hook weights when switching hook groups.
-    '''
-    MinVram = "minvram"
-    MaxSpeed = "maxspeed"
-
-class EnumHookType(enum.Enum):
-    '''
-    Hook types, each of which has different expected behavior.
-    '''
-    Weight = "weight"
-    ObjectPatch = "object_patch"
-    AdditionalModels = "add_models"
-    TransformerOptions = "transformer_options"
-    Injections = "add_injections"
-
-class EnumWeightTarget(enum.Enum):
-    Model = "model"
-    Clip = "clip"
-
-class EnumHookScope(enum.Enum):
-    '''
-    Determines if hook should be limited in its influence over sampling.
-
-    AllConditioning: hook will affect all conds used in sampling.
-    HookedOnly: hook will only affect the conds it was attached to.
-    '''
-    AllConditioning = "all_conditioning"
-    HookedOnly = "hooked_only"
-
-
-class _HookRef:
-    pass
-
-
-def default_should_register(hook: Hook, model: ModelPatcher, model_options: dict, target_dict: dict[str], registered: HookGroup):
-    '''Example for how custom_should_register function can look like.'''
-    return True
-
-
-def create_target_dict(target: EnumWeightTarget=None, **kwargs) -> dict[str]:
-    '''Creates base dictionary for use with Hooks' target param.'''
-    d = {}
-    if target is not None:
-        d['target'] = target
-    d.update(kwargs)
-    return d
-
-
-class Hook:
-    def __init__(self, hook_type: EnumHookType=None, hook_ref: _HookRef=None, hook_id: str=None,
-                 hook_keyframe: HookKeyframeGroup=None, hook_scope=EnumHookScope.AllConditioning):
-        self.hook_type = hook_type
-        '''Enum identifying the general class of this hook.'''
-        self.hook_ref = hook_ref if hook_ref else _HookRef()
-        '''Reference shared between hook clones that have the same value. Should NOT be modified.'''
-        self.hook_id = hook_id
-        '''Optional string ID to identify hook; useful if need to consolidate duplicates at registration time.'''
-        self.hook_keyframe = hook_keyframe if hook_keyframe else HookKeyframeGroup()
-        '''Keyframe storage that can be referenced to get strength for current sampling step.'''
-        self.hook_scope = hook_scope
-        '''Scope of where this hook should apply in terms of the conds used in sampling run.'''
-        self.custom_should_register = default_should_register
-        '''Can be overriden with a compatible function to decide if this hook should be registered without the need to override .should_register'''
-
-    @property
-    def strength(self):
-        return self.hook_keyframe.strength
-
-    def initialize_timesteps(self, model: BaseModel):
-        self.reset()
-        self.hook_keyframe.initialize_timesteps(model)
-
-    def reset(self):
-        self.hook_keyframe.reset()
-
-    def clone(self):
-        c: Hook = self.__class__()
-        c.hook_type = self.hook_type
-        c.hook_ref = self.hook_ref
-        c.hook_id = self.hook_id
-        c.hook_keyframe = self.hook_keyframe
-        c.hook_scope = self.hook_scope
-        c.custom_should_register = self.custom_should_register
-        return c
-
-    def should_register(self, model: ModelPatcher, model_options: dict, target_dict: dict[str], registered: HookGroup):
-        return self.custom_should_register(self, model, model_options, target_dict, registered)
-
-    def add_hook_patches(self, model: ModelPatcher, model_options: dict, target_dict: dict[str], registered: HookGroup):
-        raise NotImplementedError("add_hook_patches should be defined for Hook subclasses")
-
-    def __eq__(self, other: Hook):
-        return self.__class__ == other.__class__ and self.hook_ref == other.hook_ref
-
-    def __hash__(self):
-        return hash(self.hook_ref)
-
-class WeightHook(Hook):
-    '''
-    Hook responsible for tracking weights to be applied to some model/clip.
-
-    Note, value of hook_scope is ignored and is treated as HookedOnly.
-    '''
-    def __init__(self, strength_model=1.0, strength_clip=1.0):
-        super().__init__(hook_type=EnumHookType.Weight, hook_scope=EnumHookScope.HookedOnly)
-        self.weights: dict = None
-        self.weights_clip: dict = None
-        self.need_weight_init = True
-        self._strength_model = strength_model
-        self._strength_clip = strength_clip
-        self.hook_scope = EnumHookScope.HookedOnly # this value does not matter for WeightHooks, just for docs
-
-    @property
-    def strength_model(self):
-        return self._strength_model * self.strength
-
-    @property
-    def strength_clip(self):
-        return self._strength_clip * self.strength
-
-    def add_hook_patches(self, model: ModelPatcher, model_options: dict, target_dict: dict[str], registered: HookGroup):
-        if not self.should_register(model, model_options, target_dict, registered):
-            return False
-        weights = None
-
-        target = target_dict.get('target', None)
-        if target == EnumWeightTarget.Clip:
-            strength = self._strength_clip
-        else:
-            strength = self._strength_model
-
-        if self.need_weight_init:
-            key_map = {}
-            if target == EnumWeightTarget.Clip:
-                key_map = comfy.lora.model_lora_keys_clip(model.model, key_map)
-            else:
-                key_map = comfy.lora.model_lora_keys_unet(model.model, key_map)
-            weights = comfy.lora.load_lora(self.weights, key_map, log_missing=False)
-        else:
-            if target == EnumWeightTarget.Clip:
-                weights = self.weights_clip
-            else:
-                weights = self.weights
-        model.add_hook_patches(hook=self, patches=weights, strength_patch=strength)
-        registered.add(self)
-        return True
-        # TODO: add logs about any keys that were not applied
-
-    def clone(self):
-        c: WeightHook = super().clone()
-        c.weights = self.weights
-        c.weights_clip = self.weights_clip
-        c.need_weight_init = self.need_weight_init
-        c._strength_model = self._strength_model
-        c._strength_clip = self._strength_clip
-        return c
-
-class ObjectPatchHook(Hook):
-    def __init__(self, object_patches: dict[str]=None,
-                 hook_scope=EnumHookScope.AllConditioning):
-        super().__init__(hook_type=EnumHookType.ObjectPatch)
-        self.object_patches = object_patches
-        self.hook_scope = hook_scope
-
-    def clone(self):
-        c: ObjectPatchHook = super().clone()
-        c.object_patches = self.object_patches
-        return c
-
-    def add_hook_patches(self, model: ModelPatcher, model_options: dict, target_dict: dict[str], registered: HookGroup):
-        raise NotImplementedError("ObjectPatchHook is not supported yet in ComfyUI.")
-
-class AdditionalModelsHook(Hook):
-    '''
-    Hook responsible for telling model management any additional models that should be loaded.
-
-    Note, value of hook_scope is ignored and is treated as AllConditioning.
-    '''
-    def __init__(self, models: list[ModelPatcher]=None, key: str=None):
-        super().__init__(hook_type=EnumHookType.AdditionalModels)
-        self.models = models
-        self.key = key
-
-    def clone(self):
-        c: AdditionalModelsHook = super().clone()
-        c.models = self.models.copy() if self.models else self.models
-        c.key = self.key
-        return c
-
-    def add_hook_patches(self, model: ModelPatcher, model_options: dict, target_dict: dict[str], registered: HookGroup):
-        if not self.should_register(model, model_options, target_dict, registered):
-            return False
-        registered.add(self)
-        return True
-
-class TransformerOptionsHook(Hook):
-    '''
-    Hook responsible for adding wrappers, callbacks, patches, or anything else related to transformer_options.
-    '''
-    def __init__(self, transformers_dict: dict[str, dict[str, dict[str, list[Callable]]]]=None,
-                 hook_scope=EnumHookScope.AllConditioning):
-        super().__init__(hook_type=EnumHookType.TransformerOptions)
-        self.transformers_dict = transformers_dict
-        self.hook_scope = hook_scope
-        self._skip_adding = False
-        '''Internal value used to avoid double load of transformer_options when hook_scope is AllConditioning.'''
-
-    def clone(self):
-        c: TransformerOptionsHook = super().clone()
-        c.transformers_dict = self.transformers_dict
-        c._skip_adding = self._skip_adding
-        return c
-
-    def add_hook_patches(self, model: ModelPatcher, model_options: dict, target_dict: dict[str], registered: HookGroup):
-        if not self.should_register(model, model_options, target_dict, registered):
-            return False
-        # NOTE: to_load_options will be used to manually load patches/wrappers/callbacks from hooks
-        self._skip_adding = False
-        if self.hook_scope == EnumHookScope.AllConditioning:
-            add_model_options = {"transformer_options": self.transformers_dict,
-                                 "to_load_options": self.transformers_dict}
-            # skip_adding if included in AllConditioning to avoid double loading
-            self._skip_adding = True
-        else:
-            add_model_options = {"to_load_options": self.transformers_dict}
-        registered.add(self)
-        comfy.patcher_extension.merge_nested_dicts(model_options, add_model_options, copy_dict1=False)
-        return True
-
-    def on_apply_hooks(self, model: ModelPatcher, transformer_options: dict[str]):
-        if not self._skip_adding:
-            comfy.patcher_extension.merge_nested_dicts(transformer_options, self.transformers_dict, copy_dict1=False)
-
-WrapperHook = TransformerOptionsHook
-'''Only here for backwards compatibility, WrapperHook is identical to TransformerOptionsHook.'''
-
-class InjectionsHook(Hook):
-    def __init__(self, key: str=None, injections: list[PatcherInjection]=None,
-                 hook_scope=EnumHookScope.AllConditioning):
-        super().__init__(hook_type=EnumHookType.Injections)
-        self.key = key
-        self.injections = injections
-        self.hook_scope = hook_scope
-
-    def clone(self):
-        c: InjectionsHook = super().clone()
-        c.key = self.key
-        c.injections = self.injections.copy() if self.injections else self.injections
-        return c
-
-    def add_hook_patches(self, model: ModelPatcher, model_options: dict, target_dict: dict[str], registered: HookGroup):
-        raise NotImplementedError("InjectionsHook is not supported yet in ComfyUI.")
-
-class HookGroup:
-    '''
-    Stores groups of hooks, and allows them to be queried by type.
-
-    To prevent breaking their functionality, never modify the underlying self.hooks or self._hook_dict vars directly;
-    always use the provided functions on HookGroup.
-    '''
-    def __init__(self):
-        self.hooks: list[Hook] = []
-        self._hook_dict: dict[EnumHookType, list[Hook]] = {}
-
-    def __len__(self):
-        return len(self.hooks)
-
-    def add(self, hook: Hook):
-        if hook not in self.hooks:
-            self.hooks.append(hook)
-            self._hook_dict.setdefault(hook.hook_type, []).append(hook)
-
-    def remove(self, hook: Hook):
-        if hook in self.hooks:
-            self.hooks.remove(hook)
-            self._hook_dict[hook.hook_type].remove(hook)
-
-    def get_type(self, hook_type: EnumHookType):
-        return self._hook_dict.get(hook_type, [])
-
-    def contains(self, hook: Hook):
-        return hook in self.hooks
-
-    def is_subset_of(self, other: HookGroup):
-        self_hooks = set(self.hooks)
-        other_hooks = set(other.hooks)
-        return self_hooks.issubset(other_hooks)
-
-    def new_with_common_hooks(self, other: HookGroup):
-        c = HookGroup()
-        for hook in self.hooks:
-            if other.contains(hook):
-                c.add(hook.clone())
-        return c
-
-    def clone(self):
-        c = HookGroup()
-        for hook in self.hooks:
-            c.add(hook.clone())
-        return c
-
-    def clone_and_combine(self, other: HookGroup):
-        c = self.clone()
-        if other is not None:
-            for hook in other.hooks:
-                c.add(hook.clone())
-        return c
-
-    def set_keyframes_on_hooks(self, hook_kf: HookKeyframeGroup):
-        if hook_kf is None:
-            hook_kf = HookKeyframeGroup()
-        else:
-            hook_kf = hook_kf.clone()
-        for hook in self.hooks:
-            hook.hook_keyframe = hook_kf
-
-    def get_hooks_for_clip_schedule(self):
-        scheduled_hooks: dict[WeightHook, list[tuple[tuple[float,float], HookKeyframe]]] = {}
-        # only care about WeightHooks, for now
-        for hook in self.get_type(EnumHookType.Weight):
-            hook: WeightHook
-            hook_schedule = []
-            # if no hook keyframes, assign default value
-            if len(hook.hook_keyframe.keyframes) == 0:
-                hook_schedule.append(((0.0, 1.0), None))
-                scheduled_hooks[hook] = hook_schedule
-                continue
-            # find ranges of values
-            prev_keyframe = hook.hook_keyframe.keyframes[0]
-            for keyframe in hook.hook_keyframe.keyframes:
-                if keyframe.start_percent > prev_keyframe.start_percent and not math.isclose(keyframe.strength, prev_keyframe.strength):
-                    hook_schedule.append(((prev_keyframe.start_percent, keyframe.start_percent), prev_keyframe))
-                    prev_keyframe = keyframe
-                elif keyframe.start_percent == prev_keyframe.start_percent:
-                    prev_keyframe = keyframe
-            # create final range, assuming last start_percent was not 1.0
-            if not math.isclose(prev_keyframe.start_percent, 1.0):
-                hook_schedule.append(((prev_keyframe.start_percent, 1.0), prev_keyframe))
-            scheduled_hooks[hook] = hook_schedule
-        # hooks should not have their schedules in a list of tuples
-        all_ranges: list[tuple[float, float]] = []
-        for range_kfs in scheduled_hooks.values():
-            for t_range, keyframe in range_kfs:
-                all_ranges.append(t_range)
-        # turn list of ranges into boundaries
-        boundaries_set = set(itertools.chain.from_iterable(all_ranges))
-        boundaries_set.add(0.0)
-        boundaries = sorted(boundaries_set)
-        real_ranges = [(boundaries[i], boundaries[i + 1]) for i in range(len(boundaries) - 1)]
-        # with real ranges defined, give appropriate hooks w/ keyframes for each range
-        scheduled_keyframes: list[tuple[tuple[float,float], list[tuple[WeightHook, HookKeyframe]]]] = []
-        for t_range in real_ranges:
-            hooks_schedule = []
-            for hook, val in scheduled_hooks.items():
-                keyframe = None
-                # check if is a keyframe that works for the current t_range
-                for stored_range, stored_kf in val:
-                    # if stored start is less than current end, then fits - give it assigned keyframe
-                    if stored_range[0] < t_range[1] and stored_range[1] > t_range[0]:
-                        keyframe = stored_kf
-                        break
-                hooks_schedule.append((hook, keyframe))
-            scheduled_keyframes.append((t_range, hooks_schedule))
-        return scheduled_keyframes
-
-    def reset(self):
-        for hook in self.hooks:
-            hook.reset()
-
-    @staticmethod
-    def combine_all_hooks(hooks_list: list[HookGroup], require_count=0) -> HookGroup:
-        actual: list[HookGroup] = []
-        for group in hooks_list:
-            if group is not None:
-                actual.append(group)
-        if len(actual) < require_count:
-            raise Exception(f"Need at least {require_count} hooks to combine, but only had {len(actual)}.")
-        # if no hooks, then return None
-        if len(actual) == 0:
-            return None
-        # if only 1 hook, just return itself without cloning
-        elif len(actual) == 1:
-            return actual[0]
-        final_hook: HookGroup = None
-        for hook in actual:
-            if final_hook is None:
-                final_hook = hook.clone()
-            else:
-                final_hook = final_hook.clone_and_combine(hook)
-        return final_hook
-
-
-class HookKeyframe:
-    def __init__(self, strength: float, start_percent=0.0, guarantee_steps=1):
-        self.strength = strength
-        # scheduling
-        self.start_percent = float(start_percent)
-        self.start_t = 999999999.9
-        self.guarantee_steps = guarantee_steps
-
-    def get_effective_guarantee_steps(self, max_sigma: torch.Tensor):
-        '''If keyframe starts before current sampling range (max_sigma), treat as 0.'''
-        if self.start_t > max_sigma:
-            return 0
-        return self.guarantee_steps
-
-    def clone(self):
-        c = HookKeyframe(strength=self.strength,
-                         start_percent=self.start_percent, guarantee_steps=self.guarantee_steps)
-        c.start_t = self.start_t
-        return c
-
-class HookKeyframeGroup:
-    def __init__(self):
-        self.keyframes: list[HookKeyframe] = []
-        self._current_keyframe: HookKeyframe = None
-        self._current_used_steps = 0
-        self._current_index = 0
-        self._current_strength = None
-        self._curr_t = -1.
-
-    # properties shadow those of HookWeightsKeyframe
-    @property
-    def strength(self):
-        if self._current_keyframe is not None:
-            return self._current_keyframe.strength
-        return 1.0
-
-    def reset(self):
-        self._current_keyframe = None
-        self._current_used_steps = 0
-        self._current_index = 0
-        self._current_strength = None
-        self.curr_t = -1.
-        self._set_first_as_current()
-
-    def add(self, keyframe: HookKeyframe):
-        # add to end of list, then sort
-        self.keyframes.append(keyframe)
-        self.keyframes = get_sorted_list_via_attr(self.keyframes, "start_percent")
-        self._set_first_as_current()
-
-    def _set_first_as_current(self):
-        if len(self.keyframes) > 0:
-            self._current_keyframe = self.keyframes[0]
-        else:
-            self._current_keyframe = None
-
-    def has_guarantee_steps(self):
-        for kf in self.keyframes:
-            if kf.guarantee_steps > 0:
-                return True
-        return False
-
-    def has_index(self, index: int):
-        return index >= 0 and index < len(self.keyframes)
-
-    def is_empty(self):
-        return len(self.keyframes) == 0
-
-    def clone(self):
-        c = HookKeyframeGroup()
-        for keyframe in self.keyframes:
-            c.keyframes.append(keyframe.clone())
-        c._set_first_as_current()
-        return c
-
-    def initialize_timesteps(self, model: BaseModel):
-        for keyframe in self.keyframes:
-            keyframe.start_t = model.model_sampling.percent_to_sigma(keyframe.start_percent)
-
-    def prepare_current_keyframe(self, curr_t: float, transformer_options: dict[str, torch.Tensor]) -> bool:
-        if self.is_empty():
-            return False
-        if curr_t == self._curr_t:
-            return False
-        max_sigma = torch.max(transformer_options["sample_sigmas"])
-        prev_index = self._current_index
-        prev_strength = self._current_strength
-        # if met guaranteed steps, look for next keyframe in case need to switch
-        if self._current_used_steps >= self._current_keyframe.get_effective_guarantee_steps(max_sigma):
-            # if has next index, loop through and see if need to switch
-            if self.has_index(self._current_index+1):
-                for i in range(self._current_index+1, len(self.keyframes)):
-                    eval_c = self.keyframes[i]
-                    # check if start_t is greater or equal to curr_t
-                    # NOTE: t is in terms of sigmas, not percent, so bigger number = earlier step in sampling
-                    if eval_c.start_t >= curr_t:
-                        self._current_index = i
-                        self._current_strength = eval_c.strength
-                        self._current_keyframe = eval_c
-                        self._current_used_steps = 0
-                        # if guarantee_steps greater than zero, stop searching for other keyframes
-                        if self._current_keyframe.get_effective_guarantee_steps(max_sigma) > 0:
-                            break
-                    # if eval_c is outside the percent range, stop looking further
-                    else: break
-        # update steps current context is used
-        self._current_used_steps += 1
-        # update current timestep this was performed on
-        self._curr_t = curr_t
-        # return True if keyframe changed, False if no change
-        return prev_index != self._current_index and prev_strength != self._current_strength
-
-
-class InterpolationMethod:
-    LINEAR = "linear"
-    EASE_IN = "ease_in"
-    EASE_OUT = "ease_out"
-    EASE_IN_OUT = "ease_in_out"
-
-    _LIST = [LINEAR, EASE_IN, EASE_OUT, EASE_IN_OUT]
-
-    @classmethod
-    def get_weights(cls, num_from: float, num_to: float, length: int, method: str, reverse=False):
-        diff = num_to - num_from
-        if method == cls.LINEAR:
-            weights = torch.linspace(num_from, num_to, length)
-        elif method == cls.EASE_IN:
-            index = torch.linspace(0, 1, length)
-            weights = diff * np.power(index, 2) + num_from
-        elif method == cls.EASE_OUT:
-            index = torch.linspace(0, 1, length)
-            weights = diff * (1 - np.power(1 - index, 2)) + num_from
-        elif method == cls.EASE_IN_OUT:
-            index = torch.linspace(0, 1, length)
-            weights = diff * ((1 - np.cos(index * np.pi)) / 2) + num_from
-        else:
-            raise ValueError(f"Unrecognized interpolation method '{method}'.")
-        if reverse:
-            weights = weights.flip(dims=(0,))
-        return weights
-
-def get_sorted_list_via_attr(objects: list, attr: str) -> list:
-    if not objects:
-        return objects
-    elif len(objects) <= 1:
-        return [x for x in objects]
-    # now that we know we have to sort, do it following these rules:
-    # a) if objects have same value of attribute, maintain their relative order
-    # b) perform sorting of the groups of objects with same attributes
-    unique_attrs = {}
-    for o in objects:
-        val_attr = getattr(o, attr)
-        attr_list: list = unique_attrs.get(val_attr, list())
-        attr_list.append(o)
-        if val_attr not in unique_attrs:
-            unique_attrs[val_attr] = attr_list
-    # now that we have the unique attr values grouped together in relative order, sort them by key
-    sorted_attrs = dict(sorted(unique_attrs.items()))
-    # now flatten out the dict into a list to return
-    sorted_list = []
-    for object_list in sorted_attrs.values():
-        sorted_list.extend(object_list)
-    return sorted_list
-
-def create_transformer_options_from_hooks(model: ModelPatcher, hooks: HookGroup,  transformer_options: dict[str]=None):
-    # if no hooks or is not a ModelPatcher for sampling, return empty dict
-    if hooks is None or model.is_clip:
-        return {}
-    if transformer_options is None:
-        transformer_options = {}
-    for hook in hooks.get_type(EnumHookType.TransformerOptions):
-        hook: TransformerOptionsHook
-        hook.on_apply_hooks(model, transformer_options)
-    return transformer_options
-
-def create_hook_lora(lora: dict[str, torch.Tensor], strength_model: float, strength_clip: float):
-    hook_group = HookGroup()
-    hook = WeightHook(strength_model=strength_model, strength_clip=strength_clip)
-    hook_group.add(hook)
-    hook.weights = lora
-    return hook_group
-
-def create_hook_model_as_lora(weights_model, weights_clip, strength_model: float, strength_clip: float):
-    hook_group = HookGroup()
-    hook = WeightHook(strength_model=strength_model, strength_clip=strength_clip)
-    hook_group.add(hook)
-    patches_model = None
-    patches_clip = None
-    if weights_model is not None:
-        patches_model = {}
-        for key in weights_model:
-            patches_model[key] = ("model_as_lora", (weights_model[key],))
-    if weights_clip is not None:
-        patches_clip = {}
-        for key in weights_clip:
-            patches_clip[key] = ("model_as_lora", (weights_clip[key],))
-    hook.weights = patches_model
-    hook.weights_clip = patches_clip
-    hook.need_weight_init = False
-    return hook_group
-
-def get_patch_weights_from_model(model: ModelPatcher, discard_model_sampling=True):
-    if model is None:
-        return None
-    patches_model: dict[str, torch.Tensor] = model.model.state_dict()
-    if discard_model_sampling:
-        # do not include ANY model_sampling components of the model that should act as a patch
-        for key in list(patches_model.keys()):
-            if key.startswith("model_sampling"):
-                patches_model.pop(key, None)
-    return patches_model
-
-# NOTE: this function shows how to register weight hooks directly on the ModelPatchers
-def load_hook_lora_for_models(model: ModelPatcher, clip: CLIP, lora: dict[str, torch.Tensor],
-                              strength_model: float, strength_clip: float):
-    key_map = {}
-    if model is not None:
-        key_map = comfy.lora.model_lora_keys_unet(model.model, key_map)
-    if clip is not None:
-        key_map = comfy.lora.model_lora_keys_clip(clip.cond_stage_model, key_map)
-
-    hook_group = HookGroup()
-    hook = WeightHook()
-    hook_group.add(hook)
-    loaded: dict[str] = comfy.lora.load_lora(lora, key_map)
-    if model is not None:
-        new_modelpatcher = model.clone()
-        k = new_modelpatcher.add_hook_patches(hook=hook, patches=loaded, strength_patch=strength_model)
-    else:
-        k = ()
-        new_modelpatcher = None
-
-    if clip is not None:
-        new_clip = clip.clone()
-        k1 = new_clip.patcher.add_hook_patches(hook=hook, patches=loaded, strength_patch=strength_clip)
-    else:
-        k1 = ()
-        new_clip = None
-    k = set(k)
-    k1 = set(k1)
-    for x in loaded:
-        if (x not in k) and (x not in k1):
-            logging.warning(f"NOT LOADED {x}")
-    return (new_modelpatcher, new_clip, hook_group)
-
-def _combine_hooks_from_values(c_dict: dict[str, HookGroup], values: dict[str, HookGroup], cache: dict[tuple[HookGroup, HookGroup], HookGroup]):
-    hooks_key = 'hooks'
-    # if hooks only exist in one dict, do what's needed so that it ends up in c_dict
-    if hooks_key not in values:
-        return
-    if hooks_key not in c_dict:
-        hooks_value = values.get(hooks_key, None)
-        if hooks_value is not None:
-            c_dict[hooks_key] = hooks_value
-        return
-    # otherwise, need to combine with minimum duplication via cache
-    hooks_tuple = (c_dict[hooks_key], values[hooks_key])
-    cached_hooks = cache.get(hooks_tuple, None)
-    if cached_hooks is None:
-        new_hooks = hooks_tuple[0].clone_and_combine(hooks_tuple[1])
-        cache[hooks_tuple] = new_hooks
-        c_dict[hooks_key] = new_hooks
-    else:
-        c_dict[hooks_key] = cache[hooks_tuple]
-
-def conditioning_set_values_with_hooks(conditioning, values={}, append_hooks=True,
-                                       cache: dict[tuple[HookGroup, HookGroup], HookGroup]=None):
-    c = []
-    if cache is None:
-        cache = {}
-    for t in conditioning:
-        n = [t[0], t[1].copy()]
-        for k in values:
-            if append_hooks and k == 'hooks':
-                _combine_hooks_from_values(n[1], values, cache)
-            else:
-                n[1][k] = values[k]
-        c.append(n)
-
-    return c
-
-def set_hooks_for_conditioning(cond, hooks: HookGroup, append_hooks=True, cache: dict[tuple[HookGroup, HookGroup], HookGroup]=None):
-    if hooks is None:
-        return cond
-    return conditioning_set_values_with_hooks(cond, {'hooks': hooks}, append_hooks=append_hooks, cache=cache)
-
-def set_timesteps_for_conditioning(cond, timestep_range: tuple[float,float]):
-    if timestep_range is None:
-        return cond
-    return conditioning_set_values(cond, {"start_percent": timestep_range[0],
-                                          "end_percent": timestep_range[1]})
-
-def set_mask_for_conditioning(cond, mask: torch.Tensor, set_cond_area: str, strength: float):
-    if mask is None:
-        return cond
-    set_area_to_bounds = False
-    if set_cond_area != 'default':
-        set_area_to_bounds = True
-    if len(mask.shape) < 3:
-        mask = mask.unsqueeze(0)
-    return conditioning_set_values(cond, {'mask': mask,
-                                          'set_area_to_bounds': set_area_to_bounds,
-                                          'mask_strength': strength})
-
-def combine_conditioning(conds: list):
-    combined_conds = []
-    for cond in conds:
-        combined_conds.extend(cond)
-    return combined_conds
-
-def combine_with_new_conds(conds: list, new_conds: list):
-    combined_conds = []
-    for c, new_c in zip(conds, new_conds):
-        combined_conds.append(combine_conditioning([c, new_c]))
-    return combined_conds
-
-def set_conds_props(conds: list, strength: float, set_cond_area: str,
-                   mask: torch.Tensor=None, hooks: HookGroup=None, timesteps_range: tuple[float,float]=None, append_hooks=True):
-    final_conds = []
-    cache = {}
-    for c in conds:
-        # first, apply lora_hook to conditioning, if provided
-        c = set_hooks_for_conditioning(c, hooks, append_hooks=append_hooks, cache=cache)
-        # next, apply mask to conditioning
-        c = set_mask_for_conditioning(cond=c, mask=mask, strength=strength, set_cond_area=set_cond_area)
-        # apply timesteps, if present
-        c = set_timesteps_for_conditioning(cond=c, timestep_range=timesteps_range)
-        # finally, apply mask to conditioning and store
-        final_conds.append(c)
-    return final_conds
-
-def set_conds_props_and_combine(conds: list, new_conds: list, strength: float=1.0, set_cond_area: str="default",
-                               mask: torch.Tensor=None, hooks: HookGroup=None, timesteps_range: tuple[float,float]=None, append_hooks=True):
-    combined_conds = []
-    cache = {}
-    for c, masked_c in zip(conds, new_conds):
-        # first, apply lora_hook to new conditioning, if provided
-        masked_c = set_hooks_for_conditioning(masked_c, hooks, append_hooks=append_hooks, cache=cache)
-        # next, apply mask to new conditioning, if provided
-        masked_c = set_mask_for_conditioning(cond=masked_c, mask=mask, set_cond_area=set_cond_area, strength=strength)
-        # apply timesteps, if present
-        masked_c = set_timesteps_for_conditioning(cond=masked_c, timestep_range=timesteps_range)
-        # finally, combine with existing conditioning and store
-        combined_conds.append(combine_conditioning([c, masked_c]))
-    return combined_conds
-
-def set_default_conds_and_combine(conds: list, new_conds: list,
-                                   hooks: HookGroup=None, timesteps_range: tuple[float,float]=None, append_hooks=True):
-    combined_conds = []
-    cache = {}
-    for c, new_c in zip(conds, new_conds):
-        # first, apply lora_hook to new conditioning, if provided
-        new_c = set_hooks_for_conditioning(new_c, hooks, append_hooks=append_hooks, cache=cache)
-        # next, add default_cond key to cond so that during sampling, it can be identified
-        new_c = conditioning_set_values(new_c, {'default': True})
-        # apply timesteps, if present
-        new_c = set_timesteps_for_conditioning(cond=new_c, timestep_range=timesteps_range)
-        # finally, combine with existing conditioning and store
-        combined_conds.append(combine_conditioning([c, new_c]))
-    return combined_conds

comfy/image_encoders/dino2.py

@@ -1,141 +0,0 @@
-import torch
-from comfy.text_encoders.bert import BertAttention
-import comfy.model_management
-from comfy.ldm.modules.attention import optimized_attention_for_device
-
-
-class Dino2AttentionOutput(torch.nn.Module):
-    def __init__(self, input_dim, output_dim, layer_norm_eps, dtype, device, operations):
-        super().__init__()
-        self.dense = operations.Linear(input_dim, output_dim, dtype=dtype, device=device)
-
-    def forward(self, x):
-        return self.dense(x)
-
-
-class Dino2AttentionBlock(torch.nn.Module):
-    def __init__(self, embed_dim, heads, layer_norm_eps, dtype, device, operations):
-        super().__init__()
-        self.attention = BertAttention(embed_dim, heads, dtype, device, operations)
-        self.output = Dino2AttentionOutput(embed_dim, embed_dim, layer_norm_eps, dtype, device, operations)
-
-    def forward(self, x, mask, optimized_attention):
-        return self.output(self.attention(x, mask, optimized_attention))
-
-
-class LayerScale(torch.nn.Module):
-    def __init__(self, dim, dtype, device, operations):
-        super().__init__()
-        self.lambda1 = torch.nn.Parameter(torch.empty(dim, device=device, dtype=dtype))
-
-    def forward(self, x):
-        return x * comfy.model_management.cast_to_device(self.lambda1, x.device, x.dtype)
-
-
-class SwiGLUFFN(torch.nn.Module):
-    def __init__(self, dim, dtype, device, operations):
-        super().__init__()
-        in_features = out_features = dim
-        hidden_features = int(dim * 4)
-        hidden_features = (int(hidden_features * 2 / 3) + 7) // 8 * 8
-
-        self.weights_in = operations.Linear(in_features, 2 * hidden_features, bias=True, device=device, dtype=dtype)
-        self.weights_out = operations.Linear(hidden_features, out_features, bias=True, device=device, dtype=dtype)
-
-    def forward(self, x):
-        x = self.weights_in(x)
-        x1, x2 = x.chunk(2, dim=-1)
-        x = torch.nn.functional.silu(x1) * x2
-        return self.weights_out(x)
-
-
-class Dino2Block(torch.nn.Module):
-    def __init__(self, dim, num_heads, layer_norm_eps, dtype, device, operations):
-        super().__init__()
-        self.attention = Dino2AttentionBlock(dim, num_heads, layer_norm_eps, dtype, device, operations)
-        self.layer_scale1 = LayerScale(dim, dtype, device, operations)
-        self.layer_scale2 = LayerScale(dim, dtype, device, operations)
-        self.mlp = SwiGLUFFN(dim, dtype, device, operations)
-        self.norm1 = operations.LayerNorm(dim, eps=layer_norm_eps, dtype=dtype, device=device)
-        self.norm2 = operations.LayerNorm(dim, eps=layer_norm_eps, dtype=dtype, device=device)
-
-    def forward(self, x, optimized_attention):
-        x = x + self.layer_scale1(self.attention(self.norm1(x), None, optimized_attention))
-        x = x + self.layer_scale2(self.mlp(self.norm2(x)))
-        return x
-
-
-class Dino2Encoder(torch.nn.Module):
-    def __init__(self, dim, num_heads, layer_norm_eps, num_layers, dtype, device, operations):
-        super().__init__()
-        self.layer = torch.nn.ModuleList([Dino2Block(dim, num_heads, layer_norm_eps, dtype, device, operations) for _ in range(num_layers)])
-
-    def forward(self, x, intermediate_output=None):
-        optimized_attention = optimized_attention_for_device(x.device, False, small_input=True)
-
-        if intermediate_output is not None:
-            if intermediate_output < 0:
-                intermediate_output = len(self.layer) + intermediate_output
-
-        intermediate = None
-        for i, l in enumerate(self.layer):
-            x = l(x, optimized_attention)
-            if i == intermediate_output:
-                intermediate = x.clone()
-        return x, intermediate
-
-
-class Dino2PatchEmbeddings(torch.nn.Module):
-    def __init__(self, dim, num_channels=3, patch_size=14, image_size=518, dtype=None, device=None, operations=None):
-        super().__init__()
-        self.projection = operations.Conv2d(
-            in_channels=num_channels,
-            out_channels=dim,
-            kernel_size=patch_size,
-            stride=patch_size,
-            bias=True,
-            dtype=dtype,
-            device=device
-        )
-
-    def forward(self, pixel_values):
-        return self.projection(pixel_values).flatten(2).transpose(1, 2)
-
-
-class Dino2Embeddings(torch.nn.Module):
-    def __init__(self, dim, dtype, device, operations):
-        super().__init__()
-        patch_size = 14
-        image_size = 518
-
-        self.patch_embeddings = Dino2PatchEmbeddings(dim, patch_size=patch_size, image_size=image_size, dtype=dtype, device=device, operations=operations)
-        self.position_embeddings = torch.nn.Parameter(torch.empty(1, (image_size // patch_size) ** 2 + 1, dim, dtype=dtype, device=device))
-        self.cls_token = torch.nn.Parameter(torch.empty(1, 1, dim, dtype=dtype, device=device))
-        self.mask_token = torch.nn.Parameter(torch.empty(1, dim, dtype=dtype, device=device))
-
-    def forward(self, pixel_values):
-        x = self.patch_embeddings(pixel_values)
-        # TODO: mask_token?
-        x = torch.cat((self.cls_token.to(device=x.device, dtype=x.dtype).expand(x.shape[0], -1, -1), x), dim=1)
-        x = x + comfy.model_management.cast_to_device(self.position_embeddings, x.device, x.dtype)
-        return x
-
-
-class Dinov2Model(torch.nn.Module):
-    def __init__(self, config_dict, dtype, device, operations):
-        super().__init__()
-        num_layers = config_dict["num_hidden_layers"]
-        dim = config_dict["hidden_size"]
-        heads = config_dict["num_attention_heads"]
-        layer_norm_eps = config_dict["layer_norm_eps"]
-
-        self.embeddings = Dino2Embeddings(dim, dtype, device, operations)
-        self.encoder = Dino2Encoder(dim, heads, layer_norm_eps, num_layers, dtype, device, operations)
-        self.layernorm = operations.LayerNorm(dim, eps=layer_norm_eps, dtype=dtype, device=device)
-
-    def forward(self, pixel_values, attention_mask=None, intermediate_output=None):
-        x = self.embeddings(pixel_values)
-        x, i = self.encoder(x, intermediate_output=intermediate_output)
-        x = self.layernorm(x)
-        pooled_output = x[:, 0, :]
-        return x, i, pooled_output, None

comfy/image_encoders/dino2_giant.json

@@ -1,21 +0,0 @@
-{
-  "attention_probs_dropout_prob": 0.0,
-  "drop_path_rate": 0.0,
-  "hidden_act": "gelu",
-  "hidden_dropout_prob": 0.0,
-  "hidden_size": 1536,
-  "image_size": 518,
-  "initializer_range": 0.02,
-  "layer_norm_eps": 1e-06,
-  "layerscale_value": 1.0,
-  "mlp_ratio": 4,
-  "model_type": "dinov2",
-  "num_attention_heads": 24,
-  "num_channels": 3,
-  "num_hidden_layers": 40,
-  "patch_size": 14,
-  "qkv_bias": true,
-  "use_swiglu_ffn": true,
-  "image_mean": [0.485, 0.456, 0.406],
-  "image_std": [0.229, 0.224, 0.225]
-}

comfy/k_diffusion/deis.py

@@ -1,120 +0,0 @@
-#Taken from: https://github.com/zju-pi/diff-sampler/blob/main/gits-main/solver_utils.py
-#under Apache 2 license
-import torch
-import numpy as np
-
-# A pytorch reimplementation of DEIS (https://github.com/qsh-zh/deis).
-#############################
-### Utils for DEIS solver ###
-#############################
-#----------------------------------------------------------------------------
-# Transfer from the input time (sigma) used in EDM to that (t) used in DEIS.
-
-def edm2t(edm_steps, epsilon_s=1e-3, sigma_min=0.002, sigma_max=80):
-    vp_sigma_inv = lambda beta_d, beta_min: lambda sigma: ((beta_min ** 2 + 2 * beta_d * (sigma ** 2 + 1).log()).sqrt() - beta_min) / beta_d
-    vp_beta_d = 2 * (np.log(torch.tensor(sigma_min).cpu() ** 2 + 1) / epsilon_s - np.log(torch.tensor(sigma_max).cpu() ** 2 + 1)) / (epsilon_s - 1)
-    vp_beta_min = np.log(torch.tensor(sigma_max).cpu() ** 2 + 1) - 0.5 * vp_beta_d
-    t_steps = vp_sigma_inv(vp_beta_d.clone().detach().cpu(), vp_beta_min.clone().detach().cpu())(edm_steps.clone().detach().cpu())
-    return t_steps, vp_beta_min, vp_beta_d + vp_beta_min
-
-#----------------------------------------------------------------------------
-
-def cal_poly(prev_t, j, taus):
-    poly = 1
-    for k in range(prev_t.shape[0]):
-        if k == j:
-            continue
-        poly *= (taus - prev_t[k]) / (prev_t[j] - prev_t[k])
-    return poly
-
-#----------------------------------------------------------------------------
-# Transfer from t to alpha_t.
-
-def t2alpha_fn(beta_0, beta_1, t):
-    return torch.exp(-0.5 * t ** 2 * (beta_1 - beta_0) - t * beta_0)
-
-#----------------------------------------------------------------------------
-
-def cal_intergrand(beta_0, beta_1, taus):
-    with torch.inference_mode(mode=False):
-        taus = taus.clone()
-        beta_0 = beta_0.clone()
-        beta_1 = beta_1.clone()
-        with torch.enable_grad():
-            taus.requires_grad_(True)
-            alpha = t2alpha_fn(beta_0, beta_1, taus)
-            log_alpha = alpha.log()
-            log_alpha.sum().backward()
-            d_log_alpha_dtau = taus.grad
-    integrand = -0.5 * d_log_alpha_dtau / torch.sqrt(alpha * (1 - alpha))
-    return integrand
-
-#----------------------------------------------------------------------------
-
-def get_deis_coeff_list(t_steps, max_order, N=10000, deis_mode='tab'):
-    """
-    Get the coefficient list for DEIS sampling.
-
-    Args:
-        t_steps: A pytorch tensor. The time steps for sampling.
-        max_order: A `int`. Maximum order of the solver. 1 <= max_order <= 4
-        N: A `int`. Use how many points to perform the numerical integration when deis_mode=='tab'.
-        deis_mode: A `str`. Select between 'tab' and 'rhoab'. Type of DEIS.
-    Returns:
-        A pytorch tensor. A batch of generated samples or sampling trajectories if return_inters=True.
-    """
-    if deis_mode == 'tab':
-        t_steps, beta_0, beta_1 = edm2t(t_steps)
-        C = []
-        for i, (t_cur, t_next) in enumerate(zip(t_steps[:-1], t_steps[1:])):
-            order = min(i+1, max_order)
-            if order == 1:
-                C.append([])
-            else:
-                taus = torch.linspace(t_cur, t_next, N)   # split the interval for integral appximation
-                dtau = (t_next - t_cur) / N
-                prev_t = t_steps[[i - k for k in range(order)]]
-                coeff_temp = []
-                integrand = cal_intergrand(beta_0, beta_1, taus)
-                for j in range(order):
-                    poly = cal_poly(prev_t, j, taus)
-                    coeff_temp.append(torch.sum(integrand * poly) * dtau)
-                C.append(coeff_temp)
-
-    elif deis_mode == 'rhoab':
-        # Analytical solution, second order
-        def get_def_intergral_2(a, b, start, end, c):
-            coeff = (end**3 - start**3) / 3 - (end**2 - start**2) * (a + b) / 2 + (end - start) * a * b
-            return coeff / ((c - a) * (c - b))
-
-        # Analytical solution, third order
-        def get_def_intergral_3(a, b, c, start, end, d):
-            coeff = (end**4 - start**4) / 4 - (end**3 - start**3) * (a + b + c) / 3 \
-                    + (end**2 - start**2) * (a*b + a*c + b*c) / 2 - (end - start) * a * b * c
-            return coeff / ((d - a) * (d - b) * (d - c))
-
-        C = []
-        for i, (t_cur, t_next) in enumerate(zip(t_steps[:-1], t_steps[1:])):
-            order = min(i, max_order)
-            if order == 0:
-                C.append([])
-            else:
-                prev_t = t_steps[[i - k for k in range(order+1)]]
-                if order == 1:
-                    coeff_cur = ((t_next - prev_t[1])**2 - (t_cur - prev_t[1])**2) / (2 * (t_cur - prev_t[1]))
-                    coeff_prev1 = (t_next - t_cur)**2 / (2 * (prev_t[1] - t_cur))
-                    coeff_temp = [coeff_cur, coeff_prev1]
-                elif order == 2:
-                    coeff_cur = get_def_intergral_2(prev_t[1], prev_t[2], t_cur, t_next, t_cur)
-                    coeff_prev1 = get_def_intergral_2(t_cur, prev_t[2], t_cur, t_next, prev_t[1])
-                    coeff_prev2 = get_def_intergral_2(t_cur, prev_t[1], t_cur, t_next, prev_t[2])
-                    coeff_temp = [coeff_cur, coeff_prev1, coeff_prev2]
-                elif order == 3:
-                    coeff_cur = get_def_intergral_3(prev_t[1], prev_t[2], prev_t[3], t_cur, t_next, t_cur)
-                    coeff_prev1 = get_def_intergral_3(t_cur, prev_t[2], prev_t[3], t_cur, t_next, prev_t[1])
-                    coeff_prev2 = get_def_intergral_3(t_cur, prev_t[1], prev_t[3], t_cur, t_next, prev_t[2])
-                    coeff_prev3 = get_def_intergral_3(t_cur, prev_t[1], prev_t[2], t_cur, t_next, prev_t[3])
-                    coeff_temp = [coeff_cur, coeff_prev1, coeff_prev2, coeff_prev3]
-                C.append(coeff_temp)
-    return C
-

comfy/k_diffusion/sa_solver.py

@@ -1,121 +0,0 @@
-# SA-Solver: Stochastic Adams Solver (NeurIPS 2023, arXiv:2309.05019)
-# Conference: https://proceedings.neurips.cc/paper_files/paper/2023/file/f4a6806490d31216a3ba667eb240c897-Paper-Conference.pdf
-# Codebase ref: https://github.com/scxue/SA-Solver
-
-import math
-from typing import Union, Callable
-import torch
-
-
-def compute_exponential_coeffs(s: torch.Tensor, t: torch.Tensor, solver_order: int, tau_t: float) -> torch.Tensor:
-    """Compute (1 + tau^2) * integral of exp((1 + tau^2) * x) * x^p dx from s to t with exp((1 + tau^2) * t) factored out, using integration by parts.
-
-    Integral of exp((1 + tau^2) * x) * x^p dx
-        = product_terms[p] - (p / (1 + tau^2)) * integral of exp((1 + tau^2) * x) * x^(p-1) dx,
-    with base case p=0 where integral equals product_terms[0].
-
-    where
-        product_terms[p] = x^p * exp((1 + tau^2) * x) / (1 + tau^2).
-
-    Construct a recursive coefficient matrix following the above recursive relation to compute all integral terms up to p = (solver_order - 1).
-    Return coefficients used by the SA-Solver in data prediction mode.
-
-    Args:
-        s: Start time s.
-        t: End time t.
-        solver_order: Current order of the solver.
-        tau_t: Stochastic strength parameter in the SDE.
-
-    Returns:
-        Exponential coefficients used in data prediction, with exp((1 + tau^2) * t) factored out, ordered from p=0 to p=solver_order−1, shape (solver_order,).
-    """
-    tau_mul = 1 + tau_t ** 2
-    h = t - s
-    p = torch.arange(solver_order, dtype=s.dtype, device=s.device)
-
-    # product_terms after factoring out exp((1 + tau^2) * t)
-    # Includes (1 + tau^2) factor from outside the integral
-    product_terms_factored = (t ** p - s ** p * (-tau_mul * h).exp())
-
-    # Lower triangular recursive coefficient matrix
-    # Accumulates recursive coefficients based on p / (1 + tau^2)
-    recursive_depth_mat = p.unsqueeze(1) - p.unsqueeze(0)
-    log_factorial = (p + 1).lgamma()
-    recursive_coeff_mat = log_factorial.unsqueeze(1) - log_factorial.unsqueeze(0)
-    if tau_t > 0:
-        recursive_coeff_mat = recursive_coeff_mat - (recursive_depth_mat * math.log(tau_mul))
-    signs = torch.where(recursive_depth_mat % 2 == 0, 1.0, -1.0)
-    recursive_coeff_mat = (recursive_coeff_mat.exp() * signs).tril()
-
-    return recursive_coeff_mat @ product_terms_factored
-
-
-def compute_simple_stochastic_adams_b_coeffs(sigma_next: torch.Tensor, curr_lambdas: torch.Tensor, lambda_s: torch.Tensor, lambda_t: torch.Tensor, tau_t: float, is_corrector_step: bool = False) -> torch.Tensor:
-    """Compute simple order-2 b coefficients from SA-Solver paper (Appendix D. Implementation Details)."""
-    tau_mul = 1 + tau_t ** 2
-    h = lambda_t - lambda_s
-    alpha_t = sigma_next * lambda_t.exp()
-    if is_corrector_step:
-        # Simplified 1-step (order-2) corrector
-        b_1 = alpha_t * (0.5 * tau_mul * h)
-        b_2 = alpha_t * (-h * tau_mul).expm1().neg() - b_1
-    else:
-        # Simplified 2-step predictor
-        b_2 = alpha_t * (0.5 * tau_mul * h ** 2) / (curr_lambdas[-2] - lambda_s)
-        b_1 = alpha_t * (-h * tau_mul).expm1().neg() - b_2
-    return torch.stack([b_2, b_1])
-
-
-def compute_stochastic_adams_b_coeffs(sigma_next: torch.Tensor, curr_lambdas: torch.Tensor, lambda_s: torch.Tensor, lambda_t: torch.Tensor, tau_t: float, simple_order_2: bool = False, is_corrector_step: bool = False) -> torch.Tensor:
-    """Compute b_i coefficients for the SA-Solver (see eqs. 15 and 18).
-
-    The solver order corresponds to the number of input lambdas (half-logSNR points).
-
-    Args:
-        sigma_next: Sigma at end time t.
-        curr_lambdas: Lambda time points used to construct the Lagrange basis, shape (N,).
-        lambda_s: Lambda at start time s.
-        lambda_t: Lambda at end time t.
-        tau_t: Stochastic strength parameter in the SDE.
-        simple_order_2: Whether to enable the simple order-2 scheme.
-        is_corrector_step: Flag for corrector step in simple order-2 mode.
-
-    Returns:
-        b_i coefficients for the SA-Solver, shape (N,), where N is the solver order.
-    """
-    num_timesteps = curr_lambdas.shape[0]
-
-    if simple_order_2 and num_timesteps == 2:
-        return compute_simple_stochastic_adams_b_coeffs(sigma_next, curr_lambdas, lambda_s, lambda_t, tau_t, is_corrector_step)
-
-    # Compute coefficients by solving a linear system from Lagrange basis interpolation
-    exp_integral_coeffs = compute_exponential_coeffs(lambda_s, lambda_t, num_timesteps, tau_t)
-    vandermonde_matrix_T = torch.vander(curr_lambdas, num_timesteps, increasing=True).T
-    lagrange_integrals = torch.linalg.solve(vandermonde_matrix_T, exp_integral_coeffs)
-
-    # (sigma_t * exp(-tau^2 * lambda_t)) * exp((1 + tau^2) * lambda_t)
-    # = sigma_t * exp(lambda_t) = alpha_t
-    # exp((1 + tau^2) * lambda_t) is extracted from the integral
-    alpha_t = sigma_next * lambda_t.exp()
-    return alpha_t * lagrange_integrals
-
-
-def get_tau_interval_func(start_sigma: float, end_sigma: float, eta: float = 1.0) -> Callable[[Union[torch.Tensor, float]], float]:
-    """Return a function that controls the stochasticity of SA-Solver.
-
-    When eta = 0, SA-Solver runs as ODE. The official approach uses
-    time t to determine the SDE interval, while here we use sigma instead.
-
-    See:
-        https://github.com/scxue/SA-Solver/blob/main/README.md
-    """
-
-    def tau_func(sigma: Union[torch.Tensor, float]) -> float:
-        if eta <= 0:
-            return 0.0  # ODE
-
-        if isinstance(sigma, torch.Tensor):
-            sigma = sigma.item()
-        return eta if start_sigma >= sigma >= end_sigma else 0.0
-
-    return tau_func

comfy/k_diffusion/sampling.py

@@ -1,1761 +0,0 @@
-import math
-from functools import partial
-
-from scipy import integrate
-import torch
-from torch import nn
-import torchsde
-from tqdm.auto import trange, tqdm
-
-from . import utils
-from . import deis
-from . import sa_solver
-import comfy.model_patcher
-import comfy.model_sampling
-
-def append_zero(x):
-    return torch.cat([x, x.new_zeros([1])])
-
-
-def get_sigmas_karras(n, sigma_min, sigma_max, rho=7., device='cpu'):
-    """Constructs the noise schedule of Karras et al. (2022)."""
-    ramp = torch.linspace(0, 1, n, device=device)
-    min_inv_rho = sigma_min ** (1 / rho)
-    max_inv_rho = sigma_max ** (1 / rho)
-    sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho
-    return append_zero(sigmas).to(device)
-
-
-def get_sigmas_exponential(n, sigma_min, sigma_max, device='cpu'):
-    """Constructs an exponential noise schedule."""
-    sigmas = torch.linspace(math.log(sigma_max), math.log(sigma_min), n, device=device).exp()
-    return append_zero(sigmas)
-
-
-def get_sigmas_polyexponential(n, sigma_min, sigma_max, rho=1., device='cpu'):
-    """Constructs an polynomial in log sigma noise schedule."""
-    ramp = torch.linspace(1, 0, n, device=device) ** rho
-    sigmas = torch.exp(ramp * (math.log(sigma_max) - math.log(sigma_min)) + math.log(sigma_min))
-    return append_zero(sigmas)
-
-
-def get_sigmas_vp(n, beta_d=19.9, beta_min=0.1, eps_s=1e-3, device='cpu'):
-    """Constructs a continuous VP noise schedule."""
-    t = torch.linspace(1, eps_s, n, device=device)
-    sigmas = torch.sqrt(torch.special.expm1(beta_d * t ** 2 / 2 + beta_min * t))
-    return append_zero(sigmas)
-
-
-def get_sigmas_laplace(n, sigma_min, sigma_max, mu=0., beta=0.5, device='cpu'):
-    """Constructs the noise schedule proposed by Tiankai et al. (2024). """
-    epsilon = 1e-5 # avoid log(0)
-    x = torch.linspace(0, 1, n, device=device)
-    clamp = lambda x: torch.clamp(x, min=sigma_min, max=sigma_max)
-    lmb = mu - beta * torch.sign(0.5-x) * torch.log(1 - 2 * torch.abs(0.5-x) + epsilon)
-    sigmas = clamp(torch.exp(lmb))
-    return sigmas
-
-
-
-def to_d(x, sigma, denoised):
-    """Converts a denoiser output to a Karras ODE derivative."""
-    return (x - denoised) / utils.append_dims(sigma, x.ndim)
-
-
-def get_ancestral_step(sigma_from, sigma_to, eta=1.):
-    """Calculates the noise level (sigma_down) to step down to and the amount
-    of noise to add (sigma_up) when doing an ancestral sampling step."""
-    if not eta:
-        return sigma_to, 0.
-    sigma_up = min(sigma_to, eta * (sigma_to ** 2 * (sigma_from ** 2 - sigma_to ** 2) / sigma_from ** 2) ** 0.5)
-    sigma_down = (sigma_to ** 2 - sigma_up ** 2) ** 0.5
-    return sigma_down, sigma_up
-
-
-def default_noise_sampler(x, seed=None):
-    if seed is not None:
-        generator = torch.Generator(device=x.device)
-        generator.manual_seed(seed)
-    else:
-        generator = None
-
-    return lambda sigma, sigma_next: torch.randn(x.size(), dtype=x.dtype, layout=x.layout, device=x.device, generator=generator)
-
-
-class BatchedBrownianTree:
-    """A wrapper around torchsde.BrownianTree that enables batches of entropy."""
-
-    def __init__(self, x, t0, t1, seed=None, **kwargs):
-        self.cpu_tree = True
-        if "cpu" in kwargs:
-            self.cpu_tree = kwargs.pop("cpu")
-        t0, t1, self.sign = self.sort(t0, t1)
-        w0 = kwargs.get('w0', torch.zeros_like(x))
-        if seed is None:
-            seed = torch.randint(0, 2 ** 63 - 1, []).item()
-        self.batched = True
-        try:
-            assert len(seed) == x.shape[0]
-            w0 = w0[0]
-        except TypeError:
-            seed = [seed]
-            self.batched = False
-        if self.cpu_tree:
-            self.trees = [torchsde.BrownianTree(t0.cpu(), w0.cpu(), t1.cpu(), entropy=s, **kwargs) for s in seed]
-        else:
-            self.trees = [torchsde.BrownianTree(t0, w0, t1, entropy=s, **kwargs) for s in seed]
-
-    @staticmethod
-    def sort(a, b):
-        return (a, b, 1) if a < b else (b, a, -1)
-
-    def __call__(self, t0, t1):
-        t0, t1, sign = self.sort(t0, t1)
-        if self.cpu_tree:
-            w = torch.stack([tree(t0.cpu().float(), t1.cpu().float()).to(t0.dtype).to(t0.device) for tree in self.trees]) * (self.sign * sign)
-        else:
-            w = torch.stack([tree(t0, t1) for tree in self.trees]) * (self.sign * sign)
-
-        return w if self.batched else w[0]
-
-
-class BrownianTreeNoiseSampler:
-    """A noise sampler backed by a torchsde.BrownianTree.
-
-    Args:
-        x (Tensor): The tensor whose shape, device and dtype to use to generate
-            random samples.
-        sigma_min (float): The low end of the valid interval.
-        sigma_max (float): The high end of the valid interval.
-        seed (int or List[int]): The random seed. If a list of seeds is
-            supplied instead of a single integer, then the noise sampler will
-            use one BrownianTree per batch item, each with its own seed.
-        transform (callable): A function that maps sigma to the sampler's
-            internal timestep.
-    """
-
-    def __init__(self, x, sigma_min, sigma_max, seed=None, transform=lambda x: x, cpu=False):
-        self.transform = transform
-        t0, t1 = self.transform(torch.as_tensor(sigma_min)), self.transform(torch.as_tensor(sigma_max))
-        self.tree = BatchedBrownianTree(x, t0, t1, seed, cpu=cpu)
-
-    def __call__(self, sigma, sigma_next):
-        t0, t1 = self.transform(torch.as_tensor(sigma)), self.transform(torch.as_tensor(sigma_next))
-        return self.tree(t0, t1) / (t1 - t0).abs().sqrt()
-
-
-def sigma_to_half_log_snr(sigma, model_sampling):
-    """Convert sigma to half-logSNR log(alpha_t / sigma_t)."""
-    if isinstance(model_sampling, comfy.model_sampling.CONST):
-        # log((1 - t) / t) = log((1 - sigma) / sigma)
-        return sigma.logit().neg()
-    return sigma.log().neg()
-
-
-def half_log_snr_to_sigma(half_log_snr, model_sampling):
-    """Convert half-logSNR log(alpha_t / sigma_t) to sigma."""
-    if isinstance(model_sampling, comfy.model_sampling.CONST):
-        # 1 / (1 + exp(half_log_snr))
-        return half_log_snr.neg().sigmoid()
-    return half_log_snr.neg().exp()
-
-
-def offset_first_sigma_for_snr(sigmas, model_sampling, percent_offset=1e-4):
-    """Adjust the first sigma to avoid invalid logSNR."""
-    if len(sigmas) <= 1:
-        return sigmas
-    if isinstance(model_sampling, comfy.model_sampling.CONST):
-        if sigmas[0] >= 1:
-            sigmas = sigmas.clone()
-            sigmas[0] = model_sampling.percent_to_sigma(percent_offset)
-    return sigmas
-
-
-@torch.no_grad()
-def sample_euler(model, x, sigmas, extra_args=None, callback=None, disable=None, s_churn=0., s_tmin=0., s_tmax=float('inf'), s_noise=1.):
-    """Implements Algorithm 2 (Euler steps) from Karras et al. (2022)."""
-    extra_args = {} if extra_args is None else extra_args
-    s_in = x.new_ones([x.shape[0]])
-    for i in trange(len(sigmas) - 1, disable=disable):
-        if s_churn > 0:
-            gamma = min(s_churn / (len(sigmas) - 1), 2 ** 0.5 - 1) if s_tmin <= sigmas[i] <= s_tmax else 0.
-            sigma_hat = sigmas[i] * (gamma + 1)
-        else:
-            gamma = 0
-            sigma_hat = sigmas[i]
-
-        if gamma > 0:
-            eps = torch.randn_like(x) * s_noise
-            x = x + eps * (sigma_hat ** 2 - sigmas[i] ** 2) ** 0.5
-        denoised = model(x, sigma_hat * s_in, **extra_args)
-        d = to_d(x, sigma_hat, denoised)
-        if callback is not None:
-            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigma_hat, 'denoised': denoised})
-        dt = sigmas[i + 1] - sigma_hat
-        # Euler method
-        x = x + d * dt
-    return x
-
-
-@torch.no_grad()
-def sample_euler_ancestral(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None):
-    if isinstance(model.inner_model.inner_model.model_sampling, comfy.model_sampling.CONST):
-        return sample_euler_ancestral_RF(model, x, sigmas, extra_args, callback, disable, eta, s_noise, noise_sampler)
-    """Ancestral sampling with Euler method steps."""
-    extra_args = {} if extra_args is None else extra_args
-    seed = extra_args.get("seed", None)
-    noise_sampler = default_noise_sampler(x, seed=seed) if noise_sampler is None else noise_sampler
-    s_in = x.new_ones([x.shape[0]])
-    for i in trange(len(sigmas) - 1, disable=disable):
-        denoised = model(x, sigmas[i] * s_in, **extra_args)
-        sigma_down, sigma_up = get_ancestral_step(sigmas[i], sigmas[i + 1], eta=eta)
-        if callback is not None:
-            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
-
-        if sigma_down == 0:
-            x = denoised
-        else:
-            d = to_d(x, sigmas[i], denoised)
-            # Euler method
-            dt = sigma_down - sigmas[i]
-            x = x + d * dt + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up
-    return x
-
-@torch.no_grad()
-def sample_euler_ancestral_RF(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1.0, s_noise=1., noise_sampler=None):
-    """Ancestral sampling with Euler method steps."""
-    extra_args = {} if extra_args is None else extra_args
-    seed = extra_args.get("seed", None)
-    noise_sampler = default_noise_sampler(x, seed=seed) if noise_sampler is None else noise_sampler
-    s_in = x.new_ones([x.shape[0]])
-    for i in trange(len(sigmas) - 1, disable=disable):
-        denoised = model(x, sigmas[i] * s_in, **extra_args)
-        # sigma_down, sigma_up = get_ancestral_step(sigmas[i], sigmas[i + 1], eta=eta)
-        if callback is not None:
-            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
-
-        if sigmas[i + 1] == 0:
-            x = denoised
-        else:
-            downstep_ratio = 1 + (sigmas[i + 1] / sigmas[i] - 1) * eta
-            sigma_down = sigmas[i + 1] * downstep_ratio
-            alpha_ip1 = 1 - sigmas[i + 1]
-            alpha_down = 1 - sigma_down
-            renoise_coeff = (sigmas[i + 1]**2 - sigma_down**2 * alpha_ip1**2 / alpha_down**2)**0.5
-            # Euler method
-            sigma_down_i_ratio = sigma_down / sigmas[i]
-            x = sigma_down_i_ratio * x + (1 - sigma_down_i_ratio) * denoised
-            if eta > 0:
-                x = (alpha_ip1 / alpha_down) * x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * renoise_coeff
-    return x
-
-@torch.no_grad()
-def sample_heun(model, x, sigmas, extra_args=None, callback=None, disable=None, s_churn=0., s_tmin=0., s_tmax=float('inf'), s_noise=1.):
-    """Implements Algorithm 2 (Heun steps) from Karras et al. (2022)."""
-    extra_args = {} if extra_args is None else extra_args
-    s_in = x.new_ones([x.shape[0]])
-    for i in trange(len(sigmas) - 1, disable=disable):
-        if s_churn > 0:
-            gamma = min(s_churn / (len(sigmas) - 1), 2 ** 0.5 - 1) if s_tmin <= sigmas[i] <= s_tmax else 0.
-            sigma_hat = sigmas[i] * (gamma + 1)
-        else:
-            gamma = 0
-            sigma_hat = sigmas[i]
-
-        sigma_hat = sigmas[i] * (gamma + 1)
-        if gamma > 0:
-            eps = torch.randn_like(x) * s_noise
-            x = x + eps * (sigma_hat ** 2 - sigmas[i] ** 2) ** 0.5
-        denoised = model(x, sigma_hat * s_in, **extra_args)
-        d = to_d(x, sigma_hat, denoised)
-        if callback is not None:
-            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigma_hat, 'denoised': denoised})
-        dt = sigmas[i + 1] - sigma_hat
-        if sigmas[i + 1] == 0:
-            # Euler method
-            x = x + d * dt
-        else:
-            # Heun's method
-            x_2 = x + d * dt
-            denoised_2 = model(x_2, sigmas[i + 1] * s_in, **extra_args)
-            d_2 = to_d(x_2, sigmas[i + 1], denoised_2)
-            d_prime = (d + d_2) / 2
-            x = x + d_prime * dt
-    return x
-
-
-@torch.no_grad()
-def sample_dpm_2(model, x, sigmas, extra_args=None, callback=None, disable=None, s_churn=0., s_tmin=0., s_tmax=float('inf'), s_noise=1.):
-    """A sampler inspired by DPM-Solver-2 and Algorithm 2 from Karras et al. (2022)."""
-    extra_args = {} if extra_args is None else extra_args
-    s_in = x.new_ones([x.shape[0]])
-    for i in trange(len(sigmas) - 1, disable=disable):
-        if s_churn > 0:
-            gamma = min(s_churn / (len(sigmas) - 1), 2 ** 0.5 - 1) if s_tmin <= sigmas[i] <= s_tmax else 0.
-            sigma_hat = sigmas[i] * (gamma + 1)
-        else:
-            gamma = 0
-            sigma_hat = sigmas[i]
-
-        if gamma > 0:
-            eps = torch.randn_like(x) * s_noise
-            x = x + eps * (sigma_hat ** 2 - sigmas[i] ** 2) ** 0.5
-        denoised = model(x, sigma_hat * s_in, **extra_args)
-        d = to_d(x, sigma_hat, denoised)
-        if callback is not None:
-            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigma_hat, 'denoised': denoised})
-        if sigmas[i + 1] == 0:
-            # Euler method
-            dt = sigmas[i + 1] - sigma_hat
-            x = x + d * dt
-        else:
-            # DPM-Solver-2
-            sigma_mid = sigma_hat.log().lerp(sigmas[i + 1].log(), 0.5).exp()
-            dt_1 = sigma_mid - sigma_hat
-            dt_2 = sigmas[i + 1] - sigma_hat
-            x_2 = x + d * dt_1
-            denoised_2 = model(x_2, sigma_mid * s_in, **extra_args)
-            d_2 = to_d(x_2, sigma_mid, denoised_2)
-            x = x + d_2 * dt_2
-    return x
-
-
-@torch.no_grad()
-def sample_dpm_2_ancestral(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None):
-    if isinstance(model.inner_model.inner_model.model_sampling, comfy.model_sampling.CONST):
-        return sample_dpm_2_ancestral_RF(model, x, sigmas, extra_args, callback, disable, eta, s_noise, noise_sampler)
-
-    """Ancestral sampling with DPM-Solver second-order steps."""
-    extra_args = {} if extra_args is None else extra_args
-    seed = extra_args.get("seed", None)
-    noise_sampler = default_noise_sampler(x, seed=seed) if noise_sampler is None else noise_sampler
-    s_in = x.new_ones([x.shape[0]])
-    for i in trange(len(sigmas) - 1, disable=disable):
-        denoised = model(x, sigmas[i] * s_in, **extra_args)
-        sigma_down, sigma_up = get_ancestral_step(sigmas[i], sigmas[i + 1], eta=eta)
-        if callback is not None:
-            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
-        d = to_d(x, sigmas[i], denoised)
-        if sigma_down == 0:
-            # Euler method
-            dt = sigma_down - sigmas[i]
-            x = x + d * dt
-        else:
-            # DPM-Solver-2
-            sigma_mid = sigmas[i].log().lerp(sigma_down.log(), 0.5).exp()
-            dt_1 = sigma_mid - sigmas[i]
-            dt_2 = sigma_down - sigmas[i]
-            x_2 = x + d * dt_1
-            denoised_2 = model(x_2, sigma_mid * s_in, **extra_args)
-            d_2 = to_d(x_2, sigma_mid, denoised_2)
-            x = x + d_2 * dt_2
-            x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up
-    return x
-
-@torch.no_grad()
-def sample_dpm_2_ancestral_RF(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None):
-    """Ancestral sampling with DPM-Solver second-order steps."""
-    extra_args = {} if extra_args is None else extra_args
-    seed = extra_args.get("seed", None)
-    noise_sampler = default_noise_sampler(x, seed=seed) if noise_sampler is None else noise_sampler
-    s_in = x.new_ones([x.shape[0]])
-    for i in trange(len(sigmas) - 1, disable=disable):
-        denoised = model(x, sigmas[i] * s_in, **extra_args)
-        downstep_ratio = 1 + (sigmas[i+1]/sigmas[i] - 1) * eta
-        sigma_down = sigmas[i+1] * downstep_ratio
-        alpha_ip1 = 1 - sigmas[i+1]
-        alpha_down = 1 - sigma_down
-        renoise_coeff = (sigmas[i+1]**2 - sigma_down**2*alpha_ip1**2/alpha_down**2)**0.5
-
-        if callback is not None:
-            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
-        d = to_d(x, sigmas[i], denoised)
-        if sigma_down == 0:
-            # Euler method
-            dt = sigma_down - sigmas[i]
-            x = x + d * dt
-        else:
-            # DPM-Solver-2
-            sigma_mid = sigmas[i].log().lerp(sigma_down.log(), 0.5).exp()
-            dt_1 = sigma_mid - sigmas[i]
-            dt_2 = sigma_down - sigmas[i]
-            x_2 = x + d * dt_1
-            denoised_2 = model(x_2, sigma_mid * s_in, **extra_args)
-            d_2 = to_d(x_2, sigma_mid, denoised_2)
-            x = x + d_2 * dt_2
-            x = (alpha_ip1/alpha_down) * x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * renoise_coeff
-    return x
-
-def linear_multistep_coeff(order, t, i, j):
-    if order - 1 > i:
-        raise ValueError(f'Order {order} too high for step {i}')
-    def fn(tau):
-        prod = 1.
-        for k in range(order):
-            if j == k:
-                continue
-            prod *= (tau - t[i - k]) / (t[i - j] - t[i - k])
-        return prod
-    return integrate.quad(fn, t[i], t[i + 1], epsrel=1e-4)[0]
-
-
-@torch.no_grad()
-def sample_lms(model, x, sigmas, extra_args=None, callback=None, disable=None, order=4):
-    extra_args = {} if extra_args is None else extra_args
-    s_in = x.new_ones([x.shape[0]])
-    sigmas_cpu = sigmas.detach().cpu().numpy()
-    ds = []
-    for i in trange(len(sigmas) - 1, disable=disable):
-        denoised = model(x, sigmas[i] * s_in, **extra_args)
-        d = to_d(x, sigmas[i], denoised)
-        ds.append(d)
-        if len(ds) > order:
-            ds.pop(0)
-        if callback is not None:
-            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
-        if sigmas[i + 1] == 0:
-            # Denoising step
-            x = denoised
-        else:
-            cur_order = min(i + 1, order)
-            coeffs = [linear_multistep_coeff(cur_order, sigmas_cpu, i, j) for j in range(cur_order)]
-            x = x + sum(coeff * d for coeff, d in zip(coeffs, reversed(ds)))
-    return x
-
-
-class PIDStepSizeController:
-    """A PID controller for ODE adaptive step size control."""
-    def __init__(self, h, pcoeff, icoeff, dcoeff, order=1, accept_safety=0.81, eps=1e-8):
-        self.h = h
-        self.b1 = (pcoeff + icoeff + dcoeff) / order
-        self.b2 = -(pcoeff + 2 * dcoeff) / order
-        self.b3 = dcoeff / order
-        self.accept_safety = accept_safety
-        self.eps = eps
-        self.errs = []
-
-    def limiter(self, x):
-        return 1 + math.atan(x - 1)
-
-    def propose_step(self, error):
-        inv_error = 1 / (float(error) + self.eps)
-        if not self.errs:
-            self.errs = [inv_error, inv_error, inv_error]
-        self.errs[0] = inv_error
-        factor = self.errs[0] ** self.b1 * self.errs[1] ** self.b2 * self.errs[2] ** self.b3
-        factor = self.limiter(factor)
-        accept = factor >= self.accept_safety
-        if accept:
-            self.errs[2] = self.errs[1]
-            self.errs[1] = self.errs[0]
-        self.h *= factor
-        return accept
-
-
-class DPMSolver(nn.Module):
-    """DPM-Solver. See https://arxiv.org/abs/2206.00927."""
-
-    def __init__(self, model, extra_args=None, eps_callback=None, info_callback=None):
-        super().__init__()
-        self.model = model
-        self.extra_args = {} if extra_args is None else extra_args
-        self.eps_callback = eps_callback
-        self.info_callback = info_callback
-
-    def t(self, sigma):
-        return -sigma.log()
-
-    def sigma(self, t):
-        return t.neg().exp()
-
-    def eps(self, eps_cache, key, x, t, *args, **kwargs):
-        if key in eps_cache:
-            return eps_cache[key], eps_cache
-        sigma = self.sigma(t) * x.new_ones([x.shape[0]])
-        eps = (x - self.model(x, sigma, *args, **self.extra_args, **kwargs)) / self.sigma(t)
-        if self.eps_callback is not None:
-            self.eps_callback()
-        return eps, {key: eps, **eps_cache}
-
-    def dpm_solver_1_step(self, x, t, t_next, eps_cache=None):
-        eps_cache = {} if eps_cache is None else eps_cache
-        h = t_next - t
-        eps, eps_cache = self.eps(eps_cache, 'eps', x, t)
-        x_1 = x - self.sigma(t_next) * h.expm1() * eps
-        return x_1, eps_cache
-
-    def dpm_solver_2_step(self, x, t, t_next, r1=1 / 2, eps_cache=None):
-        eps_cache = {} if eps_cache is None else eps_cache
-        h = t_next - t
-        eps, eps_cache = self.eps(eps_cache, 'eps', x, t)
-        s1 = t + r1 * h
-        u1 = x - self.sigma(s1) * (r1 * h).expm1() * eps
-        eps_r1, eps_cache = self.eps(eps_cache, 'eps_r1', u1, s1)
-        x_2 = x - self.sigma(t_next) * h.expm1() * eps - self.sigma(t_next) / (2 * r1) * h.expm1() * (eps_r1 - eps)
-        return x_2, eps_cache
-
-    def dpm_solver_3_step(self, x, t, t_next, r1=1 / 3, r2=2 / 3, eps_cache=None):
-        eps_cache = {} if eps_cache is None else eps_cache
-        h = t_next - t
-        eps, eps_cache = self.eps(eps_cache, 'eps', x, t)
-        s1 = t + r1 * h
-        s2 = t + r2 * h
-        u1 = x - self.sigma(s1) * (r1 * h).expm1() * eps
-        eps_r1, eps_cache = self.eps(eps_cache, 'eps_r1', u1, s1)
-        u2 = x - self.sigma(s2) * (r2 * h).expm1() * eps - self.sigma(s2) * (r2 / r1) * ((r2 * h).expm1() / (r2 * h) - 1) * (eps_r1 - eps)
-        eps_r2, eps_cache = self.eps(eps_cache, 'eps_r2', u2, s2)
-        x_3 = x - self.sigma(t_next) * h.expm1() * eps - self.sigma(t_next) / r2 * (h.expm1() / h - 1) * (eps_r2 - eps)
-        return x_3, eps_cache
-
-    def dpm_solver_fast(self, x, t_start, t_end, nfe, eta=0., s_noise=1., noise_sampler=None):
-        noise_sampler = default_noise_sampler(x, seed=self.extra_args.get("seed", None)) if noise_sampler is None else noise_sampler
-        if not t_end > t_start and eta:
-            raise ValueError('eta must be 0 for reverse sampling')
-
-        m = math.floor(nfe / 3) + 1
-        ts = torch.linspace(t_start, t_end, m + 1, device=x.device)
-
-        if nfe % 3 == 0:
-            orders = [3] * (m - 2) + [2, 1]
-        else:
-            orders = [3] * (m - 1) + [nfe % 3]
-
-        for i in range(len(orders)):
-            eps_cache = {}
-            t, t_next = ts[i], ts[i + 1]
-            if eta:
-                sd, su = get_ancestral_step(self.sigma(t), self.sigma(t_next), eta)
-                t_next_ = torch.minimum(t_end, self.t(sd))
-                su = (self.sigma(t_next) ** 2 - self.sigma(t_next_) ** 2) ** 0.5
-            else:
-                t_next_, su = t_next, 0.
-
-            eps, eps_cache = self.eps(eps_cache, 'eps', x, t)
-            denoised = x - self.sigma(t) * eps
-            if self.info_callback is not None:
-                self.info_callback({'x': x, 'i': i, 't': ts[i], 't_up': t, 'denoised': denoised})
-
-            if orders[i] == 1:
-                x, eps_cache = self.dpm_solver_1_step(x, t, t_next_, eps_cache=eps_cache)
-            elif orders[i] == 2:
-                x, eps_cache = self.dpm_solver_2_step(x, t, t_next_, eps_cache=eps_cache)
-            else:
-                x, eps_cache = self.dpm_solver_3_step(x, t, t_next_, eps_cache=eps_cache)
-
-            x = x + su * s_noise * noise_sampler(self.sigma(t), self.sigma(t_next))
-
-        return x
-
-    def dpm_solver_adaptive(self, x, t_start, t_end, order=3, rtol=0.05, atol=0.0078, h_init=0.05, pcoeff=0., icoeff=1., dcoeff=0., accept_safety=0.81, eta=0., s_noise=1., noise_sampler=None):
-        noise_sampler = default_noise_sampler(x, seed=self.extra_args.get("seed", None)) if noise_sampler is None else noise_sampler
-        if order not in {2, 3}:
-            raise ValueError('order should be 2 or 3')
-        forward = t_end > t_start
-        if not forward and eta:
-            raise ValueError('eta must be 0 for reverse sampling')
-        h_init = abs(h_init) * (1 if forward else -1)
-        atol = torch.tensor(atol)
-        rtol = torch.tensor(rtol)
-        s = t_start
-        x_prev = x
-        accept = True
-        pid = PIDStepSizeController(h_init, pcoeff, icoeff, dcoeff, 1.5 if eta else order, accept_safety)
-        info = {'steps': 0, 'nfe': 0, 'n_accept': 0, 'n_reject': 0}
-
-        while s < t_end - 1e-5 if forward else s > t_end + 1e-5:
-            eps_cache = {}
-            t = torch.minimum(t_end, s + pid.h) if forward else torch.maximum(t_end, s + pid.h)
-            if eta:
-                sd, su = get_ancestral_step(self.sigma(s), self.sigma(t), eta)
-                t_ = torch.minimum(t_end, self.t(sd))
-                su = (self.sigma(t) ** 2 - self.sigma(t_) ** 2) ** 0.5
-            else:
-                t_, su = t, 0.
-
-            eps, eps_cache = self.eps(eps_cache, 'eps', x, s)
-            denoised = x - self.sigma(s) * eps
-
-            if order == 2:
-                x_low, eps_cache = self.dpm_solver_1_step(x, s, t_, eps_cache=eps_cache)
-                x_high, eps_cache = self.dpm_solver_2_step(x, s, t_, eps_cache=eps_cache)
-            else:
-                x_low, eps_cache = self.dpm_solver_2_step(x, s, t_, r1=1 / 3, eps_cache=eps_cache)
-                x_high, eps_cache = self.dpm_solver_3_step(x, s, t_, eps_cache=eps_cache)
-            delta = torch.maximum(atol, rtol * torch.maximum(x_low.abs(), x_prev.abs()))
-            error = torch.linalg.norm((x_low - x_high) / delta) / x.numel() ** 0.5
-            accept = pid.propose_step(error)
-            if accept:
-                x_prev = x_low
-                x = x_high + su * s_noise * noise_sampler(self.sigma(s), self.sigma(t))
-                s = t
-                info['n_accept'] += 1
-            else:
-                info['n_reject'] += 1
-            info['nfe'] += order
-            info['steps'] += 1
-
-            if self.info_callback is not None:
-                self.info_callback({'x': x, 'i': info['steps'] - 1, 't': s, 't_up': s, 'denoised': denoised, 'error': error, 'h': pid.h, **info})
-
-        return x, info
-
-
-@torch.no_grad()
-def sample_dpm_fast(model, x, sigma_min, sigma_max, n, extra_args=None, callback=None, disable=None, eta=0., s_noise=1., noise_sampler=None):
-    """DPM-Solver-Fast (fixed step size). See https://arxiv.org/abs/2206.00927."""
-    if sigma_min <= 0 or sigma_max <= 0:
-        raise ValueError('sigma_min and sigma_max must not be 0')
-    with tqdm(total=n, disable=disable) as pbar:
-        dpm_solver = DPMSolver(model, extra_args, eps_callback=pbar.update)
-        if callback is not None:
-            dpm_solver.info_callback = lambda info: callback({'sigma': dpm_solver.sigma(info['t']), 'sigma_hat': dpm_solver.sigma(info['t_up']), **info})
-        return dpm_solver.dpm_solver_fast(x, dpm_solver.t(torch.tensor(sigma_max)), dpm_solver.t(torch.tensor(sigma_min)), n, eta, s_noise, noise_sampler)
-
-
-@torch.no_grad()
-def sample_dpm_adaptive(model, x, sigma_min, sigma_max, extra_args=None, callback=None, disable=None, order=3, rtol=0.05, atol=0.0078, h_init=0.05, pcoeff=0., icoeff=1., dcoeff=0., accept_safety=0.81, eta=0., s_noise=1., noise_sampler=None, return_info=False):
-    """DPM-Solver-12 and 23 (adaptive step size). See https://arxiv.org/abs/2206.00927."""
-    if sigma_min <= 0 or sigma_max <= 0:
-        raise ValueError('sigma_min and sigma_max must not be 0')
-    with tqdm(disable=disable) as pbar:
-        dpm_solver = DPMSolver(model, extra_args, eps_callback=pbar.update)
-        if callback is not None:
-            dpm_solver.info_callback = lambda info: callback({'sigma': dpm_solver.sigma(info['t']), 'sigma_hat': dpm_solver.sigma(info['t_up']), **info})
-        x, info = dpm_solver.dpm_solver_adaptive(x, dpm_solver.t(torch.tensor(sigma_max)), dpm_solver.t(torch.tensor(sigma_min)), order, rtol, atol, h_init, pcoeff, icoeff, dcoeff, accept_safety, eta, s_noise, noise_sampler)
-    if return_info:
-        return x, info
-    return x
-
-
-@torch.no_grad()
-def sample_dpmpp_2s_ancestral(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None):
-    if isinstance(model.inner_model.inner_model.model_sampling, comfy.model_sampling.CONST):
-        return sample_dpmpp_2s_ancestral_RF(model, x, sigmas, extra_args, callback, disable, eta, s_noise, noise_sampler)
-
-    """Ancestral sampling with DPM-Solver++(2S) second-order steps."""
-    extra_args = {} if extra_args is None else extra_args
-    seed = extra_args.get("seed", None)
-    noise_sampler = default_noise_sampler(x, seed=seed) if noise_sampler is None else noise_sampler
-    s_in = x.new_ones([x.shape[0]])
-    sigma_fn = lambda t: t.neg().exp()
-    t_fn = lambda sigma: sigma.log().neg()
-
-    for i in trange(len(sigmas) - 1, disable=disable):
-        denoised = model(x, sigmas[i] * s_in, **extra_args)
-        sigma_down, sigma_up = get_ancestral_step(sigmas[i], sigmas[i + 1], eta=eta)
-        if callback is not None:
-            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
-        if sigma_down == 0:
-            # Euler method
-            d = to_d(x, sigmas[i], denoised)
-            dt = sigma_down - sigmas[i]
-            x = x + d * dt
-        else:
-            # DPM-Solver++(2S)
-            t, t_next = t_fn(sigmas[i]), t_fn(sigma_down)
-            r = 1 / 2
-            h = t_next - t
-            s = t + r * h
-            x_2 = (sigma_fn(s) / sigma_fn(t)) * x - (-h * r).expm1() * denoised
-            denoised_2 = model(x_2, sigma_fn(s) * s_in, **extra_args)
-            x = (sigma_fn(t_next) / sigma_fn(t)) * x - (-h).expm1() * denoised_2
-        # Noise addition
-        if sigmas[i + 1] > 0:
-            x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up
-    return x
-
-
-@torch.no_grad()
-def sample_dpmpp_2s_ancestral_RF(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None):
-    """Ancestral sampling with DPM-Solver++(2S) second-order steps."""
-    extra_args = {} if extra_args is None else extra_args
-    seed = extra_args.get("seed", None)
-    noise_sampler = default_noise_sampler(x, seed=seed) if noise_sampler is None else noise_sampler
-    s_in = x.new_ones([x.shape[0]])
-    sigma_fn = lambda lbda: (lbda.exp() + 1) ** -1
-    lambda_fn = lambda sigma: ((1-sigma)/sigma).log()
-
-    # logged_x = x.unsqueeze(0)
-
-    for i in trange(len(sigmas) - 1, disable=disable):
-        denoised = model(x, sigmas[i] * s_in, **extra_args)
-        downstep_ratio = 1 + (sigmas[i+1]/sigmas[i] - 1) * eta
-        sigma_down = sigmas[i+1] * downstep_ratio
-        alpha_ip1 = 1 - sigmas[i+1]
-        alpha_down = 1 - sigma_down
-        renoise_coeff = (sigmas[i+1]**2 - sigma_down**2*alpha_ip1**2/alpha_down**2)**0.5
-        # sigma_down, sigma_up = get_ancestral_step(sigmas[i], sigmas[i + 1], eta=eta)
-        if callback is not None:
-            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
-        if sigmas[i + 1] == 0:
-            # Euler method
-            d = to_d(x, sigmas[i], denoised)
-            dt = sigma_down - sigmas[i]
-            x = x + d * dt
-        else:
-            # DPM-Solver++(2S)
-            if sigmas[i] == 1.0:
-                sigma_s = 0.9999
-            else:
-                t_i, t_down = lambda_fn(sigmas[i]), lambda_fn(sigma_down)
-                r = 1 / 2
-                h = t_down - t_i
-                s = t_i + r * h
-                sigma_s = sigma_fn(s)
-            # sigma_s = sigmas[i+1]
-            sigma_s_i_ratio = sigma_s / sigmas[i]
-            u = sigma_s_i_ratio * x + (1 - sigma_s_i_ratio) * denoised
-            D_i = model(u, sigma_s * s_in, **extra_args)
-            sigma_down_i_ratio = sigma_down / sigmas[i]
-            x = sigma_down_i_ratio * x + (1 - sigma_down_i_ratio) * D_i
-            # print("sigma_i", sigmas[i], "sigma_ip1", sigmas[i+1],"sigma_down", sigma_down, "sigma_down_i_ratio", sigma_down_i_ratio, "sigma_s_i_ratio", sigma_s_i_ratio, "renoise_coeff", renoise_coeff)
-        # Noise addition
-        if sigmas[i + 1] > 0 and eta > 0:
-            x = (alpha_ip1/alpha_down) * x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * renoise_coeff
-        # logged_x = torch.cat((logged_x, x.unsqueeze(0)), dim=0)
-    return x
-
-
-@torch.no_grad()
-def sample_dpmpp_sde(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None, r=1 / 2):
-    """DPM-Solver++ (stochastic)."""
-    if len(sigmas) <= 1:
-        return x
-
-    extra_args = {} if extra_args is None else extra_args
-    sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max()
-    seed = extra_args.get("seed", None)
-    noise_sampler = BrownianTreeNoiseSampler(x, sigma_min, sigma_max, seed=seed, cpu=True) if noise_sampler is None else noise_sampler
-    s_in = x.new_ones([x.shape[0]])
-
-    model_sampling = model.inner_model.model_patcher.get_model_object('model_sampling')
-    sigma_fn = partial(half_log_snr_to_sigma, model_sampling=model_sampling)
-    lambda_fn = partial(sigma_to_half_log_snr, model_sampling=model_sampling)
-    sigmas = offset_first_sigma_for_snr(sigmas, model_sampling)
-
-    for i in trange(len(sigmas) - 1, disable=disable):
-        denoised = model(x, sigmas[i] * s_in, **extra_args)
-        if callback is not None:
-            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
-        if sigmas[i + 1] == 0:
-            # Denoising step
-            x = denoised
-        else:
-            # DPM-Solver++
-            lambda_s, lambda_t = lambda_fn(sigmas[i]), lambda_fn(sigmas[i + 1])
-            h = lambda_t - lambda_s
-            lambda_s_1 = lambda_s + r * h
-            fac = 1 / (2 * r)
-
-            sigma_s_1 = sigma_fn(lambda_s_1)
-
-            alpha_s = sigmas[i] * lambda_s.exp()
-            alpha_s_1 = sigma_s_1 * lambda_s_1.exp()
-            alpha_t = sigmas[i + 1] * lambda_t.exp()
-
-            # Step 1
-            sd, su = get_ancestral_step(lambda_s.neg().exp(), lambda_s_1.neg().exp(), eta)
-            lambda_s_1_ = sd.log().neg()
-            h_ = lambda_s_1_ - lambda_s
-            x_2 = (alpha_s_1 / alpha_s) * (-h_).exp() * x - alpha_s_1 * (-h_).expm1() * denoised
-            if eta > 0 and s_noise > 0:
-                x_2 = x_2 + alpha_s_1 * noise_sampler(sigmas[i], sigma_s_1) * s_noise * su
-            denoised_2 = model(x_2, sigma_s_1 * s_in, **extra_args)
-
-            # Step 2
-            sd, su = get_ancestral_step(lambda_s.neg().exp(), lambda_t.neg().exp(), eta)
-            lambda_t_ = sd.log().neg()
-            h_ = lambda_t_ - lambda_s
-            denoised_d = (1 - fac) * denoised + fac * denoised_2
-            x = (alpha_t / alpha_s) * (-h_).exp() * x - alpha_t * (-h_).expm1() * denoised_d
-            if eta > 0 and s_noise > 0:
-                x = x + alpha_t * noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * su
-    return x
-
-
-@torch.no_grad()
-def sample_dpmpp_2m(model, x, sigmas, extra_args=None, callback=None, disable=None):
-    """DPM-Solver++(2M)."""
-    extra_args = {} if extra_args is None else extra_args
-    s_in = x.new_ones([x.shape[0]])
-    sigma_fn = lambda t: t.neg().exp()
-    t_fn = lambda sigma: sigma.log().neg()
-    old_denoised = None
-
-    for i in trange(len(sigmas) - 1, disable=disable):
-        denoised = model(x, sigmas[i] * s_in, **extra_args)
-        if callback is not None:
-            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
-        t, t_next = t_fn(sigmas[i]), t_fn(sigmas[i + 1])
-        h = t_next - t
-        if old_denoised is None or sigmas[i + 1] == 0:
-            x = (sigma_fn(t_next) / sigma_fn(t)) * x - (-h).expm1() * denoised
-        else:
-            h_last = t - t_fn(sigmas[i - 1])
-            r = h_last / h
-            denoised_d = (1 + 1 / (2 * r)) * denoised - (1 / (2 * r)) * old_denoised
-            x = (sigma_fn(t_next) / sigma_fn(t)) * x - (-h).expm1() * denoised_d
-        old_denoised = denoised
-    return x
-
-
-@torch.no_grad()
-def sample_dpmpp_2m_sde(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None, solver_type='midpoint'):
-    """DPM-Solver++(2M) SDE."""
-    if len(sigmas) <= 1:
-        return x
-
-    if solver_type not in {'heun', 'midpoint'}:
-        raise ValueError('solver_type must be \'heun\' or \'midpoint\'')
-
-    extra_args = {} if extra_args is None else extra_args
-    seed = extra_args.get("seed", None)
-    sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max()
-    noise_sampler = BrownianTreeNoiseSampler(x, sigma_min, sigma_max, seed=seed, cpu=True) if noise_sampler is None else noise_sampler
-    s_in = x.new_ones([x.shape[0]])
-
-    model_sampling = model.inner_model.model_patcher.get_model_object('model_sampling')
-    lambda_fn = partial(sigma_to_half_log_snr, model_sampling=model_sampling)
-    sigmas = offset_first_sigma_for_snr(sigmas, model_sampling)
-
-    old_denoised = None
-    h, h_last = None, None
-
-    for i in trange(len(sigmas) - 1, disable=disable):
-        denoised = model(x, sigmas[i] * s_in, **extra_args)
-        if callback is not None:
-            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
-        if sigmas[i + 1] == 0:
-            # Denoising step
-            x = denoised
-        else:
-            # DPM-Solver++(2M) SDE
-            lambda_s, lambda_t = lambda_fn(sigmas[i]), lambda_fn(sigmas[i + 1])
-            h = lambda_t - lambda_s
-            h_eta = h * (eta + 1)
-
-            alpha_t = sigmas[i + 1] * lambda_t.exp()
-
-            x = sigmas[i + 1] / sigmas[i] * (-h * eta).exp() * x + alpha_t * (-h_eta).expm1().neg() * denoised
-
-            if old_denoised is not None:
-                r = h_last / h
-                if solver_type == 'heun':
-                    x = x + alpha_t * ((-h_eta).expm1().neg() / (-h_eta) + 1) * (1 / r) * (denoised - old_denoised)
-                elif solver_type == 'midpoint':
-                    x = x + 0.5 * alpha_t * (-h_eta).expm1().neg() * (1 / r) * (denoised - old_denoised)
-
-            if eta > 0 and s_noise > 0:
-                x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * sigmas[i + 1] * (-2 * h * eta).expm1().neg().sqrt() * s_noise
-
-        old_denoised = denoised
-        h_last = h
-    return x
-
-
-@torch.no_grad()
-def sample_dpmpp_3m_sde(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None):
-    """DPM-Solver++(3M) SDE."""
-
-    if len(sigmas) <= 1:
-        return x
-
-    extra_args = {} if extra_args is None else extra_args
-    seed = extra_args.get("seed", None)
-    sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max()
-    noise_sampler = BrownianTreeNoiseSampler(x, sigma_min, sigma_max, seed=seed, cpu=True) if noise_sampler is None else noise_sampler
-    s_in = x.new_ones([x.shape[0]])
-
-    model_sampling = model.inner_model.model_patcher.get_model_object('model_sampling')
-    lambda_fn = partial(sigma_to_half_log_snr, model_sampling=model_sampling)
-    sigmas = offset_first_sigma_for_snr(sigmas, model_sampling)
-
-    denoised_1, denoised_2 = None, None
-    h, h_1, h_2 = None, None, None
-
-    for i in trange(len(sigmas) - 1, disable=disable):
-        denoised = model(x, sigmas[i] * s_in, **extra_args)
-        if callback is not None:
-            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
-        if sigmas[i + 1] == 0:
-            # Denoising step
-            x = denoised
-        else:
-            lambda_s, lambda_t = lambda_fn(sigmas[i]), lambda_fn(sigmas[i + 1])
-            h = lambda_t - lambda_s
-            h_eta = h * (eta + 1)
-
-            alpha_t = sigmas[i + 1] * lambda_t.exp()
-
-            x = sigmas[i + 1] / sigmas[i] * (-h * eta).exp() * x + alpha_t * (-h_eta).expm1().neg() * denoised
-
-            if h_2 is not None:
-                # DPM-Solver++(3M) SDE
-                r0 = h_1 / h
-                r1 = h_2 / h
-                d1_0 = (denoised - denoised_1) / r0
-                d1_1 = (denoised_1 - denoised_2) / r1
-                d1 = d1_0 + (d1_0 - d1_1) * r0 / (r0 + r1)
-                d2 = (d1_0 - d1_1) / (r0 + r1)
-                phi_2 = h_eta.neg().expm1() / h_eta + 1
-                phi_3 = phi_2 / h_eta - 0.5
-                x = x + (alpha_t * phi_2) * d1 - (alpha_t * phi_3) * d2
-            elif h_1 is not None:
-                # DPM-Solver++(2M) SDE
-                r = h_1 / h
-                d = (denoised - denoised_1) / r
-                phi_2 = h_eta.neg().expm1() / h_eta + 1
-                x = x + (alpha_t * phi_2) * d
-
-            if eta > 0 and s_noise > 0:
-                x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * sigmas[i + 1] * (-2 * h * eta).expm1().neg().sqrt() * s_noise
-
-        denoised_1, denoised_2 = denoised, denoised_1
-        h_1, h_2 = h, h_1
-    return x
-
-
-@torch.no_grad()
-def sample_dpmpp_3m_sde_gpu(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None):
-    if len(sigmas) <= 1:
-        return x
-    extra_args = {} if extra_args is None else extra_args
-    sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max()
-    noise_sampler = BrownianTreeNoiseSampler(x, sigma_min, sigma_max, seed=extra_args.get("seed", None), cpu=False) if noise_sampler is None else noise_sampler
-    return sample_dpmpp_3m_sde(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, eta=eta, s_noise=s_noise, noise_sampler=noise_sampler)
-
-
-@torch.no_grad()
-def sample_dpmpp_2m_sde_gpu(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None, solver_type='midpoint'):
-    if len(sigmas) <= 1:
-        return x
-    extra_args = {} if extra_args is None else extra_args
-    sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max()
-    noise_sampler = BrownianTreeNoiseSampler(x, sigma_min, sigma_max, seed=extra_args.get("seed", None), cpu=False) if noise_sampler is None else noise_sampler
-    return sample_dpmpp_2m_sde(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, eta=eta, s_noise=s_noise, noise_sampler=noise_sampler, solver_type=solver_type)
-
-
-@torch.no_grad()
-def sample_dpmpp_sde_gpu(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None, r=1 / 2):
-    if len(sigmas) <= 1:
-        return x
-    extra_args = {} if extra_args is None else extra_args
-    sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max()
-    noise_sampler = BrownianTreeNoiseSampler(x, sigma_min, sigma_max, seed=extra_args.get("seed", None), cpu=False) if noise_sampler is None else noise_sampler
-    return sample_dpmpp_sde(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, eta=eta, s_noise=s_noise, noise_sampler=noise_sampler, r=r)
-
-
-def DDPMSampler_step(x, sigma, sigma_prev, noise, noise_sampler):
-    alpha_cumprod = 1 / ((sigma * sigma) + 1)
-    alpha_cumprod_prev = 1 / ((sigma_prev * sigma_prev) + 1)
-    alpha = (alpha_cumprod / alpha_cumprod_prev)
-
-    mu = (1.0 / alpha).sqrt() * (x - (1 - alpha) * noise / (1 - alpha_cumprod).sqrt())
-    if sigma_prev > 0:
-        mu += ((1 - alpha) * (1. - alpha_cumprod_prev) / (1. - alpha_cumprod)).sqrt() * noise_sampler(sigma, sigma_prev)
-    return mu
-
-def generic_step_sampler(model, x, sigmas, extra_args=None, callback=None, disable=None, noise_sampler=None, step_function=None):
-    extra_args = {} if extra_args is None else extra_args
-    seed = extra_args.get("seed", None)
-    noise_sampler = default_noise_sampler(x, seed=seed) if noise_sampler is None else noise_sampler
-    s_in = x.new_ones([x.shape[0]])
-
-    for i in trange(len(sigmas) - 1, disable=disable):
-        denoised = model(x, sigmas[i] * s_in, **extra_args)
-        if callback is not None:
-            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
-        x = step_function(x / torch.sqrt(1.0 + sigmas[i] ** 2.0), sigmas[i], sigmas[i + 1], (x - denoised) / sigmas[i], noise_sampler)
-        if sigmas[i + 1] != 0:
-            x *= torch.sqrt(1.0 + sigmas[i + 1] ** 2.0)
-    return x
-
-
-@torch.no_grad()
-def sample_ddpm(model, x, sigmas, extra_args=None, callback=None, disable=None, noise_sampler=None):
-    return generic_step_sampler(model, x, sigmas, extra_args, callback, disable, noise_sampler, DDPMSampler_step)
-
-@torch.no_grad()
-def sample_lcm(model, x, sigmas, extra_args=None, callback=None, disable=None, noise_sampler=None):
-    extra_args = {} if extra_args is None else extra_args
-    seed = extra_args.get("seed", None)
-    noise_sampler = default_noise_sampler(x, seed=seed) if noise_sampler is None else noise_sampler
-    s_in = x.new_ones([x.shape[0]])
-    for i in trange(len(sigmas) - 1, disable=disable):
-        denoised = model(x, sigmas[i] * s_in, **extra_args)
-        if callback is not None:
-            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
-
-        x = denoised
-        if sigmas[i + 1] > 0:
-            x = model.inner_model.inner_model.model_sampling.noise_scaling(sigmas[i + 1], noise_sampler(sigmas[i], sigmas[i + 1]), x)
-    return x
-
-
-
-@torch.no_grad()
-def sample_heunpp2(model, x, sigmas, extra_args=None, callback=None, disable=None, s_churn=0., s_tmin=0., s_tmax=float('inf'), s_noise=1.):
-    # From MIT licensed: https://github.com/Carzit/sd-webui-samplers-scheduler/
-    extra_args = {} if extra_args is None else extra_args
-    s_in = x.new_ones([x.shape[0]])
-    s_end = sigmas[-1]
-    for i in trange(len(sigmas) - 1, disable=disable):
-        gamma = min(s_churn / (len(sigmas) - 1), 2 ** 0.5 - 1) if s_tmin <= sigmas[i] <= s_tmax else 0.
-        eps = torch.randn_like(x) * s_noise
-        sigma_hat = sigmas[i] * (gamma + 1)
-        if gamma > 0:
-            x = x + eps * (sigma_hat ** 2 - sigmas[i] ** 2) ** 0.5
-        denoised = model(x, sigma_hat * s_in, **extra_args)
-        d = to_d(x, sigma_hat, denoised)
-        if callback is not None:
-            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigma_hat, 'denoised': denoised})
-        dt = sigmas[i + 1] - sigma_hat
-        if sigmas[i + 1] == s_end:
-            # Euler method
-            x = x + d * dt
-        elif sigmas[i + 2] == s_end:
-
-            # Heun's method
-            x_2 = x + d * dt
-            denoised_2 = model(x_2, sigmas[i + 1] * s_in, **extra_args)
-            d_2 = to_d(x_2, sigmas[i + 1], denoised_2)
-
-            w = 2 * sigmas[0]
-            w2 = sigmas[i+1]/w
-            w1 = 1 - w2
-
-            d_prime = d * w1 + d_2 * w2
-
-
-            x = x + d_prime * dt
-
-        else:
-            # Heun++
-            x_2 = x + d * dt
-            denoised_2 = model(x_2, sigmas[i + 1] * s_in, **extra_args)
-            d_2 = to_d(x_2, sigmas[i + 1], denoised_2)
-            dt_2 = sigmas[i + 2] - sigmas[i + 1]
-
-            x_3 = x_2 + d_2 * dt_2
-            denoised_3 = model(x_3, sigmas[i + 2] * s_in, **extra_args)
-            d_3 = to_d(x_3, sigmas[i + 2], denoised_3)
-
-            w = 3 * sigmas[0]
-            w2 = sigmas[i + 1] / w
-            w3 = sigmas[i + 2] / w
-            w1 = 1 - w2 - w3
-
-            d_prime = w1 * d + w2 * d_2 + w3 * d_3
-            x = x + d_prime * dt
-    return x
-
-
-#From https://github.com/zju-pi/diff-sampler/blob/main/diff-solvers-main/solvers.py
-#under Apache 2 license
-def sample_ipndm(model, x, sigmas, extra_args=None, callback=None, disable=None, max_order=4):
-    extra_args = {} if extra_args is None else extra_args
-    s_in = x.new_ones([x.shape[0]])
-
-    x_next = x
-
-    buffer_model = []
-    for i in trange(len(sigmas) - 1, disable=disable):
-        t_cur = sigmas[i]
-        t_next = sigmas[i + 1]
-
-        x_cur = x_next
-
-        denoised = model(x_cur, t_cur * s_in, **extra_args)
-        if callback is not None:
-            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
-
-        d_cur = (x_cur - denoised) / t_cur
-
-        order = min(max_order, i+1)
-        if t_next == 0:     # Denoising step
-            x_next = denoised
-        elif order == 1:    # First Euler step.
-            x_next = x_cur + (t_next - t_cur) * d_cur
-        elif order == 2:    # Use one history point.
-            x_next = x_cur + (t_next - t_cur) * (3 * d_cur - buffer_model[-1]) / 2
-        elif order == 3:    # Use two history points.
-            x_next = x_cur + (t_next - t_cur) * (23 * d_cur - 16 * buffer_model[-1] + 5 * buffer_model[-2]) / 12
-        elif order == 4:    # Use three history points.
-            x_next = x_cur + (t_next - t_cur) * (55 * d_cur - 59 * buffer_model[-1] + 37 * buffer_model[-2] - 9 * buffer_model[-3]) / 24
-
-        if len(buffer_model) == max_order - 1:
-            for k in range(max_order - 2):
-                buffer_model[k] = buffer_model[k+1]
-            buffer_model[-1] = d_cur
-        else:
-            buffer_model.append(d_cur)
-
-    return x_next
-
-
-#From https://github.com/zju-pi/diff-sampler/blob/main/diff-solvers-main/solvers.py
-#under Apache 2 license
-def sample_ipndm_v(model, x, sigmas, extra_args=None, callback=None, disable=None, max_order=4):
-    extra_args = {} if extra_args is None else extra_args
-    s_in = x.new_ones([x.shape[0]])
-
-    x_next = x
-    t_steps = sigmas
-
-    buffer_model = []
-    for i in trange(len(sigmas) - 1, disable=disable):
-        t_cur = sigmas[i]
-        t_next = sigmas[i + 1]
-
-        x_cur = x_next
-
-        denoised = model(x_cur, t_cur * s_in, **extra_args)
-        if callback is not None:
-            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
-
-        d_cur = (x_cur - denoised) / t_cur
-
-        order = min(max_order, i+1)
-        if t_next == 0:     # Denoising step
-            x_next = denoised
-        elif order == 1:    # First Euler step.
-            x_next = x_cur + (t_next - t_cur) * d_cur
-        elif order == 2:    # Use one history point.
-            h_n = (t_next - t_cur)
-            h_n_1 = (t_cur - t_steps[i-1])
-            coeff1 = (2 + (h_n / h_n_1)) / 2
-            coeff2 = -(h_n / h_n_1) / 2
-            x_next = x_cur + (t_next - t_cur) * (coeff1 * d_cur + coeff2 * buffer_model[-1])
-        elif order == 3:    # Use two history points.
-            h_n = (t_next - t_cur)
-            h_n_1 = (t_cur - t_steps[i-1])
-            h_n_2 = (t_steps[i-1] - t_steps[i-2])
-            temp = (1 - h_n / (3 * (h_n + h_n_1)) * (h_n * (h_n + h_n_1)) / (h_n_1 * (h_n_1 + h_n_2))) / 2
-            coeff1 = (2 + (h_n / h_n_1)) / 2 + temp
-            coeff2 = -(h_n / h_n_1) / 2 - (1 + h_n_1 / h_n_2) * temp
-            coeff3 = temp * h_n_1 / h_n_2
-            x_next = x_cur + (t_next - t_cur) * (coeff1 * d_cur + coeff2 * buffer_model[-1] + coeff3 * buffer_model[-2])
-        elif order == 4:    # Use three history points.
-            h_n = (t_next - t_cur)
-            h_n_1 = (t_cur - t_steps[i-1])
-            h_n_2 = (t_steps[i-1] - t_steps[i-2])
-            h_n_3 = (t_steps[i-2] - t_steps[i-3])
-            temp1 = (1 - h_n / (3 * (h_n + h_n_1)) * (h_n * (h_n + h_n_1)) / (h_n_1 * (h_n_1 + h_n_2))) / 2
-            temp2 = ((1 - h_n / (3 * (h_n + h_n_1))) / 2 + (1 - h_n / (2 * (h_n + h_n_1))) * h_n / (6 * (h_n + h_n_1 + h_n_2))) \
-                   * (h_n * (h_n + h_n_1) * (h_n + h_n_1 + h_n_2)) / (h_n_1 * (h_n_1 + h_n_2) * (h_n_1 + h_n_2 + h_n_3))
-            coeff1 = (2 + (h_n / h_n_1)) / 2 + temp1 + temp2
-            coeff2 = -(h_n / h_n_1) / 2 - (1 + h_n_1 / h_n_2) * temp1 - (1 + (h_n_1 / h_n_2) + (h_n_1 * (h_n_1 + h_n_2) / (h_n_2 * (h_n_2 + h_n_3)))) * temp2
-            coeff3 = temp1 * h_n_1 / h_n_2 + ((h_n_1 / h_n_2) + (h_n_1 * (h_n_1 + h_n_2) / (h_n_2 * (h_n_2 + h_n_3))) * (1 + h_n_2 / h_n_3)) * temp2
-            coeff4 = -temp2 * (h_n_1 * (h_n_1 + h_n_2) / (h_n_2 * (h_n_2 + h_n_3))) * h_n_1 / h_n_2
-            x_next = x_cur + (t_next - t_cur) * (coeff1 * d_cur + coeff2 * buffer_model[-1] + coeff3 * buffer_model[-2] + coeff4 * buffer_model[-3])
-
-        if len(buffer_model) == max_order - 1:
-            for k in range(max_order - 2):
-                buffer_model[k] = buffer_model[k+1]
-            buffer_model[-1] = d_cur.detach()
-        else:
-            buffer_model.append(d_cur.detach())
-
-    return x_next
-
-
-#From https://github.com/zju-pi/diff-sampler/blob/main/diff-solvers-main/solvers.py
-#under Apache 2 license
-@torch.no_grad()
-def sample_deis(model, x, sigmas, extra_args=None, callback=None, disable=None, max_order=3, deis_mode='tab'):
-    extra_args = {} if extra_args is None else extra_args
-    s_in = x.new_ones([x.shape[0]])
-
-    x_next = x
-    t_steps = sigmas
-
-    coeff_list = deis.get_deis_coeff_list(t_steps, max_order, deis_mode=deis_mode)
-
-    buffer_model = []
-    for i in trange(len(sigmas) - 1, disable=disable):
-        t_cur = sigmas[i]
-        t_next = sigmas[i + 1]
-
-        x_cur = x_next
-
-        denoised = model(x_cur, t_cur * s_in, **extra_args)
-        if callback is not None:
-            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
-
-        d_cur = (x_cur - denoised) / t_cur
-
-        order = min(max_order, i+1)
-        if t_next <= 0:
-            order = 1
-
-        if order == 1:          # First Euler step.
-            x_next = x_cur + (t_next - t_cur) * d_cur
-        elif order == 2:        # Use one history point.
-            coeff_cur, coeff_prev1 = coeff_list[i]
-            x_next = x_cur + coeff_cur * d_cur + coeff_prev1 * buffer_model[-1]
-        elif order == 3:        # Use two history points.
-            coeff_cur, coeff_prev1, coeff_prev2 = coeff_list[i]
-            x_next = x_cur + coeff_cur * d_cur + coeff_prev1 * buffer_model[-1] + coeff_prev2 * buffer_model[-2]
-        elif order == 4:        # Use three history points.
-            coeff_cur, coeff_prev1, coeff_prev2, coeff_prev3 = coeff_list[i]
-            x_next = x_cur + coeff_cur * d_cur + coeff_prev1 * buffer_model[-1] + coeff_prev2 * buffer_model[-2] + coeff_prev3 * buffer_model[-3]
-
-        if len(buffer_model) == max_order - 1:
-            for k in range(max_order - 2):
-                buffer_model[k] = buffer_model[k+1]
-            buffer_model[-1] = d_cur.detach()
-        else:
-            buffer_model.append(d_cur.detach())
-
-    return x_next
-
-
-@torch.no_grad()
-def sample_euler_cfg_pp(model, x, sigmas, extra_args=None, callback=None, disable=None):
-    extra_args = {} if extra_args is None else extra_args
-
-    temp = [0]
-    def post_cfg_function(args):
-        temp[0] = args["uncond_denoised"]
-        return args["denoised"]
-
-    model_options = extra_args.get("model_options", {}).copy()
-    extra_args["model_options"] = comfy.model_patcher.set_model_options_post_cfg_function(model_options, post_cfg_function, disable_cfg1_optimization=True)
-
-    s_in = x.new_ones([x.shape[0]])
-    for i in trange(len(sigmas) - 1, disable=disable):
-        sigma_hat = sigmas[i]
-        denoised = model(x, sigma_hat * s_in, **extra_args)
-        d = to_d(x, sigma_hat, temp[0])
-        if callback is not None:
-            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigma_hat, 'denoised': denoised})
-        # Euler method
-        x = denoised + d * sigmas[i + 1]
-    return x
-
-@torch.no_grad()
-def sample_euler_ancestral_cfg_pp(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None):
-    """Ancestral sampling with Euler method steps."""
-    extra_args = {} if extra_args is None else extra_args
-    seed = extra_args.get("seed", None)
-    noise_sampler = default_noise_sampler(x, seed=seed) if noise_sampler is None else noise_sampler
-
-    temp = [0]
-    def post_cfg_function(args):
-        temp[0] = args["uncond_denoised"]
-        return args["denoised"]
-
-    model_options = extra_args.get("model_options", {}).copy()
-    extra_args["model_options"] = comfy.model_patcher.set_model_options_post_cfg_function(model_options, post_cfg_function, disable_cfg1_optimization=True)
-
-    s_in = x.new_ones([x.shape[0]])
-    for i in trange(len(sigmas) - 1, disable=disable):
-        denoised = model(x, sigmas[i] * s_in, **extra_args)
-        sigma_down, sigma_up = get_ancestral_step(sigmas[i], sigmas[i + 1], eta=eta)
-        if callback is not None:
-            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
-        d = to_d(x, sigmas[i], temp[0])
-        # Euler method
-        x = denoised + d * sigma_down
-        if sigmas[i + 1] > 0:
-            x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up
-    return x
-@torch.no_grad()
-def sample_dpmpp_2s_ancestral_cfg_pp(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None):
-    """Ancestral sampling with DPM-Solver++(2S) second-order steps."""
-    extra_args = {} if extra_args is None else extra_args
-    seed = extra_args.get("seed", None)
-    noise_sampler = default_noise_sampler(x, seed=seed) if noise_sampler is None else noise_sampler
-
-    temp = [0]
-    def post_cfg_function(args):
-        temp[0] = args["uncond_denoised"]
-        return args["denoised"]
-
-    model_options = extra_args.get("model_options", {}).copy()
-    extra_args["model_options"] = comfy.model_patcher.set_model_options_post_cfg_function(model_options, post_cfg_function, disable_cfg1_optimization=True)
-
-    s_in = x.new_ones([x.shape[0]])
-    sigma_fn = lambda t: t.neg().exp()
-    t_fn = lambda sigma: sigma.log().neg()
-
-    for i in trange(len(sigmas) - 1, disable=disable):
-        denoised = model(x, sigmas[i] * s_in, **extra_args)
-        sigma_down, sigma_up = get_ancestral_step(sigmas[i], sigmas[i + 1], eta=eta)
-        if callback is not None:
-            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
-        if sigma_down == 0:
-            # Euler method
-            d = to_d(x, sigmas[i], temp[0])
-            x = denoised + d * sigma_down
-        else:
-            # DPM-Solver++(2S)
-            t, t_next = t_fn(sigmas[i]), t_fn(sigma_down)
-            # r = torch.sinh(1 + (2 - eta) * (t_next - t) / (t - t_fn(sigma_up))) works only on non-cfgpp, weird
-            r = 1 / 2
-            h = t_next - t
-            s = t + r * h
-            x_2 = (sigma_fn(s) / sigma_fn(t)) * (x + (denoised - temp[0])) - (-h * r).expm1() * denoised
-            denoised_2 = model(x_2, sigma_fn(s) * s_in, **extra_args)
-            x = (sigma_fn(t_next) / sigma_fn(t)) * (x + (denoised - temp[0])) - (-h).expm1() * denoised_2
-        # Noise addition
-        if sigmas[i + 1] > 0:
-            x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up
-    return x
-
-@torch.no_grad()
-def sample_dpmpp_2m_cfg_pp(model, x, sigmas, extra_args=None, callback=None, disable=None):
-    """DPM-Solver++(2M)."""
-    extra_args = {} if extra_args is None else extra_args
-    s_in = x.new_ones([x.shape[0]])
-    t_fn = lambda sigma: sigma.log().neg()
-
-    old_uncond_denoised = None
-    uncond_denoised = None
-    def post_cfg_function(args):
-        nonlocal uncond_denoised
-        uncond_denoised = args["uncond_denoised"]
-        return args["denoised"]
-
-    model_options = extra_args.get("model_options", {}).copy()
-    extra_args["model_options"] = comfy.model_patcher.set_model_options_post_cfg_function(model_options, post_cfg_function, disable_cfg1_optimization=True)
-
-    for i in trange(len(sigmas) - 1, disable=disable):
-        denoised = model(x, sigmas[i] * s_in, **extra_args)
-        if callback is not None:
-            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
-        t, t_next = t_fn(sigmas[i]), t_fn(sigmas[i + 1])
-        h = t_next - t
-        if old_uncond_denoised is None or sigmas[i + 1] == 0:
-            denoised_mix = -torch.exp(-h) * uncond_denoised
-        else:
-            h_last = t - t_fn(sigmas[i - 1])
-            r = h_last / h
-            denoised_mix = -torch.exp(-h) * uncond_denoised - torch.expm1(-h) * (1 / (2 * r)) * (denoised - old_uncond_denoised)
-        x = denoised + denoised_mix + torch.exp(-h) * x
-        old_uncond_denoised = uncond_denoised
-    return x
-
-@torch.no_grad()
-def res_multistep(model, x, sigmas, extra_args=None, callback=None, disable=None, s_noise=1., noise_sampler=None, eta=1., cfg_pp=False):
-    extra_args = {} if extra_args is None else extra_args
-    seed = extra_args.get("seed", None)
-    noise_sampler = default_noise_sampler(x, seed=seed) if noise_sampler is None else noise_sampler
-    s_in = x.new_ones([x.shape[0]])
-    sigma_fn = lambda t: t.neg().exp()
-    t_fn = lambda sigma: sigma.log().neg()
-    phi1_fn = lambda t: torch.expm1(t) / t
-    phi2_fn = lambda t: (phi1_fn(t) - 1.0) / t
-
-    old_sigma_down = None
-    old_denoised = None
-    uncond_denoised = None
-    def post_cfg_function(args):
-        nonlocal uncond_denoised
-        uncond_denoised = args["uncond_denoised"]
-        return args["denoised"]
-
-    if cfg_pp:
-        model_options = extra_args.get("model_options", {}).copy()
-        extra_args["model_options"] = comfy.model_patcher.set_model_options_post_cfg_function(model_options, post_cfg_function, disable_cfg1_optimization=True)
-
-    for i in trange(len(sigmas) - 1, disable=disable):
-        denoised = model(x, sigmas[i] * s_in, **extra_args)
-        sigma_down, sigma_up = get_ancestral_step(sigmas[i], sigmas[i + 1], eta=eta)
-        if callback is not None:
-            callback({"x": x, "i": i, "sigma": sigmas[i], "sigma_hat": sigmas[i], "denoised": denoised})
-        if sigma_down == 0 or old_denoised is None:
-            # Euler method
-            if cfg_pp:
-                d = to_d(x, sigmas[i], uncond_denoised)
-                x = denoised + d * sigma_down
-            else:
-                d = to_d(x, sigmas[i], denoised)
-                dt = sigma_down - sigmas[i]
-                x = x + d * dt
-        else:
-            # Second order multistep method in https://arxiv.org/pdf/2308.02157
-            t, t_old, t_next, t_prev = t_fn(sigmas[i]), t_fn(old_sigma_down), t_fn(sigma_down), t_fn(sigmas[i - 1])
-            h = t_next - t
-            c2 = (t_prev - t_old) / h
-
-            phi1_val, phi2_val = phi1_fn(-h), phi2_fn(-h)
-            b1 = torch.nan_to_num(phi1_val - phi2_val / c2, nan=0.0)
-            b2 = torch.nan_to_num(phi2_val / c2, nan=0.0)
-
-            if cfg_pp:
-                x = x + (denoised - uncond_denoised)
-                x = sigma_fn(h) * x + h * (b1 * uncond_denoised + b2 * old_denoised)
-            else:
-                x = sigma_fn(h) * x + h * (b1 * denoised + b2 * old_denoised)
-
-        # Noise addition
-        if sigmas[i + 1] > 0:
-            x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up
-
-        if cfg_pp:
-            old_denoised = uncond_denoised
-        else:
-            old_denoised = denoised
-        old_sigma_down = sigma_down
-    return x
-
-@torch.no_grad()
-def sample_res_multistep(model, x, sigmas, extra_args=None, callback=None, disable=None, s_noise=1., noise_sampler=None):
-    return res_multistep(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, s_noise=s_noise, noise_sampler=noise_sampler, eta=0., cfg_pp=False)
-
-@torch.no_grad()
-def sample_res_multistep_cfg_pp(model, x, sigmas, extra_args=None, callback=None, disable=None, s_noise=1., noise_sampler=None):
-    return res_multistep(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, s_noise=s_noise, noise_sampler=noise_sampler, eta=0., cfg_pp=True)
-
-@torch.no_grad()
-def sample_res_multistep_ancestral(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None):
-    return res_multistep(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, s_noise=s_noise, noise_sampler=noise_sampler, eta=eta, cfg_pp=False)
-
-@torch.no_grad()
-def sample_res_multistep_ancestral_cfg_pp(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None):
-    return res_multistep(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, s_noise=s_noise, noise_sampler=noise_sampler, eta=eta, cfg_pp=True)
-
-
-@torch.no_grad()
-def sample_gradient_estimation(model, x, sigmas, extra_args=None, callback=None, disable=None, ge_gamma=2., cfg_pp=False):
-    """Gradient-estimation sampler. Paper: https://openreview.net/pdf?id=o2ND9v0CeK"""
-    extra_args = {} if extra_args is None else extra_args
-    s_in = x.new_ones([x.shape[0]])
-    old_d = None
-
-    uncond_denoised = None
-    def post_cfg_function(args):
-        nonlocal uncond_denoised
-        uncond_denoised = args["uncond_denoised"]
-        return args["denoised"]
-
-    if cfg_pp:
-        model_options = extra_args.get("model_options", {}).copy()
-        extra_args["model_options"] = comfy.model_patcher.set_model_options_post_cfg_function(model_options, post_cfg_function, disable_cfg1_optimization=True)
-
-    for i in trange(len(sigmas) - 1, disable=disable):
-        denoised = model(x, sigmas[i] * s_in, **extra_args)
-        if cfg_pp:
-            d = to_d(x, sigmas[i], uncond_denoised)
-        else:
-            d = to_d(x, sigmas[i], denoised)
-        if callback is not None:
-            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
-        dt = sigmas[i + 1] - sigmas[i]
-        if sigmas[i + 1] == 0:
-            # Denoising step
-            x = denoised
-        else:
-            # Euler method
-            if cfg_pp:
-                x = denoised + d * sigmas[i + 1]
-            else:
-                x = x + d * dt
-
-            if i >= 1:
-                # Gradient estimation
-                d_bar = (ge_gamma - 1) * (d - old_d)
-                x = x + d_bar * dt
-        old_d = d
-    return x
-
-
-@torch.no_grad()
-def sample_gradient_estimation_cfg_pp(model, x, sigmas, extra_args=None, callback=None, disable=None, ge_gamma=2.):
-    return sample_gradient_estimation(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, ge_gamma=ge_gamma, cfg_pp=True)
-
-
-@torch.no_grad()
-def sample_er_sde(model, x, sigmas, extra_args=None, callback=None, disable=None, s_noise=1.0, noise_sampler=None, noise_scaler=None, max_stage=3):
-    """Extended Reverse-Time SDE solver (VP ER-SDE-Solver-3). arXiv: https://arxiv.org/abs/2309.06169.
-    Code reference: https://github.com/QinpengCui/ER-SDE-Solver/blob/main/er_sde_solver.py.
-    """
-    extra_args = {} if extra_args is None else extra_args
-    seed = extra_args.get("seed", None)
-    noise_sampler = default_noise_sampler(x, seed=seed) if noise_sampler is None else noise_sampler
-    s_in = x.new_ones([x.shape[0]])
-
-    def default_er_sde_noise_scaler(x):
-        return x * ((x ** 0.3).exp() + 10.0)
-
-    noise_scaler = default_er_sde_noise_scaler if noise_scaler is None else noise_scaler
-    num_integration_points = 200.0
-    point_indice = torch.arange(0, num_integration_points, dtype=torch.float32, device=x.device)
-
-    model_sampling = model.inner_model.model_patcher.get_model_object("model_sampling")
-    sigmas = offset_first_sigma_for_snr(sigmas, model_sampling)
-    half_log_snrs = sigma_to_half_log_snr(sigmas, model_sampling)
-    er_lambdas = half_log_snrs.neg().exp()  # er_lambda_t = sigma_t / alpha_t
-
-    old_denoised = None
-    old_denoised_d = None
-
-    for i in trange(len(sigmas) - 1, disable=disable):
-        denoised = model(x, sigmas[i] * s_in, **extra_args)
-        if callback is not None:
-            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
-        stage_used = min(max_stage, i + 1)
-        if sigmas[i + 1] == 0:
-            x = denoised
-        else:
-            er_lambda_s, er_lambda_t = er_lambdas[i], er_lambdas[i + 1]
-            alpha_s = sigmas[i] / er_lambda_s
-            alpha_t = sigmas[i + 1] / er_lambda_t
-            r_alpha = alpha_t / alpha_s
-            r = noise_scaler(er_lambda_t) / noise_scaler(er_lambda_s)
-
-            # Stage 1 Euler
-            x = r_alpha * r * x + alpha_t * (1 - r) * denoised
-
-            if stage_used >= 2:
-                dt = er_lambda_t - er_lambda_s
-                lambda_step_size = -dt / num_integration_points
-                lambda_pos = er_lambda_t + point_indice * lambda_step_size
-                scaled_pos = noise_scaler(lambda_pos)
-
-                # Stage 2
-                s = torch.sum(1 / scaled_pos) * lambda_step_size
-                denoised_d = (denoised - old_denoised) / (er_lambda_s - er_lambdas[i - 1])
-                x = x + alpha_t * (dt + s * noise_scaler(er_lambda_t)) * denoised_d
-
-                if stage_used >= 3:
-                    # Stage 3
-                    s_u = torch.sum((lambda_pos - er_lambda_s) / scaled_pos) * lambda_step_size
-                    denoised_u = (denoised_d - old_denoised_d) / ((er_lambda_s - er_lambdas[i - 2]) / 2)
-                    x = x + alpha_t * ((dt ** 2) / 2 + s_u * noise_scaler(er_lambda_t)) * denoised_u
-                old_denoised_d = denoised_d
-
-            if s_noise > 0:
-                x = x + alpha_t * noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * (er_lambda_t ** 2 - er_lambda_s ** 2 * r ** 2).sqrt().nan_to_num(nan=0.0)
-        old_denoised = denoised
-    return x
-
-
-@torch.no_grad()
-def sample_seeds_2(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None, r=0.5):
-    """SEEDS-2 - Stochastic Explicit Exponential Derivative-free Solvers (VP Data Prediction) stage 2.
-    arXiv: https://arxiv.org/abs/2305.14267
-    """
-    extra_args = {} if extra_args is None else extra_args
-    seed = extra_args.get("seed", None)
-    noise_sampler = default_noise_sampler(x, seed=seed) if noise_sampler is None else noise_sampler
-    s_in = x.new_ones([x.shape[0]])
-
-    inject_noise = eta > 0 and s_noise > 0
-
-    model_sampling = model.inner_model.model_patcher.get_model_object('model_sampling')
-    sigma_fn = partial(half_log_snr_to_sigma, model_sampling=model_sampling)
-    lambda_fn = partial(sigma_to_half_log_snr, model_sampling=model_sampling)
-    sigmas = offset_first_sigma_for_snr(sigmas, model_sampling)
-
-    for i in trange(len(sigmas) - 1, disable=disable):
-        denoised = model(x, sigmas[i] * s_in, **extra_args)
-        if callback is not None:
-            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
-        if sigmas[i + 1] == 0:
-            x = denoised
-        else:
-            lambda_s, lambda_t = lambda_fn(sigmas[i]), lambda_fn(sigmas[i + 1])
-            h = lambda_t - lambda_s
-            h_eta = h * (eta + 1)
-            lambda_s_1 = lambda_s + r * h
-            fac = 1 / (2 * r)
-            sigma_s_1 = sigma_fn(lambda_s_1)
-
-            # alpha_t = sigma_t * exp(log(alpha_t / sigma_t)) = sigma_t * exp(lambda_t)
-            alpha_s_1 = sigma_s_1 * lambda_s_1.exp()
-            alpha_t = sigmas[i + 1] * lambda_t.exp()
-
-            coeff_1, coeff_2 = (-r * h_eta).expm1(), (-h_eta).expm1()
-            if inject_noise:
-                # 0 < r < 1
-                noise_coeff_1 = (-2 * r * h * eta).expm1().neg().sqrt()
-                noise_coeff_2 = (-r * h * eta).exp() * (-2 * (1 - r) * h * eta).expm1().neg().sqrt()
-                noise_1, noise_2 = noise_sampler(sigmas[i], sigma_s_1), noise_sampler(sigma_s_1, sigmas[i + 1])
-
-            # Step 1
-            x_2 = sigma_s_1 / sigmas[i] * (-r * h * eta).exp() * x - alpha_s_1 * coeff_1 * denoised
-            if inject_noise:
-                x_2 = x_2 + sigma_s_1 * (noise_coeff_1 * noise_1) * s_noise
-            denoised_2 = model(x_2, sigma_s_1 * s_in, **extra_args)
-
-            # Step 2
-            denoised_d = (1 - fac) * denoised + fac * denoised_2
-            x = sigmas[i + 1] / sigmas[i] * (-h * eta).exp() * x - alpha_t * coeff_2 * denoised_d
-            if inject_noise:
-                x = x + sigmas[i + 1] * (noise_coeff_2 * noise_1 + noise_coeff_1 * noise_2) * s_noise
-    return x
-
-
-@torch.no_grad()
-def sample_seeds_3(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None, r_1=1./3, r_2=2./3):
-    """SEEDS-3 - Stochastic Explicit Exponential Derivative-free Solvers (VP Data Prediction) stage 3.
-    arXiv: https://arxiv.org/abs/2305.14267
-    """
-    extra_args = {} if extra_args is None else extra_args
-    seed = extra_args.get("seed", None)
-    noise_sampler = default_noise_sampler(x, seed=seed) if noise_sampler is None else noise_sampler
-    s_in = x.new_ones([x.shape[0]])
-
-    inject_noise = eta > 0 and s_noise > 0
-
-    model_sampling = model.inner_model.model_patcher.get_model_object('model_sampling')
-    sigma_fn = partial(half_log_snr_to_sigma, model_sampling=model_sampling)
-    lambda_fn = partial(sigma_to_half_log_snr, model_sampling=model_sampling)
-    sigmas = offset_first_sigma_for_snr(sigmas, model_sampling)
-
-    for i in trange(len(sigmas) - 1, disable=disable):
-        denoised = model(x, sigmas[i] * s_in, **extra_args)
-        if callback is not None:
-            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
-        if sigmas[i + 1] == 0:
-            x = denoised
-        else:
-            lambda_s, lambda_t = lambda_fn(sigmas[i]), lambda_fn(sigmas[i + 1])
-            h = lambda_t - lambda_s
-            h_eta = h * (eta + 1)
-            lambda_s_1 = lambda_s + r_1 * h
-            lambda_s_2 = lambda_s + r_2 * h
-            sigma_s_1, sigma_s_2 = sigma_fn(lambda_s_1), sigma_fn(lambda_s_2)
-
-            # alpha_t = sigma_t * exp(log(alpha_t / sigma_t)) = sigma_t * exp(lambda_t)
-            alpha_s_1 = sigma_s_1 * lambda_s_1.exp()
-            alpha_s_2 = sigma_s_2 * lambda_s_2.exp()
-            alpha_t = sigmas[i + 1] * lambda_t.exp()
-
-            coeff_1, coeff_2, coeff_3 = (-r_1 * h_eta).expm1(), (-r_2 * h_eta).expm1(), (-h_eta).expm1()
-            if inject_noise:
-                # 0 < r_1 < r_2 < 1
-                noise_coeff_1 = (-2 * r_1 * h * eta).expm1().neg().sqrt()
-                noise_coeff_2 = (-r_1 * h * eta).exp() * (-2 * (r_2 - r_1) * h * eta).expm1().neg().sqrt()
-                noise_coeff_3 = (-r_2 * h * eta).exp() * (-2 * (1 - r_2) * h * eta).expm1().neg().sqrt()
-                noise_1, noise_2, noise_3 = noise_sampler(sigmas[i], sigma_s_1), noise_sampler(sigma_s_1, sigma_s_2), noise_sampler(sigma_s_2, sigmas[i + 1])
-
-            # Step 1
-            x_2 = sigma_s_1 / sigmas[i] * (-r_1 * h * eta).exp() * x - alpha_s_1 * coeff_1 * denoised
-            if inject_noise:
-                x_2 = x_2 + sigma_s_1 * (noise_coeff_1 * noise_1) * s_noise
-            denoised_2 = model(x_2, sigma_s_1 * s_in, **extra_args)
-
-            # Step 2
-            x_3 = sigma_s_2 / sigmas[i] * (-r_2 * h * eta).exp() * x - alpha_s_2 * coeff_2 * denoised + (r_2 / r_1) * alpha_s_2 * (coeff_2 / (r_2 * h_eta) + 1) * (denoised_2 - denoised)
-            if inject_noise:
-                x_3 = x_3 + sigma_s_2 * (noise_coeff_2 * noise_1 + noise_coeff_1 * noise_2) * s_noise
-            denoised_3 = model(x_3, sigma_s_2 * s_in, **extra_args)
-
-            # Step 3
-            x = sigmas[i + 1] / sigmas[i] * (-h * eta).exp() * x - alpha_t * coeff_3 * denoised + (1. / r_2) * alpha_t * (coeff_3 / h_eta + 1) * (denoised_3 - denoised)
-            if inject_noise:
-                x = x + sigmas[i + 1] * (noise_coeff_3 * noise_1 + noise_coeff_2 * noise_2 + noise_coeff_1 * noise_3) * s_noise
-    return x
-
-
-@torch.no_grad()
-def sample_sa_solver(model, x, sigmas, extra_args=None, callback=None, disable=False, tau_func=None, s_noise=1.0, noise_sampler=None, predictor_order=3, corrector_order=4, use_pece=False, simple_order_2=False):
-    """Stochastic Adams Solver with predictor-corrector method (NeurIPS 2023)."""
-    if len(sigmas) <= 1:
-        return x
-    extra_args = {} if extra_args is None else extra_args
-    seed = extra_args.get("seed", None)
-    noise_sampler = default_noise_sampler(x, seed=seed) if noise_sampler is None else noise_sampler
-    s_in = x.new_ones([x.shape[0]])
-
-    model_sampling = model.inner_model.model_patcher.get_model_object("model_sampling")
-    sigmas = offset_first_sigma_for_snr(sigmas, model_sampling)
-    lambdas = sigma_to_half_log_snr(sigmas, model_sampling=model_sampling)
-
-    if tau_func is None:
-        # Use default interval for stochastic sampling
-        start_sigma = model_sampling.percent_to_sigma(0.2)
-        end_sigma = model_sampling.percent_to_sigma(0.8)
-        tau_func = sa_solver.get_tau_interval_func(start_sigma, end_sigma, eta=1.0)
-
-    max_used_order = max(predictor_order, corrector_order)
-    x_pred = x  # x: current state, x_pred: predicted next state
-
-    h = 0.0
-    tau_t = 0.0
-    noise = 0.0
-    pred_list = []
-
-    # Lower order near the end to improve stability
-    lower_order_to_end = sigmas[-1].item() == 0
-
-    for i in trange(len(sigmas) - 1, disable=disable):
-        # Evaluation
-        denoised = model(x_pred, sigmas[i] * s_in, **extra_args)
-        if callback is not None:
-            callback({"x": x_pred, "i": i, "sigma": sigmas[i], "sigma_hat": sigmas[i], "denoised": denoised})
-        pred_list.append(denoised)
-        pred_list = pred_list[-max_used_order:]
-
-        predictor_order_used = min(predictor_order, len(pred_list))
-        if i == 0 or (sigmas[i + 1] == 0 and not use_pece):
-            corrector_order_used = 0
-        else:
-            corrector_order_used = min(corrector_order, len(pred_list))
-
-        if lower_order_to_end:
-            predictor_order_used = min(predictor_order_used, len(sigmas) - 2 - i)
-            corrector_order_used = min(corrector_order_used, len(sigmas) - 1 - i)
-
-        # Corrector
-        if corrector_order_used == 0:
-            # Update by the predicted state
-            x = x_pred
-        else:
-            curr_lambdas = lambdas[i - corrector_order_used + 1:i + 1]
-            b_coeffs = sa_solver.compute_stochastic_adams_b_coeffs(
-                sigmas[i],
-                curr_lambdas,
-                lambdas[i - 1],
-                lambdas[i],
-                tau_t,
-                simple_order_2,
-                is_corrector_step=True,
-            )
-            pred_mat = torch.stack(pred_list[-corrector_order_used:], dim=1)    # (B, K, ...)
-            corr_res = torch.tensordot(pred_mat, b_coeffs, dims=([1], [0]))  # (B, ...)
-            x = sigmas[i] / sigmas[i - 1] * (-(tau_t ** 2) * h).exp() * x + corr_res
-
-            if tau_t > 0 and s_noise > 0:
-                # The noise from the previous predictor step
-                x = x + noise
-
-            if use_pece:
-                # Evaluate the corrected state
-                denoised = model(x, sigmas[i] * s_in, **extra_args)
-                pred_list[-1] = denoised
-
-        # Predictor
-        if sigmas[i + 1] == 0:
-            # Denoising step
-            x = denoised
-        else:
-            tau_t = tau_func(sigmas[i + 1])
-            curr_lambdas = lambdas[i - predictor_order_used + 1:i + 1]
-            b_coeffs = sa_solver.compute_stochastic_adams_b_coeffs(
-                sigmas[i + 1],
-                curr_lambdas,
-                lambdas[i],
-                lambdas[i + 1],
-                tau_t,
-                simple_order_2,
-                is_corrector_step=False,
-            )
-            pred_mat = torch.stack(pred_list[-predictor_order_used:], dim=1)    # (B, K, ...)
-            pred_res = torch.tensordot(pred_mat, b_coeffs, dims=([1], [0]))  # (B, ...)
-            h = lambdas[i + 1] - lambdas[i]
-            x_pred = sigmas[i + 1] / sigmas[i] * (-(tau_t ** 2) * h).exp() * x + pred_res
-
-            if tau_t > 0 and s_noise > 0:
-                noise = noise_sampler(sigmas[i], sigmas[i + 1]) * sigmas[i + 1] * (-2 * tau_t ** 2 * h).expm1().neg().sqrt() * s_noise
-                x_pred = x_pred + noise
-    return x
-
-
-@torch.no_grad()
-def sample_sa_solver_pece(model, x, sigmas, extra_args=None, callback=None, disable=False, tau_func=None, s_noise=1.0, noise_sampler=None, predictor_order=3, corrector_order=4, simple_order_2=False):
-    """Stochastic Adams Solver with PECE (Predict–Evaluate–Correct–Evaluate) mode (NeurIPS 2023)."""
-    return sample_sa_solver(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, tau_func=tau_func, s_noise=s_noise, noise_sampler=noise_sampler, predictor_order=predictor_order, corrector_order=corrector_order, use_pece=True, simple_order_2=simple_order_2)

comfy/k_diffusion/utils.py

@@ -1,313 +0,0 @@
-from contextlib import contextmanager
-import hashlib
-import math
-from pathlib import Path
-import shutil
-import urllib
-import warnings
-
-from PIL import Image
-import torch
-from torch import nn, optim
-from torch.utils import data
-
-
-def hf_datasets_augs_helper(examples, transform, image_key, mode='RGB'):
-    """Apply passed in transforms for HuggingFace Datasets."""
-    images = [transform(image.convert(mode)) for image in examples[image_key]]
-    return {image_key: images}
-
-
-def append_dims(x, target_dims):
-    """Appends dimensions to the end of a tensor until it has target_dims dimensions."""
-    dims_to_append = target_dims - x.ndim
-    if dims_to_append < 0:
-        raise ValueError(f'input has {x.ndim} dims but target_dims is {target_dims}, which is less')
-    expanded = x[(...,) + (None,) * dims_to_append]
-    # MPS will get inf values if it tries to index into the new axes, but detaching fixes this.
-    # https://github.com/pytorch/pytorch/issues/84364
-    return expanded.detach().clone() if expanded.device.type == 'mps' else expanded
-
-
-def n_params(module):
-    """Returns the number of trainable parameters in a module."""
-    return sum(p.numel() for p in module.parameters())
-
-
-def download_file(path, url, digest=None):
-    """Downloads a file if it does not exist, optionally checking its SHA-256 hash."""
-    path = Path(path)
-    path.parent.mkdir(parents=True, exist_ok=True)
-    if not path.exists():
-        with urllib.request.urlopen(url) as response, open(path, 'wb') as f:
-            shutil.copyfileobj(response, f)
-    if digest is not None:
-        file_digest = hashlib.sha256(open(path, 'rb').read()).hexdigest()
-        if digest != file_digest:
-            raise OSError(f'hash of {path} (url: {url}) failed to validate')
-    return path
-
-
-@contextmanager
-def train_mode(model, mode=True):
-    """A context manager that places a model into training mode and restores
-    the previous mode on exit."""
-    modes = [module.training for module in model.modules()]
-    try:
-        yield model.train(mode)
-    finally:
-        for i, module in enumerate(model.modules()):
-            module.training = modes[i]
-
-
-def eval_mode(model):
-    """A context manager that places a model into evaluation mode and restores
-    the previous mode on exit."""
-    return train_mode(model, False)
-
-
-@torch.no_grad()
-def ema_update(model, averaged_model, decay):
-    """Incorporates updated model parameters into an exponential moving averaged
-    version of a model. It should be called after each optimizer step."""
-    model_params = dict(model.named_parameters())
-    averaged_params = dict(averaged_model.named_parameters())
-    assert model_params.keys() == averaged_params.keys()
-
-    for name, param in model_params.items():
-        averaged_params[name].mul_(decay).add_(param, alpha=1 - decay)
-
-    model_buffers = dict(model.named_buffers())
-    averaged_buffers = dict(averaged_model.named_buffers())
-    assert model_buffers.keys() == averaged_buffers.keys()
-
-    for name, buf in model_buffers.items():
-        averaged_buffers[name].copy_(buf)
-
-
-class EMAWarmup:
-    """Implements an EMA warmup using an inverse decay schedule.
-    If inv_gamma=1 and power=1, implements a simple average. inv_gamma=1, power=2/3 are
-    good values for models you plan to train for a million or more steps (reaches decay
-    factor 0.999 at 31.6K steps, 0.9999 at 1M steps), inv_gamma=1, power=3/4 for models
-    you plan to train for less (reaches decay factor 0.999 at 10K steps, 0.9999 at
-    215.4k steps).
-    Args:
-        inv_gamma (float): Inverse multiplicative factor of EMA warmup. Default: 1.
-        power (float): Exponential factor of EMA warmup. Default: 1.
-        min_value (float): The minimum EMA decay rate. Default: 0.
-        max_value (float): The maximum EMA decay rate. Default: 1.
-        start_at (int): The epoch to start averaging at. Default: 0.
-        last_epoch (int): The index of last epoch. Default: 0.
-    """
-
-    def __init__(self, inv_gamma=1., power=1., min_value=0., max_value=1., start_at=0,
-                 last_epoch=0):
-        self.inv_gamma = inv_gamma
-        self.power = power
-        self.min_value = min_value
-        self.max_value = max_value
-        self.start_at = start_at
-        self.last_epoch = last_epoch
-
-    def state_dict(self):
-        """Returns the state of the class as a :class:`dict`."""
-        return dict(self.__dict__.items())
-
-    def load_state_dict(self, state_dict):
-        """Loads the class's state.
-        Args:
-            state_dict (dict): scaler state. Should be an object returned
-                from a call to :meth:`state_dict`.
-        """
-        self.__dict__.update(state_dict)
-
-    def get_value(self):
-        """Gets the current EMA decay rate."""
-        epoch = max(0, self.last_epoch - self.start_at)
-        value = 1 - (1 + epoch / self.inv_gamma) ** -self.power
-        return 0. if epoch < 0 else min(self.max_value, max(self.min_value, value))
-
-    def step(self):
-        """Updates the step count."""
-        self.last_epoch += 1
-
-
-class InverseLR(optim.lr_scheduler._LRScheduler):
-    """Implements an inverse decay learning rate schedule with an optional exponential
-    warmup. When last_epoch=-1, sets initial lr as lr.
-    inv_gamma is the number of steps/epochs required for the learning rate to decay to
-    (1 / 2)**power of its original value.
-    Args:
-        optimizer (Optimizer): Wrapped optimizer.
-        inv_gamma (float): Inverse multiplicative factor of learning rate decay. Default: 1.
-        power (float): Exponential factor of learning rate decay. Default: 1.
-        warmup (float): Exponential warmup factor (0 <= warmup < 1, 0 to disable)
-            Default: 0.
-        min_lr (float): The minimum learning rate. Default: 0.
-        last_epoch (int): The index of last epoch. Default: -1.
-        verbose (bool): If ``True``, prints a message to stdout for
-            each update. Default: ``False``.
-    """
-
-    def __init__(self, optimizer, inv_gamma=1., power=1., warmup=0., min_lr=0.,
-                 last_epoch=-1, verbose=False):
-        self.inv_gamma = inv_gamma
-        self.power = power
-        if not 0. <= warmup < 1:
-            raise ValueError('Invalid value for warmup')
-        self.warmup = warmup
-        self.min_lr = min_lr
-        super().__init__(optimizer, last_epoch, verbose)
-
-    def get_lr(self):
-        if not self._get_lr_called_within_step:
-            warnings.warn("To get the last learning rate computed by the scheduler, "
-                          "please use `get_last_lr()`.")
-
-        return self._get_closed_form_lr()
-
-    def _get_closed_form_lr(self):
-        warmup = 1 - self.warmup ** (self.last_epoch + 1)
-        lr_mult = (1 + self.last_epoch / self.inv_gamma) ** -self.power
-        return [warmup * max(self.min_lr, base_lr * lr_mult)
-                for base_lr in self.base_lrs]
-
-
-class ExponentialLR(optim.lr_scheduler._LRScheduler):
-    """Implements an exponential learning rate schedule with an optional exponential
-    warmup. When last_epoch=-1, sets initial lr as lr. Decays the learning rate
-    continuously by decay (default 0.5) every num_steps steps.
-    Args:
-        optimizer (Optimizer): Wrapped optimizer.
-        num_steps (float): The number of steps to decay the learning rate by decay in.
-        decay (float): The factor by which to decay the learning rate every num_steps
-            steps. Default: 0.5.
-        warmup (float): Exponential warmup factor (0 <= warmup < 1, 0 to disable)
-            Default: 0.
-        min_lr (float): The minimum learning rate. Default: 0.
-        last_epoch (int): The index of last epoch. Default: -1.
-        verbose (bool): If ``True``, prints a message to stdout for
-            each update. Default: ``False``.
-    """
-
-    def __init__(self, optimizer, num_steps, decay=0.5, warmup=0., min_lr=0.,
-                 last_epoch=-1, verbose=False):
-        self.num_steps = num_steps
-        self.decay = decay
-        if not 0. <= warmup < 1:
-            raise ValueError('Invalid value for warmup')
-        self.warmup = warmup
-        self.min_lr = min_lr
-        super().__init__(optimizer, last_epoch, verbose)
-
-    def get_lr(self):
-        if not self._get_lr_called_within_step:
-            warnings.warn("To get the last learning rate computed by the scheduler, "
-                          "please use `get_last_lr()`.")
-
-        return self._get_closed_form_lr()
-
-    def _get_closed_form_lr(self):
-        warmup = 1 - self.warmup ** (self.last_epoch + 1)
-        lr_mult = (self.decay ** (1 / self.num_steps)) ** self.last_epoch
-        return [warmup * max(self.min_lr, base_lr * lr_mult)
-                for base_lr in self.base_lrs]
-
-
-def rand_log_normal(shape, loc=0., scale=1., device='cpu', dtype=torch.float32):
-    """Draws samples from an lognormal distribution."""
-    return (torch.randn(shape, device=device, dtype=dtype) * scale + loc).exp()
-
-
-def rand_log_logistic(shape, loc=0., scale=1., min_value=0., max_value=float('inf'), device='cpu', dtype=torch.float32):
-    """Draws samples from an optionally truncated log-logistic distribution."""
-    min_value = torch.as_tensor(min_value, device=device, dtype=torch.float64)
-    max_value = torch.as_tensor(max_value, device=device, dtype=torch.float64)
-    min_cdf = min_value.log().sub(loc).div(scale).sigmoid()
-    max_cdf = max_value.log().sub(loc).div(scale).sigmoid()
-    u = torch.rand(shape, device=device, dtype=torch.float64) * (max_cdf - min_cdf) + min_cdf
-    return u.logit().mul(scale).add(loc).exp().to(dtype)
-
-
-def rand_log_uniform(shape, min_value, max_value, device='cpu', dtype=torch.float32):
-    """Draws samples from an log-uniform distribution."""
-    min_value = math.log(min_value)
-    max_value = math.log(max_value)
-    return (torch.rand(shape, device=device, dtype=dtype) * (max_value - min_value) + min_value).exp()
-
-
-def rand_v_diffusion(shape, sigma_data=1., min_value=0., max_value=float('inf'), device='cpu', dtype=torch.float32):
-    """Draws samples from a truncated v-diffusion training timestep distribution."""
-    min_cdf = math.atan(min_value / sigma_data) * 2 / math.pi
-    max_cdf = math.atan(max_value / sigma_data) * 2 / math.pi
-    u = torch.rand(shape, device=device, dtype=dtype) * (max_cdf - min_cdf) + min_cdf
-    return torch.tan(u * math.pi / 2) * sigma_data
-
-
-def rand_split_log_normal(shape, loc, scale_1, scale_2, device='cpu', dtype=torch.float32):
-    """Draws samples from a split lognormal distribution."""
-    n = torch.randn(shape, device=device, dtype=dtype).abs()
-    u = torch.rand(shape, device=device, dtype=dtype)
-    n_left = n * -scale_1 + loc
-    n_right = n * scale_2 + loc
-    ratio = scale_1 / (scale_1 + scale_2)
-    return torch.where(u < ratio, n_left, n_right).exp()
-
-
-class FolderOfImages(data.Dataset):
-    """Recursively finds all images in a directory. It does not support
-    classes/targets."""
-
-    IMG_EXTENSIONS = {'.jpg', '.jpeg', '.png', '.ppm', '.bmp', '.pgm', '.tif', '.tiff', '.webp'}
-
-    def __init__(self, root, transform=None):
-        super().__init__()
-        self.root = Path(root)
-        self.transform = nn.Identity() if transform is None else transform
-        self.paths = sorted(path for path in self.root.rglob('*') if path.suffix.lower() in self.IMG_EXTENSIONS)
-
-    def __repr__(self):
-        return f'FolderOfImages(root="{self.root}", len: {len(self)})'
-
-    def __len__(self):
-        return len(self.paths)
-
-    def __getitem__(self, key):
-        path = self.paths[key]
-        with open(path, 'rb') as f:
-            image = Image.open(f).convert('RGB')
-        image = self.transform(image)
-        return image,
-
-
-class CSVLogger:
-    def __init__(self, filename, columns):
-        self.filename = Path(filename)
-        self.columns = columns
-        if self.filename.exists():
-            self.file = open(self.filename, 'a')
-        else:
-            self.file = open(self.filename, 'w')
-            self.write(*self.columns)
-
-    def write(self, *args):
-        print(*args, sep=',', file=self.file, flush=True)
-
-
-@contextmanager
-def tf32_mode(cudnn=None, matmul=None):
-    """A context manager that sets whether TF32 is allowed on cuDNN or matmul."""
-    cudnn_old = torch.backends.cudnn.allow_tf32
-    matmul_old = torch.backends.cuda.matmul.allow_tf32
-    try:
-        if cudnn is not None:
-            torch.backends.cudnn.allow_tf32 = cudnn
-        if matmul is not None:
-            torch.backends.cuda.matmul.allow_tf32 = matmul
-        yield
-    finally:
-        if cudnn is not None:
-            torch.backends.cudnn.allow_tf32 = cudnn_old
-        if matmul is not None:
-            torch.backends.cuda.matmul.allow_tf32 = matmul_old

comfy/latent_formats.py

@@ -1,472 +0,0 @@
-import torch
-
-class LatentFormat:
-    scale_factor = 1.0
-    latent_channels = 4
-    latent_dimensions = 2
-    latent_rgb_factors = None
-    latent_rgb_factors_bias = None
-    taesd_decoder_name = None
-
-    def process_in(self, latent):
-        return latent * self.scale_factor
-
-    def process_out(self, latent):
-        return latent / self.scale_factor
-
-class SD15(LatentFormat):
-    def __init__(self, scale_factor=0.18215):
-        self.scale_factor = scale_factor
-        self.latent_rgb_factors = [
-                    #   R        G        B
-                    [ 0.3512,  0.2297,  0.3227],
-                    [ 0.3250,  0.4974,  0.2350],
-                    [-0.2829,  0.1762,  0.2721],
-                    [-0.2120, -0.2616, -0.7177]
-                ]
-        self.taesd_decoder_name = "taesd_decoder"
-
-class SDXL(LatentFormat):
-    scale_factor = 0.13025
-
-    def __init__(self):
-        self.latent_rgb_factors = [
-                    #   R        G        B
-                    [ 0.3651,  0.4232,  0.4341],
-                    [-0.2533, -0.0042,  0.1068],
-                    [ 0.1076,  0.1111, -0.0362],
-                    [-0.3165, -0.2492, -0.2188]
-                ]
-        self.latent_rgb_factors_bias = [ 0.1084, -0.0175, -0.0011]
-
-        self.taesd_decoder_name = "taesdxl_decoder"
-
-class SDXL_Playground_2_5(LatentFormat):
-    def __init__(self):
-        self.scale_factor = 0.5
-        self.latents_mean = torch.tensor([-1.6574, 1.886, -1.383, 2.5155]).view(1, 4, 1, 1)
-        self.latents_std = torch.tensor([8.4927, 5.9022, 6.5498, 5.2299]).view(1, 4, 1, 1)
-
-        self.latent_rgb_factors = [
-                    #   R        G        B
-                    [ 0.3920,  0.4054,  0.4549],
-                    [-0.2634, -0.0196,  0.0653],
-                    [ 0.0568,  0.1687, -0.0755],
-                    [-0.3112, -0.2359, -0.2076]
-                ]
-        self.taesd_decoder_name = "taesdxl_decoder"
-
-    def process_in(self, latent):
-        latents_mean = self.latents_mean.to(latent.device, latent.dtype)
-        latents_std = self.latents_std.to(latent.device, latent.dtype)
-        return (latent - latents_mean) * self.scale_factor / latents_std
-
-    def process_out(self, latent):
-        latents_mean = self.latents_mean.to(latent.device, latent.dtype)
-        latents_std = self.latents_std.to(latent.device, latent.dtype)
-        return latent * latents_std / self.scale_factor + latents_mean
-
-
-class SD_X4(LatentFormat):
-    def __init__(self):
-        self.scale_factor = 0.08333
-        self.latent_rgb_factors = [
-            [-0.2340, -0.3863, -0.3257],
-            [ 0.0994,  0.0885, -0.0908],
-            [-0.2833, -0.2349, -0.3741],
-            [ 0.2523, -0.0055, -0.1651]
-        ]
-
-class SC_Prior(LatentFormat):
-    latent_channels = 16
-    def __init__(self):
-        self.scale_factor = 1.0
-        self.latent_rgb_factors = [
-            [-0.0326, -0.0204, -0.0127],
-            [-0.1592, -0.0427,  0.0216],
-            [ 0.0873,  0.0638, -0.0020],
-            [-0.0602,  0.0442,  0.1304],
-            [ 0.0800, -0.0313, -0.1796],
-            [-0.0810, -0.0638, -0.1581],
-            [ 0.1791,  0.1180,  0.0967],
-            [ 0.0740,  0.1416,  0.0432],
-            [-0.1745, -0.1888, -0.1373],
-            [ 0.2412,  0.1577,  0.0928],
-            [ 0.1908,  0.0998,  0.0682],
-            [ 0.0209,  0.0365, -0.0092],
-            [ 0.0448, -0.0650, -0.1728],
-            [-0.1658, -0.1045, -0.1308],
-            [ 0.0542,  0.1545,  0.1325],
-            [-0.0352, -0.1672, -0.2541]
-        ]
-
-class SC_B(LatentFormat):
-    def __init__(self):
-        self.scale_factor = 1.0 / 0.43
-        self.latent_rgb_factors = [
-            [ 0.1121,  0.2006,  0.1023],
-            [-0.2093, -0.0222, -0.0195],
-            [-0.3087, -0.1535,  0.0366],
-            [ 0.0290, -0.1574, -0.4078]
-        ]
-
-class SD3(LatentFormat):
-    latent_channels = 16
-    def __init__(self):
-        self.scale_factor = 1.5305
-        self.shift_factor = 0.0609
-        self.latent_rgb_factors = [
-            [-0.0922, -0.0175,  0.0749],
-            [ 0.0311,  0.0633,  0.0954],
-            [ 0.1994,  0.0927,  0.0458],
-            [ 0.0856,  0.0339,  0.0902],
-            [ 0.0587,  0.0272, -0.0496],
-            [-0.0006,  0.1104,  0.0309],
-            [ 0.0978,  0.0306,  0.0427],
-            [-0.0042,  0.1038,  0.1358],
-            [-0.0194,  0.0020,  0.0669],
-            [-0.0488,  0.0130, -0.0268],
-            [ 0.0922,  0.0988,  0.0951],
-            [-0.0278,  0.0524, -0.0542],
-            [ 0.0332,  0.0456,  0.0895],
-            [-0.0069, -0.0030, -0.0810],
-            [-0.0596, -0.0465, -0.0293],
-            [-0.1448, -0.1463, -0.1189]
-        ]
-        self.latent_rgb_factors_bias = [0.2394, 0.2135, 0.1925]
-        self.taesd_decoder_name = "taesd3_decoder"
-
-    def process_in(self, latent):
-        return (latent - self.shift_factor) * self.scale_factor
-
-    def process_out(self, latent):
-        return (latent / self.scale_factor) + self.shift_factor
-
-class StableAudio1(LatentFormat):
-    latent_channels = 64
-    latent_dimensions = 1
-
-class Flux(SD3):
-    latent_channels = 16
-    def __init__(self):
-        self.scale_factor = 0.3611
-        self.shift_factor = 0.1159
-        self.latent_rgb_factors =[
-            [-0.0346,  0.0244,  0.0681],
-            [ 0.0034,  0.0210,  0.0687],
-            [ 0.0275, -0.0668, -0.0433],
-            [-0.0174,  0.0160,  0.0617],
-            [ 0.0859,  0.0721,  0.0329],
-            [ 0.0004,  0.0383,  0.0115],
-            [ 0.0405,  0.0861,  0.0915],
-            [-0.0236, -0.0185, -0.0259],
-            [-0.0245,  0.0250,  0.1180],
-            [ 0.1008,  0.0755, -0.0421],
-            [-0.0515,  0.0201,  0.0011],
-            [ 0.0428, -0.0012, -0.0036],
-            [ 0.0817,  0.0765,  0.0749],
-            [-0.1264, -0.0522, -0.1103],
-            [-0.0280, -0.0881, -0.0499],
-            [-0.1262, -0.0982, -0.0778]
-        ]
-        self.latent_rgb_factors_bias = [-0.0329, -0.0718, -0.0851]
-        self.taesd_decoder_name = "taef1_decoder"
-
-    def process_in(self, latent):
-        return (latent - self.shift_factor) * self.scale_factor
-
-    def process_out(self, latent):
-        return (latent / self.scale_factor) + self.shift_factor
-
-class Mochi(LatentFormat):
-    latent_channels = 12
-    latent_dimensions = 3
-
-    def __init__(self):
-        self.scale_factor = 1.0
-        self.latents_mean = torch.tensor([-0.06730895953510081, -0.038011381506090416, -0.07477820912866141,
-                                          -0.05565264470995561, 0.012767231469026969, -0.04703542746246419,
-                                          0.043896967884726704, -0.09346305707025976, -0.09918314763016893,
-                                          -0.008729793427399178, -0.011931556316503654, -0.0321993391887285]).view(1, self.latent_channels, 1, 1, 1)
-        self.latents_std = torch.tensor([0.9263795028493863, 0.9248894543193766, 0.9393059390890617,
-                                         0.959253732819592, 0.8244560132752793, 0.917259975397747,
-                                         0.9294154431013696, 1.3720942357788521, 0.881393668867029,
-                                         0.9168315692124348, 0.9185249279345552, 0.9274757570805041]).view(1, self.latent_channels, 1, 1, 1)
-
-        self.latent_rgb_factors =[
-            [-0.0069, -0.0045,  0.0018],
-            [ 0.0154, -0.0692, -0.0274],
-            [ 0.0333,  0.0019,  0.0206],
-            [-0.1390,  0.0628,  0.1678],
-            [-0.0725,  0.0134, -0.1898],
-            [ 0.0074, -0.0270, -0.0209],
-            [-0.0176, -0.0277, -0.0221],
-            [ 0.5294,  0.5204,  0.3852],
-            [-0.0326, -0.0446, -0.0143],
-            [-0.0659,  0.0153, -0.0153],
-            [ 0.0185, -0.0217,  0.0014],
-            [-0.0396, -0.0495, -0.0281]
-        ]
-        self.latent_rgb_factors_bias = [-0.0940, -0.1418, -0.1453]
-        self.taesd_decoder_name = None #TODO
-
-    def process_in(self, latent):
-        latents_mean = self.latents_mean.to(latent.device, latent.dtype)
-        latents_std = self.latents_std.to(latent.device, latent.dtype)
-        return (latent - latents_mean) * self.scale_factor / latents_std
-
-    def process_out(self, latent):
-        latents_mean = self.latents_mean.to(latent.device, latent.dtype)
-        latents_std = self.latents_std.to(latent.device, latent.dtype)
-        return latent * latents_std / self.scale_factor + latents_mean
-
-class LTXV(LatentFormat):
-    latent_channels = 128
-    latent_dimensions = 3
-
-    def __init__(self):
-        self.latent_rgb_factors = [
-            [ 1.1202e-02, -6.3815e-04, -1.0021e-02],
-            [ 8.6031e-02,  6.5813e-02,  9.5409e-04],
-            [-1.2576e-02, -7.5734e-03, -4.0528e-03],
-            [ 9.4063e-03, -2.1688e-03,  2.6093e-03],
-            [ 3.7636e-03,  1.2765e-02,  9.1548e-03],
-            [ 2.1024e-02, -5.2973e-03,  3.4373e-03],
-            [-8.8896e-03, -1.9703e-02, -1.8761e-02],
-            [-1.3160e-02, -1.0523e-02,  1.9709e-03],
-            [-1.5152e-03, -6.9891e-03, -7.5810e-03],
-            [-1.7247e-03,  4.6560e-04, -3.3839e-03],
-            [ 1.3617e-02,  4.7077e-03, -2.0045e-03],
-            [ 1.0256e-02,  7.7318e-03,  1.3948e-02],
-            [-1.6108e-02, -6.2151e-03,  1.1561e-03],
-            [ 7.3407e-03,  1.5628e-02,  4.4865e-04],
-            [ 9.5357e-04, -2.9518e-03, -1.4760e-02],
-            [ 1.9143e-02,  1.0868e-02,  1.2264e-02],
-            [ 4.4575e-03,  3.6682e-05, -6.8508e-03],
-            [-4.5681e-04,  3.2570e-03,  7.7929e-03],
-            [ 3.3902e-02,  3.3405e-02,  3.7454e-02],
-            [-2.3001e-02, -2.4877e-03, -3.1033e-03],
-            [ 5.0265e-02,  3.8841e-02,  3.3539e-02],
-            [-4.1018e-03, -1.1095e-03,  1.5859e-03],
-            [-1.2689e-01, -1.3107e-01, -2.1005e-01],
-            [ 2.6276e-02,  1.4189e-02, -3.5963e-03],
-            [-4.8679e-03,  8.8486e-03,  7.8029e-03],
-            [-1.6610e-03, -4.8597e-03, -5.2060e-03],
-            [-2.1010e-03,  2.3610e-03,  9.3796e-03],
-            [-2.2482e-02, -2.1305e-02, -1.5087e-02],
-            [-1.5753e-02, -1.0646e-02, -6.5083e-03],
-            [-4.6975e-03,  5.0288e-03, -6.7390e-03],
-            [ 1.1951e-02,  2.0712e-02,  1.6191e-02],
-            [-6.3704e-03, -8.4827e-03, -9.5483e-03],
-            [ 7.2610e-03, -9.9326e-03, -2.2978e-02],
-            [-9.1904e-04,  6.2882e-03,  9.5720e-03],
-            [-3.7178e-02, -3.7123e-02, -5.6713e-02],
-            [-1.3373e-01, -1.0720e-01, -5.3801e-02],
-            [-5.3702e-03,  8.1256e-03,  8.8397e-03],
-            [-1.5247e-01, -2.1437e-01, -2.1843e-01],
-            [ 3.1441e-02,  7.0335e-03, -9.7541e-03],
-            [ 2.1528e-03, -8.9817e-03, -2.1023e-02],
-            [ 3.8461e-03, -5.8957e-03, -1.5014e-02],
-            [-4.3470e-03, -1.2940e-02, -1.5972e-02],
-            [-5.4781e-03, -1.0842e-02, -3.0204e-03],
-            [-6.5347e-03,  3.0806e-03, -1.0163e-02],
-            [-5.0414e-03, -7.1503e-03, -8.9686e-04],
-            [-8.5851e-03, -2.4351e-03,  1.0674e-03],
-            [-9.0016e-03, -9.6493e-03,  1.5692e-03],
-            [ 5.0914e-03,  1.2099e-02,  1.9968e-02],
-            [ 1.3758e-02,  1.1669e-02,  8.1958e-03],
-            [-1.0518e-02, -1.1575e-02, -4.1307e-03],
-            [-2.8410e-02, -3.1266e-02, -2.2149e-02],
-            [ 2.9336e-03,  3.6511e-02,  1.8717e-02],
-            [-1.6703e-02, -1.6696e-02, -4.4529e-03],
-            [ 4.8818e-02,  4.0063e-02,  8.7410e-03],
-            [-1.5066e-02, -5.7328e-04,  2.9785e-03],
-            [-1.7613e-02, -8.1034e-03,  1.3086e-02],
-            [-9.2633e-03,  1.0803e-02, -6.3489e-03],
-            [ 3.0851e-03,  4.7750e-04,  1.2347e-02],
-            [-2.2785e-02, -2.3043e-02, -2.6005e-02],
-            [-2.4787e-02, -1.5389e-02, -2.2104e-02],
-            [-2.3572e-02,  1.0544e-03,  1.2361e-02],
-            [-7.8915e-03, -1.2271e-03, -6.0968e-03],
-            [-1.1478e-02, -1.2543e-03,  6.2679e-03],
-            [-5.4229e-02,  2.6644e-02,  6.3394e-03],
-            [ 4.4216e-03, -7.3338e-03, -1.0464e-02],
-            [-4.5013e-03,  1.6082e-03,  1.4420e-02],
-            [ 1.3673e-02,  8.8877e-03,  4.1253e-03],
-            [-1.0145e-02,  9.0072e-03,  1.5695e-02],
-            [-5.6234e-03,  1.1847e-03,  8.1261e-03],
-            [-3.7171e-03, -5.3538e-03,  1.2590e-03],
-            [ 2.9476e-02,  2.1424e-02,  3.0424e-02],
-            [-3.4925e-02, -2.4340e-02, -2.5316e-02],
-            [-3.4127e-02, -2.2406e-02, -1.0589e-02],
-            [-1.7342e-02, -1.3249e-02, -1.0719e-02],
-            [-2.1478e-03, -8.6051e-03, -2.9878e-03],
-            [ 1.2089e-03, -4.2391e-03, -6.8569e-03],
-            [ 9.0411e-04, -6.6886e-03, -6.7547e-05],
-            [ 1.6048e-02, -1.0057e-02, -2.8929e-02],
-            [ 1.2290e-03,  1.0163e-02,  1.8861e-02],
-            [ 1.7264e-02,  2.7257e-04,  1.3785e-02],
-            [-1.3482e-02, -3.6427e-03,  6.7481e-04],
-            [ 4.6782e-03, -5.2423e-03,  2.4467e-03],
-            [-5.9113e-03, -6.2244e-03, -1.8162e-03],
-            [ 1.5496e-02,  1.4582e-02,  1.9514e-03],
-            [ 7.4958e-03,  1.5886e-03, -8.2305e-03],
-            [ 1.9086e-02,  1.6360e-03, -3.9674e-03],
-            [-5.7021e-03, -2.7307e-03, -4.1066e-03],
-            [ 1.7450e-03,  1.4602e-02,  2.5794e-02],
-            [-8.2788e-04,  2.2902e-03,  4.5161e-03],
-            [ 1.1632e-02,  8.9193e-03, -7.2813e-03],
-            [ 7.5721e-03,  2.6784e-03,  1.1393e-02],
-            [ 5.1939e-03,  3.6903e-03,  1.4049e-02],
-            [-1.8383e-02, -2.2529e-02, -2.4477e-02],
-            [ 5.8842e-04, -5.7874e-03, -1.4770e-02],
-            [-1.6125e-02, -8.6101e-03, -1.4533e-02],
-            [ 2.0540e-02,  2.0729e-02,  6.4338e-03],
-            [ 3.3587e-03, -1.1226e-02, -1.6444e-02],
-            [-1.4742e-03, -1.0489e-02,  1.7097e-03],
-            [ 2.8130e-02,  2.3546e-02,  3.2791e-02],
-            [-1.8532e-02, -1.2842e-02, -8.7756e-03],
-            [-8.0533e-03, -1.0771e-02, -1.7536e-02],
-            [-3.9009e-03,  1.6150e-02,  3.3359e-02],
-            [-7.4554e-03, -1.4154e-02, -6.1910e-03],
-            [ 3.4734e-03, -1.1370e-02, -1.0581e-02],
-            [ 1.1476e-02,  3.9281e-03,  2.8231e-03],
-            [ 7.1639e-03, -1.4741e-03, -3.8066e-03],
-            [ 2.2250e-03, -8.7552e-03, -9.5719e-03],
-            [ 2.4146e-02,  2.1696e-02,  2.8056e-02],
-            [-5.4365e-03, -2.4291e-02, -1.7802e-02],
-            [ 7.4263e-03,  1.0510e-02,  1.2705e-02],
-            [ 6.2669e-03,  6.2658e-03,  1.9211e-02],
-            [ 1.6378e-02,  9.4933e-03,  6.6971e-03],
-            [ 1.7173e-02,  2.3601e-02,  2.3296e-02],
-            [-1.4568e-02, -9.8279e-03, -1.1556e-02],
-            [ 1.4431e-02,  1.4430e-02,  6.6362e-03],
-            [-6.8230e-03,  1.8863e-02,  1.4555e-02],
-            [ 6.1156e-03,  3.4700e-03, -2.6662e-03],
-            [-2.6983e-03, -5.9402e-03, -9.2276e-03],
-            [ 1.0235e-02,  7.4173e-03, -7.6243e-03],
-            [-1.3255e-02,  1.9322e-02, -9.2153e-04],
-            [ 2.4222e-03, -4.8039e-03, -1.5759e-02],
-            [ 2.6244e-02,  2.5951e-02,  2.0249e-02],
-            [ 1.5711e-02,  1.8498e-02,  2.7407e-03],
-            [-2.1714e-03,  4.7214e-03, -2.2443e-02],
-            [-7.4747e-03,  7.4166e-03,  1.4430e-02],
-            [-8.3906e-03, -7.9776e-03,  9.7927e-03],
-            [ 3.8321e-02,  9.6622e-03, -1.9268e-02],
-            [-1.4605e-02, -6.7032e-03,  3.9675e-03]
-        ]
-
-        self.latent_rgb_factors_bias = [-0.0571, -0.1657, -0.2512]
-
-class HunyuanVideo(LatentFormat):
-    latent_channels = 16
-    latent_dimensions = 3
-    scale_factor = 0.476986
-    latent_rgb_factors = [
-        [-0.0395, -0.0331,  0.0445],
-        [ 0.0696,  0.0795,  0.0518],
-        [ 0.0135, -0.0945, -0.0282],
-        [ 0.0108, -0.0250, -0.0765],
-        [-0.0209,  0.0032,  0.0224],
-        [-0.0804, -0.0254, -0.0639],
-        [-0.0991,  0.0271, -0.0669],
-        [-0.0646, -0.0422, -0.0400],
-        [-0.0696, -0.0595, -0.0894],
-        [-0.0799, -0.0208, -0.0375],
-        [ 0.1166,  0.1627,  0.0962],
-        [ 0.1165,  0.0432,  0.0407],
-        [-0.2315, -0.1920, -0.1355],
-        [-0.0270,  0.0401, -0.0821],
-        [-0.0616, -0.0997, -0.0727],
-        [ 0.0249, -0.0469, -0.1703]
-    ]
-
-    latent_rgb_factors_bias = [ 0.0259, -0.0192, -0.0761]
-
-class Cosmos1CV8x8x8(LatentFormat):
-    latent_channels = 16
-    latent_dimensions = 3
-
-    latent_rgb_factors = [
-        [ 0.1817,  0.2284,  0.2423],
-        [-0.0586, -0.0862, -0.3108],
-        [-0.4703, -0.4255, -0.3995],
-        [ 0.0803,  0.1963,  0.1001],
-        [-0.0820, -0.1050,  0.0400],
-        [ 0.2511,  0.3098,  0.2787],
-        [-0.1830, -0.2117, -0.0040],
-        [-0.0621, -0.2187, -0.0939],
-        [ 0.3619,  0.1082,  0.1455],
-        [ 0.3164,  0.3922,  0.2575],
-        [ 0.1152,  0.0231, -0.0462],
-        [-0.1434, -0.3609, -0.3665],
-        [ 0.0635,  0.1471,  0.1680],
-        [-0.3635, -0.1963, -0.3248],
-        [-0.1865,  0.0365,  0.2346],
-        [ 0.0447,  0.0994,  0.0881]
-    ]
-
-    latent_rgb_factors_bias = [-0.1223, -0.1889, -0.1976]
-
-class Wan21(LatentFormat):
-    latent_channels = 16
-    latent_dimensions = 3
-
-    latent_rgb_factors = [
-            [-0.1299, -0.1692,  0.2932],
-            [ 0.0671,  0.0406,  0.0442],
-            [ 0.3568,  0.2548,  0.1747],
-            [ 0.0372,  0.2344,  0.1420],
-            [ 0.0313,  0.0189, -0.0328],
-            [ 0.0296, -0.0956, -0.0665],
-            [-0.3477, -0.4059, -0.2925],
-            [ 0.0166,  0.1902,  0.1975],
-            [-0.0412,  0.0267, -0.1364],
-            [-0.1293,  0.0740,  0.1636],
-            [ 0.0680,  0.3019,  0.1128],
-            [ 0.0032,  0.0581,  0.0639],
-            [-0.1251,  0.0927,  0.1699],
-            [ 0.0060, -0.0633,  0.0005],
-            [ 0.3477,  0.2275,  0.2950],
-            [ 0.1984,  0.0913,  0.1861]
-        ]
-
-    latent_rgb_factors_bias = [-0.1835, -0.0868, -0.3360]
-
-    def __init__(self):
-        self.scale_factor = 1.0
-        self.latents_mean = torch.tensor([
-            -0.7571, -0.7089, -0.9113, 0.1075, -0.1745, 0.9653, -0.1517, 1.5508,
-            0.4134, -0.0715, 0.5517, -0.3632, -0.1922, -0.9497, 0.2503, -0.2921
-        ]).view(1, self.latent_channels, 1, 1, 1)
-        self.latents_std = torch.tensor([
-            2.8184, 1.4541, 2.3275, 2.6558, 1.2196, 1.7708, 2.6052, 2.0743,
-            3.2687, 2.1526, 2.8652, 1.5579, 1.6382, 1.1253, 2.8251, 1.9160
-        ]).view(1, self.latent_channels, 1, 1, 1)
-
-
-        self.taesd_decoder_name = None #TODO
-
-    def process_in(self, latent):
-        latents_mean = self.latents_mean.to(latent.device, latent.dtype)
-        latents_std = self.latents_std.to(latent.device, latent.dtype)
-        return (latent - latents_mean) * self.scale_factor / latents_std
-
-    def process_out(self, latent):
-        latents_mean = self.latents_mean.to(latent.device, latent.dtype)
-        latents_std = self.latents_std.to(latent.device, latent.dtype)
-        return latent * latents_std / self.scale_factor + latents_mean
-
-class Hunyuan3Dv2(LatentFormat):
-    latent_channels = 64
-    latent_dimensions = 1
-    scale_factor = 0.9990943042622529
-
-class Hunyuan3Dv2mini(LatentFormat):
-    latent_channels = 64
-    latent_dimensions = 1
-    scale_factor = 1.0188137142395404
-
-class ACEAudio(LatentFormat):
-    latent_channels = 8
-    latent_dimensions = 2

comfy/ldm/ace/attention.py

@@ -1,761 +0,0 @@
-# Original from: https://github.com/ace-step/ACE-Step/blob/main/models/attention.py
-# Copyright 2024 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-from typing import Tuple, Union, Optional
-
-import torch
-import torch.nn.functional as F
-from torch import nn
-
-import comfy.model_management
-from comfy.ldm.modules.attention import optimized_attention
-
-class Attention(nn.Module):
-    def __init__(
-        self,
-        query_dim: int,
-        cross_attention_dim: Optional[int] = None,
-        heads: int = 8,
-        kv_heads: Optional[int] = None,
-        dim_head: int = 64,
-        dropout: float = 0.0,
-        bias: bool = False,
-        qk_norm: Optional[str] = None,
-        added_kv_proj_dim: Optional[int] = None,
-        added_proj_bias: Optional[bool] = True,
-        out_bias: bool = True,
-        scale_qk: bool = True,
-        only_cross_attention: bool = False,
-        eps: float = 1e-5,
-        rescale_output_factor: float = 1.0,
-        residual_connection: bool = False,
-        processor=None,
-        out_dim: int = None,
-        out_context_dim: int = None,
-        context_pre_only=None,
-        pre_only=False,
-        elementwise_affine: bool = True,
-        is_causal: bool = False,
-        dtype=None, device=None, operations=None
-    ):
-        super().__init__()
-
-        self.inner_dim = out_dim if out_dim is not None else dim_head * heads
-        self.inner_kv_dim = self.inner_dim if kv_heads is None else dim_head * kv_heads
-        self.query_dim = query_dim
-        self.use_bias = bias
-        self.is_cross_attention = cross_attention_dim is not None
-        self.cross_attention_dim = cross_attention_dim if cross_attention_dim is not None else query_dim
-        self.rescale_output_factor = rescale_output_factor
-        self.residual_connection = residual_connection
-        self.dropout = dropout
-        self.fused_projections = False
-        self.out_dim = out_dim if out_dim is not None else query_dim
-        self.out_context_dim = out_context_dim if out_context_dim is not None else query_dim
-        self.context_pre_only = context_pre_only
-        self.pre_only = pre_only
-        self.is_causal = is_causal
-
-        self.scale_qk = scale_qk
-        self.scale = dim_head**-0.5 if self.scale_qk else 1.0
-
-        self.heads = out_dim // dim_head if out_dim is not None else heads
-        # for slice_size > 0 the attention score computation
-        # is split across the batch axis to save memory
-        # You can set slice_size with `set_attention_slice`
-        self.sliceable_head_dim = heads
-
-        self.added_kv_proj_dim = added_kv_proj_dim
-        self.only_cross_attention = only_cross_attention
-
-        if self.added_kv_proj_dim is None and self.only_cross_attention:
-            raise ValueError(
-                "`only_cross_attention` can only be set to True if `added_kv_proj_dim` is not None. Make sure to set either `only_cross_attention=False` or define `added_kv_proj_dim`."
-            )
-
-        self.group_norm = None
-        self.spatial_norm = None
-
-        self.norm_q = None
-        self.norm_k = None
-
-        self.norm_cross = None
-        self.to_q = operations.Linear(query_dim, self.inner_dim, bias=bias, dtype=dtype, device=device)
-
-        if not self.only_cross_attention:
-            # only relevant for the `AddedKVProcessor` classes
-            self.to_k = operations.Linear(self.cross_attention_dim, self.inner_kv_dim, bias=bias, dtype=dtype, device=device)
-            self.to_v = operations.Linear(self.cross_attention_dim, self.inner_kv_dim, bias=bias, dtype=dtype, device=device)
-        else:
-            self.to_k = None
-            self.to_v = None
-
-        self.added_proj_bias = added_proj_bias
-        if self.added_kv_proj_dim is not None:
-            self.add_k_proj = operations.Linear(added_kv_proj_dim, self.inner_kv_dim, bias=added_proj_bias, dtype=dtype, device=device)
-            self.add_v_proj = operations.Linear(added_kv_proj_dim, self.inner_kv_dim, bias=added_proj_bias, dtype=dtype, device=device)
-            if self.context_pre_only is not None:
-                self.add_q_proj = operations.Linear(added_kv_proj_dim, self.inner_dim, bias=added_proj_bias, dtype=dtype, device=device)
-        else:
-            self.add_q_proj = None
-            self.add_k_proj = None
-            self.add_v_proj = None
-
-        if not self.pre_only:
-            self.to_out = nn.ModuleList([])
-            self.to_out.append(operations.Linear(self.inner_dim, self.out_dim, bias=out_bias, dtype=dtype, device=device))
-            self.to_out.append(nn.Dropout(dropout))
-        else:
-            self.to_out = None
-
-        if self.context_pre_only is not None and not self.context_pre_only:
-            self.to_add_out = operations.Linear(self.inner_dim, self.out_context_dim, bias=out_bias, dtype=dtype, device=device)
-        else:
-            self.to_add_out = None
-
-        self.norm_added_q = None
-        self.norm_added_k = None
-        self.processor = processor
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        encoder_hidden_states: Optional[torch.Tensor] = None,
-        attention_mask: Optional[torch.Tensor] = None,
-        **cross_attention_kwargs,
-    ) -> torch.Tensor:
-        return self.processor(
-            self,
-            hidden_states,
-            encoder_hidden_states=encoder_hidden_states,
-            attention_mask=attention_mask,
-            **cross_attention_kwargs,
-        )
-
-
-class CustomLiteLAProcessor2_0:
-    """Attention processor used typically in processing the SD3-like self-attention projections. add rms norm for query and key and apply RoPE"""
-
-    def __init__(self):
-        self.kernel_func = nn.ReLU(inplace=False)
-        self.eps = 1e-15
-        self.pad_val = 1.0
-
-    def apply_rotary_emb(
-        self,
-        x: torch.Tensor,
-        freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor]],
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
-        """
-        Apply rotary embeddings to input tensors using the given frequency tensor. This function applies rotary embeddings
-        to the given query or key 'x' tensors using the provided frequency tensor 'freqs_cis'. The input tensors are
-        reshaped as complex numbers, and the frequency tensor is reshaped for broadcasting compatibility. The resulting
-        tensors contain rotary embeddings and are returned as real tensors.
-
-        Args:
-            x (`torch.Tensor`):
-                Query or key tensor to apply rotary embeddings. [B, H, S, D] xk (torch.Tensor): Key tensor to apply
-            freqs_cis (`Tuple[torch.Tensor]`): Precomputed frequency tensor for complex exponentials. ([S, D], [S, D],)
-
-        Returns:
-            Tuple[torch.Tensor, torch.Tensor]: Tuple of modified query tensor and key tensor with rotary embeddings.
-        """
-        cos, sin = freqs_cis  # [S, D]
-        cos = cos[None, None]
-        sin = sin[None, None]
-        cos, sin = cos.to(x.device), sin.to(x.device)
-
-        x_real, x_imag = x.reshape(*x.shape[:-1], -1, 2).unbind(-1)  # [B, S, H, D//2]
-        x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(3)
-        out = (x.float() * cos + x_rotated.float() * sin).to(x.dtype)
-
-        return out
-
-    def __call__(
-        self,
-        attn: Attention,
-        hidden_states: torch.FloatTensor,
-        encoder_hidden_states: torch.FloatTensor = None,
-        attention_mask: Optional[torch.FloatTensor] = None,
-        encoder_attention_mask: Optional[torch.FloatTensor] = None,
-        rotary_freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor]] = None,
-        rotary_freqs_cis_cross: Union[torch.Tensor, Tuple[torch.Tensor]] = None,
-        *args,
-        **kwargs,
-    ) -> torch.FloatTensor:
-        hidden_states_len = hidden_states.shape[1]
-
-        input_ndim = hidden_states.ndim
-        if input_ndim == 4:
-            batch_size, channel, height, width = hidden_states.shape
-            hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
-        if encoder_hidden_states is not None:
-            context_input_ndim = encoder_hidden_states.ndim
-            if context_input_ndim == 4:
-                batch_size, channel, height, width = encoder_hidden_states.shape
-                encoder_hidden_states = encoder_hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
-
-        batch_size = hidden_states.shape[0]
-
-        # `sample` projections.
-        dtype = hidden_states.dtype
-        query = attn.to_q(hidden_states)
-        key = attn.to_k(hidden_states)
-        value = attn.to_v(hidden_states)
-
-        # `context` projections.
-        has_encoder_hidden_state_proj = hasattr(attn, "add_q_proj") and hasattr(attn, "add_k_proj") and hasattr(attn, "add_v_proj")
-        if encoder_hidden_states is not None and has_encoder_hidden_state_proj:
-            encoder_hidden_states_query_proj = attn.add_q_proj(encoder_hidden_states)
-            encoder_hidden_states_key_proj = attn.add_k_proj(encoder_hidden_states)
-            encoder_hidden_states_value_proj = attn.add_v_proj(encoder_hidden_states)
-
-            # attention
-            if not attn.is_cross_attention:
-                query = torch.cat([query, encoder_hidden_states_query_proj], dim=1)
-                key = torch.cat([key, encoder_hidden_states_key_proj], dim=1)
-                value = torch.cat([value, encoder_hidden_states_value_proj], dim=1)
-            else:
-                query = hidden_states
-                key = encoder_hidden_states
-                value = encoder_hidden_states
-
-        inner_dim = key.shape[-1]
-        head_dim = inner_dim // attn.heads
-
-        query = query.transpose(-1, -2).reshape(batch_size, attn.heads, head_dim, -1)
-        key = key.transpose(-1, -2).reshape(batch_size, attn.heads, head_dim, -1).transpose(-1, -2)
-        value = value.transpose(-1, -2).reshape(batch_size, attn.heads, head_dim, -1)
-
-        # RoPE需要 [B, H, S, D] 输入
-        # 此时 query是 [B, H, D, S], 需要转成 [B, H, S, D] 才能应用RoPE
-        query = query.permute(0, 1, 3, 2)  # [B, H, S, D]  (从 [B, H, D, S])
-
-        # Apply query and key normalization if needed
-        if attn.norm_q is not None:
-            query = attn.norm_q(query)
-        if attn.norm_k is not None:
-            key = attn.norm_k(key)
-
-        # Apply RoPE if needed
-        if rotary_freqs_cis is not None:
-            query = self.apply_rotary_emb(query, rotary_freqs_cis)
-            if not attn.is_cross_attention:
-                key = self.apply_rotary_emb(key, rotary_freqs_cis)
-            elif rotary_freqs_cis_cross is not None and has_encoder_hidden_state_proj:
-                key = self.apply_rotary_emb(key, rotary_freqs_cis_cross)
-
-        # 此时 query是 [B, H, S, D]，需要还原成 [B, H, D, S]
-        query = query.permute(0, 1, 3, 2)  # [B, H, D, S]
-
-        if attention_mask is not None:
-            # attention_mask: [B, S] -> [B, 1, S, 1]
-            attention_mask = attention_mask[:, None, :, None].to(key.dtype)  # [B, 1, S, 1]
-            query = query * attention_mask.permute(0, 1, 3, 2)  # [B, H, S, D] * [B, 1, S, 1]
-            if not attn.is_cross_attention:
-                key = key * attention_mask  # key: [B, h, S, D] 与 mask [B, 1, S, 1] 相乘
-                value = value * attention_mask.permute(0, 1, 3, 2)  # 如果 value 是 [B, h, D, S]，那么需调整mask以匹配S维度
-
-        if attn.is_cross_attention and encoder_attention_mask is not None and has_encoder_hidden_state_proj:
-            encoder_attention_mask = encoder_attention_mask[:, None, :, None].to(key.dtype)  # [B, 1, S_enc, 1]
-            # 此时 key: [B, h, S_enc, D], value: [B, h, D, S_enc]
-            key = key * encoder_attention_mask  # [B, h, S_enc, D] * [B, 1, S_enc, 1]
-            value = value * encoder_attention_mask.permute(0, 1, 3, 2)  # [B, h, D, S_enc] * [B, 1, 1, S_enc]
-
-        query = self.kernel_func(query)
-        key = self.kernel_func(key)
-
-        query, key, value = query.float(), key.float(), value.float()
-
-        value = F.pad(value, (0, 0, 0, 1), mode="constant", value=self.pad_val)
-
-        vk = torch.matmul(value, key)
-
-        hidden_states = torch.matmul(vk, query)
-
-        if hidden_states.dtype in [torch.float16, torch.bfloat16]:
-            hidden_states = hidden_states.float()
-
-        hidden_states = hidden_states[:, :, :-1] / (hidden_states[:, :, -1:] + self.eps)
-
-        hidden_states = hidden_states.view(batch_size, attn.heads * head_dim, -1).permute(0, 2, 1)
-
-        hidden_states = hidden_states.to(dtype)
-        if encoder_hidden_states is not None:
-            encoder_hidden_states = encoder_hidden_states.to(dtype)
-
-        # Split the attention outputs.
-        if encoder_hidden_states is not None and not attn.is_cross_attention and has_encoder_hidden_state_proj:
-            hidden_states, encoder_hidden_states = (
-                hidden_states[:, : hidden_states_len],
-                hidden_states[:, hidden_states_len:],
-            )
-
-        # linear proj
-        hidden_states = attn.to_out[0](hidden_states)
-        # dropout
-        hidden_states = attn.to_out[1](hidden_states)
-        if encoder_hidden_states is not None and not attn.context_pre_only and not attn.is_cross_attention and hasattr(attn, "to_add_out"):
-            encoder_hidden_states = attn.to_add_out(encoder_hidden_states)
-
-        if input_ndim == 4:
-            hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
-        if encoder_hidden_states is not None and context_input_ndim == 4:
-            encoder_hidden_states = encoder_hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
-
-        if torch.get_autocast_gpu_dtype() == torch.float16:
-            hidden_states = hidden_states.clip(-65504, 65504)
-            if encoder_hidden_states is not None:
-                encoder_hidden_states = encoder_hidden_states.clip(-65504, 65504)
-
-        return hidden_states, encoder_hidden_states
-
-
-class CustomerAttnProcessor2_0:
-    r"""
-    Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0).
-    """
-
-    def apply_rotary_emb(
-        self,
-        x: torch.Tensor,
-        freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor]],
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
-        """
-        Apply rotary embeddings to input tensors using the given frequency tensor. This function applies rotary embeddings
-        to the given query or key 'x' tensors using the provided frequency tensor 'freqs_cis'. The input tensors are
-        reshaped as complex numbers, and the frequency tensor is reshaped for broadcasting compatibility. The resulting
-        tensors contain rotary embeddings and are returned as real tensors.
-
-        Args:
-            x (`torch.Tensor`):
-                Query or key tensor to apply rotary embeddings. [B, H, S, D] xk (torch.Tensor): Key tensor to apply
-            freqs_cis (`Tuple[torch.Tensor]`): Precomputed frequency tensor for complex exponentials. ([S, D], [S, D],)
-
-        Returns:
-            Tuple[torch.Tensor, torch.Tensor]: Tuple of modified query tensor and key tensor with rotary embeddings.
-        """
-        cos, sin = freqs_cis  # [S, D]
-        cos = cos[None, None]
-        sin = sin[None, None]
-        cos, sin = cos.to(x.device), sin.to(x.device)
-
-        x_real, x_imag = x.reshape(*x.shape[:-1], -1, 2).unbind(-1)  # [B, S, H, D//2]
-        x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(3)
-        out = (x.float() * cos + x_rotated.float() * sin).to(x.dtype)
-
-        return out
-
-    def __call__(
-        self,
-        attn: Attention,
-        hidden_states: torch.FloatTensor,
-        encoder_hidden_states: torch.FloatTensor = None,
-        attention_mask: Optional[torch.FloatTensor] = None,
-        encoder_attention_mask: Optional[torch.FloatTensor] = None,
-        rotary_freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor]] = None,
-        rotary_freqs_cis_cross: Union[torch.Tensor, Tuple[torch.Tensor]] = None,
-        *args,
-        **kwargs,
-    ) -> torch.Tensor:
-
-        residual = hidden_states
-        input_ndim = hidden_states.ndim
-
-        if input_ndim == 4:
-            batch_size, channel, height, width = hidden_states.shape
-            hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
-
-        batch_size, sequence_length, _ = (
-            hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
-        )
-
-        has_encoder_hidden_state_proj = hasattr(attn, "add_q_proj") and hasattr(attn, "add_k_proj") and hasattr(attn, "add_v_proj")
-
-        if attn.group_norm is not None:
-            hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
-
-        query = attn.to_q(hidden_states)
-
-        if encoder_hidden_states is None:
-            encoder_hidden_states = hidden_states
-        elif attn.norm_cross:
-            encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
-
-        key = attn.to_k(encoder_hidden_states)
-        value = attn.to_v(encoder_hidden_states)
-
-        inner_dim = key.shape[-1]
-        head_dim = inner_dim // attn.heads
-
-        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
-
-        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
-        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
-
-        if attn.norm_q is not None:
-            query = attn.norm_q(query)
-        if attn.norm_k is not None:
-            key = attn.norm_k(key)
-
-        # Apply RoPE if needed
-        if rotary_freqs_cis is not None:
-            query = self.apply_rotary_emb(query, rotary_freqs_cis)
-            if not attn.is_cross_attention:
-                key = self.apply_rotary_emb(key, rotary_freqs_cis)
-            elif rotary_freqs_cis_cross is not None and has_encoder_hidden_state_proj:
-                key = self.apply_rotary_emb(key, rotary_freqs_cis_cross)
-
-        if attn.is_cross_attention and encoder_attention_mask is not None and has_encoder_hidden_state_proj:
-            # attention_mask: N x S1
-            # encoder_attention_mask: N x S2
-            # cross attention 整合attention_mask和encoder_attention_mask
-            combined_mask = attention_mask[:, :, None] * encoder_attention_mask[:, None, :]
-            attention_mask = torch.where(combined_mask == 1, 0.0, -torch.inf)
-            attention_mask = attention_mask[:, None, :, :].expand(-1, attn.heads, -1, -1).to(query.dtype)
-
-        elif not attn.is_cross_attention and attention_mask is not None:
-            attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
-            # scaled_dot_product_attention expects attention_mask shape to be
-            # (batch, heads, source_length, target_length)
-            attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
-
-        # the output of sdp = (batch, num_heads, seq_len, head_dim)
-        hidden_states = optimized_attention(
-            query, key, value, heads=query.shape[1], mask=attention_mask, skip_reshape=True,
-        ).to(query.dtype)
-
-        # linear proj
-        hidden_states = attn.to_out[0](hidden_states)
-        # dropout
-        hidden_states = attn.to_out[1](hidden_states)
-
-        if input_ndim == 4:
-            hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
-
-        if attn.residual_connection:
-            hidden_states = hidden_states + residual
-
-        hidden_states = hidden_states / attn.rescale_output_factor
-
-        return hidden_states
-
-def val2list(x: list or tuple or any, repeat_time=1) -> list:  # type: ignore
-    """Repeat `val` for `repeat_time` times and return the list or val if list/tuple."""
-    if isinstance(x, (list, tuple)):
-        return list(x)
-    return [x for _ in range(repeat_time)]
-
-
-def val2tuple(x: list or tuple or any, min_len: int = 1, idx_repeat: int = -1) -> tuple:  # type: ignore
-    """Return tuple with min_len by repeating element at idx_repeat."""
-    # convert to list first
-    x = val2list(x)
-
-    # repeat elements if necessary
-    if len(x) > 0:
-        x[idx_repeat:idx_repeat] = [x[idx_repeat] for _ in range(min_len - len(x))]
-
-    return tuple(x)
-
-
-def t2i_modulate(x, shift, scale):
-    return x * (1 + scale) + shift
-
-
-def get_same_padding(kernel_size: Union[int, Tuple[int, ...]]) -> Union[int, Tuple[int, ...]]:
-    if isinstance(kernel_size, tuple):
-        return tuple([get_same_padding(ks) for ks in kernel_size])
-    else:
-        assert kernel_size % 2 > 0, f"kernel size {kernel_size} should be odd number"
-        return kernel_size // 2
-
-class ConvLayer(nn.Module):
-    def __init__(
-        self,
-        in_dim: int,
-        out_dim: int,
-        kernel_size=3,
-        stride=1,
-        dilation=1,
-        groups=1,
-        padding: Union[int, None] = None,
-        use_bias=False,
-        norm=None,
-        act=None,
-        dtype=None, device=None, operations=None
-    ):
-        super().__init__()
-        if padding is None:
-            padding = get_same_padding(kernel_size)
-            padding *= dilation
-
-        self.in_dim = in_dim
-        self.out_dim = out_dim
-        self.kernel_size = kernel_size
-        self.stride = stride
-        self.dilation = dilation
-        self.groups = groups
-        self.padding = padding
-        self.use_bias = use_bias
-
-        self.conv = operations.Conv1d(
-            in_dim,
-            out_dim,
-            kernel_size=kernel_size,
-            stride=stride,
-            padding=padding,
-            dilation=dilation,
-            groups=groups,
-            bias=use_bias,
-            device=device,
-            dtype=dtype
-        )
-        if norm is not None:
-            self.norm = operations.RMSNorm(out_dim, elementwise_affine=False, dtype=dtype, device=device)
-        else:
-            self.norm = None
-        if act is not None:
-            self.act = nn.SiLU(inplace=True)
-        else:
-            self.act = None
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        x = self.conv(x)
-        if self.norm:
-            x = self.norm(x)
-        if self.act:
-            x = self.act(x)
-        return x
-
-
-class GLUMBConv(nn.Module):
-    def __init__(
-        self,
-        in_features: int,
-        hidden_features: int,
-        out_feature=None,
-        kernel_size=3,
-        stride=1,
-        padding: Union[int, None] = None,
-        use_bias=False,
-        norm=(None, None, None),
-        act=("silu", "silu", None),
-        dilation=1,
-        dtype=None, device=None, operations=None
-    ):
-        out_feature = out_feature or in_features
-        super().__init__()
-        use_bias = val2tuple(use_bias, 3)
-        norm = val2tuple(norm, 3)
-        act = val2tuple(act, 3)
-
-        self.glu_act = nn.SiLU(inplace=False)
-        self.inverted_conv = ConvLayer(
-            in_features,
-            hidden_features * 2,
-            1,
-            use_bias=use_bias[0],
-            norm=norm[0],
-            act=act[0],
-            dtype=dtype,
-            device=device,
-            operations=operations,
-        )
-        self.depth_conv = ConvLayer(
-            hidden_features * 2,
-            hidden_features * 2,
-            kernel_size,
-            stride=stride,
-            groups=hidden_features * 2,
-            padding=padding,
-            use_bias=use_bias[1],
-            norm=norm[1],
-            act=None,
-            dilation=dilation,
-            dtype=dtype,
-            device=device,
-            operations=operations,
-        )
-        self.point_conv = ConvLayer(
-            hidden_features,
-            out_feature,
-            1,
-            use_bias=use_bias[2],
-            norm=norm[2],
-            act=act[2],
-            dtype=dtype,
-            device=device,
-            operations=operations,
-        )
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        x = x.transpose(1, 2)
-        x = self.inverted_conv(x)
-        x = self.depth_conv(x)
-
-        x, gate = torch.chunk(x, 2, dim=1)
-        gate = self.glu_act(gate)
-        x = x * gate
-
-        x = self.point_conv(x)
-        x = x.transpose(1, 2)
-
-        return x
-
-
-class LinearTransformerBlock(nn.Module):
-    """
-    A Sana block with global shared adaptive layer norm (adaLN-single) conditioning.
-    """
-    def __init__(
-        self,
-        dim,
-        num_attention_heads,
-        attention_head_dim,
-        use_adaln_single=True,
-        cross_attention_dim=None,
-        added_kv_proj_dim=None,
-        context_pre_only=False,
-        mlp_ratio=4.0,
-        add_cross_attention=False,
-        add_cross_attention_dim=None,
-        qk_norm=None,
-        dtype=None, device=None, operations=None
-    ):
-        super().__init__()
-
-        self.norm1 = operations.RMSNorm(dim, elementwise_affine=False, eps=1e-6)
-        self.attn = Attention(
-            query_dim=dim,
-            cross_attention_dim=cross_attention_dim,
-            added_kv_proj_dim=added_kv_proj_dim,
-            dim_head=attention_head_dim,
-            heads=num_attention_heads,
-            out_dim=dim,
-            bias=True,
-            qk_norm=qk_norm,
-            processor=CustomLiteLAProcessor2_0(),
-            dtype=dtype,
-            device=device,
-            operations=operations,
-        )
-
-        self.add_cross_attention = add_cross_attention
-        self.context_pre_only = context_pre_only
-
-        if add_cross_attention and add_cross_attention_dim is not None:
-            self.cross_attn = Attention(
-                query_dim=dim,
-                cross_attention_dim=add_cross_attention_dim,
-                added_kv_proj_dim=add_cross_attention_dim,
-                dim_head=attention_head_dim,
-                heads=num_attention_heads,
-                out_dim=dim,
-                context_pre_only=context_pre_only,
-                bias=True,
-                qk_norm=qk_norm,
-                processor=CustomerAttnProcessor2_0(),
-                dtype=dtype,
-                device=device,
-                operations=operations,
-            )
-
-        self.norm2 = operations.RMSNorm(dim, 1e-06, elementwise_affine=False)
-
-        self.ff = GLUMBConv(
-            in_features=dim,
-            hidden_features=int(dim * mlp_ratio),
-            use_bias=(True, True, False),
-            norm=(None, None, None),
-            act=("silu", "silu", None),
-            dtype=dtype,
-            device=device,
-            operations=operations,
-        )
-        self.use_adaln_single = use_adaln_single
-        if use_adaln_single:
-            self.scale_shift_table = nn.Parameter(torch.empty(6, dim, dtype=dtype, device=device))
-
-    def forward(
-        self,
-        hidden_states: torch.FloatTensor,
-        encoder_hidden_states: torch.FloatTensor = None,
-        attention_mask: torch.FloatTensor = None,
-        encoder_attention_mask: torch.FloatTensor = None,
-        rotary_freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor]] = None,
-        rotary_freqs_cis_cross: Union[torch.Tensor, Tuple[torch.Tensor]] = None,
-        temb: torch.FloatTensor = None,
-    ):
-
-        N = hidden_states.shape[0]
-
-        # step 1: AdaLN single
-        if self.use_adaln_single:
-            shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
-                comfy.model_management.cast_to(self.scale_shift_table[None], dtype=temb.dtype, device=temb.device) + temb.reshape(N, 6, -1)
-            ).chunk(6, dim=1)
-
-        norm_hidden_states = self.norm1(hidden_states)
-        if self.use_adaln_single:
-            norm_hidden_states = norm_hidden_states * (1 + scale_msa) + shift_msa
-
-        # step 2: attention
-        if not self.add_cross_attention:
-            attn_output, encoder_hidden_states = self.attn(
-                hidden_states=norm_hidden_states,
-                attention_mask=attention_mask,
-                encoder_hidden_states=encoder_hidden_states,
-                encoder_attention_mask=encoder_attention_mask,
-                rotary_freqs_cis=rotary_freqs_cis,
-                rotary_freqs_cis_cross=rotary_freqs_cis_cross,
-            )
-        else:
-            attn_output, _ = self.attn(
-                hidden_states=norm_hidden_states,
-                attention_mask=attention_mask,
-                encoder_hidden_states=None,
-                encoder_attention_mask=None,
-                rotary_freqs_cis=rotary_freqs_cis,
-                rotary_freqs_cis_cross=None,
-            )
-
-        if self.use_adaln_single:
-            attn_output = gate_msa * attn_output
-        hidden_states = attn_output + hidden_states
-
-        if self.add_cross_attention:
-            attn_output = self.cross_attn(
-                hidden_states=hidden_states,
-                attention_mask=attention_mask,
-                encoder_hidden_states=encoder_hidden_states,
-                encoder_attention_mask=encoder_attention_mask,
-                rotary_freqs_cis=rotary_freqs_cis,
-                rotary_freqs_cis_cross=rotary_freqs_cis_cross,
-            )
-            hidden_states = attn_output + hidden_states
-
-        # step 3: add norm
-        norm_hidden_states = self.norm2(hidden_states)
-        if self.use_adaln_single:
-            norm_hidden_states = norm_hidden_states * (1 + scale_mlp) + shift_mlp
-
-        # step 4: feed forward
-        ff_output = self.ff(norm_hidden_states)
-        if self.use_adaln_single:
-            ff_output = gate_mlp * ff_output
-
-        hidden_states = hidden_states + ff_output
-
-        return hidden_states

comfy/ldm/ace/lyric_encoder.py

@@ -1,1067 +0,0 @@
-# Original from: https://github.com/ace-step/ACE-Step/blob/main/models/lyrics_utils/lyric_encoder.py
-from typing import Optional, Tuple, Union
-import math
-import torch
-from torch import nn
-
-import comfy.model_management
-
-class ConvolutionModule(nn.Module):
-    """ConvolutionModule in Conformer model."""
-
-    def __init__(self,
-                 channels: int,
-                 kernel_size: int = 15,
-                 activation: nn.Module = nn.ReLU(),
-                 norm: str = "batch_norm",
-                 causal: bool = False,
-                 bias: bool = True,
-                 dtype=None, device=None, operations=None):
-        """Construct an ConvolutionModule object.
-        Args:
-            channels (int): The number of channels of conv layers.
-            kernel_size (int): Kernel size of conv layers.
-            causal (int): Whether use causal convolution or not
-        """
-        super().__init__()
-
-        self.pointwise_conv1 = operations.Conv1d(
-            channels,
-            2 * channels,
-            kernel_size=1,
-            stride=1,
-            padding=0,
-            bias=bias,
-            dtype=dtype, device=device
-        )
-        # self.lorder is used to distinguish if it's a causal convolution,
-        # if self.lorder > 0: it's a causal convolution, the input will be
-        #    padded with self.lorder frames on the left in forward.
-        # else: it's a symmetrical convolution
-        if causal:
-            padding = 0
-            self.lorder = kernel_size - 1
-        else:
-            # kernel_size should be an odd number for none causal convolution
-            assert (kernel_size - 1) % 2 == 0
-            padding = (kernel_size - 1) // 2
-            self.lorder = 0
-        self.depthwise_conv = operations.Conv1d(
-            channels,
-            channels,
-            kernel_size,
-            stride=1,
-            padding=padding,
-            groups=channels,
-            bias=bias,
-            dtype=dtype, device=device
-        )
-
-        assert norm in ['batch_norm', 'layer_norm']
-        if norm == "batch_norm":
-            self.use_layer_norm = False
-            self.norm = nn.BatchNorm1d(channels)
-        else:
-            self.use_layer_norm = True
-            self.norm = operations.LayerNorm(channels, dtype=dtype, device=device)
-
-        self.pointwise_conv2 = operations.Conv1d(
-            channels,
-            channels,
-            kernel_size=1,
-            stride=1,
-            padding=0,
-            bias=bias,
-            dtype=dtype, device=device
-        )
-        self.activation = activation
-
-    def forward(
-        self,
-        x: torch.Tensor,
-        mask_pad: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool),
-        cache: torch.Tensor = torch.zeros((0, 0, 0)),
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
-        """Compute convolution module.
-        Args:
-            x (torch.Tensor): Input tensor (#batch, time, channels).
-            mask_pad (torch.Tensor): used for batch padding (#batch, 1, time),
-                (0, 0, 0) means fake mask.
-            cache (torch.Tensor): left context cache, it is only
-                used in causal convolution (#batch, channels, cache_t),
-                (0, 0, 0) meas fake cache.
-        Returns:
-            torch.Tensor: Output tensor (#batch, time, channels).
-        """
-        # exchange the temporal dimension and the feature dimension
-        x = x.transpose(1, 2)  # (#batch, channels, time)
-
-        # mask batch padding
-        if mask_pad.size(2) > 0:  # time > 0
-            x.masked_fill_(~mask_pad, 0.0)
-
-        if self.lorder > 0:
-            if cache.size(2) == 0:  # cache_t == 0
-                x = nn.functional.pad(x, (self.lorder, 0), 'constant', 0.0)
-            else:
-                assert cache.size(0) == x.size(0)  # equal batch
-                assert cache.size(1) == x.size(1)  # equal channel
-                x = torch.cat((cache, x), dim=2)
-            assert (x.size(2) > self.lorder)
-            new_cache = x[:, :, -self.lorder:]
-        else:
-            # It's better we just return None if no cache is required,
-            # However, for JIT export, here we just fake one tensor instead of
-            # None.
-            new_cache = torch.zeros((0, 0, 0), dtype=x.dtype, device=x.device)
-
-        # GLU mechanism
-        x = self.pointwise_conv1(x)  # (batch, 2*channel, dim)
-        x = nn.functional.glu(x, dim=1)  # (batch, channel, dim)
-
-        # 1D Depthwise Conv
-        x = self.depthwise_conv(x)
-        if self.use_layer_norm:
-            x = x.transpose(1, 2)
-        x = self.activation(self.norm(x))
-        if self.use_layer_norm:
-            x = x.transpose(1, 2)
-        x = self.pointwise_conv2(x)
-        # mask batch padding
-        if mask_pad.size(2) > 0:  # time > 0
-            x.masked_fill_(~mask_pad, 0.0)
-
-        return x.transpose(1, 2), new_cache
-
-class PositionwiseFeedForward(torch.nn.Module):
-    """Positionwise feed forward layer.
-
-    FeedForward are appied on each position of the sequence.
-    The output dim is same with the input dim.
-
-    Args:
-        idim (int): Input dimenstion.
-        hidden_units (int): The number of hidden units.
-        dropout_rate (float): Dropout rate.
-        activation (torch.nn.Module): Activation function
-    """
-
-    def __init__(
-            self,
-            idim: int,
-            hidden_units: int,
-            dropout_rate: float,
-            activation: torch.nn.Module = torch.nn.ReLU(),
-            dtype=None, device=None, operations=None
-    ):
-        """Construct a PositionwiseFeedForward object."""
-        super(PositionwiseFeedForward, self).__init__()
-        self.w_1 = operations.Linear(idim, hidden_units, dtype=dtype, device=device)
-        self.activation = activation
-        self.dropout = torch.nn.Dropout(dropout_rate)
-        self.w_2 = operations.Linear(hidden_units, idim, dtype=dtype, device=device)
-
-    def forward(self, xs: torch.Tensor) -> torch.Tensor:
-        """Forward function.
-
-        Args:
-            xs: input tensor (B, L, D)
-        Returns:
-            output tensor, (B, L, D)
-        """
-        return self.w_2(self.dropout(self.activation(self.w_1(xs))))
-
-class Swish(torch.nn.Module):
-    """Construct an Swish object."""
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        """Return Swish activation function."""
-        return x * torch.sigmoid(x)
-
-class MultiHeadedAttention(nn.Module):
-    """Multi-Head Attention layer.
-
-    Args:
-        n_head (int): The number of heads.
-        n_feat (int): The number of features.
-        dropout_rate (float): Dropout rate.
-
-    """
-
-    def __init__(self,
-                 n_head: int,
-                 n_feat: int,
-                 dropout_rate: float,
-                 key_bias: bool = True,
-                 dtype=None, device=None, operations=None):
-        """Construct an MultiHeadedAttention object."""
-        super().__init__()
-        assert n_feat % n_head == 0
-        # We assume d_v always equals d_k
-        self.d_k = n_feat // n_head
-        self.h = n_head
-        self.linear_q = operations.Linear(n_feat, n_feat, dtype=dtype, device=device)
-        self.linear_k = operations.Linear(n_feat, n_feat, bias=key_bias, dtype=dtype, device=device)
-        self.linear_v = operations.Linear(n_feat, n_feat, dtype=dtype, device=device)
-        self.linear_out = operations.Linear(n_feat, n_feat, dtype=dtype, device=device)
-        self.dropout = nn.Dropout(p=dropout_rate)
-
-    def forward_qkv(
-        self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor
-    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        """Transform query, key and value.
-
-        Args:
-            query (torch.Tensor): Query tensor (#batch, time1, size).
-            key (torch.Tensor): Key tensor (#batch, time2, size).
-            value (torch.Tensor): Value tensor (#batch, time2, size).
-
-        Returns:
-            torch.Tensor: Transformed query tensor, size
-                (#batch, n_head, time1, d_k).
-            torch.Tensor: Transformed key tensor, size
-                (#batch, n_head, time2, d_k).
-            torch.Tensor: Transformed value tensor, size
-                (#batch, n_head, time2, d_k).
-
-        """
-        n_batch = query.size(0)
-        q = self.linear_q(query).view(n_batch, -1, self.h, self.d_k)
-        k = self.linear_k(key).view(n_batch, -1, self.h, self.d_k)
-        v = self.linear_v(value).view(n_batch, -1, self.h, self.d_k)
-        q = q.transpose(1, 2)  # (batch, head, time1, d_k)
-        k = k.transpose(1, 2)  # (batch, head, time2, d_k)
-        v = v.transpose(1, 2)  # (batch, head, time2, d_k)
-        return q, k, v
-
-    def forward_attention(
-        self,
-        value: torch.Tensor,
-        scores: torch.Tensor,
-        mask: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool)
-    ) -> torch.Tensor:
-        """Compute attention context vector.
-
-        Args:
-            value (torch.Tensor): Transformed value, size
-                (#batch, n_head, time2, d_k).
-            scores (torch.Tensor): Attention score, size
-                (#batch, n_head, time1, time2).
-            mask (torch.Tensor): Mask, size (#batch, 1, time2) or
-                (#batch, time1, time2), (0, 0, 0) means fake mask.
-
-        Returns:
-            torch.Tensor: Transformed value (#batch, time1, d_model)
-                weighted by the attention score (#batch, time1, time2).
-
-        """
-        n_batch = value.size(0)
-
-        if mask is not None and mask.size(2) > 0:  # time2 > 0
-            mask = mask.unsqueeze(1).eq(0)  # (batch, 1, *, time2)
-            # For last chunk, time2 might be larger than scores.size(-1)
-            mask = mask[:, :, :, :scores.size(-1)]  # (batch, 1, *, time2)
-            scores = scores.masked_fill(mask, -float('inf'))
-            attn = torch.softmax(scores, dim=-1).masked_fill(
-                mask, 0.0)  # (batch, head, time1, time2)
-
-        else:
-            attn = torch.softmax(scores, dim=-1)  # (batch, head, time1, time2)
-
-        p_attn = self.dropout(attn)
-        x = torch.matmul(p_attn, value)  # (batch, head, time1, d_k)
-        x = (x.transpose(1, 2).contiguous().view(n_batch, -1,
-                                                 self.h * self.d_k)
-             )  # (batch, time1, d_model)
-
-        return self.linear_out(x)  # (batch, time1, d_model)
-
-    def forward(
-        self,
-        query: torch.Tensor,
-        key: torch.Tensor,
-        value: torch.Tensor,
-        mask: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool),
-        pos_emb: torch.Tensor = torch.empty(0),
-        cache: torch.Tensor = torch.zeros((0, 0, 0, 0))
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
-        """Compute scaled dot product attention.
-
-        Args:
-            query (torch.Tensor): Query tensor (#batch, time1, size).
-            key (torch.Tensor): Key tensor (#batch, time2, size).
-            value (torch.Tensor): Value tensor (#batch, time2, size).
-            mask (torch.Tensor): Mask tensor (#batch, 1, time2) or
-                (#batch, time1, time2).
-                1.When applying cross attention between decoder and encoder,
-                the batch padding mask for input is in (#batch, 1, T) shape.
-                2.When applying self attention of encoder,
-                the mask is in (#batch, T, T)  shape.
-                3.When applying self attention of decoder,
-                the mask is in (#batch, L, L)  shape.
-                4.If the different position in decoder see different block
-                of the encoder, such as Mocha, the passed in mask could be
-                in (#batch, L, T) shape. But there is no such case in current
-                CosyVoice.
-            cache (torch.Tensor): Cache tensor (1, head, cache_t, d_k * 2),
-                where `cache_t == chunk_size * num_decoding_left_chunks`
-                and `head * d_k == size`
-
-
-        Returns:
-            torch.Tensor: Output tensor (#batch, time1, d_model).
-            torch.Tensor: Cache tensor (1, head, cache_t + time1, d_k * 2)
-                where `cache_t == chunk_size * num_decoding_left_chunks`
-                and `head * d_k == size`
-
-        """
-        q, k, v = self.forward_qkv(query, key, value)
-        if cache.size(0) > 0:
-            key_cache, value_cache = torch.split(cache,
-                                                 cache.size(-1) // 2,
-                                                 dim=-1)
-            k = torch.cat([key_cache, k], dim=2)
-            v = torch.cat([value_cache, v], dim=2)
-        new_cache = torch.cat((k, v), dim=-1)
-
-        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)
-        return self.forward_attention(v, scores, mask), new_cache
-
-
-class RelPositionMultiHeadedAttention(MultiHeadedAttention):
-    """Multi-Head Attention layer with relative position encoding.
-    Paper: https://arxiv.org/abs/1901.02860
-    Args:
-        n_head (int): The number of heads.
-        n_feat (int): The number of features.
-        dropout_rate (float): Dropout rate.
-    """
-
-    def __init__(self,
-                 n_head: int,
-                 n_feat: int,
-                 dropout_rate: float,
-                 key_bias: bool = True,
-                 dtype=None, device=None, operations=None):
-        """Construct an RelPositionMultiHeadedAttention object."""
-        super().__init__(n_head, n_feat, dropout_rate, key_bias, dtype=dtype, device=device, operations=operations)
-        # linear transformation for positional encoding
-        self.linear_pos = operations.Linear(n_feat, n_feat, bias=False, dtype=dtype, device=device)
-        # these two learnable bias are used in matrix c and matrix d
-        # as described in https://arxiv.org/abs/1901.02860 Section 3.3
-        self.pos_bias_u = nn.Parameter(torch.empty(self.h, self.d_k, dtype=dtype, device=device))
-        self.pos_bias_v = nn.Parameter(torch.empty(self.h, self.d_k, dtype=dtype, device=device))
-        # torch.nn.init.xavier_uniform_(self.pos_bias_u)
-        # torch.nn.init.xavier_uniform_(self.pos_bias_v)
-
-    def rel_shift(self, x: torch.Tensor) -> torch.Tensor:
-        """Compute relative positional encoding.
-
-        Args:
-            x (torch.Tensor): Input tensor (batch, head, time1, 2*time1-1).
-            time1 means the length of query vector.
-
-        Returns:
-            torch.Tensor: Output tensor.
-
-        """
-        zero_pad = torch.zeros((x.size()[0], x.size()[1], x.size()[2], 1),
-                               device=x.device,
-                               dtype=x.dtype)
-        x_padded = torch.cat([zero_pad, x], dim=-1)
-
-        x_padded = x_padded.view(x.size()[0],
-                                 x.size()[1],
-                                 x.size(3) + 1, x.size(2))
-        x = x_padded[:, :, 1:].view_as(x)[
-            :, :, :, : x.size(-1) // 2 + 1
-        ]  # only keep the positions from 0 to time2
-        return x
-
-    def forward(
-        self,
-        query: torch.Tensor,
-        key: torch.Tensor,
-        value: torch.Tensor,
-        mask: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool),
-        pos_emb: torch.Tensor = torch.empty(0),
-        cache: torch.Tensor = torch.zeros((0, 0, 0, 0))
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
-        """Compute 'Scaled Dot Product Attention' with rel. positional encoding.
-        Args:
-            query (torch.Tensor): Query tensor (#batch, time1, size).
-            key (torch.Tensor): Key tensor (#batch, time2, size).
-            value (torch.Tensor): Value tensor (#batch, time2, size).
-            mask (torch.Tensor): Mask tensor (#batch, 1, time2) or
-                (#batch, time1, time2), (0, 0, 0) means fake mask.
-            pos_emb (torch.Tensor): Positional embedding tensor
-                (#batch, time2, size).
-            cache (torch.Tensor): Cache tensor (1, head, cache_t, d_k * 2),
-                where `cache_t == chunk_size * num_decoding_left_chunks`
-                and `head * d_k == size`
-        Returns:
-            torch.Tensor: Output tensor (#batch, time1, d_model).
-            torch.Tensor: Cache tensor (1, head, cache_t + time1, d_k * 2)
-                where `cache_t == chunk_size * num_decoding_left_chunks`
-                and `head * d_k == size`
-        """
-        q, k, v = self.forward_qkv(query, key, value)
-        q = q.transpose(1, 2)  # (batch, time1, head, d_k)
-
-        if cache.size(0) > 0:
-            key_cache, value_cache = torch.split(cache,
-                                                 cache.size(-1) // 2,
-                                                 dim=-1)
-            k = torch.cat([key_cache, k], dim=2)
-            v = torch.cat([value_cache, v], dim=2)
-        # NOTE(xcsong): We do cache slicing in encoder.forward_chunk, since it's
-        #   non-trivial to calculate `next_cache_start` here.
-        new_cache = torch.cat((k, v), dim=-1)
-
-        n_batch_pos = pos_emb.size(0)
-        p = self.linear_pos(pos_emb).view(n_batch_pos, -1, self.h, self.d_k)
-        p = p.transpose(1, 2)  # (batch, head, time1, d_k)
-
-        # (batch, head, time1, d_k)
-        q_with_bias_u = (q + comfy.model_management.cast_to(self.pos_bias_u, dtype=q.dtype, device=q.device)).transpose(1, 2)
-        # (batch, head, time1, d_k)
-        q_with_bias_v = (q + comfy.model_management.cast_to(self.pos_bias_v, dtype=q.dtype, device=q.device)).transpose(1, 2)
-
-        # compute attention score
-        # first compute matrix a and matrix c
-        # as described in https://arxiv.org/abs/1901.02860 Section 3.3
-        # (batch, head, time1, time2)
-        matrix_ac = torch.matmul(q_with_bias_u, k.transpose(-2, -1))
-
-        # compute matrix b and matrix d
-        # (batch, head, time1, time2)
-        matrix_bd = torch.matmul(q_with_bias_v, p.transpose(-2, -1))
-        # NOTE(Xiang Lyu): Keep rel_shift since espnet rel_pos_emb is used
-        if matrix_ac.shape != matrix_bd.shape:
-            matrix_bd = self.rel_shift(matrix_bd)
-
-        scores = (matrix_ac + matrix_bd) / math.sqrt(
-            self.d_k)  # (batch, head, time1, time2)
-
-        return self.forward_attention(v, scores, mask), new_cache
-
-
-
-def subsequent_mask(
-        size: int,
-        device: torch.device = torch.device("cpu"),
-) -> torch.Tensor:
-    """Create mask for subsequent steps (size, size).
-
-    This mask is used only in decoder which works in an auto-regressive mode.
-    This means the current step could only do attention with its left steps.
-
-    In encoder, fully attention is used when streaming is not necessary and
-    the sequence is not long. In this  case, no attention mask is needed.
-
-    When streaming is need, chunk-based attention is used in encoder. See
-    subsequent_chunk_mask for the chunk-based attention mask.
-
-    Args:
-        size (int): size of mask
-        str device (str): "cpu" or "cuda" or torch.Tensor.device
-        dtype (torch.device): result dtype
-
-    Returns:
-        torch.Tensor: mask
-
-    Examples:
-        >>> subsequent_mask(3)
-        [[1, 0, 0],
-         [1, 1, 0],
-         [1, 1, 1]]
-    """
-    arange = torch.arange(size, device=device)
-    mask = arange.expand(size, size)
-    arange = arange.unsqueeze(-1)
-    mask = mask <= arange
-    return mask
-
-
-def subsequent_chunk_mask(
-        size: int,
-        chunk_size: int,
-        num_left_chunks: int = -1,
-        device: torch.device = torch.device("cpu"),
-    ) -> torch.Tensor:
-    """Create mask for subsequent steps (size, size) with chunk size,
-       this is for streaming encoder
-
-    Args:
-        size (int): size of mask
-        chunk_size (int): size of chunk
-        num_left_chunks (int): number of left chunks
-            <0: use full chunk
-            >=0: use num_left_chunks
-        device (torch.device): "cpu" or "cuda" or torch.Tensor.device
-
-    Returns:
-        torch.Tensor: mask
-
-    Examples:
-        >>> subsequent_chunk_mask(4, 2)
-        [[1, 1, 0, 0],
-         [1, 1, 0, 0],
-         [1, 1, 1, 1],
-         [1, 1, 1, 1]]
-    """
-    ret = torch.zeros(size, size, device=device, dtype=torch.bool)
-    for i in range(size):
-        if num_left_chunks < 0:
-            start = 0
-        else:
-            start = max((i // chunk_size - num_left_chunks) * chunk_size, 0)
-        ending = min((i // chunk_size + 1) * chunk_size, size)
-        ret[i, start:ending] = True
-    return ret
-
-def add_optional_chunk_mask(xs: torch.Tensor,
-                            masks: torch.Tensor,
-                            use_dynamic_chunk: bool,
-                            use_dynamic_left_chunk: bool,
-                            decoding_chunk_size: int,
-                            static_chunk_size: int,
-                            num_decoding_left_chunks: int,
-                            enable_full_context: bool = True):
-    """ Apply optional mask for encoder.
-
-    Args:
-        xs (torch.Tensor): padded input, (B, L, D), L for max length
-        mask (torch.Tensor): mask for xs, (B, 1, L)
-        use_dynamic_chunk (bool): whether to use dynamic chunk or not
-        use_dynamic_left_chunk (bool): whether to use dynamic left chunk for
-            training.
-        decoding_chunk_size (int): decoding chunk size for dynamic chunk, it's
-            0: default for training, use random dynamic chunk.
-            <0: for decoding, use full chunk.
-            >0: for decoding, use fixed chunk size as set.
-        static_chunk_size (int): chunk size for static chunk training/decoding
-            if it's greater than 0, if use_dynamic_chunk is true,
-            this parameter will be ignored
-        num_decoding_left_chunks: number of left chunks, this is for decoding,
-            the chunk size is decoding_chunk_size.
-            >=0: use num_decoding_left_chunks
-            <0: use all left chunks
-        enable_full_context (bool):
-            True: chunk size is either [1, 25] or full context(max_len)
-            False: chunk size ~ U[1, 25]
-
-    Returns:
-        torch.Tensor: chunk mask of the input xs.
-    """
-    # Whether to use chunk mask or not
-    if use_dynamic_chunk:
-        max_len = xs.size(1)
-        if decoding_chunk_size < 0:
-            chunk_size = max_len
-            num_left_chunks = -1
-        elif decoding_chunk_size > 0:
-            chunk_size = decoding_chunk_size
-            num_left_chunks = num_decoding_left_chunks
-        else:
-            # chunk size is either [1, 25] or full context(max_len).
-            # Since we use 4 times subsampling and allow up to 1s(100 frames)
-            # delay, the maximum frame is 100 / 4 = 25.
-            chunk_size = torch.randint(1, max_len, (1, )).item()
-            num_left_chunks = -1
-            if chunk_size > max_len // 2 and enable_full_context:
-                chunk_size = max_len
-            else:
-                chunk_size = chunk_size % 25 + 1
-                if use_dynamic_left_chunk:
-                    max_left_chunks = (max_len - 1) // chunk_size
-                    num_left_chunks = torch.randint(0, max_left_chunks,
-                                                    (1, )).item()
-        chunk_masks = subsequent_chunk_mask(xs.size(1), chunk_size,
-                                            num_left_chunks,
-                                            xs.device)  # (L, L)
-        chunk_masks = chunk_masks.unsqueeze(0)  # (1, L, L)
-        chunk_masks = masks & chunk_masks  # (B, L, L)
-    elif static_chunk_size > 0:
-        num_left_chunks = num_decoding_left_chunks
-        chunk_masks = subsequent_chunk_mask(xs.size(1), static_chunk_size,
-                                            num_left_chunks,
-                                            xs.device)  # (L, L)
-        chunk_masks = chunk_masks.unsqueeze(0)  # (1, L, L)
-        chunk_masks = masks & chunk_masks  # (B, L, L)
-    else:
-        chunk_masks = masks
-    return chunk_masks
-
-
-class ConformerEncoderLayer(nn.Module):
-    """Encoder layer module.
-    Args:
-        size (int): Input dimension.
-        self_attn (torch.nn.Module): Self-attention module instance.
-            `MultiHeadedAttention` or `RelPositionMultiHeadedAttention`
-            instance can be used as the argument.
-        feed_forward (torch.nn.Module): Feed-forward module instance.
-            `PositionwiseFeedForward` instance can be used as the argument.
-        feed_forward_macaron (torch.nn.Module): Additional feed-forward module
-             instance.
-            `PositionwiseFeedForward` instance can be used as the argument.
-        conv_module (torch.nn.Module): Convolution module instance.
-            `ConvlutionModule` instance can be used as the argument.
-        dropout_rate (float): Dropout rate.
-        normalize_before (bool):
-            True: use layer_norm before each sub-block.
-            False: use layer_norm after each sub-block.
-    """
-
-    def __init__(
-        self,
-        size: int,
-        self_attn: torch.nn.Module,
-        feed_forward: Optional[nn.Module] = None,
-        feed_forward_macaron: Optional[nn.Module] = None,
-        conv_module: Optional[nn.Module] = None,
-        dropout_rate: float = 0.1,
-        normalize_before: bool = True,
-        dtype=None, device=None, operations=None
-    ):
-        """Construct an EncoderLayer object."""
-        super().__init__()
-        self.self_attn = self_attn
-        self.feed_forward = feed_forward
-        self.feed_forward_macaron = feed_forward_macaron
-        self.conv_module = conv_module
-        self.norm_ff = operations.LayerNorm(size, eps=1e-5, dtype=dtype, device=device)  # for the FNN module
-        self.norm_mha = operations.LayerNorm(size, eps=1e-5, dtype=dtype, device=device)  # for the MHA module
-        if feed_forward_macaron is not None:
-            self.norm_ff_macaron = operations.LayerNorm(size, eps=1e-5, dtype=dtype, device=device)
-            self.ff_scale = 0.5
-        else:
-            self.ff_scale = 1.0
-        if self.conv_module is not None:
-            self.norm_conv = operations.LayerNorm(size, eps=1e-5, dtype=dtype, device=device)  # for the CNN module
-            self.norm_final = operations.LayerNorm(
-                size, eps=1e-5, dtype=dtype, device=device)  # for the final output of the block
-        self.dropout = nn.Dropout(dropout_rate)
-        self.size = size
-        self.normalize_before = normalize_before
-
-    def forward(
-        self,
-        x: torch.Tensor,
-        mask: torch.Tensor,
-        pos_emb: torch.Tensor,
-        mask_pad: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool),
-        att_cache: torch.Tensor = torch.zeros((0, 0, 0, 0)),
-        cnn_cache: torch.Tensor = torch.zeros((0, 0, 0, 0)),
-    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
-        """Compute encoded features.
-
-        Args:
-            x (torch.Tensor): (#batch, time, size)
-            mask (torch.Tensor): Mask tensor for the input (#batch, time，time),
-                (0, 0, 0) means fake mask.
-            pos_emb (torch.Tensor): positional encoding, must not be None
-                for ConformerEncoderLayer.
-            mask_pad (torch.Tensor): batch padding mask used for conv module.
-                (#batch, 1，time), (0, 0, 0) means fake mask.
-            att_cache (torch.Tensor): Cache tensor of the KEY & VALUE
-                (#batch=1, head, cache_t1, d_k * 2), head * d_k == size.
-            cnn_cache (torch.Tensor): Convolution cache in conformer layer
-                (#batch=1, size, cache_t2)
-        Returns:
-            torch.Tensor: Output tensor (#batch, time, size).
-            torch.Tensor: Mask tensor (#batch, time, time).
-            torch.Tensor: att_cache tensor,
-                (#batch=1, head, cache_t1 + time, d_k * 2).
-            torch.Tensor: cnn_cahce tensor (#batch, size, cache_t2).
-        """
-
-        # whether to use macaron style
-        if self.feed_forward_macaron is not None:
-            residual = x
-            if self.normalize_before:
-                x = self.norm_ff_macaron(x)
-            x = residual + self.ff_scale * self.dropout(
-                self.feed_forward_macaron(x))
-            if not self.normalize_before:
-                x = self.norm_ff_macaron(x)
-
-        # multi-headed self-attention module
-        residual = x
-        if self.normalize_before:
-            x = self.norm_mha(x)
-        x_att, new_att_cache = self.self_attn(x, x, x, mask, pos_emb,
-                                              att_cache)
-        x = residual + self.dropout(x_att)
-        if not self.normalize_before:
-            x = self.norm_mha(x)
-
-        # convolution module
-        # Fake new cnn cache here, and then change it in conv_module
-        new_cnn_cache = torch.zeros((0, 0, 0), dtype=x.dtype, device=x.device)
-        if self.conv_module is not None:
-            residual = x
-            if self.normalize_before:
-                x = self.norm_conv(x)
-            x, new_cnn_cache = self.conv_module(x, mask_pad, cnn_cache)
-            x = residual + self.dropout(x)
-
-            if not self.normalize_before:
-                x = self.norm_conv(x)
-
-        # feed forward module
-        residual = x
-        if self.normalize_before:
-            x = self.norm_ff(x)
-
-        x = residual + self.ff_scale * self.dropout(self.feed_forward(x))
-        if not self.normalize_before:
-            x = self.norm_ff(x)
-
-        if self.conv_module is not None:
-            x = self.norm_final(x)
-
-        return x, mask, new_att_cache, new_cnn_cache
-
-
-
-class EspnetRelPositionalEncoding(torch.nn.Module):
-    """Relative positional encoding module (new implementation).
-
-    Details can be found in https://github.com/espnet/espnet/pull/2816.
-
-    See : Appendix B in https://arxiv.org/abs/1901.02860
-
-    Args:
-        d_model (int): Embedding dimension.
-        dropout_rate (float): Dropout rate.
-        max_len (int): Maximum input length.
-
-    """
-
-    def __init__(self, d_model: int, dropout_rate: float, max_len: int = 5000):
-        """Construct an PositionalEncoding object."""
-        super(EspnetRelPositionalEncoding, self).__init__()
-        self.d_model = d_model
-        self.xscale = math.sqrt(self.d_model)
-        self.dropout = torch.nn.Dropout(p=dropout_rate)
-        self.pe = None
-        self.extend_pe(torch.tensor(0.0).expand(1, max_len))
-
-    def extend_pe(self, x: torch.Tensor):
-        """Reset the positional encodings."""
-        if self.pe is not None:
-            # self.pe contains both positive and negative parts
-            # the length of self.pe is 2 * input_len - 1
-            if self.pe.size(1) >= x.size(1) * 2 - 1:
-                if self.pe.dtype != x.dtype or self.pe.device != x.device:
-                    self.pe = self.pe.to(dtype=x.dtype, device=x.device)
-                return
-        # Suppose `i` means to the position of query vecotr and `j` means the
-        # position of key vector. We use position relative positions when keys
-        # are to the left (i>j) and negative relative positions otherwise (i<j).
-        pe_positive = torch.zeros(x.size(1), self.d_model)
-        pe_negative = torch.zeros(x.size(1), self.d_model)
-        position = torch.arange(0, x.size(1), dtype=torch.float32).unsqueeze(1)
-        div_term = torch.exp(
-            torch.arange(0, self.d_model, 2, dtype=torch.float32)
-            * -(math.log(10000.0) / self.d_model)
-        )
-        pe_positive[:, 0::2] = torch.sin(position * div_term)
-        pe_positive[:, 1::2] = torch.cos(position * div_term)
-        pe_negative[:, 0::2] = torch.sin(-1 * position * div_term)
-        pe_negative[:, 1::2] = torch.cos(-1 * position * div_term)
-
-        # Reserve the order of positive indices and concat both positive and
-        # negative indices. This is used to support the shifting trick
-        # as in https://arxiv.org/abs/1901.02860
-        pe_positive = torch.flip(pe_positive, [0]).unsqueeze(0)
-        pe_negative = pe_negative[1:].unsqueeze(0)
-        pe = torch.cat([pe_positive, pe_negative], dim=1)
-        self.pe = pe.to(device=x.device, dtype=x.dtype)
-
-    def forward(self, x: torch.Tensor, offset: Union[int, torch.Tensor] = 0) \
-            -> Tuple[torch.Tensor, torch.Tensor]:
-        """Add positional encoding.
-
-        Args:
-            x (torch.Tensor): Input tensor (batch, time, `*`).
-
-        Returns:
-            torch.Tensor: Encoded tensor (batch, time, `*`).
-
-        """
-        self.extend_pe(x)
-        x = x * self.xscale
-        pos_emb = self.position_encoding(size=x.size(1), offset=offset)
-        return self.dropout(x), self.dropout(pos_emb)
-
-    def position_encoding(self,
-                          offset: Union[int, torch.Tensor],
-                          size: int) -> torch.Tensor:
-        """ For getting encoding in a streaming fashion
-
-        Attention!!!!!
-        we apply dropout only once at the whole utterance level in a none
-        streaming way, but will call this function several times with
-        increasing input size in a streaming scenario, so the dropout will
-        be applied several times.
-
-        Args:
-            offset (int or torch.tensor): start offset
-            size (int): required size of position encoding
-
-        Returns:
-            torch.Tensor: Corresponding encoding
-        """
-        pos_emb = self.pe[
-            :,
-            self.pe.size(1) // 2 - size + 1: self.pe.size(1) // 2 + size,
-        ]
-        return pos_emb
-
-
-
-class LinearEmbed(torch.nn.Module):
-    """Linear transform the input without subsampling
-
-    Args:
-        idim (int): Input dimension.
-        odim (int): Output dimension.
-        dropout_rate (float): Dropout rate.
-
-    """
-
-    def __init__(self, idim: int, odim: int, dropout_rate: float,
-                 pos_enc_class: torch.nn.Module, dtype=None, device=None, operations=None):
-        """Construct an linear object."""
-        super().__init__()
-        self.out = torch.nn.Sequential(
-            operations.Linear(idim, odim, dtype=dtype, device=device),
-            operations.LayerNorm(odim, eps=1e-5, dtype=dtype, device=device),
-            torch.nn.Dropout(dropout_rate),
-        )
-        self.pos_enc = pos_enc_class #rel_pos_espnet
-
-    def position_encoding(self, offset: Union[int, torch.Tensor],
-                          size: int) -> torch.Tensor:
-        return self.pos_enc.position_encoding(offset, size)
-
-    def forward(
-        self,
-        x: torch.Tensor,
-        offset: Union[int, torch.Tensor] = 0
-    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        """Input x.
-
-        Args:
-            x (torch.Tensor): Input tensor (#batch, time, idim).
-            x_mask (torch.Tensor): Input mask (#batch, 1, time).
-
-        Returns:
-            torch.Tensor: linear input tensor (#batch, time', odim),
-                where time' = time .
-            torch.Tensor: linear input mask (#batch, 1, time'),
-                where time' = time .
-
-        """
-        x = self.out(x)
-        x, pos_emb = self.pos_enc(x, offset)
-        return x, pos_emb
-
-
-ATTENTION_CLASSES = {
-    "selfattn": MultiHeadedAttention,
-    "rel_selfattn": RelPositionMultiHeadedAttention,
-}
-
-ACTIVATION_CLASSES = {
-    "hardtanh": torch.nn.Hardtanh,
-    "tanh": torch.nn.Tanh,
-    "relu": torch.nn.ReLU,
-    "selu": torch.nn.SELU,
-    "swish": getattr(torch.nn, "SiLU", Swish),
-    "gelu": torch.nn.GELU,
-}
-
-
-def make_pad_mask(lengths: torch.Tensor, max_len: int = 0) -> torch.Tensor:
-    """Make mask tensor containing indices of padded part.
-
-    See description of make_non_pad_mask.
-
-    Args:
-        lengths (torch.Tensor): Batch of lengths (B,).
-    Returns:
-        torch.Tensor: Mask tensor containing indices of padded part.
-
-    Examples:
-        >>> lengths = [5, 3, 2]
-        >>> make_pad_mask(lengths)
-        masks = [[0, 0, 0, 0 ,0],
-                 [0, 0, 0, 1, 1],
-                 [0, 0, 1, 1, 1]]
-    """
-    batch_size = lengths.size(0)
-    max_len = max_len if max_len > 0 else lengths.max().item()
-    seq_range = torch.arange(0,
-                             max_len,
-                             dtype=torch.int64,
-                             device=lengths.device)
-    seq_range_expand = seq_range.unsqueeze(0).expand(batch_size, max_len)
-    seq_length_expand = lengths.unsqueeze(-1)
-    mask = seq_range_expand >= seq_length_expand
-    return mask
-
-#https://github.com/FunAudioLLM/CosyVoice/blob/main/examples/magicdata-read/cosyvoice/conf/cosyvoice.yaml
-class ConformerEncoder(torch.nn.Module):
-    """Conformer encoder module."""
-
-    def __init__(
-        self,
-        input_size: int,
-        output_size: int = 1024,
-        attention_heads: int = 16,
-        linear_units: int = 4096,
-        num_blocks: int = 6,
-        dropout_rate: float = 0.1,
-        positional_dropout_rate: float = 0.1,
-        attention_dropout_rate: float = 0.0,
-        input_layer: str = 'linear',
-        pos_enc_layer_type: str = 'rel_pos_espnet',
-        normalize_before: bool = True,
-        static_chunk_size: int = 1, # 1: causal_mask; 0: full_mask
-        use_dynamic_chunk: bool = False,
-        use_dynamic_left_chunk: bool = False,
-        positionwise_conv_kernel_size: int = 1,
-        macaron_style: bool =False,
-        selfattention_layer_type: str = "rel_selfattn",
-        activation_type: str = "swish",
-        use_cnn_module: bool = False,
-        cnn_module_kernel: int = 15,
-        causal: bool = False,
-        cnn_module_norm: str = "batch_norm",
-        key_bias: bool = True,
-        dtype=None, device=None, operations=None
-    ):
-        """Construct ConformerEncoder
-
-        Args:
-            input_size to use_dynamic_chunk, see in BaseEncoder
-            positionwise_conv_kernel_size (int): Kernel size of positionwise
-                conv1d layer.
-            macaron_style (bool): Whether to use macaron style for
-                positionwise layer.
-            selfattention_layer_type (str): Encoder attention layer type,
-                the parameter has no effect now, it's just for configure
-                compatibility. #'rel_selfattn'
-            activation_type (str): Encoder activation function type.
-            use_cnn_module (bool): Whether to use convolution module.
-            cnn_module_kernel (int): Kernel size of convolution module.
-            causal (bool): whether to use causal convolution or not.
-            key_bias: whether use bias in attention.linear_k, False for whisper models.
-        """
-        super().__init__()
-        self.output_size = output_size
-        self.embed = LinearEmbed(input_size, output_size, dropout_rate,
-                                        EspnetRelPositionalEncoding(output_size, positional_dropout_rate), dtype=dtype, device=device, operations=operations)
-        self.normalize_before = normalize_before
-        self.after_norm = operations.LayerNorm(output_size, eps=1e-5, dtype=dtype, device=device)
-        self.use_dynamic_chunk = use_dynamic_chunk
-
-        self.static_chunk_size = static_chunk_size
-        self.use_dynamic_chunk = use_dynamic_chunk
-        self.use_dynamic_left_chunk = use_dynamic_left_chunk
-        activation = ACTIVATION_CLASSES[activation_type]()
-
-        # self-attention module definition
-        encoder_selfattn_layer_args = (
-            attention_heads,
-            output_size,
-            attention_dropout_rate,
-            key_bias,
-        )
-        # feed-forward module definition
-        positionwise_layer_args = (
-            output_size,
-            linear_units,
-            dropout_rate,
-            activation,
-        )
-        # convolution module definition
-        convolution_layer_args = (output_size, cnn_module_kernel, activation,
-                                  cnn_module_norm, causal)
-
-        self.encoders = torch.nn.ModuleList([
-            ConformerEncoderLayer(
-                output_size,
-                RelPositionMultiHeadedAttention(
-                    *encoder_selfattn_layer_args, dtype=dtype, device=device, operations=operations),
-                PositionwiseFeedForward(*positionwise_layer_args, dtype=dtype, device=device, operations=operations),
-                PositionwiseFeedForward(
-                    *positionwise_layer_args, dtype=dtype, device=device, operations=operations) if macaron_style else None,
-                ConvolutionModule(
-                    *convolution_layer_args, dtype=dtype, device=device, operations=operations) if use_cnn_module else None,
-                dropout_rate,
-                normalize_before, dtype=dtype, device=device, operations=operations
-            ) for _ in range(num_blocks)
-        ])
-
-    def forward_layers(self, xs: torch.Tensor, chunk_masks: torch.Tensor,
-        pos_emb: torch.Tensor,
-        mask_pad: torch.Tensor) -> torch.Tensor:
-        for layer in self.encoders:
-            xs, chunk_masks, _, _ = layer(xs, chunk_masks, pos_emb, mask_pad)
-        return xs
-
-    def forward(
-        self,
-        xs: torch.Tensor,
-        pad_mask: torch.Tensor,
-        decoding_chunk_size: int = 0,
-        num_decoding_left_chunks: int = -1,
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
-        """Embed positions in tensor.
-
-        Args:
-            xs: padded input tensor (B, T, D)
-            xs_lens: input length (B)
-            decoding_chunk_size: decoding chunk size for dynamic chunk
-                0: default for training, use random dynamic chunk.
-                <0: for decoding, use full chunk.
-                >0: for decoding, use fixed chunk size as set.
-            num_decoding_left_chunks: number of left chunks, this is for decoding,
-            the chunk size is decoding_chunk_size.
-                >=0: use num_decoding_left_chunks
-                <0: use all left chunks
-        Returns:
-            encoder output tensor xs, and subsampled masks
-            xs: padded output tensor (B, T' ~= T/subsample_rate, D)
-            masks: torch.Tensor batch padding mask after subsample
-                (B, 1, T' ~= T/subsample_rate)
-        NOTE(xcsong):
-            We pass the `__call__` method of the modules instead of `forward` to the
-            checkpointing API because `__call__` attaches all the hooks of the module.
-            https://discuss.pytorch.org/t/any-different-between-model-input-and-model-forward-input/3690/2
-        """
-        masks = None
-        if pad_mask is not None:
-            masks = pad_mask.to(torch.bool).unsqueeze(1)  # (B, 1, T)
-        xs, pos_emb = self.embed(xs)
-        mask_pad = masks  # (B, 1, T/subsample_rate)
-        chunk_masks = add_optional_chunk_mask(xs, masks,
-                                              self.use_dynamic_chunk,
-                                              self.use_dynamic_left_chunk,
-                                              decoding_chunk_size,
-                                              self.static_chunk_size,
-                                              num_decoding_left_chunks)
-
-        xs = self.forward_layers(xs, chunk_masks, pos_emb, mask_pad)
-        if self.normalize_before:
-            xs = self.after_norm(xs)
-        # Here we assume the mask is not changed in encoder layers, so just
-        # return the masks before encoder layers, and the masks will be used
-        # for cross attention with decoder later
-        return xs, masks
-

comfy/ldm/ace/model.py

@@ -1,385 +0,0 @@
-# Original from: https://github.com/ace-step/ACE-Step/blob/main/models/ace_step_transformer.py
-
-# Copyright 2024 The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-from typing import Optional, List, Union
-
-import torch
-from torch import nn
-
-import comfy.model_management
-
-from comfy.ldm.lightricks.model import TimestepEmbedding, Timesteps
-from .attention import LinearTransformerBlock, t2i_modulate
-from .lyric_encoder import ConformerEncoder as LyricEncoder
-
-
-def cross_norm(hidden_states, controlnet_input):
-    # input N x T x c
-    mean_hidden_states, std_hidden_states = hidden_states.mean(dim=(1,2), keepdim=True), hidden_states.std(dim=(1,2), keepdim=True)
-    mean_controlnet_input, std_controlnet_input = controlnet_input.mean(dim=(1,2), keepdim=True), controlnet_input.std(dim=(1,2), keepdim=True)
-    controlnet_input = (controlnet_input - mean_controlnet_input) * (std_hidden_states / (std_controlnet_input + 1e-12)) + mean_hidden_states
-    return controlnet_input
-
-
-# Copied from transformers.models.mixtral.modeling_mixtral.MixtralRotaryEmbedding with Mixtral->Qwen2
-class Qwen2RotaryEmbedding(nn.Module):
-    def __init__(self, dim, max_position_embeddings=2048, base=10000, dtype=None, device=None):
-        super().__init__()
-
-        self.dim = dim
-        self.max_position_embeddings = max_position_embeddings
-        self.base = base
-        inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, dtype=torch.int64, device=device).float() / self.dim))
-        self.register_buffer("inv_freq", inv_freq, persistent=False)
-
-        # Build here to make `torch.jit.trace` work.
-        self._set_cos_sin_cache(
-            seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.float32
-        )
-
-    def _set_cos_sin_cache(self, seq_len, device, dtype):
-        self.max_seq_len_cached = seq_len
-        t = torch.arange(self.max_seq_len_cached, device=device, dtype=torch.int64).type_as(self.inv_freq)
-
-        freqs = torch.outer(t, self.inv_freq)
-        # Different from paper, but it uses a different permutation in order to obtain the same calculation
-        emb = torch.cat((freqs, freqs), dim=-1)
-        self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
-        self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
-
-    def forward(self, x, seq_len=None):
-        # x: [bs, num_attention_heads, seq_len, head_size]
-        if seq_len > self.max_seq_len_cached:
-            self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype)
-
-        return (
-            self.cos_cached[:seq_len].to(dtype=x.dtype),
-            self.sin_cached[:seq_len].to(dtype=x.dtype),
-        )
-
-
-class T2IFinalLayer(nn.Module):
-    """
-    The final layer of Sana.
-    """
-
-    def __init__(self, hidden_size, patch_size=[16, 1], out_channels=256, dtype=None, device=None, operations=None):
-        super().__init__()
-        self.norm_final = operations.RMSNorm(hidden_size, elementwise_affine=False, eps=1e-6, dtype=dtype, device=device)
-        self.linear = operations.Linear(hidden_size, patch_size[0] * patch_size[1] * out_channels, bias=True, dtype=dtype, device=device)
-        self.scale_shift_table = nn.Parameter(torch.empty(2, hidden_size, dtype=dtype, device=device))
-        self.out_channels = out_channels
-        self.patch_size = patch_size
-
-    def unpatchfy(
-        self,
-        hidden_states: torch.Tensor,
-        width: int,
-    ):
-        # 4 unpatchify
-        new_height, new_width = 1, hidden_states.size(1)
-        hidden_states = hidden_states.reshape(
-            shape=(hidden_states.shape[0], new_height, new_width, self.patch_size[0], self.patch_size[1], self.out_channels)
-        ).contiguous()
-        hidden_states = torch.einsum("nhwpqc->nchpwq", hidden_states)
-        output = hidden_states.reshape(
-            shape=(hidden_states.shape[0], self.out_channels, new_height * self.patch_size[0], new_width * self.patch_size[1])
-        ).contiguous()
-        if width > new_width:
-            output = torch.nn.functional.pad(output, (0, width - new_width, 0, 0), 'constant', 0)
-        elif width < new_width:
-            output = output[:, :, :, :width]
-        return output
-
-    def forward(self, x, t, output_length):
-        shift, scale = (comfy.model_management.cast_to(self.scale_shift_table[None], device=t.device, dtype=t.dtype) + t[:, None]).chunk(2, dim=1)
-        x = t2i_modulate(self.norm_final(x), shift, scale)
-        x = self.linear(x)
-        # unpatchify
-        output = self.unpatchfy(x, output_length)
-        return output
-
-
-class PatchEmbed(nn.Module):
-    """2D Image to Patch Embedding"""
-
-    def __init__(
-        self,
-        height=16,
-        width=4096,
-        patch_size=(16, 1),
-        in_channels=8,
-        embed_dim=1152,
-        bias=True,
-        dtype=None, device=None, operations=None
-    ):
-        super().__init__()
-        patch_size_h, patch_size_w = patch_size
-        self.early_conv_layers = nn.Sequential(
-            operations.Conv2d(in_channels, in_channels*256, kernel_size=patch_size, stride=patch_size, padding=0, bias=bias, dtype=dtype, device=device),
-            operations.GroupNorm(num_groups=32, num_channels=in_channels*256, eps=1e-6, affine=True, dtype=dtype, device=device),
-            operations.Conv2d(in_channels*256, embed_dim, kernel_size=1, stride=1, padding=0, bias=bias, dtype=dtype, device=device)
-        )
-        self.patch_size = patch_size
-        self.height, self.width = height // patch_size_h, width // patch_size_w
-        self.base_size = self.width
-
-    def forward(self, latent):
-        # early convolutions, N x C x H x W -> N x 256 * sqrt(patch_size) x H/patch_size x W/patch_size
-        latent = self.early_conv_layers(latent)
-        latent = latent.flatten(2).transpose(1, 2)  # BCHW -> BNC
-        return latent
-
-
-class ACEStepTransformer2DModel(nn.Module):
-    # _supports_gradient_checkpointing = True
-
-    def __init__(
-        self,
-        in_channels: Optional[int] = 8,
-        num_layers: int = 28,
-        inner_dim: int = 1536,
-        attention_head_dim: int = 64,
-        num_attention_heads: int = 24,
-        mlp_ratio: float = 4.0,
-        out_channels: int = 8,
-        max_position: int = 32768,
-        rope_theta: float = 1000000.0,
-        speaker_embedding_dim: int = 512,
-        text_embedding_dim: int = 768,
-        ssl_encoder_depths: List[int] = [9, 9],
-        ssl_names: List[str] = ["mert", "m-hubert"],
-        ssl_latent_dims: List[int] = [1024, 768],
-        lyric_encoder_vocab_size: int = 6681,
-        lyric_hidden_size: int = 1024,
-        patch_size: List[int] = [16, 1],
-        max_height: int = 16,
-        max_width: int = 4096,
-        audio_model=None,
-        dtype=None, device=None, operations=None
-
-    ):
-        super().__init__()
-
-        self.dtype = dtype
-        self.num_attention_heads = num_attention_heads
-        self.attention_head_dim = attention_head_dim
-        inner_dim = num_attention_heads * attention_head_dim
-        self.inner_dim = inner_dim
-        self.out_channels = out_channels
-        self.max_position = max_position
-        self.patch_size = patch_size
-
-        self.rope_theta = rope_theta
-
-        self.rotary_emb = Qwen2RotaryEmbedding(
-            dim=self.attention_head_dim,
-            max_position_embeddings=self.max_position,
-            base=self.rope_theta,
-            dtype=dtype,
-            device=device,
-        )
-
-        # 2. Define input layers
-        self.in_channels = in_channels
-
-        self.num_layers = num_layers
-        # 3. Define transformers blocks
-        self.transformer_blocks = nn.ModuleList(
-            [
-                LinearTransformerBlock(
-                    dim=self.inner_dim,
-                    num_attention_heads=self.num_attention_heads,
-                    attention_head_dim=attention_head_dim,
-                    mlp_ratio=mlp_ratio,
-                    add_cross_attention=True,
-                    add_cross_attention_dim=self.inner_dim,
-                    dtype=dtype,
-                    device=device,
-                    operations=operations,
-                )
-                for i in range(self.num_layers)
-            ]
-        )
-
-        self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0)
-        self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=self.inner_dim, dtype=dtype, device=device, operations=operations)
-        self.t_block = nn.Sequential(nn.SiLU(), operations.Linear(self.inner_dim, 6 * self.inner_dim, bias=True, dtype=dtype, device=device))
-
-        # speaker
-        self.speaker_embedder = operations.Linear(speaker_embedding_dim, self.inner_dim, dtype=dtype, device=device)
-
-        # genre
-        self.genre_embedder = operations.Linear(text_embedding_dim, self.inner_dim, dtype=dtype, device=device)
-
-        # lyric
-        self.lyric_embs = operations.Embedding(lyric_encoder_vocab_size, lyric_hidden_size, dtype=dtype, device=device)
-        self.lyric_encoder = LyricEncoder(input_size=lyric_hidden_size, static_chunk_size=0, dtype=dtype, device=device, operations=operations)
-        self.lyric_proj = operations.Linear(lyric_hidden_size, self.inner_dim, dtype=dtype, device=device)
-
-        projector_dim = 2 * self.inner_dim
-
-        self.projectors = nn.ModuleList([
-            nn.Sequential(
-                operations.Linear(self.inner_dim, projector_dim, dtype=dtype, device=device),
-                nn.SiLU(),
-                operations.Linear(projector_dim, projector_dim, dtype=dtype, device=device),
-                nn.SiLU(),
-                operations.Linear(projector_dim, ssl_dim, dtype=dtype, device=device),
-            ) for ssl_dim in ssl_latent_dims
-        ])
-
-        self.proj_in = PatchEmbed(
-            height=max_height,
-            width=max_width,
-            patch_size=patch_size,
-            embed_dim=self.inner_dim,
-            bias=True,
-            dtype=dtype,
-            device=device,
-            operations=operations,
-        )
-
-        self.final_layer = T2IFinalLayer(self.inner_dim, patch_size=patch_size, out_channels=out_channels, dtype=dtype, device=device, operations=operations)
-
-    def forward_lyric_encoder(
-        self,
-        lyric_token_idx: Optional[torch.LongTensor] = None,
-        lyric_mask: Optional[torch.LongTensor] = None,
-        out_dtype=None,
-    ):
-        # N x T x D
-        lyric_embs = self.lyric_embs(lyric_token_idx, out_dtype=out_dtype)
-        prompt_prenet_out, _mask = self.lyric_encoder(lyric_embs, lyric_mask, decoding_chunk_size=1, num_decoding_left_chunks=-1)
-        prompt_prenet_out = self.lyric_proj(prompt_prenet_out)
-        return prompt_prenet_out
-
-    def encode(
-        self,
-        encoder_text_hidden_states: Optional[torch.Tensor] = None,
-        text_attention_mask: Optional[torch.LongTensor] = None,
-        speaker_embeds: Optional[torch.FloatTensor] = None,
-        lyric_token_idx: Optional[torch.LongTensor] = None,
-        lyric_mask: Optional[torch.LongTensor] = None,
-        lyrics_strength=1.0,
-    ):
-
-        bs = encoder_text_hidden_states.shape[0]
-        device = encoder_text_hidden_states.device
-
-        # speaker embedding
-        encoder_spk_hidden_states = self.speaker_embedder(speaker_embeds).unsqueeze(1)
-
-        # genre embedding
-        encoder_text_hidden_states = self.genre_embedder(encoder_text_hidden_states)
-
-        # lyric
-        encoder_lyric_hidden_states = self.forward_lyric_encoder(
-            lyric_token_idx=lyric_token_idx,
-            lyric_mask=lyric_mask,
-            out_dtype=encoder_text_hidden_states.dtype,
-        )
-
-        encoder_lyric_hidden_states *= lyrics_strength
-
-        encoder_hidden_states = torch.cat([encoder_spk_hidden_states, encoder_text_hidden_states, encoder_lyric_hidden_states], dim=1)
-
-        encoder_hidden_mask = None
-        if text_attention_mask is not None:
-            speaker_mask = torch.ones(bs, 1, device=device)
-            encoder_hidden_mask = torch.cat([speaker_mask, text_attention_mask, lyric_mask], dim=1)
-
-        return encoder_hidden_states, encoder_hidden_mask
-
-    def decode(
-        self,
-        hidden_states: torch.Tensor,
-        attention_mask: torch.Tensor,
-        encoder_hidden_states: torch.Tensor,
-        encoder_hidden_mask: torch.Tensor,
-        timestep: Optional[torch.Tensor],
-        output_length: int = 0,
-        block_controlnet_hidden_states: Optional[Union[List[torch.Tensor], torch.Tensor]] = None,
-        controlnet_scale: Union[float, torch.Tensor] = 1.0,
-    ):
-        embedded_timestep = self.timestep_embedder(self.time_proj(timestep).to(dtype=hidden_states.dtype))
-        temb = self.t_block(embedded_timestep)
-
-        hidden_states = self.proj_in(hidden_states)
-
-        # controlnet logic
-        if block_controlnet_hidden_states is not None:
-            control_condi = cross_norm(hidden_states, block_controlnet_hidden_states)
-            hidden_states = hidden_states + control_condi * controlnet_scale
-
-        # inner_hidden_states = []
-
-        rotary_freqs_cis = self.rotary_emb(hidden_states, seq_len=hidden_states.shape[1])
-        encoder_rotary_freqs_cis = self.rotary_emb(encoder_hidden_states, seq_len=encoder_hidden_states.shape[1])
-
-        for index_block, block in enumerate(self.transformer_blocks):
-            hidden_states = block(
-                hidden_states=hidden_states,
-                attention_mask=attention_mask,
-                encoder_hidden_states=encoder_hidden_states,
-                encoder_attention_mask=encoder_hidden_mask,
-                rotary_freqs_cis=rotary_freqs_cis,
-                rotary_freqs_cis_cross=encoder_rotary_freqs_cis,
-                temb=temb,
-            )
-
-        output = self.final_layer(hidden_states, embedded_timestep, output_length)
-        return output
-
-    def forward(
-        self,
-        x,
-        timestep,
-        attention_mask=None,
-        context: Optional[torch.Tensor] = None,
-        text_attention_mask: Optional[torch.LongTensor] = None,
-        speaker_embeds: Optional[torch.FloatTensor] = None,
-        lyric_token_idx: Optional[torch.LongTensor] = None,
-        lyric_mask: Optional[torch.LongTensor] = None,
-        block_controlnet_hidden_states: Optional[Union[List[torch.Tensor], torch.Tensor]] = None,
-        controlnet_scale: Union[float, torch.Tensor] = 1.0,
-        lyrics_strength=1.0,
-        **kwargs
-    ):
-        hidden_states = x
-        encoder_text_hidden_states = context
-        encoder_hidden_states, encoder_hidden_mask = self.encode(
-            encoder_text_hidden_states=encoder_text_hidden_states,
-            text_attention_mask=text_attention_mask,
-            speaker_embeds=speaker_embeds,
-            lyric_token_idx=lyric_token_idx,
-            lyric_mask=lyric_mask,
-            lyrics_strength=lyrics_strength,
-        )
-
-        output_length = hidden_states.shape[-1]
-
-        output = self.decode(
-            hidden_states=hidden_states,
-            attention_mask=attention_mask,
-            encoder_hidden_states=encoder_hidden_states,
-            encoder_hidden_mask=encoder_hidden_mask,
-            timestep=timestep,
-            output_length=output_length,
-            block_controlnet_hidden_states=block_controlnet_hidden_states,
-            controlnet_scale=controlnet_scale,
-        )
-
-        return output

comfy/ldm/ace/vae/autoencoder_dc.py

@@ -1,644 +0,0 @@
-# Rewritten from diffusers
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from typing import Tuple, Union
-
-import comfy.model_management
-import comfy.ops
-ops = comfy.ops.disable_weight_init
-
-
-class RMSNorm(ops.RMSNorm):
-    def __init__(self, dim, eps=1e-5, elementwise_affine=True, bias=False):
-        super().__init__(dim, eps=eps, elementwise_affine=elementwise_affine)
-        if elementwise_affine:
-            self.bias = nn.Parameter(torch.empty(dim)) if bias else None
-
-    def forward(self, x):
-        x = super().forward(x)
-        if self.elementwise_affine:
-            if self.bias is not None:
-                x = x + comfy.model_management.cast_to(self.bias, dtype=x.dtype, device=x.device)
-        return x
-
-
-def get_normalization(norm_type, num_features, num_groups=32, eps=1e-5):
-    if norm_type == "batch_norm":
-        return nn.BatchNorm2d(num_features)
-    elif norm_type == "group_norm":
-        return ops.GroupNorm(num_groups, num_features)
-    elif norm_type == "layer_norm":
-        return ops.LayerNorm(num_features)
-    elif norm_type == "rms_norm":
-        return RMSNorm(num_features, eps=eps, elementwise_affine=True, bias=True)
-    else:
-        raise ValueError(f"Unknown normalization type: {norm_type}")
-
-
-def get_activation(activation_type):
-    if activation_type == "relu":
-        return nn.ReLU()
-    elif activation_type == "relu6":
-        return nn.ReLU6()
-    elif activation_type == "silu":
-        return nn.SiLU()
-    elif activation_type == "leaky_relu":
-        return nn.LeakyReLU(0.2)
-    else:
-        raise ValueError(f"Unknown activation type: {activation_type}")
-
-
-class ResBlock(nn.Module):
-    def __init__(
-        self,
-        in_channels: int,
-        out_channels: int,
-        norm_type: str = "batch_norm",
-        act_fn: str = "relu6",
-    ) -> None:
-        super().__init__()
-
-        self.norm_type = norm_type
-        self.nonlinearity = get_activation(act_fn) if act_fn is not None else nn.Identity()
-        self.conv1 = ops.Conv2d(in_channels, in_channels, 3, 1, 1)
-        self.conv2 = ops.Conv2d(in_channels, out_channels, 3, 1, 1, bias=False)
-        self.norm = get_normalization(norm_type, out_channels)
-
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        residual = hidden_states
-        hidden_states = self.conv1(hidden_states)
-        hidden_states = self.nonlinearity(hidden_states)
-        hidden_states = self.conv2(hidden_states)
-
-        if self.norm_type == "rms_norm":
-            # move channel to the last dimension so we apply RMSnorm across channel dimension
-            hidden_states = self.norm(hidden_states.movedim(1, -1)).movedim(-1, 1)
-        else:
-            hidden_states = self.norm(hidden_states)
-
-        return hidden_states + residual
-
-class SanaMultiscaleAttentionProjection(nn.Module):
-    def __init__(
-        self,
-        in_channels: int,
-        num_attention_heads: int,
-        kernel_size: int,
-    ) -> None:
-        super().__init__()
-
-        channels = 3 * in_channels
-        self.proj_in = ops.Conv2d(
-            channels,
-            channels,
-            kernel_size,
-            padding=kernel_size // 2,
-            groups=channels,
-            bias=False,
-        )
-        self.proj_out = ops.Conv2d(channels, channels, 1, 1, 0, groups=3 * num_attention_heads, bias=False)
-
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        hidden_states = self.proj_in(hidden_states)
-        hidden_states = self.proj_out(hidden_states)
-        return hidden_states
-
-class SanaMultiscaleLinearAttention(nn.Module):
-    def __init__(
-        self,
-        in_channels: int,
-        out_channels: int,
-        num_attention_heads: int = None,
-        attention_head_dim: int = 8,
-        mult: float = 1.0,
-        norm_type: str = "batch_norm",
-        kernel_sizes: tuple = (5,),
-        eps: float = 1e-15,
-        residual_connection: bool = False,
-    ):
-        super().__init__()
-
-        self.eps = eps
-        self.attention_head_dim = attention_head_dim
-        self.norm_type = norm_type
-        self.residual_connection = residual_connection
-
-        num_attention_heads = (
-            int(in_channels // attention_head_dim * mult)
-            if num_attention_heads is None
-            else num_attention_heads
-        )
-        inner_dim = num_attention_heads * attention_head_dim
-
-        self.to_q = ops.Linear(in_channels, inner_dim, bias=False)
-        self.to_k = ops.Linear(in_channels, inner_dim, bias=False)
-        self.to_v = ops.Linear(in_channels, inner_dim, bias=False)
-
-        self.to_qkv_multiscale = nn.ModuleList()
-        for kernel_size in kernel_sizes:
-            self.to_qkv_multiscale.append(
-                SanaMultiscaleAttentionProjection(inner_dim, num_attention_heads, kernel_size)
-            )
-
-        self.nonlinearity = nn.ReLU()
-        self.to_out = ops.Linear(inner_dim * (1 + len(kernel_sizes)), out_channels, bias=False)
-        self.norm_out = get_normalization(norm_type, out_channels)
-
-    def apply_linear_attention(self, query, key, value):
-        value = F.pad(value, (0, 0, 0, 1), mode="constant", value=1)
-        scores = torch.matmul(value, key.transpose(-1, -2))
-        hidden_states = torch.matmul(scores, query)
-
-        hidden_states = hidden_states.to(dtype=torch.float32)
-        hidden_states = hidden_states[:, :, :-1] / (hidden_states[:, :, -1:] + self.eps)
-        return hidden_states
-
-    def apply_quadratic_attention(self, query, key, value):
-        scores = torch.matmul(key.transpose(-1, -2), query)
-        scores = scores.to(dtype=torch.float32)
-        scores = scores / (torch.sum(scores, dim=2, keepdim=True) + self.eps)
-        hidden_states = torch.matmul(value, scores.to(value.dtype))
-        return hidden_states
-
-    def forward(self, hidden_states):
-        height, width = hidden_states.shape[-2:]
-        if height * width > self.attention_head_dim:
-            use_linear_attention = True
-        else:
-            use_linear_attention = False
-
-        residual = hidden_states
-
-        batch_size, _, height, width = list(hidden_states.size())
-        original_dtype = hidden_states.dtype
-
-        hidden_states = hidden_states.movedim(1, -1)
-        query = self.to_q(hidden_states)
-        key = self.to_k(hidden_states)
-        value = self.to_v(hidden_states)
-        hidden_states = torch.cat([query, key, value], dim=3)
-        hidden_states = hidden_states.movedim(-1, 1)
-
-        multi_scale_qkv = [hidden_states]
-        for block in self.to_qkv_multiscale:
-            multi_scale_qkv.append(block(hidden_states))
-
-        hidden_states = torch.cat(multi_scale_qkv, dim=1)
-
-        if use_linear_attention:
-            # for linear attention upcast hidden_states to float32
-            hidden_states = hidden_states.to(dtype=torch.float32)
-
-        hidden_states = hidden_states.reshape(batch_size, -1, 3 * self.attention_head_dim, height * width)
-
-        query, key, value = hidden_states.chunk(3, dim=2)
-        query = self.nonlinearity(query)
-        key = self.nonlinearity(key)
-
-        if use_linear_attention:
-            hidden_states = self.apply_linear_attention(query, key, value)
-            hidden_states = hidden_states.to(dtype=original_dtype)
-        else:
-            hidden_states = self.apply_quadratic_attention(query, key, value)
-
-        hidden_states = torch.reshape(hidden_states, (batch_size, -1, height, width))
-        hidden_states = self.to_out(hidden_states.movedim(1, -1)).movedim(-1, 1)
-
-        if self.norm_type == "rms_norm":
-            hidden_states = self.norm_out(hidden_states.movedim(1, -1)).movedim(-1, 1)
-        else:
-            hidden_states = self.norm_out(hidden_states)
-
-        if self.residual_connection:
-            hidden_states = hidden_states + residual
-
-        return hidden_states
-
-
-class EfficientViTBlock(nn.Module):
-    def __init__(
-        self,
-        in_channels: int,
-        mult: float = 1.0,
-        attention_head_dim: int = 32,
-        qkv_multiscales: tuple = (5,),
-        norm_type: str = "batch_norm",
-    ) -> None:
-        super().__init__()
-
-        self.attn = SanaMultiscaleLinearAttention(
-            in_channels=in_channels,
-            out_channels=in_channels,
-            mult=mult,
-            attention_head_dim=attention_head_dim,
-            norm_type=norm_type,
-            kernel_sizes=qkv_multiscales,
-            residual_connection=True,
-        )
-
-        self.conv_out = GLUMBConv(
-            in_channels=in_channels,
-            out_channels=in_channels,
-            norm_type="rms_norm",
-        )
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        x = self.attn(x)
-        x = self.conv_out(x)
-        return x
-
-
-class GLUMBConv(nn.Module):
-    def __init__(
-        self,
-        in_channels: int,
-        out_channels: int,
-        expand_ratio: float = 4,
-        norm_type: str = None,
-        residual_connection: bool = True,
-    ) -> None:
-        super().__init__()
-
-        hidden_channels = int(expand_ratio * in_channels)
-        self.norm_type = norm_type
-        self.residual_connection = residual_connection
-
-        self.nonlinearity = nn.SiLU()
-        self.conv_inverted = ops.Conv2d(in_channels, hidden_channels * 2, 1, 1, 0)
-        self.conv_depth = ops.Conv2d(hidden_channels * 2, hidden_channels * 2, 3, 1, 1, groups=hidden_channels * 2)
-        self.conv_point = ops.Conv2d(hidden_channels, out_channels, 1, 1, 0, bias=False)
-
-        self.norm = None
-        if norm_type == "rms_norm":
-            self.norm = RMSNorm(out_channels, eps=1e-5, elementwise_affine=True, bias=True)
-
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        if self.residual_connection:
-            residual = hidden_states
-
-        hidden_states = self.conv_inverted(hidden_states)
-        hidden_states = self.nonlinearity(hidden_states)
-
-        hidden_states = self.conv_depth(hidden_states)
-        hidden_states, gate = torch.chunk(hidden_states, 2, dim=1)
-        hidden_states = hidden_states * self.nonlinearity(gate)
-
-        hidden_states = self.conv_point(hidden_states)
-
-        if self.norm_type == "rms_norm":
-            # move channel to the last dimension so we apply RMSnorm across channel dimension
-            hidden_states = self.norm(hidden_states.movedim(1, -1)).movedim(-1, 1)
-
-        if self.residual_connection:
-            hidden_states = hidden_states + residual
-
-        return hidden_states
-
-
-def get_block(
-    block_type: str,
-    in_channels: int,
-    out_channels: int,
-    attention_head_dim: int,
-    norm_type: str,
-    act_fn: str,
-    qkv_mutliscales: tuple = (),
-):
-    if block_type == "ResBlock":
-        block = ResBlock(in_channels, out_channels, norm_type, act_fn)
-    elif block_type == "EfficientViTBlock":
-        block = EfficientViTBlock(
-            in_channels,
-            attention_head_dim=attention_head_dim,
-            norm_type=norm_type,
-            qkv_multiscales=qkv_mutliscales
-        )
-    else:
-        raise ValueError(f"Block with {block_type=} is not supported.")
-
-    return block
-
-
-class DCDownBlock2d(nn.Module):
-    def __init__(self, in_channels: int, out_channels: int, downsample: bool = False, shortcut: bool = True) -> None:
-        super().__init__()
-
-        self.downsample = downsample
-        self.factor = 2
-        self.stride = 1 if downsample else 2
-        self.group_size = in_channels * self.factor**2 // out_channels
-        self.shortcut = shortcut
-
-        out_ratio = self.factor**2
-        if downsample:
-            assert out_channels % out_ratio == 0
-            out_channels = out_channels // out_ratio
-
-        self.conv = ops.Conv2d(
-            in_channels,
-            out_channels,
-            kernel_size=3,
-            stride=self.stride,
-            padding=1,
-        )
-
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        x = self.conv(hidden_states)
-        if self.downsample:
-            x = F.pixel_unshuffle(x, self.factor)
-
-        if self.shortcut:
-            y = F.pixel_unshuffle(hidden_states, self.factor)
-            y = y.unflatten(1, (-1, self.group_size))
-            y = y.mean(dim=2)
-            hidden_states = x + y
-        else:
-            hidden_states = x
-
-        return hidden_states
-
-
-class DCUpBlock2d(nn.Module):
-    def __init__(
-        self,
-        in_channels: int,
-        out_channels: int,
-        interpolate: bool = False,
-        shortcut: bool = True,
-        interpolation_mode: str = "nearest",
-    ) -> None:
-        super().__init__()
-
-        self.interpolate = interpolate
-        self.interpolation_mode = interpolation_mode
-        self.shortcut = shortcut
-        self.factor = 2
-        self.repeats = out_channels * self.factor**2 // in_channels
-
-        out_ratio = self.factor**2
-        if not interpolate:
-            out_channels = out_channels * out_ratio
-
-        self.conv = ops.Conv2d(in_channels, out_channels, 3, 1, 1)
-
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        if self.interpolate:
-            x = F.interpolate(hidden_states, scale_factor=self.factor, mode=self.interpolation_mode)
-            x = self.conv(x)
-        else:
-            x = self.conv(hidden_states)
-            x = F.pixel_shuffle(x, self.factor)
-
-        if self.shortcut:
-            y = hidden_states.repeat_interleave(self.repeats, dim=1, output_size=hidden_states.shape[1] * self.repeats)
-            y = F.pixel_shuffle(y, self.factor)
-            hidden_states = x + y
-        else:
-            hidden_states = x
-
-        return hidden_states
-
-
-class Encoder(nn.Module):
-    def __init__(
-        self,
-        in_channels: int,
-        latent_channels: int,
-        attention_head_dim: int = 32,
-        block_type: str or tuple = "ResBlock",
-        block_out_channels: tuple = (128, 256, 512, 512, 1024, 1024),
-        layers_per_block: tuple = (2, 2, 2, 2, 2, 2),
-        qkv_multiscales: tuple = ((), (), (), (5,), (5,), (5,)),
-        downsample_block_type: str = "pixel_unshuffle",
-        out_shortcut: bool = True,
-    ):
-        super().__init__()
-
-        num_blocks = len(block_out_channels)
-
-        if isinstance(block_type, str):
-            block_type = (block_type,) * num_blocks
-
-        if layers_per_block[0] > 0:
-            self.conv_in = ops.Conv2d(
-                in_channels,
-                block_out_channels[0] if layers_per_block[0] > 0 else block_out_channels[1],
-                kernel_size=3,
-                stride=1,
-                padding=1,
-            )
-        else:
-            self.conv_in = DCDownBlock2d(
-                in_channels=in_channels,
-                out_channels=block_out_channels[0] if layers_per_block[0] > 0 else block_out_channels[1],
-                downsample=downsample_block_type == "pixel_unshuffle",
-                shortcut=False,
-            )
-
-        down_blocks = []
-        for i, (out_channel, num_layers) in enumerate(zip(block_out_channels, layers_per_block)):
-            down_block_list = []
-
-            for _ in range(num_layers):
-                block = get_block(
-                    block_type[i],
-                    out_channel,
-                    out_channel,
-                    attention_head_dim=attention_head_dim,
-                    norm_type="rms_norm",
-                    act_fn="silu",
-                    qkv_mutliscales=qkv_multiscales[i],
-                )
-                down_block_list.append(block)
-
-            if i < num_blocks - 1 and num_layers > 0:
-                downsample_block = DCDownBlock2d(
-                    in_channels=out_channel,
-                    out_channels=block_out_channels[i + 1],
-                    downsample=downsample_block_type == "pixel_unshuffle",
-                    shortcut=True,
-                )
-                down_block_list.append(downsample_block)
-
-            down_blocks.append(nn.Sequential(*down_block_list))
-
-        self.down_blocks = nn.ModuleList(down_blocks)
-
-        self.conv_out = ops.Conv2d(block_out_channels[-1], latent_channels, 3, 1, 1)
-
-        self.out_shortcut = out_shortcut
-        if out_shortcut:
-            self.out_shortcut_average_group_size = block_out_channels[-1] // latent_channels
-
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        hidden_states = self.conv_in(hidden_states)
-        for down_block in self.down_blocks:
-            hidden_states = down_block(hidden_states)
-
-        if self.out_shortcut:
-            x = hidden_states.unflatten(1, (-1, self.out_shortcut_average_group_size))
-            x = x.mean(dim=2)
-            hidden_states = self.conv_out(hidden_states) + x
-        else:
-            hidden_states = self.conv_out(hidden_states)
-
-        return hidden_states
-
-
-class Decoder(nn.Module):
-    def __init__(
-        self,
-        in_channels: int,
-        latent_channels: int,
-        attention_head_dim: int = 32,
-        block_type: str or tuple = "ResBlock",
-        block_out_channels: tuple = (128, 256, 512, 512, 1024, 1024),
-        layers_per_block: tuple = (2, 2, 2, 2, 2, 2),
-        qkv_multiscales: tuple = ((), (), (), (5,), (5,), (5,)),
-        norm_type: str or tuple = "rms_norm",
-        act_fn: str or tuple = "silu",
-        upsample_block_type: str = "pixel_shuffle",
-        in_shortcut: bool = True,
-    ):
-        super().__init__()
-
-        num_blocks = len(block_out_channels)
-
-        if isinstance(block_type, str):
-            block_type = (block_type,) * num_blocks
-        if isinstance(norm_type, str):
-            norm_type = (norm_type,) * num_blocks
-        if isinstance(act_fn, str):
-            act_fn = (act_fn,) * num_blocks
-
-        self.conv_in = ops.Conv2d(latent_channels, block_out_channels[-1], 3, 1, 1)
-
-        self.in_shortcut = in_shortcut
-        if in_shortcut:
-            self.in_shortcut_repeats = block_out_channels[-1] // latent_channels
-
-        up_blocks = []
-        for i, (out_channel, num_layers) in reversed(list(enumerate(zip(block_out_channels, layers_per_block)))):
-            up_block_list = []
-
-            if i < num_blocks - 1 and num_layers > 0:
-                upsample_block = DCUpBlock2d(
-                    block_out_channels[i + 1],
-                    out_channel,
-                    interpolate=upsample_block_type == "interpolate",
-                    shortcut=True,
-                )
-                up_block_list.append(upsample_block)
-
-            for _ in range(num_layers):
-                block = get_block(
-                    block_type[i],
-                    out_channel,
-                    out_channel,
-                    attention_head_dim=attention_head_dim,
-                    norm_type=norm_type[i],
-                    act_fn=act_fn[i],
-                    qkv_mutliscales=qkv_multiscales[i],
-                )
-                up_block_list.append(block)
-
-            up_blocks.insert(0, nn.Sequential(*up_block_list))
-
-        self.up_blocks = nn.ModuleList(up_blocks)
-
-        channels = block_out_channels[0] if layers_per_block[0] > 0 else block_out_channels[1]
-
-        self.norm_out = RMSNorm(channels, 1e-5, elementwise_affine=True, bias=True)
-        self.conv_act = nn.ReLU()
-        self.conv_out = None
-
-        if layers_per_block[0] > 0:
-            self.conv_out = ops.Conv2d(channels, in_channels, 3, 1, 1)
-        else:
-            self.conv_out = DCUpBlock2d(
-                channels, in_channels, interpolate=upsample_block_type == "interpolate", shortcut=False
-            )
-
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        if self.in_shortcut:
-            x = hidden_states.repeat_interleave(
-                self.in_shortcut_repeats, dim=1, output_size=hidden_states.shape[1] * self.in_shortcut_repeats
-            )
-            hidden_states = self.conv_in(hidden_states) + x
-        else:
-            hidden_states = self.conv_in(hidden_states)
-
-        for up_block in reversed(self.up_blocks):
-            hidden_states = up_block(hidden_states)
-
-        hidden_states = self.norm_out(hidden_states.movedim(1, -1)).movedim(-1, 1)
-        hidden_states = self.conv_act(hidden_states)
-        hidden_states = self.conv_out(hidden_states)
-        return hidden_states
-
-
-class AutoencoderDC(nn.Module):
-    def __init__(
-        self,
-        in_channels: int = 2,
-        latent_channels: int = 8,
-        attention_head_dim: int = 32,
-        encoder_block_types: Union[str, Tuple[str]] = ["ResBlock", "ResBlock", "ResBlock", "EfficientViTBlock"],
-        decoder_block_types: Union[str, Tuple[str]] = ["ResBlock", "ResBlock", "ResBlock", "EfficientViTBlock"],
-        encoder_block_out_channels: Tuple[int, ...] = (128, 256, 512, 1024),
-        decoder_block_out_channels: Tuple[int, ...] = (128, 256, 512, 1024),
-        encoder_layers_per_block: Tuple[int] = (2, 2, 3, 3),
-        decoder_layers_per_block: Tuple[int] = (3, 3, 3, 3),
-        encoder_qkv_multiscales: Tuple[Tuple[int, ...], ...] = ((), (), (5,), (5,)),
-        decoder_qkv_multiscales: Tuple[Tuple[int, ...], ...] = ((), (), (5,), (5,)),
-        upsample_block_type: str = "interpolate",
-        downsample_block_type: str = "Conv",
-        decoder_norm_types: Union[str, Tuple[str]] = "rms_norm",
-        decoder_act_fns: Union[str, Tuple[str]] = "silu",
-        scaling_factor: float = 0.41407,
-    ) -> None:
-        super().__init__()
-
-        self.encoder = Encoder(
-            in_channels=in_channels,
-            latent_channels=latent_channels,
-            attention_head_dim=attention_head_dim,
-            block_type=encoder_block_types,
-            block_out_channels=encoder_block_out_channels,
-            layers_per_block=encoder_layers_per_block,
-            qkv_multiscales=encoder_qkv_multiscales,
-            downsample_block_type=downsample_block_type,
-        )
-
-        self.decoder = Decoder(
-            in_channels=in_channels,
-            latent_channels=latent_channels,
-            attention_head_dim=attention_head_dim,
-            block_type=decoder_block_types,
-            block_out_channels=decoder_block_out_channels,
-            layers_per_block=decoder_layers_per_block,
-            qkv_multiscales=decoder_qkv_multiscales,
-            norm_type=decoder_norm_types,
-            act_fn=decoder_act_fns,
-            upsample_block_type=upsample_block_type,
-        )
-
-        self.scaling_factor = scaling_factor
-        self.spatial_compression_ratio = 2 ** (len(encoder_block_out_channels) - 1)
-
-    def encode(self, x: torch.Tensor) -> torch.Tensor:
-        """Internal encoding function."""
-        encoded = self.encoder(x)
-        return encoded * self.scaling_factor
-
-    def decode(self, z: torch.Tensor) -> torch.Tensor:
-        # Scale the latents back
-        z = z / self.scaling_factor
-        decoded = self.decoder(z)
-        return decoded
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        z = self.encode(x)
-        return self.decode(z)
-

comfy/ldm/ace/vae/music_dcae_pipeline.py

@@ -1,109 +0,0 @@
-# Original from: https://github.com/ace-step/ACE-Step/blob/main/music_dcae/music_dcae_pipeline.py
-import torch
-from .autoencoder_dc import AutoencoderDC
-import logging
-try:
-    import torchaudio
-except:
-    logging.warning("torchaudio missing, ACE model will be broken")
-
-import torchvision.transforms as transforms
-from .music_vocoder import ADaMoSHiFiGANV1
-
-
-class MusicDCAE(torch.nn.Module):
-    def __init__(self, source_sample_rate=None, dcae_config={}, vocoder_config={}):
-        super(MusicDCAE, self).__init__()
-
-        self.dcae = AutoencoderDC(**dcae_config)
-        self.vocoder = ADaMoSHiFiGANV1(**vocoder_config)
-
-        if source_sample_rate is None:
-            self.source_sample_rate = 48000
-        else:
-            self.source_sample_rate = source_sample_rate
-
-        # self.resampler = torchaudio.transforms.Resample(source_sample_rate, 44100)
-
-        self.transform = transforms.Compose([
-            transforms.Normalize(0.5, 0.5),
-        ])
-        self.min_mel_value = -11.0
-        self.max_mel_value = 3.0
-        self.audio_chunk_size = int(round((1024 * 512 / 44100 * 48000)))
-        self.mel_chunk_size = 1024
-        self.time_dimention_multiple = 8
-        self.latent_chunk_size = self.mel_chunk_size // self.time_dimention_multiple
-        self.scale_factor = 0.1786
-        self.shift_factor = -1.9091
-
-    def load_audio(self, audio_path):
-        audio, sr = torchaudio.load(audio_path)
-        return audio, sr
-
-    def forward_mel(self, audios):
-        mels = []
-        for i in range(len(audios)):
-            image = self.vocoder.mel_transform(audios[i])
-            mels.append(image)
-        mels = torch.stack(mels)
-        return mels
-
-    @torch.no_grad()
-    def encode(self, audios, audio_lengths=None, sr=None):
-        if audio_lengths is None:
-            audio_lengths = torch.tensor([audios.shape[2]] * audios.shape[0])
-            audio_lengths = audio_lengths.to(audios.device)
-
-        if sr is None:
-            sr = self.source_sample_rate
-
-        if sr != 44100:
-            audios = torchaudio.functional.resample(audios, sr, 44100)
-
-        max_audio_len = audios.shape[-1]
-        if max_audio_len % (8 * 512) != 0:
-            audios = torch.nn.functional.pad(audios, (0, 8 * 512 - max_audio_len % (8 * 512)))
-
-        mels = self.forward_mel(audios)
-        mels = (mels - self.min_mel_value) / (self.max_mel_value - self.min_mel_value)
-        mels = self.transform(mels)
-        latents = []
-        for mel in mels:
-            latent = self.dcae.encoder(mel.unsqueeze(0))
-            latents.append(latent)
-        latents = torch.cat(latents, dim=0)
-        # latent_lengths = (audio_lengths / sr * 44100 / 512 / self.time_dimention_multiple).long()
-        latents = (latents - self.shift_factor) * self.scale_factor
-        return latents
-        # return latents, latent_lengths
-
-    @torch.no_grad()
-    def decode(self, latents, audio_lengths=None, sr=None):
-        latents = latents / self.scale_factor + self.shift_factor
-
-        pred_wavs = []
-
-        for latent in latents:
-            mels = self.dcae.decoder(latent.unsqueeze(0))
-            mels = mels * 0.5 + 0.5
-            mels = mels * (self.max_mel_value - self.min_mel_value) + self.min_mel_value
-            wav = self.vocoder.decode(mels[0]).squeeze(1)
-
-            if sr is not None:
-                # resampler = torchaudio.transforms.Resample(44100, sr).to(latents.device).to(latents.dtype)
-                wav = torchaudio.functional.resample(wav, 44100, sr)
-                # wav = resampler(wav)
-            else:
-                sr = 44100
-            pred_wavs.append(wav)
-
-        if audio_lengths is not None:
-            pred_wavs = [wav[:, :length].cpu() for wav, length in zip(pred_wavs, audio_lengths)]
-        return torch.stack(pred_wavs)
-        # return sr, pred_wavs
-
-    def forward(self, audios, audio_lengths=None, sr=None):
-        latents, latent_lengths = self.encode(audios=audios, audio_lengths=audio_lengths, sr=sr)
-        sr, pred_wavs = self.decode(latents=latents, audio_lengths=audio_lengths, sr=sr)
-        return sr, pred_wavs, latents, latent_lengths

comfy/ldm/ace/vae/music_log_mel.py

@@ -1,113 +0,0 @@
-# Original from: https://github.com/ace-step/ACE-Step/blob/main/music_dcae/music_log_mel.py
-import torch
-import torch.nn as nn
-from torch import Tensor
-import logging
-try:
-    from torchaudio.transforms import MelScale
-except:
-    logging.warning("torchaudio missing, ACE model will be broken")
-
-import comfy.model_management
-
-class LinearSpectrogram(nn.Module):
-    def __init__(
-        self,
-        n_fft=2048,
-        win_length=2048,
-        hop_length=512,
-        center=False,
-        mode="pow2_sqrt",
-    ):
-        super().__init__()
-
-        self.n_fft = n_fft
-        self.win_length = win_length
-        self.hop_length = hop_length
-        self.center = center
-        self.mode = mode
-
-        self.register_buffer("window", torch.hann_window(win_length))
-
-    def forward(self, y: Tensor) -> Tensor:
-        if y.ndim == 3:
-            y = y.squeeze(1)
-
-        y = torch.nn.functional.pad(
-            y.unsqueeze(1),
-            (
-                (self.win_length - self.hop_length) // 2,
-                (self.win_length - self.hop_length + 1) // 2,
-            ),
-            mode="reflect",
-        ).squeeze(1)
-        dtype = y.dtype
-        spec = torch.stft(
-            y.float(),
-            self.n_fft,
-            hop_length=self.hop_length,
-            win_length=self.win_length,
-            window=comfy.model_management.cast_to(self.window, dtype=torch.float32, device=y.device),
-            center=self.center,
-            pad_mode="reflect",
-            normalized=False,
-            onesided=True,
-            return_complex=True,
-        )
-        spec = torch.view_as_real(spec)
-
-        if self.mode == "pow2_sqrt":
-            spec = torch.sqrt(spec.pow(2).sum(-1) + 1e-6)
-        spec = spec.to(dtype)
-        return spec
-
-
-class LogMelSpectrogram(nn.Module):
-    def __init__(
-        self,
-        sample_rate=44100,
-        n_fft=2048,
-        win_length=2048,
-        hop_length=512,
-        n_mels=128,
-        center=False,
-        f_min=0.0,
-        f_max=None,
-    ):
-        super().__init__()
-
-        self.sample_rate = sample_rate
-        self.n_fft = n_fft
-        self.win_length = win_length
-        self.hop_length = hop_length
-        self.center = center
-        self.n_mels = n_mels
-        self.f_min = f_min
-        self.f_max = f_max or sample_rate // 2
-
-        self.spectrogram = LinearSpectrogram(n_fft, win_length, hop_length, center)
-        self.mel_scale = MelScale(
-            self.n_mels,
-            self.sample_rate,
-            self.f_min,
-            self.f_max,
-            self.n_fft // 2 + 1,
-            "slaney",
-            "slaney",
-        )
-
-    def compress(self, x: Tensor) -> Tensor:
-        return torch.log(torch.clamp(x, min=1e-5))
-
-    def decompress(self, x: Tensor) -> Tensor:
-        return torch.exp(x)
-
-    def forward(self, x: Tensor, return_linear: bool = False) -> Tensor:
-        linear = self.spectrogram(x)
-        x = self.mel_scale(linear)
-        x = self.compress(x)
-        # print(x.shape)
-        if return_linear:
-            return x, self.compress(linear)
-
-        return x

comfy/ldm/ace/vae/music_vocoder.py

@@ -1,538 +0,0 @@
-# Original from: https://github.com/ace-step/ACE-Step/blob/main/music_dcae/music_vocoder.py
-import torch
-from torch import nn
-
-from functools import partial
-from math import prod
-from typing import Callable, Tuple, List
-
-import numpy as np
-import torch.nn.functional as F
-from torch.nn.utils.parametrize import remove_parametrizations as remove_weight_norm
-
-from .music_log_mel import LogMelSpectrogram
-
-import comfy.model_management
-import comfy.ops
-ops = comfy.ops.disable_weight_init
-
-
-def drop_path(
-    x, drop_prob: float = 0.0, training: bool = False, scale_by_keep: bool = True
-):
-    """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
-
-    This is the same as the DropConnect impl I created for EfficientNet, etc networks, however,
-    the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper...
-    See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for
-    changing the layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use
-    'survival rate' as the argument.
-
-    """  # noqa: E501
-
-    if drop_prob == 0.0 or not training:
-        return x
-    keep_prob = 1 - drop_prob
-    shape = (x.shape[0],) + (1,) * (
-        x.ndim - 1
-    )  # work with diff dim tensors, not just 2D ConvNets
-    random_tensor = x.new_empty(shape).bernoulli_(keep_prob)
-    if keep_prob > 0.0 and scale_by_keep:
-        random_tensor.div_(keep_prob)
-    return x * random_tensor
-
-
-class DropPath(nn.Module):
-    """Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks)."""  # noqa: E501
-
-    def __init__(self, drop_prob: float = 0.0, scale_by_keep: bool = True):
-        super(DropPath, self).__init__()
-        self.drop_prob = drop_prob
-        self.scale_by_keep = scale_by_keep
-
-    def forward(self, x):
-        return drop_path(x, self.drop_prob, self.training, self.scale_by_keep)
-
-    def extra_repr(self):
-        return f"drop_prob={round(self.drop_prob,3):0.3f}"
-
-
-class LayerNorm(nn.Module):
-    r"""LayerNorm that supports two data formats: channels_last (default) or channels_first.
-    The ordering of the dimensions in the inputs. channels_last corresponds to inputs with
-    shape (batch_size, height, width, channels) while channels_first corresponds to inputs
-    with shape (batch_size, channels, height, width).
-    """  # noqa: E501
-
-    def __init__(self, normalized_shape, eps=1e-6, data_format="channels_last"):
-        super().__init__()
-        self.weight = nn.Parameter(torch.ones(normalized_shape))
-        self.bias = nn.Parameter(torch.zeros(normalized_shape))
-        self.eps = eps
-        self.data_format = data_format
-        if self.data_format not in ["channels_last", "channels_first"]:
-            raise NotImplementedError
-        self.normalized_shape = (normalized_shape,)
-
-    def forward(self, x):
-        if self.data_format == "channels_last":
-            return F.layer_norm(
-                x, self.normalized_shape, comfy.model_management.cast_to(self.weight, dtype=x.dtype, device=x.device), comfy.model_management.cast_to(self.bias, dtype=x.dtype, device=x.device), self.eps
-            )
-        elif self.data_format == "channels_first":
-            u = x.mean(1, keepdim=True)
-            s = (x - u).pow(2).mean(1, keepdim=True)
-            x = (x - u) / torch.sqrt(s + self.eps)
-            x = comfy.model_management.cast_to(self.weight[:, None], dtype=x.dtype, device=x.device) * x + comfy.model_management.cast_to(self.bias[:, None], dtype=x.dtype, device=x.device)
-            return x
-
-
-class ConvNeXtBlock(nn.Module):
-    r"""ConvNeXt Block. There are two equivalent implementations:
-    (1) DwConv -> LayerNorm (channels_first) -> 1x1 Conv -> GELU -> 1x1 Conv; all in (N, C, H, W)
-    (2) DwConv -> Permute to (N, H, W, C); LayerNorm (channels_last) -> Linear -> GELU -> Linear; Permute back
-    We use (2) as we find it slightly faster in PyTorch
-
-    Args:
-        dim (int): Number of input channels.
-        drop_path (float): Stochastic depth rate. Default: 0.0
-        layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6.
-        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4.0.
-        kernel_size (int): Kernel size for depthwise conv. Default: 7.
-        dilation (int): Dilation for depthwise conv. Default: 1.
-    """  # noqa: E501
-
-    def __init__(
-        self,
-        dim: int,
-        drop_path: float = 0.0,
-        layer_scale_init_value: float = 1e-6,
-        mlp_ratio: float = 4.0,
-        kernel_size: int = 7,
-        dilation: int = 1,
-    ):
-        super().__init__()
-
-        self.dwconv = ops.Conv1d(
-            dim,
-            dim,
-            kernel_size=kernel_size,
-            padding=int(dilation * (kernel_size - 1) / 2),
-            groups=dim,
-        )  # depthwise conv
-        self.norm = LayerNorm(dim, eps=1e-6)
-        self.pwconv1 = ops.Linear(
-            dim, int(mlp_ratio * dim)
-        )  # pointwise/1x1 convs, implemented with linear layers
-        self.act = nn.GELU()
-        self.pwconv2 = ops.Linear(int(mlp_ratio * dim), dim)
-        self.gamma = (
-            nn.Parameter(torch.empty((dim)), requires_grad=False)
-            if layer_scale_init_value > 0
-            else None
-        )
-        self.drop_path = DropPath(
-            drop_path) if drop_path > 0.0 else nn.Identity()
-
-    def forward(self, x, apply_residual: bool = True):
-        input = x
-
-        x = self.dwconv(x)
-        x = x.permute(0, 2, 1)  # (N, C, L) -> (N, L, C)
-        x = self.norm(x)
-        x = self.pwconv1(x)
-        x = self.act(x)
-        x = self.pwconv2(x)
-
-        if self.gamma is not None:
-            x = comfy.model_management.cast_to(self.gamma, dtype=x.dtype, device=x.device) * x
-
-        x = x.permute(0, 2, 1)  # (N, L, C) -> (N, C, L)
-        x = self.drop_path(x)
-
-        if apply_residual:
-            x = input + x
-
-        return x
-
-
-class ParallelConvNeXtBlock(nn.Module):
-    def __init__(self, kernel_sizes: List[int], *args, **kwargs):
-        super().__init__()
-        self.blocks = nn.ModuleList(
-            [
-                ConvNeXtBlock(kernel_size=kernel_size, *args, **kwargs)
-                for kernel_size in kernel_sizes
-            ]
-        )
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        return torch.stack(
-            [block(x, apply_residual=False) for block in self.blocks] + [x],
-            dim=1,
-        ).sum(dim=1)
-
-
-class ConvNeXtEncoder(nn.Module):
-    def __init__(
-        self,
-        input_channels=3,
-        depths=[3, 3, 9, 3],
-        dims=[96, 192, 384, 768],
-        drop_path_rate=0.0,
-        layer_scale_init_value=1e-6,
-        kernel_sizes: Tuple[int] = (7,),
-    ):
-        super().__init__()
-        assert len(depths) == len(dims)
-
-        self.channel_layers = nn.ModuleList()
-        stem = nn.Sequential(
-            ops.Conv1d(
-                input_channels,
-                dims[0],
-                kernel_size=7,
-                padding=3,
-                padding_mode="replicate",
-            ),
-            LayerNorm(dims[0], eps=1e-6, data_format="channels_first"),
-        )
-        self.channel_layers.append(stem)
-
-        for i in range(len(depths) - 1):
-            mid_layer = nn.Sequential(
-                LayerNorm(dims[i], eps=1e-6, data_format="channels_first"),
-                ops.Conv1d(dims[i], dims[i + 1], kernel_size=1),
-            )
-            self.channel_layers.append(mid_layer)
-
-        block_fn = (
-            partial(ConvNeXtBlock, kernel_size=kernel_sizes[0])
-            if len(kernel_sizes) == 1
-            else partial(ParallelConvNeXtBlock, kernel_sizes=kernel_sizes)
-        )
-
-        self.stages = nn.ModuleList()
-        drop_path_rates = [
-            x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))
-        ]
-
-        cur = 0
-        for i in range(len(depths)):
-            stage = nn.Sequential(
-                *[
-                    block_fn(
-                        dim=dims[i],
-                        drop_path=drop_path_rates[cur + j],
-                        layer_scale_init_value=layer_scale_init_value,
-                    )
-                    for j in range(depths[i])
-                ]
-            )
-            self.stages.append(stage)
-            cur += depths[i]
-
-        self.norm = LayerNorm(dims[-1], eps=1e-6, data_format="channels_first")
-
-    def forward(
-        self,
-        x: torch.Tensor,
-    ) -> torch.Tensor:
-        for channel_layer, stage in zip(self.channel_layers, self.stages):
-            x = channel_layer(x)
-            x = stage(x)
-
-        return self.norm(x)
-
-
-def get_padding(kernel_size, dilation=1):
-    return (kernel_size * dilation - dilation) // 2
-
-
-class ResBlock1(torch.nn.Module):
-    def __init__(self, channels, kernel_size=3, dilation=(1, 3, 5)):
-        super().__init__()
-
-        self.convs1 = nn.ModuleList(
-            [
-                torch.nn.utils.parametrizations.weight_norm(
-                    ops.Conv1d(
-                        channels,
-                        channels,
-                        kernel_size,
-                        1,
-                        dilation=dilation[0],
-                        padding=get_padding(kernel_size, dilation[0]),
-                    )
-                ),
-                torch.nn.utils.parametrizations.weight_norm(
-                    ops.Conv1d(
-                        channels,
-                        channels,
-                        kernel_size,
-                        1,
-                        dilation=dilation[1],
-                        padding=get_padding(kernel_size, dilation[1]),
-                    )
-                ),
-                torch.nn.utils.parametrizations.weight_norm(
-                    ops.Conv1d(
-                        channels,
-                        channels,
-                        kernel_size,
-                        1,
-                        dilation=dilation[2],
-                        padding=get_padding(kernel_size, dilation[2]),
-                    )
-                ),
-            ]
-        )
-
-        self.convs2 = nn.ModuleList(
-            [
-                torch.nn.utils.parametrizations.weight_norm(
-                    ops.Conv1d(
-                        channels,
-                        channels,
-                        kernel_size,
-                        1,
-                        dilation=1,
-                        padding=get_padding(kernel_size, 1),
-                    )
-                ),
-                torch.nn.utils.parametrizations.weight_norm(
-                    ops.Conv1d(
-                        channels,
-                        channels,
-                        kernel_size,
-                        1,
-                        dilation=1,
-                        padding=get_padding(kernel_size, 1),
-                    )
-                ),
-                torch.nn.utils.parametrizations.weight_norm(
-                    ops.Conv1d(
-                        channels,
-                        channels,
-                        kernel_size,
-                        1,
-                        dilation=1,
-                        padding=get_padding(kernel_size, 1),
-                    )
-                ),
-            ]
-        )
-
-    def forward(self, x):
-        for c1, c2 in zip(self.convs1, self.convs2):
-            xt = F.silu(x)
-            xt = c1(xt)
-            xt = F.silu(xt)
-            xt = c2(xt)
-            x = xt + x
-        return x
-
-    def remove_weight_norm(self):
-        for conv in self.convs1:
-            remove_weight_norm(conv)
-        for conv in self.convs2:
-            remove_weight_norm(conv)
-
-
-class HiFiGANGenerator(nn.Module):
-    def __init__(
-        self,
-        *,
-        hop_length: int = 512,
-        upsample_rates: Tuple[int] = (8, 8, 2, 2, 2),
-        upsample_kernel_sizes: Tuple[int] = (16, 16, 8, 2, 2),
-        resblock_kernel_sizes: Tuple[int] = (3, 7, 11),
-        resblock_dilation_sizes: Tuple[Tuple[int]] = (
-            (1, 3, 5), (1, 3, 5), (1, 3, 5)),
-        num_mels: int = 128,
-        upsample_initial_channel: int = 512,
-        use_template: bool = True,
-        pre_conv_kernel_size: int = 7,
-        post_conv_kernel_size: int = 7,
-        post_activation: Callable = partial(nn.SiLU, inplace=True),
-    ):
-        super().__init__()
-
-        assert (
-            prod(upsample_rates) == hop_length
-        ), f"hop_length must be {prod(upsample_rates)}"
-
-        self.conv_pre = torch.nn.utils.parametrizations.weight_norm(
-            ops.Conv1d(
-                num_mels,
-                upsample_initial_channel,
-                pre_conv_kernel_size,
-                1,
-                padding=get_padding(pre_conv_kernel_size),
-            )
-        )
-
-        self.num_upsamples = len(upsample_rates)
-        self.num_kernels = len(resblock_kernel_sizes)
-
-        self.noise_convs = nn.ModuleList()
-        self.use_template = use_template
-        self.ups = nn.ModuleList()
-
-        for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)):
-            c_cur = upsample_initial_channel // (2 ** (i + 1))
-            self.ups.append(
-                torch.nn.utils.parametrizations.weight_norm(
-                    ops.ConvTranspose1d(
-                        upsample_initial_channel // (2**i),
-                        upsample_initial_channel // (2 ** (i + 1)),
-                        k,
-                        u,
-                        padding=(k - u) // 2,
-                    )
-                )
-            )
-
-            if not use_template:
-                continue
-
-            if i + 1 < len(upsample_rates):
-                stride_f0 = np.prod(upsample_rates[i + 1:])
-                self.noise_convs.append(
-                    ops.Conv1d(
-                        1,
-                        c_cur,
-                        kernel_size=stride_f0 * 2,
-                        stride=stride_f0,
-                        padding=stride_f0 // 2,
-                    )
-                )
-            else:
-                self.noise_convs.append(ops.Conv1d(1, c_cur, kernel_size=1))
-
-        self.resblocks = nn.ModuleList()
-        for i in range(len(self.ups)):
-            ch = upsample_initial_channel // (2 ** (i + 1))
-            for k, d in zip(resblock_kernel_sizes, resblock_dilation_sizes):
-                self.resblocks.append(ResBlock1(ch, k, d))
-
-        self.activation_post = post_activation()
-        self.conv_post = torch.nn.utils.parametrizations.weight_norm(
-            ops.Conv1d(
-                ch,
-                1,
-                post_conv_kernel_size,
-                1,
-                padding=get_padding(post_conv_kernel_size),
-            )
-        )
-
-    def forward(self, x, template=None):
-        x = self.conv_pre(x)
-
-        for i in range(self.num_upsamples):
-            x = F.silu(x, inplace=True)
-            x = self.ups[i](x)
-
-            if self.use_template:
-                x = x + self.noise_convs[i](template)
-
-            xs = None
-
-            for j in range(self.num_kernels):
-                if xs is None:
-                    xs = self.resblocks[i * self.num_kernels + j](x)
-                else:
-                    xs += self.resblocks[i * self.num_kernels + j](x)
-
-            x = xs / self.num_kernels
-
-        x = self.activation_post(x)
-        x = self.conv_post(x)
-        x = torch.tanh(x)
-
-        return x
-
-    def remove_weight_norm(self):
-        for up in self.ups:
-            remove_weight_norm(up)
-        for block in self.resblocks:
-            block.remove_weight_norm()
-        remove_weight_norm(self.conv_pre)
-        remove_weight_norm(self.conv_post)
-
-
-class ADaMoSHiFiGANV1(nn.Module):
-    def __init__(
-        self,
-        input_channels: int = 128,
-        depths: List[int] = [3, 3, 9, 3],
-        dims: List[int] = [128, 256, 384, 512],
-        drop_path_rate: float = 0.0,
-        kernel_sizes: Tuple[int] = (7,),
-        upsample_rates: Tuple[int] = (4, 4, 2, 2, 2, 2, 2),
-        upsample_kernel_sizes: Tuple[int] = (8, 8, 4, 4, 4, 4, 4),
-        resblock_kernel_sizes: Tuple[int] = (3, 7, 11, 13),
-        resblock_dilation_sizes: Tuple[Tuple[int]] = (
-            (1, 3, 5), (1, 3, 5), (1, 3, 5), (1, 3, 5)),
-        num_mels: int = 512,
-        upsample_initial_channel: int = 1024,
-        use_template: bool = False,
-        pre_conv_kernel_size: int = 13,
-        post_conv_kernel_size: int = 13,
-        sampling_rate: int = 44100,
-        n_fft: int = 2048,
-        win_length: int = 2048,
-        hop_length: int = 512,
-        f_min: int = 40,
-        f_max: int = 16000,
-        n_mels: int = 128,
-    ):
-        super().__init__()
-
-        self.backbone = ConvNeXtEncoder(
-            input_channels=input_channels,
-            depths=depths,
-            dims=dims,
-            drop_path_rate=drop_path_rate,
-            kernel_sizes=kernel_sizes,
-        )
-
-        self.head = HiFiGANGenerator(
-            hop_length=hop_length,
-            upsample_rates=upsample_rates,
-            upsample_kernel_sizes=upsample_kernel_sizes,
-            resblock_kernel_sizes=resblock_kernel_sizes,
-            resblock_dilation_sizes=resblock_dilation_sizes,
-            num_mels=num_mels,
-            upsample_initial_channel=upsample_initial_channel,
-            use_template=use_template,
-            pre_conv_kernel_size=pre_conv_kernel_size,
-            post_conv_kernel_size=post_conv_kernel_size,
-        )
-        self.sampling_rate = sampling_rate
-        self.mel_transform = LogMelSpectrogram(
-            sample_rate=sampling_rate,
-            n_fft=n_fft,
-            win_length=win_length,
-            hop_length=hop_length,
-            f_min=f_min,
-            f_max=f_max,
-            n_mels=n_mels,
-        )
-        self.eval()
-
-    @torch.no_grad()
-    def decode(self, mel):
-        y = self.backbone(mel)
-        y = self.head(y)
-        return y
-
-    @torch.no_grad()
-    def encode(self, x):
-        return self.mel_transform(x)
-
-    def forward(self, mel):
-        y = self.backbone(mel)
-        y = self.head(y)
-        return y

comfy/ldm/audio/autoencoder.py

@@ -1,276 +0,0 @@
-# code adapted from: https://github.com/Stability-AI/stable-audio-tools
-
-import torch
-from torch import nn
-from typing import Literal
-import math
-import comfy.ops
-ops = comfy.ops.disable_weight_init
-
-def vae_sample(mean, scale):
-        stdev = nn.functional.softplus(scale) + 1e-4
-        var = stdev * stdev
-        logvar = torch.log(var)
-        latents = torch.randn_like(mean) * stdev + mean
-
-        kl = (mean * mean + var - logvar - 1).sum(1).mean()
-
-        return latents, kl
-
-class VAEBottleneck(nn.Module):
-    def __init__(self):
-        super().__init__()
-        self.is_discrete = False
-
-    def encode(self, x, return_info=False, **kwargs):
-        info = {}
-
-        mean, scale = x.chunk(2, dim=1)
-
-        x, kl = vae_sample(mean, scale)
-
-        info["kl"] = kl
-
-        if return_info:
-            return x, info
-        else:
-            return x
-
-    def decode(self, x):
-        return x
-
-
-def snake_beta(x, alpha, beta):
-    return x + (1.0 / (beta + 0.000000001)) * pow(torch.sin(x * alpha), 2)
-
-# Adapted from https://github.com/NVIDIA/BigVGAN/blob/main/activations.py under MIT license
-class SnakeBeta(nn.Module):
-
-    def __init__(self, in_features, alpha=1.0, alpha_trainable=True, alpha_logscale=True):
-        super(SnakeBeta, self).__init__()
-        self.in_features = in_features
-
-        # initialize alpha
-        self.alpha_logscale = alpha_logscale
-        if self.alpha_logscale: # log scale alphas initialized to zeros
-            self.alpha = nn.Parameter(torch.zeros(in_features) * alpha)
-            self.beta = nn.Parameter(torch.zeros(in_features) * alpha)
-        else: # linear scale alphas initialized to ones
-            self.alpha = nn.Parameter(torch.ones(in_features) * alpha)
-            self.beta = nn.Parameter(torch.ones(in_features) * alpha)
-
-        # self.alpha.requires_grad = alpha_trainable
-        # self.beta.requires_grad = alpha_trainable
-
-        self.no_div_by_zero = 0.000000001
-
-    def forward(self, x):
-        alpha = self.alpha.unsqueeze(0).unsqueeze(-1).to(x.device) # line up with x to [B, C, T]
-        beta = self.beta.unsqueeze(0).unsqueeze(-1).to(x.device)
-        if self.alpha_logscale:
-            alpha = torch.exp(alpha)
-            beta = torch.exp(beta)
-        x = snake_beta(x, alpha, beta)
-
-        return x
-
-def WNConv1d(*args, **kwargs):
-    return torch.nn.utils.parametrizations.weight_norm(ops.Conv1d(*args, **kwargs))
-
-def WNConvTranspose1d(*args, **kwargs):
-    return torch.nn.utils.parametrizations.weight_norm(ops.ConvTranspose1d(*args, **kwargs))
-
-def get_activation(activation: Literal["elu", "snake", "none"], antialias=False, channels=None) -> nn.Module:
-    if activation == "elu":
-        act = torch.nn.ELU()
-    elif activation == "snake":
-        act = SnakeBeta(channels)
-    elif activation == "none":
-        act = torch.nn.Identity()
-    else:
-        raise ValueError(f"Unknown activation {activation}")
-
-    if antialias:
-        act = Activation1d(act)  # noqa: F821 Activation1d is not defined
-
-    return act
-
-
-class ResidualUnit(nn.Module):
-    def __init__(self, in_channels, out_channels, dilation, use_snake=False, antialias_activation=False):
-        super().__init__()
-
-        self.dilation = dilation
-
-        padding = (dilation * (7-1)) // 2
-
-        self.layers = nn.Sequential(
-            get_activation("snake" if use_snake else "elu", antialias=antialias_activation, channels=out_channels),
-            WNConv1d(in_channels=in_channels, out_channels=out_channels,
-                      kernel_size=7, dilation=dilation, padding=padding),
-            get_activation("snake" if use_snake else "elu", antialias=antialias_activation, channels=out_channels),
-            WNConv1d(in_channels=out_channels, out_channels=out_channels,
-                      kernel_size=1)
-        )
-
-    def forward(self, x):
-        res = x
-
-        #x = checkpoint(self.layers, x)
-        x = self.layers(x)
-
-        return x + res
-
-class EncoderBlock(nn.Module):
-    def __init__(self, in_channels, out_channels, stride, use_snake=False, antialias_activation=False):
-        super().__init__()
-
-        self.layers = nn.Sequential(
-            ResidualUnit(in_channels=in_channels,
-                         out_channels=in_channels, dilation=1, use_snake=use_snake),
-            ResidualUnit(in_channels=in_channels,
-                         out_channels=in_channels, dilation=3, use_snake=use_snake),
-            ResidualUnit(in_channels=in_channels,
-                         out_channels=in_channels, dilation=9, use_snake=use_snake),
-            get_activation("snake" if use_snake else "elu", antialias=antialias_activation, channels=in_channels),
-            WNConv1d(in_channels=in_channels, out_channels=out_channels,
-                      kernel_size=2*stride, stride=stride, padding=math.ceil(stride/2)),
-        )
-
-    def forward(self, x):
-        return self.layers(x)
-
-class DecoderBlock(nn.Module):
-    def __init__(self, in_channels, out_channels, stride, use_snake=False, antialias_activation=False, use_nearest_upsample=False):
-        super().__init__()
-
-        if use_nearest_upsample:
-            upsample_layer = nn.Sequential(
-                nn.Upsample(scale_factor=stride, mode="nearest"),
-                WNConv1d(in_channels=in_channels,
-                        out_channels=out_channels,
-                        kernel_size=2*stride,
-                        stride=1,
-                        bias=False,
-                        padding='same')
-            )
-        else:
-            upsample_layer = WNConvTranspose1d(in_channels=in_channels,
-                               out_channels=out_channels,
-                               kernel_size=2*stride, stride=stride, padding=math.ceil(stride/2))
-
-        self.layers = nn.Sequential(
-            get_activation("snake" if use_snake else "elu", antialias=antialias_activation, channels=in_channels),
-            upsample_layer,
-            ResidualUnit(in_channels=out_channels, out_channels=out_channels,
-                         dilation=1, use_snake=use_snake),
-            ResidualUnit(in_channels=out_channels, out_channels=out_channels,
-                         dilation=3, use_snake=use_snake),
-            ResidualUnit(in_channels=out_channels, out_channels=out_channels,
-                         dilation=9, use_snake=use_snake),
-        )
-
-    def forward(self, x):
-        return self.layers(x)
-
-class OobleckEncoder(nn.Module):
-    def __init__(self,
-                 in_channels=2,
-                 channels=128,
-                 latent_dim=32,
-                 c_mults = [1, 2, 4, 8],
-                 strides = [2, 4, 8, 8],
-                 use_snake=False,
-                 antialias_activation=False
-        ):
-        super().__init__()
-
-        c_mults = [1] + c_mults
-
-        self.depth = len(c_mults)
-
-        layers = [
-            WNConv1d(in_channels=in_channels, out_channels=c_mults[0] * channels, kernel_size=7, padding=3)
-        ]
-
-        for i in range(self.depth-1):
-            layers += [EncoderBlock(in_channels=c_mults[i]*channels, out_channels=c_mults[i+1]*channels, stride=strides[i], use_snake=use_snake)]
-
-        layers += [
-            get_activation("snake" if use_snake else "elu", antialias=antialias_activation, channels=c_mults[-1] * channels),
-            WNConv1d(in_channels=c_mults[-1]*channels, out_channels=latent_dim, kernel_size=3, padding=1)
-        ]
-
-        self.layers = nn.Sequential(*layers)
-
-    def forward(self, x):
-        return self.layers(x)
-
-
-class OobleckDecoder(nn.Module):
-    def __init__(self,
-                 out_channels=2,
-                 channels=128,
-                 latent_dim=32,
-                 c_mults = [1, 2, 4, 8],
-                 strides = [2, 4, 8, 8],
-                 use_snake=False,
-                 antialias_activation=False,
-                 use_nearest_upsample=False,
-                 final_tanh=True):
-        super().__init__()
-
-        c_mults = [1] + c_mults
-
-        self.depth = len(c_mults)
-
-        layers = [
-            WNConv1d(in_channels=latent_dim, out_channels=c_mults[-1]*channels, kernel_size=7, padding=3),
-        ]
-
-        for i in range(self.depth-1, 0, -1):
-            layers += [DecoderBlock(
-                in_channels=c_mults[i]*channels,
-                out_channels=c_mults[i-1]*channels,
-                stride=strides[i-1],
-                use_snake=use_snake,
-                antialias_activation=antialias_activation,
-                use_nearest_upsample=use_nearest_upsample
-                )
-            ]
-
-        layers += [
-            get_activation("snake" if use_snake else "elu", antialias=antialias_activation, channels=c_mults[0] * channels),
-            WNConv1d(in_channels=c_mults[0] * channels, out_channels=out_channels, kernel_size=7, padding=3, bias=False),
-            nn.Tanh() if final_tanh else nn.Identity()
-        ]
-
-        self.layers = nn.Sequential(*layers)
-
-    def forward(self, x):
-        return self.layers(x)
-
-
-class AudioOobleckVAE(nn.Module):
-    def __init__(self,
-                 in_channels=2,
-                 channels=128,
-                 latent_dim=64,
-                 c_mults = [1, 2, 4, 8, 16],
-                 strides = [2, 4, 4, 8, 8],
-                 use_snake=True,
-                 antialias_activation=False,
-                 use_nearest_upsample=False,
-                 final_tanh=False):
-        super().__init__()
-        self.encoder = OobleckEncoder(in_channels, channels, latent_dim * 2, c_mults, strides, use_snake, antialias_activation)
-        self.decoder = OobleckDecoder(in_channels, channels, latent_dim, c_mults, strides, use_snake, antialias_activation,
-                                      use_nearest_upsample=use_nearest_upsample, final_tanh=final_tanh)
-        self.bottleneck = VAEBottleneck()
-
-    def encode(self, x):
-        return self.bottleneck.encode(self.encoder(x))
-
-    def decode(self, x):
-        return self.decoder(self.bottleneck.decode(x))
-

comfy/ldm/audio/dit.py

@@ -1,896 +0,0 @@
-# code adapted from: https://github.com/Stability-AI/stable-audio-tools
-
-from comfy.ldm.modules.attention import optimized_attention
-import typing as tp
-
-import torch
-
-from einops import rearrange
-from torch import nn
-from torch.nn import functional as F
-import math
-import comfy.ops
-
-class FourierFeatures(nn.Module):
-    def __init__(self, in_features, out_features, std=1., dtype=None, device=None):
-        super().__init__()
-        assert out_features % 2 == 0
-        self.weight = nn.Parameter(torch.empty(
-            [out_features // 2, in_features], dtype=dtype, device=device))
-
-    def forward(self, input):
-        f = 2 * math.pi * input @ comfy.ops.cast_to_input(self.weight.T, input)
-        return torch.cat([f.cos(), f.sin()], dim=-1)
-
-# norms
-class LayerNorm(nn.Module):
-    def __init__(self, dim, bias=False, fix_scale=False, dtype=None, device=None):
-        """
-        bias-less layernorm has been shown to be more stable. most newer models have moved towards rmsnorm, also bias-less
-        """
-        super().__init__()
-
-        self.gamma = nn.Parameter(torch.empty(dim, dtype=dtype, device=device))
-
-        if bias:
-            self.beta = nn.Parameter(torch.empty(dim, dtype=dtype, device=device))
-        else:
-            self.beta = None
-
-    def forward(self, x):
-        beta = self.beta
-        if beta is not None:
-            beta = comfy.ops.cast_to_input(beta, x)
-        return F.layer_norm(x, x.shape[-1:], weight=comfy.ops.cast_to_input(self.gamma, x), bias=beta)
-
-class GLU(nn.Module):
-    def __init__(
-        self,
-        dim_in,
-        dim_out,
-        activation,
-        use_conv = False,
-        conv_kernel_size = 3,
-        dtype=None,
-        device=None,
-        operations=None,
-    ):
-        super().__init__()
-        self.act = activation
-        self.proj = operations.Linear(dim_in, dim_out * 2, dtype=dtype, device=device) if not use_conv else operations.Conv1d(dim_in, dim_out * 2, conv_kernel_size, padding = (conv_kernel_size // 2), dtype=dtype, device=device)
-        self.use_conv = use_conv
-
-    def forward(self, x):
-        if self.use_conv:
-            x = rearrange(x, 'b n d -> b d n')
-            x = self.proj(x)
-            x = rearrange(x, 'b d n -> b n d')
-        else:
-            x = self.proj(x)
-
-        x, gate = x.chunk(2, dim = -1)
-        return x * self.act(gate)
-
-class AbsolutePositionalEmbedding(nn.Module):
-    def __init__(self, dim, max_seq_len):
-        super().__init__()
-        self.scale = dim ** -0.5
-        self.max_seq_len = max_seq_len
-        self.emb = nn.Embedding(max_seq_len, dim)
-
-    def forward(self, x, pos = None, seq_start_pos = None):
-        seq_len, device = x.shape[1], x.device
-        assert seq_len <= self.max_seq_len, f'you are passing in a sequence length of {seq_len} but your absolute positional embedding has a max sequence length of {self.max_seq_len}'
-
-        if pos is None:
-            pos = torch.arange(seq_len, device = device)
-
-        if seq_start_pos is not None:
-            pos = (pos - seq_start_pos[..., None]).clamp(min = 0)
-
-        pos_emb = self.emb(pos)
-        pos_emb = pos_emb * self.scale
-        return pos_emb
-
-class ScaledSinusoidalEmbedding(nn.Module):
-    def __init__(self, dim, theta = 10000):
-        super().__init__()
-        assert (dim % 2) == 0, 'dimension must be divisible by 2'
-        self.scale = nn.Parameter(torch.ones(1) * dim ** -0.5)
-
-        half_dim = dim // 2
-        freq_seq = torch.arange(half_dim).float() / half_dim
-        inv_freq = theta ** -freq_seq
-        self.register_buffer('inv_freq', inv_freq, persistent = False)
-
-    def forward(self, x, pos = None, seq_start_pos = None):
-        seq_len, device = x.shape[1], x.device
-
-        if pos is None:
-            pos = torch.arange(seq_len, device = device)
-
-        if seq_start_pos is not None:
-            pos = pos - seq_start_pos[..., None]
-
-        emb = torch.einsum('i, j -> i j', pos, self.inv_freq)
-        emb = torch.cat((emb.sin(), emb.cos()), dim = -1)
-        return emb * self.scale
-
-class RotaryEmbedding(nn.Module):
-    def __init__(
-        self,
-        dim,
-        use_xpos = False,
-        scale_base = 512,
-        interpolation_factor = 1.,
-        base = 10000,
-        base_rescale_factor = 1.,
-        dtype=None,
-        device=None,
-    ):
-        super().__init__()
-        # proposed by reddit user bloc97, to rescale rotary embeddings to longer sequence length without fine-tuning
-        # has some connection to NTK literature
-        # https://www.reddit.com/r/LocalLLaMA/comments/14lz7j5/ntkaware_scaled_rope_allows_llama_models_to_have/
-        base *= base_rescale_factor ** (dim / (dim - 2))
-
-        # inv_freq = 1. / (base ** (torch.arange(0, dim, 2).float() / dim))
-        self.register_buffer('inv_freq', torch.empty((dim // 2,), device=device, dtype=dtype))
-
-        assert interpolation_factor >= 1.
-        self.interpolation_factor = interpolation_factor
-
-        if not use_xpos:
-            self.register_buffer('scale', None)
-            return
-
-        scale = (torch.arange(0, dim, 2) + 0.4 * dim) / (1.4 * dim)
-
-        self.scale_base = scale_base
-        self.register_buffer('scale', scale)
-
-    def forward_from_seq_len(self, seq_len, device, dtype):
-        # device = self.inv_freq.device
-
-        t = torch.arange(seq_len, device=device, dtype=dtype)
-        return self.forward(t)
-
-    def forward(self, t):
-        # device = self.inv_freq.device
-        device = t.device
-
-        # t = t.to(torch.float32)
-
-        t = t / self.interpolation_factor
-
-        freqs = torch.einsum('i , j -> i j', t, comfy.ops.cast_to_input(self.inv_freq, t))
-        freqs = torch.cat((freqs, freqs), dim = -1)
-
-        if self.scale is None:
-            return freqs, 1.
-
-        power = (torch.arange(seq_len, device = device) - (seq_len // 2)) / self.scale_base  # noqa: F821 seq_len is not defined
-        scale = comfy.ops.cast_to_input(self.scale, t) ** rearrange(power, 'n -> n 1')
-        scale = torch.cat((scale, scale), dim = -1)
-
-        return freqs, scale
-
-def rotate_half(x):
-    x = rearrange(x, '... (j d) -> ... j d', j = 2)
-    x1, x2 = x.unbind(dim = -2)
-    return torch.cat((-x2, x1), dim = -1)
-
-def apply_rotary_pos_emb(t, freqs, scale = 1):
-    out_dtype = t.dtype
-
-    # cast to float32 if necessary for numerical stability
-    dtype = t.dtype #reduce(torch.promote_types, (t.dtype, freqs.dtype, torch.float32))
-    rot_dim, seq_len = freqs.shape[-1], t.shape[-2]
-    freqs, t = freqs.to(dtype), t.to(dtype)
-    freqs = freqs[-seq_len:, :]
-
-    if t.ndim == 4 and freqs.ndim == 3:
-        freqs = rearrange(freqs, 'b n d -> b 1 n d')
-
-    # partial rotary embeddings, Wang et al. GPT-J
-    t, t_unrotated = t[..., :rot_dim], t[..., rot_dim:]
-    t = (t * freqs.cos() * scale) + (rotate_half(t) * freqs.sin() * scale)
-
-    t, t_unrotated = t.to(out_dtype), t_unrotated.to(out_dtype)
-
-    return torch.cat((t, t_unrotated), dim = -1)
-
-class FeedForward(nn.Module):
-    def __init__(
-        self,
-        dim,
-        dim_out = None,
-        mult = 4,
-        no_bias = False,
-        glu = True,
-        use_conv = False,
-        conv_kernel_size = 3,
-        zero_init_output = True,
-        dtype=None,
-        device=None,
-        operations=None,
-    ):
-        super().__init__()
-        inner_dim = int(dim * mult)
-
-        # Default to SwiGLU
-
-        activation = nn.SiLU()
-
-        dim_out = dim if dim_out is None else dim_out
-
-        if glu:
-            linear_in = GLU(dim, inner_dim, activation, dtype=dtype, device=device, operations=operations)
-        else:
-            linear_in = nn.Sequential(
-                rearrange('b n d -> b d n') if use_conv else nn.Identity(),
-                operations.Linear(dim, inner_dim, bias = not no_bias, dtype=dtype, device=device) if not use_conv else operations.Conv1d(dim, inner_dim, conv_kernel_size, padding = (conv_kernel_size // 2), bias = not no_bias, dtype=dtype, device=device),
-                rearrange('b n d -> b d n') if use_conv else nn.Identity(),
-                activation
-            )
-
-        linear_out = operations.Linear(inner_dim, dim_out, bias = not no_bias, dtype=dtype, device=device) if not use_conv else operations.Conv1d(inner_dim, dim_out, conv_kernel_size, padding = (conv_kernel_size // 2), bias = not no_bias, dtype=dtype, device=device)
-
-        # # init last linear layer to 0
-        # if zero_init_output:
-        #     nn.init.zeros_(linear_out.weight)
-        #     if not no_bias:
-        #         nn.init.zeros_(linear_out.bias)
-
-
-        self.ff = nn.Sequential(
-            linear_in,
-            rearrange('b d n -> b n d') if use_conv else nn.Identity(),
-            linear_out,
-            rearrange('b n d -> b d n') if use_conv else nn.Identity(),
-        )
-
-    def forward(self, x):
-        return self.ff(x)
-
-class Attention(nn.Module):
-    def __init__(
-        self,
-        dim,
-        dim_heads = 64,
-        dim_context = None,
-        causal = False,
-        zero_init_output=True,
-        qk_norm = False,
-        natten_kernel_size = None,
-        dtype=None,
-        device=None,
-        operations=None,
-    ):
-        super().__init__()
-        self.dim = dim
-        self.dim_heads = dim_heads
-        self.causal = causal
-
-        dim_kv = dim_context if dim_context is not None else dim
-
-        self.num_heads = dim // dim_heads
-        self.kv_heads = dim_kv // dim_heads
-
-        if dim_context is not None:
-            self.to_q = operations.Linear(dim, dim, bias=False, dtype=dtype, device=device)
-            self.to_kv = operations.Linear(dim_kv, dim_kv * 2, bias=False, dtype=dtype, device=device)
-        else:
-            self.to_qkv = operations.Linear(dim, dim * 3, bias=False, dtype=dtype, device=device)
-
-        self.to_out = operations.Linear(dim, dim, bias=False, dtype=dtype, device=device)
-
-        # if zero_init_output:
-        #     nn.init.zeros_(self.to_out.weight)
-
-        self.qk_norm = qk_norm
-
-
-    def forward(
-        self,
-        x,
-        context = None,
-        mask = None,
-        context_mask = None,
-        rotary_pos_emb = None,
-        causal = None
-    ):
-        h, kv_h, has_context = self.num_heads, self.kv_heads, context is not None
-
-        kv_input = context if has_context else x
-
-        if hasattr(self, 'to_q'):
-            # Use separate linear projections for q and k/v
-            q = self.to_q(x)
-            q = rearrange(q, 'b n (h d) -> b h n d', h = h)
-
-            k, v = self.to_kv(kv_input).chunk(2, dim=-1)
-
-            k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = kv_h), (k, v))
-        else:
-            # Use fused linear projection
-            q, k, v = self.to_qkv(x).chunk(3, dim=-1)
-            q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = h), (q, k, v))
-
-        # Normalize q and k for cosine sim attention
-        if self.qk_norm:
-            q = F.normalize(q, dim=-1)
-            k = F.normalize(k, dim=-1)
-
-        if rotary_pos_emb is not None and not has_context:
-            freqs, _ = rotary_pos_emb
-
-            q_dtype = q.dtype
-            k_dtype = k.dtype
-
-            q = q.to(torch.float32)
-            k = k.to(torch.float32)
-            freqs = freqs.to(torch.float32)
-
-            q = apply_rotary_pos_emb(q, freqs)
-            k = apply_rotary_pos_emb(k, freqs)
-
-            q = q.to(q_dtype)
-            k = k.to(k_dtype)
-
-        input_mask = context_mask
-
-        if input_mask is None and not has_context:
-            input_mask = mask
-
-        # determine masking
-        masks = []
-
-        if input_mask is not None:
-            input_mask = rearrange(input_mask, 'b j -> b 1 1 j')
-            masks.append(~input_mask)
-
-        # Other masks will be added here later
-        n = q.shape[-2]
-
-        causal = self.causal if causal is None else causal
-
-        if n == 1 and causal:
-            causal = False
-
-        if h != kv_h:
-            # Repeat interleave kv_heads to match q_heads
-            heads_per_kv_head = h // kv_h
-            k, v = map(lambda t: t.repeat_interleave(heads_per_kv_head, dim = 1), (k, v))
-
-        out = optimized_attention(q, k, v, h, skip_reshape=True)
-        out = self.to_out(out)
-
-        if mask is not None:
-            mask = rearrange(mask, 'b n -> b n 1')
-            out = out.masked_fill(~mask, 0.)
-
-        return out
-
-class ConformerModule(nn.Module):
-    def __init__(
-        self,
-        dim,
-        norm_kwargs = {},
-    ):
-
-        super().__init__()
-
-        self.dim = dim
-
-        self.in_norm = LayerNorm(dim, **norm_kwargs)
-        self.pointwise_conv = nn.Conv1d(dim, dim, kernel_size=1, bias=False)
-        self.glu = GLU(dim, dim, nn.SiLU())
-        self.depthwise_conv = nn.Conv1d(dim, dim, kernel_size=17, groups=dim, padding=8, bias=False)
-        self.mid_norm = LayerNorm(dim, **norm_kwargs) # This is a batch norm in the original but I don't like batch norm
-        self.swish = nn.SiLU()
-        self.pointwise_conv_2 = nn.Conv1d(dim, dim, kernel_size=1, bias=False)
-
-    def forward(self, x):
-        x = self.in_norm(x)
-        x = rearrange(x, 'b n d -> b d n')
-        x = self.pointwise_conv(x)
-        x = rearrange(x, 'b d n -> b n d')
-        x = self.glu(x)
-        x = rearrange(x, 'b n d -> b d n')
-        x = self.depthwise_conv(x)
-        x = rearrange(x, 'b d n -> b n d')
-        x = self.mid_norm(x)
-        x = self.swish(x)
-        x = rearrange(x, 'b n d -> b d n')
-        x = self.pointwise_conv_2(x)
-        x = rearrange(x, 'b d n -> b n d')
-
-        return x
-
-class TransformerBlock(nn.Module):
-    def __init__(
-            self,
-            dim,
-            dim_heads = 64,
-            cross_attend = False,
-            dim_context = None,
-            global_cond_dim = None,
-            causal = False,
-            zero_init_branch_outputs = True,
-            conformer = False,
-            layer_ix = -1,
-            remove_norms = False,
-            attn_kwargs = {},
-            ff_kwargs = {},
-            norm_kwargs = {},
-            dtype=None,
-            device=None,
-            operations=None,
-    ):
-
-        super().__init__()
-        self.dim = dim
-        self.dim_heads = dim_heads
-        self.cross_attend = cross_attend
-        self.dim_context = dim_context
-        self.causal = causal
-
-        self.pre_norm = LayerNorm(dim, dtype=dtype, device=device, **norm_kwargs) if not remove_norms else nn.Identity()
-
-        self.self_attn = Attention(
-            dim,
-            dim_heads = dim_heads,
-            causal = causal,
-            zero_init_output=zero_init_branch_outputs,
-            dtype=dtype,
-            device=device,
-            operations=operations,
-            **attn_kwargs
-        )
-
-        if cross_attend:
-            self.cross_attend_norm = LayerNorm(dim, dtype=dtype, device=device, **norm_kwargs) if not remove_norms else nn.Identity()
-            self.cross_attn = Attention(
-                dim,
-                dim_heads = dim_heads,
-                dim_context=dim_context,
-                causal = causal,
-                zero_init_output=zero_init_branch_outputs,
-                dtype=dtype,
-                device=device,
-                operations=operations,
-                **attn_kwargs
-            )
-
-        self.ff_norm = LayerNorm(dim, dtype=dtype, device=device, **norm_kwargs) if not remove_norms else nn.Identity()
-        self.ff = FeedForward(dim, zero_init_output=zero_init_branch_outputs, dtype=dtype, device=device, operations=operations,**ff_kwargs)
-
-        self.layer_ix = layer_ix
-
-        self.conformer = ConformerModule(dim, norm_kwargs=norm_kwargs) if conformer else None
-
-        self.global_cond_dim = global_cond_dim
-
-        if global_cond_dim is not None:
-            self.to_scale_shift_gate = nn.Sequential(
-                nn.SiLU(),
-                nn.Linear(global_cond_dim, dim * 6, bias=False)
-            )
-
-            nn.init.zeros_(self.to_scale_shift_gate[1].weight)
-            #nn.init.zeros_(self.to_scale_shift_gate_self[1].bias)
-
-    def forward(
-        self,
-        x,
-        context = None,
-        global_cond=None,
-        mask = None,
-        context_mask = None,
-        rotary_pos_emb = None
-    ):
-        if self.global_cond_dim is not None and self.global_cond_dim > 0 and global_cond is not None:
-
-            scale_self, shift_self, gate_self, scale_ff, shift_ff, gate_ff = self.to_scale_shift_gate(global_cond).unsqueeze(1).chunk(6, dim = -1)
-
-            # self-attention with adaLN
-            residual = x
-            x = self.pre_norm(x)
-            x = x * (1 + scale_self) + shift_self
-            x = self.self_attn(x, mask = mask, rotary_pos_emb = rotary_pos_emb)
-            x = x * torch.sigmoid(1 - gate_self)
-            x = x + residual
-
-            if context is not None:
-                x = x + self.cross_attn(self.cross_attend_norm(x), context = context, context_mask = context_mask)
-
-            if self.conformer is not None:
-                x = x + self.conformer(x)
-
-            # feedforward with adaLN
-            residual = x
-            x = self.ff_norm(x)
-            x = x * (1 + scale_ff) + shift_ff
-            x = self.ff(x)
-            x = x * torch.sigmoid(1 - gate_ff)
-            x = x + residual
-
-        else:
-            x = x + self.self_attn(self.pre_norm(x), mask = mask, rotary_pos_emb = rotary_pos_emb)
-
-            if context is not None:
-                x = x + self.cross_attn(self.cross_attend_norm(x), context = context, context_mask = context_mask)
-
-            if self.conformer is not None:
-                x = x + self.conformer(x)
-
-            x = x + self.ff(self.ff_norm(x))
-
-        return x
-
-class ContinuousTransformer(nn.Module):
-    def __init__(
-        self,
-        dim,
-        depth,
-        *,
-        dim_in = None,
-        dim_out = None,
-        dim_heads = 64,
-        cross_attend=False,
-        cond_token_dim=None,
-        global_cond_dim=None,
-        causal=False,
-        rotary_pos_emb=True,
-        zero_init_branch_outputs=True,
-        conformer=False,
-        use_sinusoidal_emb=False,
-        use_abs_pos_emb=False,
-        abs_pos_emb_max_length=10000,
-        dtype=None,
-        device=None,
-        operations=None,
-        **kwargs
-        ):
-
-        super().__init__()
-
-        self.dim = dim
-        self.depth = depth
-        self.causal = causal
-        self.layers = nn.ModuleList([])
-
-        self.project_in = operations.Linear(dim_in, dim, bias=False, dtype=dtype, device=device) if dim_in is not None else nn.Identity()
-        self.project_out = operations.Linear(dim, dim_out, bias=False, dtype=dtype, device=device) if dim_out is not None else nn.Identity()
-
-        if rotary_pos_emb:
-            self.rotary_pos_emb = RotaryEmbedding(max(dim_heads // 2, 32), device=device, dtype=dtype)
-        else:
-            self.rotary_pos_emb = None
-
-        self.use_sinusoidal_emb = use_sinusoidal_emb
-        if use_sinusoidal_emb:
-            self.pos_emb = ScaledSinusoidalEmbedding(dim)
-
-        self.use_abs_pos_emb = use_abs_pos_emb
-        if use_abs_pos_emb:
-            self.pos_emb = AbsolutePositionalEmbedding(dim, abs_pos_emb_max_length)
-
-        for i in range(depth):
-            self.layers.append(
-                TransformerBlock(
-                    dim,
-                    dim_heads = dim_heads,
-                    cross_attend = cross_attend,
-                    dim_context = cond_token_dim,
-                    global_cond_dim = global_cond_dim,
-                    causal = causal,
-                    zero_init_branch_outputs = zero_init_branch_outputs,
-                    conformer=conformer,
-                    layer_ix=i,
-                    dtype=dtype,
-                    device=device,
-                    operations=operations,
-                    **kwargs
-                )
-            )
-
-    def forward(
-        self,
-        x,
-        mask = None,
-        prepend_embeds = None,
-        prepend_mask = None,
-        global_cond = None,
-        return_info = False,
-        **kwargs
-    ):
-        patches_replace = kwargs.get("transformer_options", {}).get("patches_replace", {})
-        batch, seq, device = *x.shape[:2], x.device
-        context = kwargs["context"]
-
-        info = {
-            "hidden_states": [],
-        }
-
-        x = self.project_in(x)
-
-        if prepend_embeds is not None:
-            prepend_length, prepend_dim = prepend_embeds.shape[1:]
-
-            assert prepend_dim == x.shape[-1], 'prepend dimension must match sequence dimension'
-
-            x = torch.cat((prepend_embeds, x), dim = -2)
-
-            if prepend_mask is not None or mask is not None:
-                mask = mask if mask is not None else torch.ones((batch, seq), device = device, dtype = torch.bool)
-                prepend_mask = prepend_mask if prepend_mask is not None else torch.ones((batch, prepend_length), device = device, dtype = torch.bool)
-
-                mask = torch.cat((prepend_mask, mask), dim = -1)
-
-        # Attention layers
-
-        if self.rotary_pos_emb is not None:
-            rotary_pos_emb = self.rotary_pos_emb.forward_from_seq_len(x.shape[1], dtype=x.dtype, device=x.device)
-        else:
-            rotary_pos_emb = None
-
-        if self.use_sinusoidal_emb or self.use_abs_pos_emb:
-            x = x + self.pos_emb(x)
-
-        blocks_replace = patches_replace.get("dit", {})
-        # Iterate over the transformer layers
-        for i, layer in enumerate(self.layers):
-            if ("double_block", i) in blocks_replace:
-                def block_wrap(args):
-                    out = {}
-                    out["img"] = layer(args["img"], rotary_pos_emb=args["pe"], global_cond=args["vec"], context=args["txt"])
-                    return out
-
-                out = blocks_replace[("double_block", i)]({"img": x, "txt": context, "vec": global_cond, "pe": rotary_pos_emb}, {"original_block": block_wrap})
-                x = out["img"]
-            else:
-                x = layer(x, rotary_pos_emb = rotary_pos_emb, global_cond=global_cond, context=context)
-            # x = checkpoint(layer, x, rotary_pos_emb = rotary_pos_emb, global_cond=global_cond, **kwargs)
-
-            if return_info:
-                info["hidden_states"].append(x)
-
-        x = self.project_out(x)
-
-        if return_info:
-            return x, info
-
-        return x
-
-class AudioDiffusionTransformer(nn.Module):
-    def __init__(self,
-        io_channels=64,
-        patch_size=1,
-        embed_dim=1536,
-        cond_token_dim=768,
-        project_cond_tokens=False,
-        global_cond_dim=1536,
-        project_global_cond=True,
-        input_concat_dim=0,
-        prepend_cond_dim=0,
-        depth=24,
-        num_heads=24,
-        transformer_type: tp.Literal["continuous_transformer"] = "continuous_transformer",
-        global_cond_type: tp.Literal["prepend", "adaLN"] = "prepend",
-        audio_model="",
-        dtype=None,
-        device=None,
-        operations=None,
-        **kwargs):
-
-        super().__init__()
-
-        self.dtype = dtype
-        self.cond_token_dim = cond_token_dim
-
-        # Timestep embeddings
-        timestep_features_dim = 256
-
-        self.timestep_features = FourierFeatures(1, timestep_features_dim, dtype=dtype, device=device)
-
-        self.to_timestep_embed = nn.Sequential(
-            operations.Linear(timestep_features_dim, embed_dim, bias=True, dtype=dtype, device=device),
-            nn.SiLU(),
-            operations.Linear(embed_dim, embed_dim, bias=True, dtype=dtype, device=device),
-        )
-
-        if cond_token_dim > 0:
-            # Conditioning tokens
-
-            cond_embed_dim = cond_token_dim if not project_cond_tokens else embed_dim
-            self.to_cond_embed = nn.Sequential(
-                operations.Linear(cond_token_dim, cond_embed_dim, bias=False, dtype=dtype, device=device),
-                nn.SiLU(),
-                operations.Linear(cond_embed_dim, cond_embed_dim, bias=False, dtype=dtype, device=device)
-            )
-        else:
-            cond_embed_dim = 0
-
-        if global_cond_dim > 0:
-            # Global conditioning
-            global_embed_dim = global_cond_dim if not project_global_cond else embed_dim
-            self.to_global_embed = nn.Sequential(
-                operations.Linear(global_cond_dim, global_embed_dim, bias=False, dtype=dtype, device=device),
-                nn.SiLU(),
-                operations.Linear(global_embed_dim, global_embed_dim, bias=False, dtype=dtype, device=device)
-            )
-
-        if prepend_cond_dim > 0:
-            # Prepend conditioning
-            self.to_prepend_embed = nn.Sequential(
-                operations.Linear(prepend_cond_dim, embed_dim, bias=False, dtype=dtype, device=device),
-                nn.SiLU(),
-                operations.Linear(embed_dim, embed_dim, bias=False, dtype=dtype, device=device)
-            )
-
-        self.input_concat_dim = input_concat_dim
-
-        dim_in = io_channels + self.input_concat_dim
-
-        self.patch_size = patch_size
-
-        # Transformer
-
-        self.transformer_type = transformer_type
-
-        self.global_cond_type = global_cond_type
-
-        if self.transformer_type == "continuous_transformer":
-
-            global_dim = None
-
-            if self.global_cond_type == "adaLN":
-                # The global conditioning is projected to the embed_dim already at this point
-                global_dim = embed_dim
-
-            self.transformer = ContinuousTransformer(
-                dim=embed_dim,
-                depth=depth,
-                dim_heads=embed_dim // num_heads,
-                dim_in=dim_in * patch_size,
-                dim_out=io_channels * patch_size,
-                cross_attend = cond_token_dim > 0,
-                cond_token_dim = cond_embed_dim,
-                global_cond_dim=global_dim,
-                dtype=dtype,
-                device=device,
-                operations=operations,
-                **kwargs
-            )
-        else:
-            raise ValueError(f"Unknown transformer type: {self.transformer_type}")
-
-        self.preprocess_conv = operations.Conv1d(dim_in, dim_in, 1, bias=False, dtype=dtype, device=device)
-        self.postprocess_conv = operations.Conv1d(io_channels, io_channels, 1, bias=False, dtype=dtype, device=device)
-
-    def _forward(
-        self,
-        x,
-        t,
-        mask=None,
-        cross_attn_cond=None,
-        cross_attn_cond_mask=None,
-        input_concat_cond=None,
-        global_embed=None,
-        prepend_cond=None,
-        prepend_cond_mask=None,
-        return_info=False,
-        **kwargs):
-
-        if cross_attn_cond is not None:
-            cross_attn_cond = self.to_cond_embed(cross_attn_cond)
-
-        if global_embed is not None:
-            # Project the global conditioning to the embedding dimension
-            global_embed = self.to_global_embed(global_embed)
-
-        prepend_inputs = None
-        prepend_mask = None
-        prepend_length = 0
-        if prepend_cond is not None:
-            # Project the prepend conditioning to the embedding dimension
-            prepend_cond = self.to_prepend_embed(prepend_cond)
-
-            prepend_inputs = prepend_cond
-            if prepend_cond_mask is not None:
-                prepend_mask = prepend_cond_mask
-
-        if input_concat_cond is not None:
-
-            # Interpolate input_concat_cond to the same length as x
-            if input_concat_cond.shape[2] != x.shape[2]:
-                input_concat_cond = F.interpolate(input_concat_cond, (x.shape[2], ), mode='nearest')
-
-            x = torch.cat([x, input_concat_cond], dim=1)
-
-        # Get the batch of timestep embeddings
-        timestep_embed = self.to_timestep_embed(self.timestep_features(t[:, None]).to(x.dtype)) # (b, embed_dim)
-
-        # Timestep embedding is considered a global embedding. Add to the global conditioning if it exists
-        if global_embed is not None:
-            global_embed = global_embed + timestep_embed
-        else:
-            global_embed = timestep_embed
-
-        # Add the global_embed to the prepend inputs if there is no global conditioning support in the transformer
-        if self.global_cond_type == "prepend":
-            if prepend_inputs is None:
-                # Prepend inputs are just the global embed, and the mask is all ones
-                prepend_inputs = global_embed.unsqueeze(1)
-                prepend_mask = torch.ones((x.shape[0], 1), device=x.device, dtype=torch.bool)
-            else:
-                # Prepend inputs are the prepend conditioning + the global embed
-                prepend_inputs = torch.cat([prepend_inputs, global_embed.unsqueeze(1)], dim=1)
-                prepend_mask = torch.cat([prepend_mask, torch.ones((x.shape[0], 1), device=x.device, dtype=torch.bool)], dim=1)
-
-            prepend_length = prepend_inputs.shape[1]
-
-        x = self.preprocess_conv(x) + x
-
-        x = rearrange(x, "b c t -> b t c")
-
-        extra_args = {}
-
-        if self.global_cond_type == "adaLN":
-            extra_args["global_cond"] = global_embed
-
-        if self.patch_size > 1:
-            x = rearrange(x, "b (t p) c -> b t (c p)", p=self.patch_size)
-
-        if self.transformer_type == "x-transformers":
-            output = self.transformer(x, prepend_embeds=prepend_inputs, context=cross_attn_cond, context_mask=cross_attn_cond_mask, mask=mask, prepend_mask=prepend_mask, **extra_args, **kwargs)
-        elif self.transformer_type == "continuous_transformer":
-            output = self.transformer(x, prepend_embeds=prepend_inputs, context=cross_attn_cond, context_mask=cross_attn_cond_mask, mask=mask, prepend_mask=prepend_mask, return_info=return_info, **extra_args, **kwargs)
-
-            if return_info:
-                output, info = output
-        elif self.transformer_type == "mm_transformer":
-            output = self.transformer(x, context=cross_attn_cond, mask=mask, context_mask=cross_attn_cond_mask, **extra_args, **kwargs)
-
-        output = rearrange(output, "b t c -> b c t")[:,:,prepend_length:]
-
-        if self.patch_size > 1:
-            output = rearrange(output, "b (c p) t -> b c (t p)", p=self.patch_size)
-
-        output = self.postprocess_conv(output) + output
-
-        if return_info:
-            return output, info
-
-        return output
-
-    def forward(
-        self,
-        x,
-        timestep,
-        context=None,
-        context_mask=None,
-        input_concat_cond=None,
-        global_embed=None,
-        negative_global_embed=None,
-        prepend_cond=None,
-        prepend_cond_mask=None,
-        mask=None,
-        return_info=False,
-        control=None,
-        **kwargs):
-            return self._forward(
-                x,
-                timestep,
-                cross_attn_cond=context,
-                cross_attn_cond_mask=context_mask,
-                input_concat_cond=input_concat_cond,
-                global_embed=global_embed,
-                prepend_cond=prepend_cond,
-                prepend_cond_mask=prepend_cond_mask,
-                mask=mask,
-                return_info=return_info,
-                **kwargs
-            )

comfy/ldm/audio/embedders.py

@@ -1,108 +0,0 @@
-# code adapted from: https://github.com/Stability-AI/stable-audio-tools
-
-import torch
-import torch.nn as nn
-from torch import Tensor
-from typing import List, Union
-from einops import rearrange
-import math
-import comfy.ops
-
-class LearnedPositionalEmbedding(nn.Module):
-    """Used for continuous time"""
-
-    def __init__(self, dim: int):
-        super().__init__()
-        assert (dim % 2) == 0
-        half_dim = dim // 2
-        self.weights = nn.Parameter(torch.empty(half_dim))
-
-    def forward(self, x: Tensor) -> Tensor:
-        x = rearrange(x, "b -> b 1")
-        freqs = x * rearrange(self.weights, "d -> 1 d") * 2 * math.pi
-        fouriered = torch.cat((freqs.sin(), freqs.cos()), dim=-1)
-        fouriered = torch.cat((x, fouriered), dim=-1)
-        return fouriered
-
-def TimePositionalEmbedding(dim: int, out_features: int) -> nn.Module:
-    return nn.Sequential(
-        LearnedPositionalEmbedding(dim),
-        comfy.ops.manual_cast.Linear(in_features=dim + 1, out_features=out_features),
-    )
-
-
-class NumberEmbedder(nn.Module):
-    def __init__(
-        self,
-        features: int,
-        dim: int = 256,
-    ):
-        super().__init__()
-        self.features = features
-        self.embedding = TimePositionalEmbedding(dim=dim, out_features=features)
-
-    def forward(self, x: Union[List[float], Tensor]) -> Tensor:
-        if not torch.is_tensor(x):
-            device = next(self.embedding.parameters()).device
-            x = torch.tensor(x, device=device)
-        assert isinstance(x, Tensor)
-        shape = x.shape
-        x = rearrange(x, "... -> (...)")
-        embedding = self.embedding(x)
-        x = embedding.view(*shape, self.features)
-        return x  # type: ignore
-
-
-class Conditioner(nn.Module):
-    def __init__(
-            self,
-            dim: int,
-            output_dim: int,
-            project_out: bool = False
-            ):
-
-        super().__init__()
-
-        self.dim = dim
-        self.output_dim = output_dim
-        self.proj_out = nn.Linear(dim, output_dim) if (dim != output_dim or project_out) else nn.Identity()
-
-    def forward(self, x):
-        raise NotImplementedError()
-
-class NumberConditioner(Conditioner):
-    '''
-        Conditioner that takes a list of floats, normalizes them for a given range, and returns a list of embeddings
-    '''
-    def __init__(self,
-                output_dim: int,
-                min_val: float=0,
-                max_val: float=1
-                ):
-        super().__init__(output_dim, output_dim)
-
-        self.min_val = min_val
-        self.max_val = max_val
-
-        self.embedder = NumberEmbedder(features=output_dim)
-
-    def forward(self, floats, device=None):
-            # Cast the inputs to floats
-            floats = [float(x) for x in floats]
-
-            if device is None:
-                device = next(self.embedder.parameters()).device
-
-            floats = torch.tensor(floats).to(device)
-
-            floats = floats.clamp(self.min_val, self.max_val)
-
-            normalized_floats = (floats - self.min_val) / (self.max_val - self.min_val)
-
-            # Cast floats to same type as embedder
-            embedder_dtype = next(self.embedder.parameters()).dtype
-            normalized_floats = normalized_floats.to(embedder_dtype)
-
-            float_embeds = self.embedder(normalized_floats).unsqueeze(1)
-
-            return [float_embeds, torch.ones(float_embeds.shape[0], 1).to(device)]

comfy/ldm/aura/mmdit.py

@@ -1,498 +0,0 @@
-#AuraFlow MMDiT
-#Originally written by the AuraFlow Authors
-
-import math
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-from comfy.ldm.modules.attention import optimized_attention
-import comfy.ops
-import comfy.ldm.common_dit
-
-def modulate(x, shift, scale):
-    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
-
-
-def find_multiple(n: int, k: int) -> int:
-    if n % k == 0:
-        return n
-    return n + k - (n % k)
-
-
-class MLP(nn.Module):
-    def __init__(self, dim, hidden_dim=None, dtype=None, device=None, operations=None) -> None:
-        super().__init__()
-        if hidden_dim is None:
-            hidden_dim = 4 * dim
-
-        n_hidden = int(2 * hidden_dim / 3)
-        n_hidden = find_multiple(n_hidden, 256)
-
-        self.c_fc1 = operations.Linear(dim, n_hidden, bias=False, dtype=dtype, device=device)
-        self.c_fc2 = operations.Linear(dim, n_hidden, bias=False, dtype=dtype, device=device)
-        self.c_proj = operations.Linear(n_hidden, dim, bias=False, dtype=dtype, device=device)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        x = F.silu(self.c_fc1(x)) * self.c_fc2(x)
-        x = self.c_proj(x)
-        return x
-
-
-class MultiHeadLayerNorm(nn.Module):
-    def __init__(self, hidden_size=None, eps=1e-5, dtype=None, device=None):
-        # Copy pasta from https://github.com/huggingface/transformers/blob/e5f71ecaae50ea476d1e12351003790273c4b2ed/src/transformers/models/cohere/modeling_cohere.py#L78
-
-        super().__init__()
-        self.weight = nn.Parameter(torch.empty(hidden_size, dtype=dtype, device=device))
-        self.variance_epsilon = eps
-
-    def forward(self, hidden_states):
-        input_dtype = hidden_states.dtype
-        hidden_states = hidden_states.to(torch.float32)
-        mean = hidden_states.mean(-1, keepdim=True)
-        variance = (hidden_states - mean).pow(2).mean(-1, keepdim=True)
-        hidden_states = (hidden_states - mean) * torch.rsqrt(
-            variance + self.variance_epsilon
-        )
-        hidden_states = self.weight.to(torch.float32) * hidden_states
-        return hidden_states.to(input_dtype)
-
-class SingleAttention(nn.Module):
-    def __init__(self, dim, n_heads, mh_qknorm=False, dtype=None, device=None, operations=None):
-        super().__init__()
-
-        self.n_heads = n_heads
-        self.head_dim = dim // n_heads
-
-        # this is for cond
-        self.w1q = operations.Linear(dim, dim, bias=False, dtype=dtype, device=device)
-        self.w1k = operations.Linear(dim, dim, bias=False, dtype=dtype, device=device)
-        self.w1v = operations.Linear(dim, dim, bias=False, dtype=dtype, device=device)
-        self.w1o = operations.Linear(dim, dim, bias=False, dtype=dtype, device=device)
-
-        self.q_norm1 = (
-            MultiHeadLayerNorm((self.n_heads, self.head_dim), dtype=dtype, device=device)
-            if mh_qknorm
-            else operations.LayerNorm(self.head_dim, elementwise_affine=False, dtype=dtype, device=device)
-        )
-        self.k_norm1 = (
-            MultiHeadLayerNorm((self.n_heads, self.head_dim), dtype=dtype, device=device)
-            if mh_qknorm
-            else operations.LayerNorm(self.head_dim, elementwise_affine=False, dtype=dtype, device=device)
-        )
-
-    #@torch.compile()
-    def forward(self, c):
-
-        bsz, seqlen1, _ = c.shape
-
-        q, k, v = self.w1q(c), self.w1k(c), self.w1v(c)
-        q = q.view(bsz, seqlen1, self.n_heads, self.head_dim)
-        k = k.view(bsz, seqlen1, self.n_heads, self.head_dim)
-        v = v.view(bsz, seqlen1, self.n_heads, self.head_dim)
-        q, k = self.q_norm1(q), self.k_norm1(k)
-
-        output = optimized_attention(q.permute(0, 2, 1, 3), k.permute(0, 2, 1, 3), v.permute(0, 2, 1, 3), self.n_heads, skip_reshape=True)
-        c = self.w1o(output)
-        return c
-
-
-
-class DoubleAttention(nn.Module):
-    def __init__(self, dim, n_heads, mh_qknorm=False, dtype=None, device=None, operations=None):
-        super().__init__()
-
-        self.n_heads = n_heads
-        self.head_dim = dim // n_heads
-
-        # this is for cond
-        self.w1q = operations.Linear(dim, dim, bias=False, dtype=dtype, device=device)
-        self.w1k = operations.Linear(dim, dim, bias=False, dtype=dtype, device=device)
-        self.w1v = operations.Linear(dim, dim, bias=False, dtype=dtype, device=device)
-        self.w1o = operations.Linear(dim, dim, bias=False, dtype=dtype, device=device)
-
-        # this is for x
-        self.w2q = operations.Linear(dim, dim, bias=False, dtype=dtype, device=device)
-        self.w2k = operations.Linear(dim, dim, bias=False, dtype=dtype, device=device)
-        self.w2v = operations.Linear(dim, dim, bias=False, dtype=dtype, device=device)
-        self.w2o = operations.Linear(dim, dim, bias=False, dtype=dtype, device=device)
-
-        self.q_norm1 = (
-            MultiHeadLayerNorm((self.n_heads, self.head_dim), dtype=dtype, device=device)
-            if mh_qknorm
-            else operations.LayerNorm(self.head_dim, elementwise_affine=False, dtype=dtype, device=device)
-        )
-        self.k_norm1 = (
-            MultiHeadLayerNorm((self.n_heads, self.head_dim), dtype=dtype, device=device)
-            if mh_qknorm
-            else operations.LayerNorm(self.head_dim, elementwise_affine=False, dtype=dtype, device=device)
-        )
-
-        self.q_norm2 = (
-            MultiHeadLayerNorm((self.n_heads, self.head_dim), dtype=dtype, device=device)
-            if mh_qknorm
-            else operations.LayerNorm(self.head_dim, elementwise_affine=False, dtype=dtype, device=device)
-        )
-        self.k_norm2 = (
-            MultiHeadLayerNorm((self.n_heads, self.head_dim), dtype=dtype, device=device)
-            if mh_qknorm
-            else operations.LayerNorm(self.head_dim, elementwise_affine=False, dtype=dtype, device=device)
-        )
-
-
-    #@torch.compile()
-    def forward(self, c, x):
-
-        bsz, seqlen1, _ = c.shape
-        bsz, seqlen2, _ = x.shape
-
-        cq, ck, cv = self.w1q(c), self.w1k(c), self.w1v(c)
-        cq = cq.view(bsz, seqlen1, self.n_heads, self.head_dim)
-        ck = ck.view(bsz, seqlen1, self.n_heads, self.head_dim)
-        cv = cv.view(bsz, seqlen1, self.n_heads, self.head_dim)
-        cq, ck = self.q_norm1(cq), self.k_norm1(ck)
-
-        xq, xk, xv = self.w2q(x), self.w2k(x), self.w2v(x)
-        xq = xq.view(bsz, seqlen2, self.n_heads, self.head_dim)
-        xk = xk.view(bsz, seqlen2, self.n_heads, self.head_dim)
-        xv = xv.view(bsz, seqlen2, self.n_heads, self.head_dim)
-        xq, xk = self.q_norm2(xq), self.k_norm2(xk)
-
-        # concat all
-        q, k, v = (
-            torch.cat([cq, xq], dim=1),
-            torch.cat([ck, xk], dim=1),
-            torch.cat([cv, xv], dim=1),
-        )
-
-        output = optimized_attention(q.permute(0, 2, 1, 3), k.permute(0, 2, 1, 3), v.permute(0, 2, 1, 3), self.n_heads, skip_reshape=True)
-
-        c, x = output.split([seqlen1, seqlen2], dim=1)
-        c = self.w1o(c)
-        x = self.w2o(x)
-
-        return c, x
-
-
-class MMDiTBlock(nn.Module):
-    def __init__(self, dim, heads=8, global_conddim=1024, is_last=False, dtype=None, device=None, operations=None):
-        super().__init__()
-
-        self.normC1 = operations.LayerNorm(dim, elementwise_affine=False, dtype=dtype, device=device)
-        self.normC2 = operations.LayerNorm(dim, elementwise_affine=False, dtype=dtype, device=device)
-        if not is_last:
-            self.mlpC = MLP(dim, hidden_dim=dim * 4, dtype=dtype, device=device, operations=operations)
-            self.modC = nn.Sequential(
-                nn.SiLU(),
-                operations.Linear(global_conddim, 6 * dim, bias=False, dtype=dtype, device=device),
-            )
-        else:
-            self.modC = nn.Sequential(
-                nn.SiLU(),
-                operations.Linear(global_conddim, 2 * dim, bias=False, dtype=dtype, device=device),
-            )
-
-        self.normX1 = operations.LayerNorm(dim, elementwise_affine=False, dtype=dtype, device=device)
-        self.normX2 = operations.LayerNorm(dim, elementwise_affine=False, dtype=dtype, device=device)
-        self.mlpX = MLP(dim, hidden_dim=dim * 4, dtype=dtype, device=device, operations=operations)
-        self.modX = nn.Sequential(
-            nn.SiLU(),
-            operations.Linear(global_conddim, 6 * dim, bias=False, dtype=dtype, device=device),
-        )
-
-        self.attn = DoubleAttention(dim, heads, dtype=dtype, device=device, operations=operations)
-        self.is_last = is_last
-
-    #@torch.compile()
-    def forward(self, c, x, global_cond, **kwargs):
-
-        cres, xres = c, x
-
-        cshift_msa, cscale_msa, cgate_msa, cshift_mlp, cscale_mlp, cgate_mlp = (
-            self.modC(global_cond).chunk(6, dim=1)
-        )
-
-        c = modulate(self.normC1(c), cshift_msa, cscale_msa)
-
-        # xpath
-        xshift_msa, xscale_msa, xgate_msa, xshift_mlp, xscale_mlp, xgate_mlp = (
-            self.modX(global_cond).chunk(6, dim=1)
-        )
-
-        x = modulate(self.normX1(x), xshift_msa, xscale_msa)
-
-        # attention
-        c, x = self.attn(c, x)
-
-
-        c = self.normC2(cres + cgate_msa.unsqueeze(1) * c)
-        c = cgate_mlp.unsqueeze(1) * self.mlpC(modulate(c, cshift_mlp, cscale_mlp))
-        c = cres + c
-
-        x = self.normX2(xres + xgate_msa.unsqueeze(1) * x)
-        x = xgate_mlp.unsqueeze(1) * self.mlpX(modulate(x, xshift_mlp, xscale_mlp))
-        x = xres + x
-
-        return c, x
-
-class DiTBlock(nn.Module):
-    # like MMDiTBlock, but it only has X
-    def __init__(self, dim, heads=8, global_conddim=1024, dtype=None, device=None, operations=None):
-        super().__init__()
-
-        self.norm1 = operations.LayerNorm(dim, elementwise_affine=False, dtype=dtype, device=device)
-        self.norm2 = operations.LayerNorm(dim, elementwise_affine=False, dtype=dtype, device=device)
-
-        self.modCX = nn.Sequential(
-            nn.SiLU(),
-            operations.Linear(global_conddim, 6 * dim, bias=False, dtype=dtype, device=device),
-        )
-
-        self.attn = SingleAttention(dim, heads, dtype=dtype, device=device, operations=operations)
-        self.mlp = MLP(dim, hidden_dim=dim * 4, dtype=dtype, device=device, operations=operations)
-
-    #@torch.compile()
-    def forward(self, cx, global_cond, **kwargs):
-        cxres = cx
-        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.modCX(
-            global_cond
-        ).chunk(6, dim=1)
-        cx = modulate(self.norm1(cx), shift_msa, scale_msa)
-        cx = self.attn(cx)
-        cx = self.norm2(cxres + gate_msa.unsqueeze(1) * cx)
-        mlpout = self.mlp(modulate(cx, shift_mlp, scale_mlp))
-        cx = gate_mlp.unsqueeze(1) * mlpout
-
-        cx = cxres + cx
-
-        return cx
-
-
-
-class TimestepEmbedder(nn.Module):
-    def __init__(self, hidden_size, frequency_embedding_size=256, dtype=None, device=None, operations=None):
-        super().__init__()
-        self.mlp = nn.Sequential(
-            operations.Linear(frequency_embedding_size, hidden_size, dtype=dtype, device=device),
-            nn.SiLU(),
-            operations.Linear(hidden_size, hidden_size, dtype=dtype, device=device),
-        )
-        self.frequency_embedding_size = frequency_embedding_size
-
-    @staticmethod
-    def timestep_embedding(t, dim, max_period=10000):
-        half = dim // 2
-        freqs = 1000 * torch.exp(
-            -math.log(max_period) * torch.arange(start=0, end=half) / half
-        ).to(t.device)
-        args = t[:, None] * freqs[None]
-        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
-        if dim % 2:
-            embedding = torch.cat(
-                [embedding, torch.zeros_like(embedding[:, :1])], dim=-1
-            )
-        return embedding
-
-    #@torch.compile()
-    def forward(self, t, dtype):
-        t_freq = self.timestep_embedding(t, self.frequency_embedding_size).to(dtype)
-        t_emb = self.mlp(t_freq)
-        return t_emb
-
-
-class MMDiT(nn.Module):
-    def __init__(
-        self,
-        in_channels=4,
-        out_channels=4,
-        patch_size=2,
-        dim=3072,
-        n_layers=36,
-        n_double_layers=4,
-        n_heads=12,
-        global_conddim=3072,
-        cond_seq_dim=2048,
-        max_seq=32 * 32,
-        device=None,
-        dtype=None,
-        operations=None,
-    ):
-        super().__init__()
-        self.dtype = dtype
-
-        self.t_embedder = TimestepEmbedder(global_conddim, dtype=dtype, device=device, operations=operations)
-
-        self.cond_seq_linear = operations.Linear(
-            cond_seq_dim, dim, bias=False, dtype=dtype, device=device
-        )  # linear for something like text sequence.
-        self.init_x_linear = operations.Linear(
-            patch_size * patch_size * in_channels, dim, dtype=dtype, device=device
-        )  # init linear for patchified image.
-
-        self.positional_encoding = nn.Parameter(torch.empty(1, max_seq, dim, dtype=dtype, device=device))
-        self.register_tokens = nn.Parameter(torch.empty(1, 8, dim, dtype=dtype, device=device))
-
-        self.double_layers = nn.ModuleList([])
-        self.single_layers = nn.ModuleList([])
-
-
-        for idx in range(n_double_layers):
-            self.double_layers.append(
-                MMDiTBlock(dim, n_heads, global_conddim, is_last=(idx == n_layers - 1), dtype=dtype, device=device, operations=operations)
-            )
-
-        for idx in range(n_double_layers, n_layers):
-            self.single_layers.append(
-                DiTBlock(dim, n_heads, global_conddim, dtype=dtype, device=device, operations=operations)
-            )
-
-
-        self.final_linear = operations.Linear(
-            dim, patch_size * patch_size * out_channels, bias=False, dtype=dtype, device=device
-        )
-
-        self.modF = nn.Sequential(
-            nn.SiLU(),
-            operations.Linear(global_conddim, 2 * dim, bias=False, dtype=dtype, device=device),
-        )
-
-        self.out_channels = out_channels
-        self.patch_size = patch_size
-        self.n_double_layers = n_double_layers
-        self.n_layers = n_layers
-
-        self.h_max = round(max_seq**0.5)
-        self.w_max = round(max_seq**0.5)
-
-    @torch.no_grad()
-    def extend_pe(self, init_dim=(16, 16), target_dim=(64, 64)):
-        # extend pe
-        pe_data = self.positional_encoding.data.squeeze(0)[: init_dim[0] * init_dim[1]]
-
-        pe_as_2d = pe_data.view(init_dim[0], init_dim[1], -1).permute(2, 0, 1)
-
-        # now we need to extend this to target_dim. for this we will use interpolation.
-        # we will use torch.nn.functional.interpolate
-        pe_as_2d = F.interpolate(
-            pe_as_2d.unsqueeze(0), size=target_dim, mode="bilinear"
-        )
-        pe_new = pe_as_2d.squeeze(0).permute(1, 2, 0).flatten(0, 1)
-        self.positional_encoding.data = pe_new.unsqueeze(0).contiguous()
-        self.h_max, self.w_max = target_dim
-
-    def pe_selection_index_based_on_dim(self, h, w):
-        h_p, w_p = h // self.patch_size, w // self.patch_size
-        original_pe_indexes = torch.arange(self.positional_encoding.shape[1])
-        original_pe_indexes = original_pe_indexes.view(self.h_max, self.w_max)
-        starth =  self.h_max // 2 - h_p // 2
-        endh =starth + h_p
-        startw = self.w_max // 2 - w_p // 2
-        endw = startw + w_p
-        original_pe_indexes = original_pe_indexes[
-            starth:endh, startw:endw
-        ]
-        return original_pe_indexes.flatten()
-
-    def unpatchify(self, x, h, w):
-        c = self.out_channels
-        p = self.patch_size
-
-        x = x.reshape(shape=(x.shape[0], h, w, p, p, c))
-        x = torch.einsum("nhwpqc->nchpwq", x)
-        imgs = x.reshape(shape=(x.shape[0], c, h * p, w * p))
-        return imgs
-
-    def patchify(self, x):
-        B, C, H, W = x.size()
-        x = comfy.ldm.common_dit.pad_to_patch_size(x, (self.patch_size, self.patch_size))
-        x = x.view(
-            B,
-            C,
-            (H + 1) // self.patch_size,
-            self.patch_size,
-            (W + 1) // self.patch_size,
-            self.patch_size,
-        )
-        x = x.permute(0, 2, 4, 1, 3, 5).flatten(-3).flatten(1, 2)
-        return x
-
-    def apply_pos_embeds(self, x, h, w):
-        h = (h + 1) // self.patch_size
-        w = (w + 1) // self.patch_size
-        max_dim = max(h, w)
-
-        cur_dim = self.h_max
-        pos_encoding = comfy.ops.cast_to_input(self.positional_encoding.reshape(1, cur_dim, cur_dim, -1), x)
-
-        if max_dim > cur_dim:
-            pos_encoding = F.interpolate(pos_encoding.movedim(-1, 1), (max_dim, max_dim), mode="bilinear").movedim(1, -1)
-            cur_dim = max_dim
-
-        from_h = (cur_dim - h) // 2
-        from_w = (cur_dim - w) // 2
-        pos_encoding = pos_encoding[:,from_h:from_h+h,from_w:from_w+w]
-        return x + pos_encoding.reshape(1, -1, self.positional_encoding.shape[-1])
-
-    def forward(self, x, timestep, context, transformer_options={}, **kwargs):
-        patches_replace = transformer_options.get("patches_replace", {})
-        # patchify x, add PE
-        b, c, h, w = x.shape
-
-        # pe_indexes = self.pe_selection_index_based_on_dim(h, w)
-        # print(pe_indexes, pe_indexes.shape)
-
-        x = self.init_x_linear(self.patchify(x))  # B, T_x, D
-        x = self.apply_pos_embeds(x, h, w)
-        # x = x + self.positional_encoding[:, : x.size(1)].to(device=x.device, dtype=x.dtype)
-        # x = x + self.positional_encoding[:, pe_indexes].to(device=x.device, dtype=x.dtype)
-
-        # process conditions for MMDiT Blocks
-        c_seq = context  # B, T_c, D_c
-        t = timestep
-
-        c = self.cond_seq_linear(c_seq)  # B, T_c, D
-        c = torch.cat([comfy.ops.cast_to_input(self.register_tokens, c).repeat(c.size(0), 1, 1), c], dim=1)
-
-        global_cond = self.t_embedder(t, x.dtype)  # B, D
-
-        blocks_replace = patches_replace.get("dit", {})
-        if len(self.double_layers) > 0:
-            for i, layer in enumerate(self.double_layers):
-                if ("double_block", i) in blocks_replace:
-                    def block_wrap(args):
-                        out = {}
-                        out["txt"], out["img"] = layer(args["txt"],
-                                                       args["img"],
-                                                       args["vec"])
-                        return out
-                    out = blocks_replace[("double_block", i)]({"img": x, "txt": c, "vec": global_cond}, {"original_block": block_wrap})
-                    c = out["txt"]
-                    x = out["img"]
-                else:
-                    c, x = layer(c, x, global_cond, **kwargs)
-
-        if len(self.single_layers) > 0:
-            c_len = c.size(1)
-            cx = torch.cat([c, x], dim=1)
-            for i, layer in enumerate(self.single_layers):
-                if ("single_block", i) in blocks_replace:
-                    def block_wrap(args):
-                        out = {}
-                        out["img"] = layer(args["img"], args["vec"])
-                        return out
-
-                    out = blocks_replace[("single_block", i)]({"img": cx, "vec": global_cond}, {"original_block": block_wrap})
-                    cx = out["img"]
-                else:
-                    cx = layer(cx, global_cond, **kwargs)
-
-            x = cx[:, c_len:]
-
-        fshift, fscale = self.modF(global_cond).chunk(2, dim=1)
-
-        x = modulate(x, fshift, fscale)
-        x = self.final_linear(x)
-        x = self.unpatchify(x, (h + 1) // self.patch_size, (w + 1) // self.patch_size)[:,:,:h,:w]
-        return x

comfy/ldm/cascade/common.py

@@ -1,154 +0,0 @@
-"""
-    This file is part of ComfyUI.
-    Copyright (C) 2024 Stability AI
-
-    This program is free software: you can redistribute it and/or modify
-    it under the terms of the GNU General Public License as published by
-    the Free Software Foundation, either version 3 of the License, or
-    (at your option) any later version.
-
-    This program is distributed in the hope that it will be useful,
-    but WITHOUT ANY WARRANTY; without even the implied warranty of
-    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
-    GNU General Public License for more details.
-
-    You should have received a copy of the GNU General Public License
-    along with this program.  If not, see <https://www.gnu.org/licenses/>.
-"""
-
-import torch
-import torch.nn as nn
-from comfy.ldm.modules.attention import optimized_attention
-import comfy.ops
-
-class OptimizedAttention(nn.Module):
-    def __init__(self, c, nhead, dropout=0.0, dtype=None, device=None, operations=None):
-        super().__init__()
-        self.heads = nhead
-
-        self.to_q = operations.Linear(c, c, bias=True, dtype=dtype, device=device)
-        self.to_k = operations.Linear(c, c, bias=True, dtype=dtype, device=device)
-        self.to_v = operations.Linear(c, c, bias=True, dtype=dtype, device=device)
-
-        self.out_proj = operations.Linear(c, c, bias=True, dtype=dtype, device=device)
-
-    def forward(self, q, k, v):
-        q = self.to_q(q)
-        k = self.to_k(k)
-        v = self.to_v(v)
-
-        out = optimized_attention(q, k, v, self.heads)
-
-        return self.out_proj(out)
-
-class Attention2D(nn.Module):
-    def __init__(self, c, nhead, dropout=0.0, dtype=None, device=None, operations=None):
-        super().__init__()
-        self.attn = OptimizedAttention(c, nhead, dtype=dtype, device=device, operations=operations)
-        # self.attn = nn.MultiheadAttention(c, nhead, dropout=dropout, bias=True, batch_first=True, dtype=dtype, device=device)
-
-    def forward(self, x, kv, self_attn=False):
-        orig_shape = x.shape
-        x = x.view(x.size(0), x.size(1), -1).permute(0, 2, 1)  # Bx4xHxW -> Bx(HxW)x4
-        if self_attn:
-            kv = torch.cat([x, kv], dim=1)
-        # x = self.attn(x, kv, kv, need_weights=False)[0]
-        x = self.attn(x, kv, kv)
-        x = x.permute(0, 2, 1).view(*orig_shape)
-        return x
-
-
-def LayerNorm2d_op(operations):
-    class LayerNorm2d(operations.LayerNorm):
-        def __init__(self, *args, **kwargs):
-            super().__init__(*args, **kwargs)
-
-        def forward(self, x):
-            return super().forward(x.permute(0, 2, 3, 1)).permute(0, 3, 1, 2)
-    return LayerNorm2d
-
-class GlobalResponseNorm(nn.Module):
-    "from https://github.com/facebookresearch/ConvNeXt-V2/blob/3608f67cc1dae164790c5d0aead7bf2d73d9719b/models/utils.py#L105"
-    def __init__(self, dim, dtype=None, device=None):
-        super().__init__()
-        self.gamma = nn.Parameter(torch.empty(1, 1, 1, dim, dtype=dtype, device=device))
-        self.beta = nn.Parameter(torch.empty(1, 1, 1, dim, dtype=dtype, device=device))
-
-    def forward(self, x):
-        Gx = torch.norm(x, p=2, dim=(1, 2), keepdim=True)
-        Nx = Gx / (Gx.mean(dim=-1, keepdim=True) + 1e-6)
-        return comfy.ops.cast_to_input(self.gamma, x) * (x * Nx) + comfy.ops.cast_to_input(self.beta, x) + x
-
-
-class ResBlock(nn.Module):
-    def __init__(self, c, c_skip=0, kernel_size=3, dropout=0.0, dtype=None, device=None, operations=None):  # , num_heads=4, expansion=2):
-        super().__init__()
-        self.depthwise = operations.Conv2d(c, c, kernel_size=kernel_size, padding=kernel_size // 2, groups=c, dtype=dtype, device=device)
-        #         self.depthwise = SAMBlock(c, num_heads, expansion)
-        self.norm = LayerNorm2d_op(operations)(c, elementwise_affine=False, eps=1e-6, dtype=dtype, device=device)
-        self.channelwise = nn.Sequential(
-            operations.Linear(c + c_skip, c * 4, dtype=dtype, device=device),
-            nn.GELU(),
-            GlobalResponseNorm(c * 4, dtype=dtype, device=device),
-            nn.Dropout(dropout),
-            operations.Linear(c * 4, c, dtype=dtype, device=device)
-        )
-
-    def forward(self, x, x_skip=None):
-        x_res = x
-        x = self.norm(self.depthwise(x))
-        if x_skip is not None:
-            x = torch.cat([x, x_skip], dim=1)
-        x = self.channelwise(x.permute(0, 2, 3, 1)).permute(0, 3, 1, 2)
-        return x + x_res
-
-
-class AttnBlock(nn.Module):
-    def __init__(self, c, c_cond, nhead, self_attn=True, dropout=0.0, dtype=None, device=None, operations=None):
-        super().__init__()
-        self.self_attn = self_attn
-        self.norm = LayerNorm2d_op(operations)(c, elementwise_affine=False, eps=1e-6, dtype=dtype, device=device)
-        self.attention = Attention2D(c, nhead, dropout, dtype=dtype, device=device, operations=operations)
-        self.kv_mapper = nn.Sequential(
-            nn.SiLU(),
-            operations.Linear(c_cond, c, dtype=dtype, device=device)
-        )
-
-    def forward(self, x, kv):
-        kv = self.kv_mapper(kv)
-        x = x + self.attention(self.norm(x), kv, self_attn=self.self_attn)
-        return x
-
-
-class FeedForwardBlock(nn.Module):
-    def __init__(self, c, dropout=0.0, dtype=None, device=None, operations=None):
-        super().__init__()
-        self.norm = LayerNorm2d_op(operations)(c, elementwise_affine=False, eps=1e-6, dtype=dtype, device=device)
-        self.channelwise = nn.Sequential(
-            operations.Linear(c, c * 4, dtype=dtype, device=device),
-            nn.GELU(),
-            GlobalResponseNorm(c * 4, dtype=dtype, device=device),
-            nn.Dropout(dropout),
-            operations.Linear(c * 4, c, dtype=dtype, device=device)
-        )
-
-    def forward(self, x):
-        x = x + self.channelwise(self.norm(x).permute(0, 2, 3, 1)).permute(0, 3, 1, 2)
-        return x
-
-
-class TimestepBlock(nn.Module):
-    def __init__(self, c, c_timestep, conds=['sca'], dtype=None, device=None, operations=None):
-        super().__init__()
-        self.mapper = operations.Linear(c_timestep, c * 2, dtype=dtype, device=device)
-        self.conds = conds
-        for cname in conds:
-            setattr(self, f"mapper_{cname}", operations.Linear(c_timestep, c * 2, dtype=dtype, device=device))
-
-    def forward(self, x, t):
-        t = t.chunk(len(self.conds) + 1, dim=1)
-        a, b = self.mapper(t[0])[:, :, None, None].chunk(2, dim=1)
-        for i, c in enumerate(self.conds):
-            ac, bc = getattr(self, f"mapper_{c}")(t[i + 1])[:, :, None, None].chunk(2, dim=1)
-            a, b = a + ac, b + bc
-        return x * (1 + a) + b