From e21ba60d05cb91e041df2ec04cd76901d39d66f6 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Wed, 1 Jul 2026 22:47:07 +0200 Subject: [PATCH] update hf hijack, simplify anima loader, add few more models Signed-off-by: Vladimir Mandic --- CHANGELOG.md | 8 +- data/previews.json | 4 +- data/reference-community.json | 9 +- data/reference-quantized.json | 8 + modules/sd_hijack_hfhub.py | 11 +- pipelines/anima/modeling_llm_adapter.py | 215 ++++++++++++++ pipelines/anima/pipeline.py | 372 ++++++++++++++++++++++++ pipelines/model_anima.py | 17 +- 8 files changed, 634 insertions(+), 10 deletions(-) create mode 100644 pipelines/anima/modeling_llm_adapter.py create mode 100644 pipelines/anima/pipeline.py diff --git a/CHANGELOG.md b/CHANGELOG.md index e107e2fc7..ffb652f6f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,7 +5,7 @@ ### Highlights for 2026-07-01 Service-pack update with several fixes and quality-of-life improvements -Plus few new models: **Krea 2** +Plus few new models: **Krea 2**, **Photoroom PRXPixel**, **FLUX.2 Klein 9B KV** And **SDNQ** improvements: now with *NPU* support and its own native *attention* kernels! [Home](https://vladmandic.github.io/sdnext/) | [ChangeLog](https://github.com/vladmandic/automatic/blob/master/CHANGELOG.md) | [Docs](https://vladmandic.github.io/sdnext-docs/) | [Discord](https://discord.com/invite/sd-next-federal-batch-inspectors-1101998836328697867) | [Sponsor](https://github.com/sponsors/vladmandic) @@ -18,7 +18,9 @@ And **SDNQ** improvements: now with *NPU* support and its own native *attention* K2 is a 12.9B single-stream flow-matching DiT and using a Qwen3-VL-4B text encoder - [Photoroom PRXPixel](https://huggingface.co/Photoroom/prxpixel-t2i) pixel-space PRX variant using a Qwen3-VL text encoder and flow-matching scheduler supports *direct RGB* generation without a VAE and uses a *1024px* default sample size - - **Microsoft Lens** got unpublished, but we still got a mirror + - [Microsoft Lens](https://huggingface.co/Jinstudio/Lens) got unpublished, but we still got a mirror + - [FLUX.2 Klein 9B KV](https://huggingface.co/black-forest-labs/FLUX.2-klein-9b-kv) a bit late, but finally here in both *bf16* and *sdnq* pre-quantized variants + - plus several new community models... - **Features** - **SDNQ-Attention** modelled after *sage-attention*, but modified to support AMD and Intel GPUs in addition to nVidia @@ -41,6 +43,7 @@ And **SDNQ** improvements: now with *NPU* support and its own native *attention* - **Internal** - delay init of video models - **Fixes** + - anima: simplify loader - amd: hipBLASLt improved detection, thanks @0xDELUXA - api: add missing endpoint registration - api: openapi schema exposure @@ -48,6 +51,7 @@ And **SDNQ** improvements: now with *NPU* support and its own native *attention* - caption: button in standard-ui - embeddings: handle textual-inversion with new transformers - extensions: handle extension without remote + - huggingface: strip corrupt headers on download - insightface: missing dependencies - live preview: configurable pause when not in focus, thanks @Artheriax - log: strip ansi sequences from ring buffer and client side logging diff --git a/data/previews.json b/data/previews.json index 846aa42c5..d9f65721e 100644 --- a/data/previews.json +++ b/data/previews.json @@ -35,5 +35,7 @@ "Jinstudio--Lens": "microsoft--Lens.jpg", "Jinstudio--Lens-Base": "microsoft--Lens-Base.jpg", "Jinstudio--Lens-Turbo": "microsoft--Lens-Turbo.jpg", - "SahilCarterr--BRIA-3.2": "briaai--BRIA-3.2.jpg" + "SahilCarterr--BRIA-3.2": "briaai--BRIA-3.2.jpg", + "vladmandic--Flux.2-Klein-9B-KV-sdnq-hadamard-uint4": "black-forest-labs--FLUX.2-klein-9b-kv.jpg", + "vladmandic--Anima-1.0-Base-Merge-sdnq-hadamard-uint4": "vladmandic--Anima-1.0-Base.jpg" } diff --git a/data/reference-community.json b/data/reference-community.json index 04954dba2..50b582cf5 100644 --- a/data/reference-community.json +++ b/data/reference-community.json @@ -156,11 +156,18 @@ "date": "2026 February", "size": 57.7 }, - "Skywork/Unipic3-DMD": { + "Skywork Unipic3-DMD": { "path": "Skywork/Unipic3-DMD", "preview": "Skywork--Unipic3-DMD.jpg", "desc": "UniPic3-DMD-Model is a few-step image editing and multi-image composition model trained using Distribution Matching Distillation (DMD) and is a fine-tune of Qwen-Image-Edit.", "date": "2026 February", "size": 57.7 + }, + "Anima 1.0 Base V-Merge sdnq-hadamard-uint4": { + "path": "vladmandic/Anima-1.0-Base-Merge-sdnq-hadamard-uint4", + "preview": "vladmandic--Anima-1.0-Base.jpg", + "desc": "Anima 1.0 Base pre-merged with several LoRAs and quantized to uint4 using SDNQ with Hadamard. Flexible as it can be used with and without guidance.", + "date": "2026 July", + "size": 2.22 } } diff --git a/data/reference-quantized.json b/data/reference-quantized.json index 2c8942afd..cc58dbb8c 100644 --- a/data/reference-quantized.json +++ b/data/reference-quantized.json @@ -50,6 +50,14 @@ "size": 12.59, "date": "2026 January" }, + "Black Forest Labs FLUX.2 Klein 9B KV sdnq-uint4-dynamic-svd": { + "path": "vladmandic/Flux.2-Klein-9B-KV-sdnq-hadamard-uint4", + "preview": "black-forest-labs--FLUX.2-klein-9b-kv.jpg", + "desc": "Dynamic 4-bit quantization of black-forest-labs/FLUX.2-klein-9B-KV using SDNQ with Hadamard.", + "extras": "sampler: Default, cfg_scale: 1.0, steps: 4", + "size": 11.67, + "date": "2026 March" + }, "Chroma1-HD sdnq-svd-uint4": { "path": "Disty0/Chroma1-HD-SDNQ-uint4-svd-r32", "preview": "Disty0--Chroma1-HD-SDNQ-uint4-svd-r32.jpg", diff --git a/modules/sd_hijack_hfhub.py b/modules/sd_hijack_hfhub.py index f90ba94dc..e2819b5ef 100644 --- a/modules/sd_hijack_hfhub.py +++ b/modules/sd_hijack_hfhub.py @@ -6,6 +6,7 @@ from modules.logger import log debug = log.trace if os.environ.get('SD_DOWNLOAD_DEBUG', None) is not None else lambda *args, **kwargs: None orig_http_get = None orig_xet_get = None +orig_build_hf_headers = None def clean_user_agent(headers): @@ -59,11 +60,19 @@ def xet_get_hijack(*args, **kwargs): return res +def build_hf_headers_hijack(*args, **kwargs): + headers = orig_build_hf_headers(*args, **kwargs) + headers = clean_user_agent(headers) + return headers + + def init_hijack(): from huggingface_hub import file_download - global orig_http_get, orig_xet_get # pylint: disable=global-statement + global orig_http_get, orig_xet_get, orig_build_hf_headers # pylint: disable=global-statement if orig_http_get is None or orig_xet_get is None: orig_http_get = file_download.http_get orig_xet_get = file_download.xet_get + orig_build_hf_headers = file_download.build_hf_headers # pylint: disable=protected-access file_download.http_get = http_get_hijack file_download.xet_get = xet_get_hijack + file_download.build_hf_headers = build_hf_headers_hijack # pylint: disable=protected-access diff --git a/pipelines/anima/modeling_llm_adapter.py b/pipelines/anima/modeling_llm_adapter.py new file mode 100644 index 000000000..ac7bce1d4 --- /dev/null +++ b/pipelines/anima/modeling_llm_adapter.py @@ -0,0 +1,215 @@ +import torch +from torch import nn +import torch.nn.functional as F +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.models.modeling_utils import ModelMixin + + +def rotate_half(x): + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +def apply_rotary_pos_emb(x, cos, sin, unsqueeze_dim=1): + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + return (x * cos) + (rotate_half(x) * sin) + + +class RotaryEmbedding(nn.Module): + def __init__(self, head_dim): + super().__init__() + self.rope_theta = 10000 + inv_freq = 1.0 / ( + self.rope_theta + ** (torch.arange(0, head_dim, 2, dtype=torch.int64).to(dtype=torch.float) / head_dim) + ) + self.register_buffer("inv_freq", inv_freq, persistent=False) + + @torch.no_grad() + def forward(self, x, position_ids): + inv_freq_expanded = ( + self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device) + ) + position_ids_expanded = position_ids[:, None, :].float() + + device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu" + with torch.autocast(device_type=device_type, enabled=False): + freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2) + emb = torch.cat((freqs, freqs), dim=-1) + cos = emb.cos() + sin = emb.sin() + return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) + + +class Attention(nn.Module): + def __init__(self, query_dim, context_dim, n_heads, head_dim): + super().__init__() + inner_dim = head_dim * n_heads + self.n_heads = n_heads + self.head_dim = head_dim + + self.q_proj = nn.Linear(query_dim, inner_dim, bias=False) + self.q_norm = nn.RMSNorm(head_dim, eps=1e-6) + self.k_proj = nn.Linear(context_dim, inner_dim, bias=False) + self.k_norm = nn.RMSNorm(head_dim, eps=1e-6) + self.v_proj = nn.Linear(context_dim, inner_dim, bias=False) + self.o_proj = nn.Linear(inner_dim, query_dim, bias=False) + + def forward(self, x, mask=None, context=None, position_embeddings=None, position_embeddings_context=None): + context = x if context is None else context + input_shape = x.shape[:-1] + q_shape = (*input_shape, self.n_heads, self.head_dim) + context_shape = context.shape[:-1] + kv_shape = (*context_shape, self.n_heads, self.head_dim) + + query_states = self.q_norm(self.q_proj(x).view(q_shape)).transpose(1, 2) + key_states = self.k_norm(self.k_proj(context).view(kv_shape)).transpose(1, 2) + value_states = self.v_proj(context).view(kv_shape).transpose(1, 2) + + if position_embeddings is not None: + assert position_embeddings_context is not None + cos, sin = position_embeddings + query_states = apply_rotary_pos_emb(query_states, cos, sin) + cos, sin = position_embeddings_context + key_states = apply_rotary_pos_emb(key_states, cos, sin) + + attn_output = F.scaled_dot_product_attention(query_states, key_states, value_states, attn_mask=mask) + attn_output = attn_output.transpose(1, 2).reshape(*input_shape, -1).contiguous() + return self.o_proj(attn_output) + + +class TransformerBlock(nn.Module): + def __init__(self, source_dim, model_dim, num_heads=16, mlp_ratio=4.0, use_self_attn=True): + super().__init__() + self.use_self_attn = use_self_attn + + if self.use_self_attn: + self.norm_self_attn = nn.RMSNorm(model_dim, eps=1e-6) + self.self_attn = Attention( + query_dim=model_dim, + context_dim=model_dim, + n_heads=num_heads, + head_dim=model_dim // num_heads, + ) + + self.norm_cross_attn = nn.RMSNorm(model_dim, eps=1e-6) + self.cross_attn = Attention( + query_dim=model_dim, + context_dim=source_dim, + n_heads=num_heads, + head_dim=model_dim // num_heads, + ) + + self.norm_mlp = nn.RMSNorm(model_dim, eps=1e-6) + self.mlp = nn.Sequential( + nn.Linear(model_dim, int(model_dim * mlp_ratio)), + nn.GELU(), + nn.Linear(int(model_dim * mlp_ratio), model_dim), + ) + + def forward( + self, + x, + context, + target_attention_mask=None, + source_attention_mask=None, + position_embeddings=None, + position_embeddings_context=None, + ): + if self.use_self_attn: + normed = self.norm_self_attn(x) + attn_out = self.self_attn( + normed, + mask=target_attention_mask, + position_embeddings=position_embeddings, + position_embeddings_context=position_embeddings, + ) + x = x + attn_out + + normed = self.norm_cross_attn(x) + attn_out = self.cross_attn( + normed, + mask=source_attention_mask, + context=context, + position_embeddings=position_embeddings, + position_embeddings_context=position_embeddings_context, + ) + x = x + attn_out + x = x + self.mlp(self.norm_mlp(x)) + return x + + +class AnimaLLMAdapter(ModelMixin, ConfigMixin): + @register_to_config + def __init__( + self, + source_dim: int = 1024, + target_dim: int = 1024, + model_dim: int = 1024, + num_layers: int = 6, + num_heads: int = 16, + mlp_ratio: float = 4.0, + vocab_size: int = 32128, + use_self_attn: bool = True, + ): + super().__init__() + + self.embed = nn.Embedding(vocab_size, target_dim) + if model_dim != target_dim: + self.in_proj = nn.Linear(target_dim, model_dim) + else: + self.in_proj = nn.Identity() + self.rotary_emb = RotaryEmbedding(model_dim // num_heads) + self.blocks = nn.ModuleList( + [ + TransformerBlock( + source_dim, + model_dim, + num_heads=num_heads, + mlp_ratio=mlp_ratio, + use_self_attn=use_self_attn, + ) + for _ in range(num_layers) + ] + ) + self.out_proj = nn.Linear(model_dim, target_dim) + self.norm = nn.RMSNorm(target_dim, eps=1e-6) + + def forward( + self, + source_hidden_states: torch.Tensor, + target_input_ids: torch.Tensor, + target_attention_mask: torch.Tensor = None, + source_attention_mask: torch.Tensor = None, + ) -> torch.Tensor: + if target_attention_mask is not None: + target_attention_mask = target_attention_mask.to(torch.bool) + if target_attention_mask.ndim == 2: + target_attention_mask = target_attention_mask.unsqueeze(1).unsqueeze(1) + + if source_attention_mask is not None: + source_attention_mask = source_attention_mask.to(torch.bool) + if source_attention_mask.ndim == 2: + source_attention_mask = source_attention_mask.unsqueeze(1).unsqueeze(1) + + x = self.in_proj(self.embed(target_input_ids)) + context = source_hidden_states + + position_ids = torch.arange(x.shape[1], device=x.device).unsqueeze(0) + position_ids_context = torch.arange(context.shape[1], device=x.device).unsqueeze(0) + position_embeddings = self.rotary_emb(x, position_ids) + position_embeddings_context = self.rotary_emb(x, position_ids_context) + + for block in self.blocks: + x = block( + x, + context, + target_attention_mask=target_attention_mask, + source_attention_mask=source_attention_mask, + position_embeddings=position_embeddings, + position_embeddings_context=position_embeddings_context, + ) + + return self.norm(self.out_proj(x)) diff --git a/pipelines/anima/pipeline.py b/pipelines/anima/pipeline.py new file mode 100644 index 000000000..aa10b165d --- /dev/null +++ b/pipelines/anima/pipeline.py @@ -0,0 +1,372 @@ +# pylint: disable + +from typing import Callable, Dict, List, Optional, Union + +import numpy as np +import torch +from transformers import PreTrainedModel, PreTrainedTokenizerFast + +from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback +from diffusers.models import AutoencoderKLWan, CosmosTransformer3DModel +from diffusers.schedulers import FlowMatchEulerDiscreteScheduler +from diffusers.utils import logging +from diffusers.utils.torch_utils import randn_tensor +from diffusers.video_processor import VideoProcessor +from diffusers.pipelines.pipeline_utils import DiffusionPipeline +from diffusers.pipelines.cosmos.pipeline_output import CosmosImagePipelineOutput + +logger = logging.get_logger(__name__) + + +def retrieve_timesteps(scheduler, num_inference_steps=None, device=None, timesteps=None, sigmas=None, **kwargs): + if timesteps is not None and sigmas is not None: + raise ValueError("Only one of `timesteps` or `sigmas` can be passed.") + if timesteps is not None: + scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + elif sigmas is not None: + scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + else: + scheduler.set_timesteps(num_inference_steps, device=device, **kwargs) + timesteps = scheduler.timesteps + return timesteps, num_inference_steps + + +class AnimaTextToImagePipeline(DiffusionPipeline): + """Pipeline for text-to-image generation using the Anima model. + + Anima uses a Cosmos Predict2 backbone with a Qwen3 text encoder and an LLM adapter + that cross-attends T5 token embeddings to Qwen3 hidden states. + """ + + model_cpu_offload_seq = "text_encoder->llm_adapter->transformer->vae" + _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"] + + def __init__( + self, + text_encoder: PreTrainedModel, + tokenizer: PreTrainedTokenizerFast, + t5_tokenizer: PreTrainedTokenizerFast, + llm_adapter, + transformer: CosmosTransformer3DModel, + vae: AutoencoderKLWan, + scheduler: FlowMatchEulerDiscreteScheduler, + ): + super().__init__() + + self.register_modules( + text_encoder=text_encoder, + tokenizer=tokenizer, + t5_tokenizer=t5_tokenizer, + llm_adapter=llm_adapter, + transformer=transformer, + vae=vae, + scheduler=scheduler, + ) + + self.vae_scale_factor_temporal = 2 ** sum(self.vae.temperal_downsample) if getattr(self, "vae", None) else 4 + self.vae_scale_factor_spatial = 2 ** len(self.vae.temperal_downsample) if getattr(self, "vae", None) else 8 + self.video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial) + + def _encode_prompt( + self, + prompt: Union[str, List[str]], + device: torch.device, + dtype: torch.dtype, + max_sequence_length: int = 512, + ): + """Encode prompt through Qwen3 and run LLM adapter with T5 token IDs.""" + prompt = [prompt] if isinstance(prompt, str) else prompt + batch_size = len(prompt) + + # Check for empty prompts - return zero embeddings directly + all_empty = all(p.strip() == "" for p in prompt) + if all_empty: + return torch.zeros(batch_size, 512, self.llm_adapter.config.target_dim, device=device, dtype=dtype) + + # Tokenize with Qwen3 tokenizer + qwen_inputs = self.tokenizer( + prompt, + padding=True, + truncation=True, + max_length=max_sequence_length, + return_tensors="pt", + ) + qwen_input_ids = qwen_inputs.input_ids.to(device) + qwen_attention_mask = qwen_inputs.attention_mask.to(device) + + # Get Qwen3 hidden states + qwen_outputs = self.text_encoder( + input_ids=qwen_input_ids, + attention_mask=qwen_attention_mask, + ) + qwen_hidden_states = qwen_outputs.last_hidden_state.to(dtype=dtype) + + # Tokenize with T5 tokenizer (we only need the IDs for the adapter embedding) + t5_inputs = self.t5_tokenizer( + prompt, + padding=True, + truncation=True, + max_length=max_sequence_length, + return_tensors="pt", + ) + t5_input_ids = t5_inputs.input_ids.to(device) + + # Run LLM adapter: T5 token embeddings attend to Qwen3 hidden states + adapted_embeds = self.llm_adapter( + source_hidden_states=qwen_hidden_states, + target_input_ids=t5_input_ids, + ) + + # Pad to 512 sequence length if shorter + if adapted_embeds.shape[1] < 512: + adapted_embeds = torch.nn.functional.pad( + adapted_embeds, (0, 0, 0, 512 - adapted_embeds.shape[1]) + ) + + return adapted_embeds + + def encode_prompt( + self, + prompt: Union[str, List[str]], + negative_prompt: Optional[Union[str, List[str]]] = None, + do_classifier_free_guidance: bool = True, + num_images_per_prompt: int = 1, + prompt_embeds: Optional[torch.Tensor] = None, + negative_prompt_embeds: Optional[torch.Tensor] = None, + max_sequence_length: int = 512, + device: Optional[torch.device] = None, + dtype: Optional[torch.dtype] = None, + ): + device = device or self._execution_device + dtype = dtype or self.text_encoder.dtype + prompt = [prompt] if isinstance(prompt, str) else prompt + + if prompt is not None: + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + if prompt_embeds is None: + prompt_embeds = self._encode_prompt(prompt, device, dtype, max_sequence_length) + _, seq_len, _ = prompt_embeds.shape + prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1) + prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1) + + if do_classifier_free_guidance and negative_prompt_embeds is None: + negative_prompt = negative_prompt or "" + negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt + negative_prompt_embeds = self._encode_prompt(negative_prompt, device, dtype, max_sequence_length) + _, seq_len, _ = negative_prompt_embeds.shape + negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1) + negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1) + + return prompt_embeds, negative_prompt_embeds + + def prepare_latents( + self, + batch_size: int, + num_channels_latents: int, + height: int, + width: int, + num_frames: int = 1, + dtype: torch.dtype = None, + device: torch.device = None, + generator=None, + latents: torch.Tensor = None, + ): + num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1 + latent_height = height // self.vae_scale_factor_spatial + latent_width = width // self.vae_scale_factor_spatial + + if latents is not None: + return latents.to(device=device, dtype=dtype) + + shape = (batch_size, num_channels_latents, num_latent_frames, latent_height, latent_width) + latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) + return latents + + def check_inputs(self, prompt, height, width, prompt_embeds=None): + if height % 16 != 0 or width % 16 != 0: + raise ValueError(f"`height` and `width` have to be divisible by 16 but are {height} and {width}.") + if prompt is not None and prompt_embeds is not None: + raise ValueError("Cannot forward both `prompt` and `prompt_embeds`.") + elif prompt is None and prompt_embeds is None: + raise ValueError("Provide either `prompt` or `prompt_embeds`.") + + @property + def guidance_scale(self): + return self._guidance_scale + + @property + def do_classifier_free_guidance(self): + return self._guidance_scale > 1.0 + + @property + def num_timesteps(self): + return self._num_timesteps + + @property + def interrupt(self): + return self._interrupt + + @torch.no_grad() + def __call__( + self, + prompt: str | list[str] | None = None, + negative_prompt: Optional[str | list[str]] = None, + height: int = 768, + width: int = 1360, + num_inference_steps: int = 35, + guidance_scale: float = 7.0, + num_images_per_prompt: Optional[int] = 1, + generator: Optional[torch.Generator | list[torch.Generator]] = None, + latents: Optional[torch.Tensor] = None, + prompt_embeds: Optional[torch.Tensor] = None, + negative_prompt_embeds: Optional[torch.Tensor] = None, + output_type: Optional[str] = "pil", + return_dict: bool = True, + callback_on_step_end: Optional[ + Callable[[int, int, Dict], None] | PipelineCallback | MultiPipelineCallbacks + ] = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], + max_sequence_length: int = 512, + ): + if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)): + callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs + + num_frames = 1 + + self.check_inputs(prompt, height, width, prompt_embeds) + self._guidance_scale = guidance_scale # pylint: disable=attribute-defined-outside-init + self._current_timestep = None # pylint: disable=attribute-defined-outside-init + self._interrupt = False # pylint: disable=attribute-defined-outside-init + + device = self._execution_device + + if prompt is not None and isinstance(prompt, str): + batch_size = 1 + elif prompt is not None and isinstance(prompt, list): + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + # Encode prompt + prompt_embeds, negative_prompt_embeds = self.encode_prompt( + prompt=prompt, + negative_prompt=negative_prompt, + do_classifier_free_guidance=self.do_classifier_free_guidance, + num_images_per_prompt=num_images_per_prompt, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + device=device, + max_sequence_length=max_sequence_length, + ) + + # Prepare timesteps - use default descending schedule (1→0) + timesteps, num_inference_steps = retrieve_timesteps( + self.scheduler, num_inference_steps=num_inference_steps, device=device + ) + + # Prepare latents + transformer_dtype = self.transformer.dtype + num_channels_latents = self.transformer.config.in_channels + latents = self.prepare_latents( + batch_size * num_images_per_prompt, + num_channels_latents, + height, + width, + num_frames, + torch.float32, + device, + generator, + latents, + ) + + padding_mask = latents.new_zeros(1, 1, height, width, dtype=transformer_dtype) + + # Denoising loop using CONST preconditioning (flow matching velocity model): + # - c_in = 1.0 (no input scaling) + # - timestep = sigma (passed directly) + # - model output is the velocity: denoised = x - velocity * sigma + # - CFG applied to velocity (equivalent to applying to denoised for linear preconditioning) + num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order + self._num_timesteps = len(timesteps) # pylint: disable=attribute-defined-outside-init + + with self.progress_bar(total=num_inference_steps) as progress_bar: + for i, t in enumerate(timesteps): + if self.interrupt: + continue + + self._current_timestep = t # pylint: disable=attribute-defined-outside-init + sigma = self.scheduler.sigmas[i] + + # Pass sigma directly as timestep (CONST preconditioning) + timestep = sigma.expand(latents.shape[0]).to(transformer_dtype) + latent_model_input = latents.to(transformer_dtype) + + # Model predicts velocity (raw output IS the velocity for CONST) + velocity = self.transformer( + hidden_states=latent_model_input, + timestep=timestep, + encoder_hidden_states=prompt_embeds, + padding_mask=padding_mask, + return_dict=False, + )[0].float() + + if self.do_classifier_free_guidance: + velocity_uncond = self.transformer( + hidden_states=latent_model_input, + timestep=timestep, + encoder_hidden_states=negative_prompt_embeds, + padding_mask=padding_mask, + return_dict=False, + )[0].float() + velocity = velocity_uncond + self.guidance_scale * (velocity - velocity_uncond) + + # Euler step: scheduler computes x_next = x + (sigma_next - sigma) * velocity + latents = self.scheduler.step(velocity, t, latents, return_dict=False)[0] + + if callback_on_step_end is not None: + callback_kwargs = {} + for k in callback_on_step_end_tensor_inputs: + callback_kwargs[k] = locals()[k] + callback_outputs = callback_on_step_end(self, i, t, callback_kwargs) + latents = callback_outputs.pop("latents", latents) + prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds) + negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds) + + if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): + progress_bar.update() + + self._current_timestep = None # pylint: disable=attribute-defined-outside-init + + if output_type != "latent": + latents_mean = ( + torch.tensor(self.vae.config.latents_mean) + .view(1, self.vae.config.z_dim, 1, 1, 1) + .to(latents.device, latents.dtype) + ) + latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(1, self.vae.config.z_dim, 1, 1, 1).to( + latents.device, latents.dtype + ) + latents = latents / latents_std + latents_mean + video = self.vae.decode(latents.to(self.vae.dtype), return_dict=False)[0] + video = self.video_processor.postprocess_video(video, output_type=output_type) + image = [batch[0] for batch in video] + if isinstance(video, torch.Tensor): + image = torch.stack(image) + elif isinstance(video, np.ndarray): + image = np.stack(image) + else: + image = latents[:, :, 0] + + self.maybe_free_model_hooks() + + if not return_dict: + return (image,) + + return CosmosImagePipelineOutput(images=image) diff --git a/pipelines/model_anima.py b/pipelines/model_anima.py index 5a88a240b..a54d0e8df 100644 --- a/pipelines/model_anima.py +++ b/pipelines/model_anima.py @@ -1,9 +1,6 @@ -import os -import sys import importlib.util import transformers import diffusers -import huggingface_hub as hf from modules import shared, devices, sd_models, model_quant, sd_hijack_te, sd_hijack_vae, errors from modules.logger import log from pipelines import generic @@ -62,6 +59,11 @@ def load_anima(checkpoint_info, diffusers_load_config=None): return None # load-or-download custom pipeline modules from repo + """ + import os + import sys + import huggingface_hub as hf + if os.path.exists(os.path.join(repo_id, 'pipeline.py')): pipeline_file = os.path.join(repo_id, 'pipeline.py') else: @@ -96,7 +98,12 @@ def load_anima(checkpoint_info, diffusers_load_config=None): sys.modules['pipeline'] = pipeline_mod AnimaTextToImagePipeline = pipeline_mod.AnimaTextToImagePipeline AnimaLLMAdapter = adapter_mod.AnimaLLMAdapter + """ + import sys + from pipelines.anima import modeling_llm_adapter + sys.modules['modeling_llm_adapter'] = modeling_llm_adapter + from pipelines.anima.pipeline import AnimaTextToImagePipeline from pipelines.anima.anima_image import build_anima_pipeline_classes AnimaImageToImagePipeline, AnimaInpaintPipeline = build_anima_pipeline_classes(AnimaTextToImagePipeline) diffusers.pipelines.auto_pipeline.AUTO_TEXT2IMAGE_PIPELINES_MAPPING["anima"] = AnimaTextToImagePipeline @@ -106,7 +113,7 @@ def load_anima(checkpoint_info, diffusers_load_config=None): # UNET dropdown (shared.opts.sd_unet) may redirect the transformer to a # community file that bundles both the transformer and the llm_adapter. - transformer, llm_adapter = init_transformer_component(repo_id, diffusers_load_config, AnimaLLMAdapter) + transformer, llm_adapter = init_transformer_component(repo_id, diffusers_load_config, modeling_llm_adapter.AnimaLLMAdapter) if transformer is None: return None text_encoder = generic.load_text_encoder( @@ -120,7 +127,7 @@ def load_anima(checkpoint_info, diffusers_load_config=None): if llm_adapter is None: shared.state.begin('Load adapter') try: - llm_adapter = AnimaLLMAdapter.from_pretrained( + llm_adapter = modeling_llm_adapter.AnimaLLMAdapter.from_pretrained( repo_id, subfolder="llm_adapter", cache_dir=shared.opts.hfcache_dir,