from __future__ import annotations from typing import Any import types import torch def _rope_encode_with_source_id( wan_model, t, h, w, *, t_start=0, steps_t=None, steps_h=None, steps_w=None, device=None, dtype=None, transformer_options=None, source_id=0, ): freqs = wan_model.rope_encode( t, h, w, t_start=t_start, steps_t=steps_t, steps_h=steps_h, steps_w=steps_w, device=device, dtype=dtype, transformer_options=transformer_options or {}, ) if source_id: from comfy.ldm.flux.math import rope d = wan_model.dim // wan_model.num_heads pos = torch.tensor([[float(source_id)]], device=freqs.device, dtype=torch.float32) id_rot = rope(pos, d, wan_model.rope_embedder.theta).reshape(1, 1, 1, d // 2, 2, 2).to(freqs.dtype) freqs = torch.einsum("...ij,...jk->...ik", freqs, id_rot) return freqs def _bernini_forward_orig( wan_model, x, t, context, *, clip_fea=None, freqs=None, transformer_options=None, **kwargs, ): import comfy.ldm.wan.model as wan_model_module if transformer_options is None: transformer_options = {} x = wan_model.patch_embedding(x.float()).to(x.dtype) grid_sizes = x.shape[2:] transformer_options["grid_sizes"] = grid_sizes x = x.flatten(2).transpose(1, 2) e = wan_model.time_embedding( wan_model_module.sinusoidal_embedding_1d(wan_model.freq_dim, t.flatten()).to(dtype=x[0].dtype) ) e = e.reshape(t.shape[0], -1, e.shape[-1]) e0 = wan_model.time_projection(e).unflatten(2, (6, wan_model.dim)) full_ref = None if wan_model.ref_conv is not None: full_ref = kwargs.get("reference_latent", None) if full_ref is not None: full_ref = wan_model.ref_conv(full_ref).flatten(2).transpose(1, 2) x = torch.concat((full_ref, x), dim=1) context_latents = kwargs.get("context_latents", None) main_len = x.shape[1] if context_latents is not None: for lat in context_latents: cl = wan_model.patch_embedding(lat.float().to(x.device)).to(x.dtype).flatten(2).transpose(1, 2) x = torch.cat([x, cl], dim=1) context = wan_model.text_embedding(context) context_img_len = None if clip_fea is not None: if wan_model.img_emb is not None: context_clip = wan_model.img_emb(clip_fea) context = torch.concat([context_clip, context], dim=1) context_img_len = clip_fea.shape[-2] patches_replace = transformer_options.get("patches_replace", {}) blocks_replace = patches_replace.get("dit", {}) transformer_options["total_blocks"] = len(wan_model.blocks) transformer_options["block_type"] = "double" for i, block in enumerate(wan_model.blocks): transformer_options["block_index"] = i if ("double_block", i) in blocks_replace: def block_wrap(args): return { "img": block( args["img"], context=args["txt"], e=args["vec"], freqs=args["pe"], context_img_len=context_img_len, transformer_options=args["transformer_options"], ) } out = blocks_replace[("double_block", i)]( { "img": x, "txt": context, "vec": e0, "pe": freqs, "transformer_options": transformer_options, }, {"original_block": block_wrap}, ) x = out["img"] else: x = block( x, e=e0, freqs=freqs, context=context, context_img_len=context_img_len, transformer_options=transformer_options, ) x = wan_model.head(x, e) if context_latents is not None: x = x[:, :main_len] if full_ref is not None: x = x[:, full_ref.shape[1] :] return wan_model.unpatchify(x, grid_sizes) def bind_bernini_conds(model): import comfy.conds base_model = model.model if getattr(base_model, "_bernini_conds", False): return orig_extra_conds = base_model.extra_conds orig_resize_cond = base_model.resize_cond_for_context_window def extra_conds(self, **kwargs): out = orig_extra_conds(**kwargs) context_latents = kwargs.get("context_latents", None) if context_latents is not None: out["context_latents"] = comfy.conds.CONDList([self.process_latent_in(lat) for lat in context_latents]) return out def resize_cond_for_context_window(self, cond_key, cond_value, window, x_in, device, retain_index_list=[]): if cond_key == "context_latents" and isinstance(getattr(cond_value, "cond", None), list): dim = window.dim out = [] for lat in cond_value.cond: if lat.ndim > dim and lat.shape[dim] > 1 and lat.shape[dim] == x_in.shape[dim]: idx = tuple([slice(None)] * dim + [window.index_list]) out.append(lat[idx].to(device)) else: out.append(lat.to(device)) return cond_value._copy_with(out) return orig_resize_cond(cond_key, cond_value, window, x_in, device, retain_index_list=retain_index_list) base_model.extra_conds = types.MethodType(extra_conds, base_model) base_model.resize_cond_for_context_window = types.MethodType(resize_cond_for_context_window, base_model) base_model._bernini_conds = True def bernini_forward_wrapper( executor, x, timestep, context, clip_fea=None, time_dim_concat=None, transformer_options=None, **kwargs, ): import comfy.ldm.common_dit context_latents = kwargs.get("context_latents", None) if context_latents is None: return executor(x, timestep, context, clip_fea, time_dim_concat, transformer_options, **kwargs) if transformer_options is None: transformer_options = {} wan_model = executor.class_obj _, _, t, h, w = x.shape x = comfy.ldm.common_dit.pad_to_patch_size(x, wan_model.patch_size) t_len = t if time_dim_concat is not None: time_dim_concat = comfy.ldm.common_dit.pad_to_patch_size(time_dim_concat, wan_model.patch_size) x = torch.cat([x, time_dim_concat], dim=2) t_len = x.shape[2] if wan_model.ref_conv is not None and "reference_latent" in kwargs: t_len += 1 freqs = wan_model.rope_encode(t_len, h, w, device=x.device, dtype=x.dtype, transformer_options=transformer_options) padded_context_latents = [ comfy.ldm.common_dit.pad_to_patch_size(lat, wan_model.patch_size) for lat in context_latents ] for i, lat in enumerate(padded_context_latents): freqs = torch.cat( [ freqs, _rope_encode_with_source_id( wan_model, lat.shape[-3], lat.shape[-2], lat.shape[-1], device=x.device, dtype=x.dtype, transformer_options=transformer_options, source_id=i + 1, ), ], dim=1, ) kwargs = {**kwargs, "context_latents": padded_context_latents} out = _bernini_forward_orig( wan_model, x, timestep, context, clip_fea=clip_fea, freqs=freqs, transformer_options=transformer_options, **kwargs, ) return out[:, :, :t, :h, :w]