Spaces:
Running on Zero
Running on Zero
| """Cosmos-Reference (Mirumo0u0) for DiffSynth Anima — Anima Edit support. | |
| Implements the same inference contract as | |
| https://github.com/Mirumo0u0/ComfyUI-Cosmos-Reference : | |
| * VAE-encode a source image to a reference latent | |
| * concatenate the reference latent as an extra temporal frame on every DiT step | |
| * slice the reference frame(s) off the DiT output before the sampler advances | |
| * no attention bias / strength gating (unlike In-Context Character); control | |
| strength comes from the Edit LoRA and the denoising strength | |
| This is the mechanism required by Anima Edit v1 / v2 and other edit LoRAs that | |
| expect Cosmos-style reference latents. | |
| """ | |
| from __future__ import annotations | |
| from contextlib import contextmanager | |
| from dataclasses import dataclass, field | |
| from pathlib import Path | |
| from typing import Any, Iterator, Optional | |
| import torch | |
| from PIL import Image | |
| # Official ComfyUI-Cosmos-Reference constants (kept for clarity / debugging). | |
| MAX_REF_LATENTS = 1 | |
| COND_REF_LATENTS_KEY = "ref_latents" | |
| class CosmosReferenceState: | |
| """Request-scoped state for one Cosmos-Reference generation.""" | |
| ref_latents: list[torch.Tensor] = field(default_factory=list) # each (C, H, W) | |
| enabled: bool = False | |
| active: bool = False | |
| _STATE = CosmosReferenceState() | |
| def get_state() -> CosmosReferenceState: | |
| return _STATE | |
| def clear_state() -> None: | |
| _STATE.ref_latents = [] | |
| _STATE.enabled = False | |
| _STATE.active = False | |
| def encode_reference_image( | |
| pipe: Any, | |
| image: Image.Image | None, | |
| width: int, | |
| height: int, | |
| ) -> Optional[torch.Tensor]: | |
| """Encode one RGB image with the Anima VAE to a single-frame latent (C, h, w). | |
| Matches the resize + centre-pad behaviour used by typical Anima Edit | |
| workflows (letterbox into the generation canvas, then VAE-encode). | |
| """ | |
| if image is None: | |
| return None | |
| if not isinstance(image, Image.Image): | |
| try: | |
| image = Image.fromarray(image) | |
| except Exception: | |
| return None | |
| image = image.convert("RGB") | |
| src_w, src_h = image.size | |
| if src_w < 1 or src_h < 1: | |
| return None | |
| # Letterbox into target canvas (same idea as the sample workflow resize node). | |
| scale = min(width / src_w, height / src_h) | |
| new_w = max(1, int(round(src_w * scale))) | |
| new_h = max(1, int(round(src_h * scale))) | |
| resized = image.resize((new_w, new_h), Image.Resampling.LANCZOS) | |
| canvas = Image.new("RGB", (width, height), (0, 0, 0)) | |
| canvas.paste(resized, ((width - new_w) // 2, (height - new_h) // 2)) | |
| device = pipe.device | |
| dtype = pipe.torch_dtype | |
| vae = pipe.vae | |
| if hasattr(pipe, "load_models_to_device"): | |
| pipe.load_models_to_device(["vae"]) | |
| with torch.no_grad(): | |
| if hasattr(pipe, "preprocess_image"): | |
| tensor = pipe.preprocess_image(canvas).to(device=device, dtype=dtype) | |
| else: | |
| import numpy as np | |
| arr = np.asarray(canvas, dtype="float32") / 127.5 - 1.0 | |
| tensor = ( | |
| torch.from_numpy(arr) | |
| .to(device=device, dtype=dtype) | |
| .permute(2, 0, 1) | |
| ) | |
| # Match DiffSynth / Anima VAE encode paths used elsewhere in the app. | |
| try: | |
| if tensor.ndim == 3: | |
| enc = vae.encode(tensor.unsqueeze(2), device=device) | |
| else: | |
| enc = vae.encode(tensor.unsqueeze(2), device=device) | |
| except TypeError: | |
| video = tensor if tensor.ndim == 3 else tensor[0] | |
| if video.ndim == 3: | |
| video = video.unsqueeze(1) | |
| enc = vae.encode([video], device=device) | |
| except Exception: | |
| video = tensor if tensor.ndim == 3 else tensor[0] | |
| if video.ndim == 3: | |
| video = video.unsqueeze(1) | |
| enc = vae.encode([video], device=device) | |
| if isinstance(enc, (list, tuple)): | |
| enc = torch.stack( | |
| [e if isinstance(e, torch.Tensor) else e for e in enc] | |
| ) | |
| if enc.ndim == 5: | |
| frame = enc[0, :, 0] if enc.shape[0] >= 1 else enc.squeeze(0)[:, 0] | |
| elif enc.ndim == 4: | |
| if enc.shape[0] == 1: | |
| frame = enc[0] | |
| else: | |
| frame = enc[:, 0] if enc.shape[1] < enc.shape[0] else enc[0] | |
| elif enc.ndim == 3: | |
| frame = enc | |
| else: | |
| raise RuntimeError(f"Unexpected VAE encode shape: {tuple(enc.shape)}") | |
| return frame.detach().to(device="cpu", dtype=torch.bfloat16).contiguous() | |
| def cosmos_model_fn( | |
| dit=None, | |
| latents=None, | |
| timestep=None, | |
| prompt_emb=None, | |
| t5xxl_ids=None, | |
| t5xxl_weights=None, | |
| use_gradient_checkpointing=False, | |
| use_gradient_checkpointing_offload=False, | |
| **kwargs, | |
| ): | |
| """DiT forward with optional Cosmos-Reference temporal concatenation. | |
| Matches the official ``cosmos_diffusion_reference_wrapper`` exactly: | |
| * concat ref frames on the temporal axis of ``x`` only | |
| * do **not** rewrite timesteps (Comfy leaves the original ``t`` intact) | |
| * slice the output back to the original temporal length | |
| """ | |
| state = get_state() | |
| # Pipeline passes (B, C, H, W); DiT expects a temporal axis. | |
| if latents.ndim == 4: | |
| x = latents.unsqueeze(2) # (B, C, 1, H, W) | |
| squeezed = True | |
| else: | |
| x = latents | |
| squeezed = False | |
| t = timestep / 1000 if timestep is not None else timestep | |
| use_ref = state.enabled and bool(state.ref_latents) | |
| if not use_ref: | |
| model_output = dit( | |
| x=x, | |
| timesteps=t, | |
| context=prompt_emb, | |
| t5xxl_ids=t5xxl_ids, | |
| t5xxl_weights=t5xxl_weights, | |
| use_gradient_checkpointing=use_gradient_checkpointing, | |
| use_gradient_checkpointing_offload=use_gradient_checkpointing_offload, | |
| ) | |
| if squeezed and getattr(model_output, "ndim", 0) == 5: | |
| return model_output.squeeze(2) | |
| return model_output | |
| # Official node: only extend x; leave timesteps unchanged. | |
| x_temporal_dim = int(x.shape[2]) | |
| B, C, _, H, W = x.shape | |
| pieces = [x] | |
| for ref in state.ref_latents: | |
| r = ref.to(device=x.device, dtype=x.dtype) | |
| if r.ndim == 3: # (C, H, W) | |
| r = r.unsqueeze(0).unsqueeze(2) # (1, C, 1, H, W) | |
| elif r.ndim == 4: # (B|1, C, H, W) | |
| r = r.unsqueeze(2) | |
| if r.shape[0] == 1 and B > 1: | |
| r = r.expand(B, -1, -1, -1, -1) | |
| if r.shape[-2] != H or r.shape[-1] != W: | |
| b, c, tt = r.shape[0], r.shape[1], r.shape[2] | |
| r = r.reshape(b * tt, c, r.shape[-2], r.shape[-1]) | |
| r = torch.nn.functional.interpolate( | |
| r, size=(H, W), mode="bilinear", align_corners=False | |
| ) | |
| r = r.reshape(b, c, tt, H, W) | |
| pieces.append(r) | |
| x_cat = torch.cat(pieces, dim=2) | |
| state.active = True | |
| try: | |
| model_output = dit( | |
| x=x_cat, | |
| timesteps=t, # unchanged — matches ComfyUI-Cosmos-Reference | |
| context=prompt_emb, | |
| t5xxl_ids=t5xxl_ids, | |
| t5xxl_weights=t5xxl_weights, | |
| use_gradient_checkpointing=use_gradient_checkpointing, | |
| use_gradient_checkpointing_offload=use_gradient_checkpointing_offload, | |
| ) | |
| finally: | |
| state.active = False | |
| if getattr(model_output, "ndim", 0) == 5: | |
| model_output = model_output[:, :, :x_temporal_dim] | |
| if squeezed: | |
| model_output = model_output.squeeze(2) | |
| return model_output | |
| def temporary_cosmos_reference( | |
| pipe: Any, | |
| source_image: Image.Image | None, | |
| width: int, | |
| height: int, | |
| ) -> Iterator[bool]: | |
| """Encode the source image and enable Cosmos-Reference for the duration of the block.""" | |
| clear_state() | |
| if source_image is None: | |
| yield False | |
| return | |
| try: | |
| latent = encode_reference_image(pipe, source_image, width, height) | |
| if latent is None: | |
| yield False | |
| return | |
| st = get_state() | |
| st.ref_latents = [latent] | |
| st.enabled = True | |
| yield True | |
| finally: | |
| clear_state() | |