vladmandic commited on
Commit
c78b2bb
·
verified ·
1 Parent(s): b245a37

Upload folder using huggingface_hub

Browse files
.gitattributes CHANGED
@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
33
  *.zip filter=lfs diff=lfs merge=lfs -text
34
  *.zst filter=lfs diff=lfs merge=lfs -text
35
  *tfevents* filter=lfs diff=lfs merge=lfs -text
 
 
33
  *.zip filter=lfs diff=lfs merge=lfs -text
34
  *.zst filter=lfs diff=lfs merge=lfs -text
35
  *tfevents* filter=lfs diff=lfs merge=lfs -text
36
+ tokenizer/tokenizer.json filter=lfs diff=lfs merge=lfs -text
__pycache__/pipeline.cpython-313.pyc ADDED
Binary file (16.5 kB). View file
 
llm_adapter/__pycache__/modeling_llm_adapter.cpython-313.pyc ADDED
Binary file (11.8 kB). View file
 
llm_adapter/config.json ADDED
@@ -0,0 +1,12 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "_class_name": "AnimaLLMAdapter",
3
+ "_diffusers_version": "0.39.0.dev0",
4
+ "mlp_ratio": 4.0,
5
+ "model_dim": 1024,
6
+ "num_heads": 16,
7
+ "num_layers": 6,
8
+ "source_dim": 1024,
9
+ "target_dim": 1024,
10
+ "use_self_attn": true,
11
+ "vocab_size": 32128
12
+ }
llm_adapter/diffusion_pytorch_model.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:fdcf606524892e35b270b96fc8bbd8d77e615ab882675a049f6e7c2ff6b71da3
3
+ size 269339400
llm_adapter/modeling_llm_adapter.py ADDED
@@ -0,0 +1,215 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ from torch import nn
3
+ import torch.nn.functional as F
4
+ from diffusers.configuration_utils import ConfigMixin, register_to_config
5
+ from diffusers.models.modeling_utils import ModelMixin
6
+
7
+
8
+ def rotate_half(x):
9
+ x1 = x[..., : x.shape[-1] // 2]
10
+ x2 = x[..., x.shape[-1] // 2 :]
11
+ return torch.cat((-x2, x1), dim=-1)
12
+
13
+
14
+ def apply_rotary_pos_emb(x, cos, sin, unsqueeze_dim=1):
15
+ cos = cos.unsqueeze(unsqueeze_dim)
16
+ sin = sin.unsqueeze(unsqueeze_dim)
17
+ return (x * cos) + (rotate_half(x) * sin)
18
+
19
+
20
+ class RotaryEmbedding(nn.Module):
21
+ def __init__(self, head_dim):
22
+ super().__init__()
23
+ self.rope_theta = 10000
24
+ inv_freq = 1.0 / (
25
+ self.rope_theta
26
+ ** (torch.arange(0, head_dim, 2, dtype=torch.int64).to(dtype=torch.float) / head_dim)
27
+ )
28
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
29
+
30
+ @torch.no_grad()
31
+ def forward(self, x, position_ids):
32
+ inv_freq_expanded = (
33
+ self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
34
+ )
35
+ position_ids_expanded = position_ids[:, None, :].float()
36
+
37
+ device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
38
+ with torch.autocast(device_type=device_type, enabled=False):
39
+ freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
40
+ emb = torch.cat((freqs, freqs), dim=-1)
41
+ cos = emb.cos()
42
+ sin = emb.sin()
43
+ return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
44
+
45
+
46
+ class Attention(nn.Module):
47
+ def __init__(self, query_dim, context_dim, n_heads, head_dim):
48
+ super().__init__()
49
+ inner_dim = head_dim * n_heads
50
+ self.n_heads = n_heads
51
+ self.head_dim = head_dim
52
+
53
+ self.q_proj = nn.Linear(query_dim, inner_dim, bias=False)
54
+ self.q_norm = nn.RMSNorm(head_dim, eps=1e-6)
55
+ self.k_proj = nn.Linear(context_dim, inner_dim, bias=False)
56
+ self.k_norm = nn.RMSNorm(head_dim, eps=1e-6)
57
+ self.v_proj = nn.Linear(context_dim, inner_dim, bias=False)
58
+ self.o_proj = nn.Linear(inner_dim, query_dim, bias=False)
59
+
60
+ def forward(self, x, mask=None, context=None, position_embeddings=None, position_embeddings_context=None):
61
+ context = x if context is None else context
62
+ input_shape = x.shape[:-1]
63
+ q_shape = (*input_shape, self.n_heads, self.head_dim)
64
+ context_shape = context.shape[:-1]
65
+ kv_shape = (*context_shape, self.n_heads, self.head_dim)
66
+
67
+ query_states = self.q_norm(self.q_proj(x).view(q_shape)).transpose(1, 2)
68
+ key_states = self.k_norm(self.k_proj(context).view(kv_shape)).transpose(1, 2)
69
+ value_states = self.v_proj(context).view(kv_shape).transpose(1, 2)
70
+
71
+ if position_embeddings is not None:
72
+ assert position_embeddings_context is not None
73
+ cos, sin = position_embeddings
74
+ query_states = apply_rotary_pos_emb(query_states, cos, sin)
75
+ cos, sin = position_embeddings_context
76
+ key_states = apply_rotary_pos_emb(key_states, cos, sin)
77
+
78
+ attn_output = F.scaled_dot_product_attention(query_states, key_states, value_states, attn_mask=mask)
79
+ attn_output = attn_output.transpose(1, 2).reshape(*input_shape, -1).contiguous()
80
+ return self.o_proj(attn_output)
81
+
82
+
83
+ class TransformerBlock(nn.Module):
84
+ def __init__(self, source_dim, model_dim, num_heads=16, mlp_ratio=4.0, use_self_attn=True):
85
+ super().__init__()
86
+ self.use_self_attn = use_self_attn
87
+
88
+ if self.use_self_attn:
89
+ self.norm_self_attn = nn.RMSNorm(model_dim, eps=1e-6)
90
+ self.self_attn = Attention(
91
+ query_dim=model_dim,
92
+ context_dim=model_dim,
93
+ n_heads=num_heads,
94
+ head_dim=model_dim // num_heads,
95
+ )
96
+
97
+ self.norm_cross_attn = nn.RMSNorm(model_dim, eps=1e-6)
98
+ self.cross_attn = Attention(
99
+ query_dim=model_dim,
100
+ context_dim=source_dim,
101
+ n_heads=num_heads,
102
+ head_dim=model_dim // num_heads,
103
+ )
104
+
105
+ self.norm_mlp = nn.RMSNorm(model_dim, eps=1e-6)
106
+ self.mlp = nn.Sequential(
107
+ nn.Linear(model_dim, int(model_dim * mlp_ratio)),
108
+ nn.GELU(),
109
+ nn.Linear(int(model_dim * mlp_ratio), model_dim),
110
+ )
111
+
112
+ def forward(
113
+ self,
114
+ x,
115
+ context,
116
+ target_attention_mask=None,
117
+ source_attention_mask=None,
118
+ position_embeddings=None,
119
+ position_embeddings_context=None,
120
+ ):
121
+ if self.use_self_attn:
122
+ normed = self.norm_self_attn(x)
123
+ attn_out = self.self_attn(
124
+ normed,
125
+ mask=target_attention_mask,
126
+ position_embeddings=position_embeddings,
127
+ position_embeddings_context=position_embeddings,
128
+ )
129
+ x = x + attn_out
130
+
131
+ normed = self.norm_cross_attn(x)
132
+ attn_out = self.cross_attn(
133
+ normed,
134
+ mask=source_attention_mask,
135
+ context=context,
136
+ position_embeddings=position_embeddings,
137
+ position_embeddings_context=position_embeddings_context,
138
+ )
139
+ x = x + attn_out
140
+ x = x + self.mlp(self.norm_mlp(x))
141
+ return x
142
+
143
+
144
+ class AnimaLLMAdapter(ModelMixin, ConfigMixin):
145
+ @register_to_config
146
+ def __init__(
147
+ self,
148
+ source_dim: int = 1024,
149
+ target_dim: int = 1024,
150
+ model_dim: int = 1024,
151
+ num_layers: int = 6,
152
+ num_heads: int = 16,
153
+ mlp_ratio: float = 4.0,
154
+ vocab_size: int = 32128,
155
+ use_self_attn: bool = True,
156
+ ):
157
+ super().__init__()
158
+
159
+ self.embed = nn.Embedding(vocab_size, target_dim)
160
+ if model_dim != target_dim:
161
+ self.in_proj = nn.Linear(target_dim, model_dim)
162
+ else:
163
+ self.in_proj = nn.Identity()
164
+ self.rotary_emb = RotaryEmbedding(model_dim // num_heads)
165
+ self.blocks = nn.ModuleList(
166
+ [
167
+ TransformerBlock(
168
+ source_dim,
169
+ model_dim,
170
+ num_heads=num_heads,
171
+ mlp_ratio=mlp_ratio,
172
+ use_self_attn=use_self_attn,
173
+ )
174
+ for _ in range(num_layers)
175
+ ]
176
+ )
177
+ self.out_proj = nn.Linear(model_dim, target_dim)
178
+ self.norm = nn.RMSNorm(target_dim, eps=1e-6)
179
+
180
+ def forward(
181
+ self,
182
+ source_hidden_states: torch.Tensor,
183
+ target_input_ids: torch.Tensor,
184
+ target_attention_mask: torch.Tensor = None,
185
+ source_attention_mask: torch.Tensor = None,
186
+ ) -> torch.Tensor:
187
+ if target_attention_mask is not None:
188
+ target_attention_mask = target_attention_mask.to(torch.bool)
189
+ if target_attention_mask.ndim == 2:
190
+ target_attention_mask = target_attention_mask.unsqueeze(1).unsqueeze(1)
191
+
192
+ if source_attention_mask is not None:
193
+ source_attention_mask = source_attention_mask.to(torch.bool)
194
+ if source_attention_mask.ndim == 2:
195
+ source_attention_mask = source_attention_mask.unsqueeze(1).unsqueeze(1)
196
+
197
+ x = self.in_proj(self.embed(target_input_ids))
198
+ context = source_hidden_states
199
+
200
+ position_ids = torch.arange(x.shape[1], device=x.device).unsqueeze(0)
201
+ position_ids_context = torch.arange(context.shape[1], device=x.device).unsqueeze(0)
202
+ position_embeddings = self.rotary_emb(x, position_ids)
203
+ position_embeddings_context = self.rotary_emb(x, position_ids_context)
204
+
205
+ for block in self.blocks:
206
+ x = block(
207
+ x,
208
+ context,
209
+ target_attention_mask=target_attention_mask,
210
+ source_attention_mask=source_attention_mask,
211
+ position_embeddings=position_embeddings,
212
+ position_embeddings_context=position_embeddings_context,
213
+ )
214
+
215
+ return self.norm(self.out_proj(x))
model_index.json ADDED
@@ -0,0 +1,32 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "_class_name": "AnimaTextToImagePipeline",
3
+ "_diffusers_version": "0.39.0.dev0",
4
+ "llm_adapter": [
5
+ "modeling_llm_adapter",
6
+ "AnimaLLMAdapter"
7
+ ],
8
+ "scheduler": [
9
+ "diffusers",
10
+ "FlowMatchEulerDiscreteScheduler"
11
+ ],
12
+ "t5_tokenizer": [
13
+ "transformers",
14
+ "T5Tokenizer"
15
+ ],
16
+ "text_encoder": [
17
+ "transformers",
18
+ "Qwen3Model"
19
+ ],
20
+ "tokenizer": [
21
+ "transformers",
22
+ "Qwen2Tokenizer"
23
+ ],
24
+ "transformer": [
25
+ "diffusers",
26
+ "CosmosTransformer3DModel"
27
+ ],
28
+ "vae": [
29
+ "diffusers",
30
+ "AutoencoderKLWan"
31
+ ]
32
+ }
pipeline.py ADDED
@@ -0,0 +1,371 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import inspect
2
+ from typing import Callable, Dict, List, Optional, Union
3
+
4
+ import numpy as np
5
+ import torch
6
+ from transformers import PreTrainedModel, PreTrainedTokenizerFast
7
+
8
+ from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback
9
+ from diffusers.models import AutoencoderKLWan, CosmosTransformer3DModel
10
+ from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
11
+ from diffusers.utils import logging
12
+ from diffusers.utils.torch_utils import randn_tensor
13
+ from diffusers.video_processor import VideoProcessor
14
+ from diffusers.pipelines.pipeline_utils import DiffusionPipeline
15
+ from diffusers.pipelines.cosmos.pipeline_output import CosmosImagePipelineOutput
16
+
17
+ logger = logging.get_logger(__name__)
18
+
19
+
20
+ def retrieve_timesteps(scheduler, num_inference_steps=None, device=None, timesteps=None, sigmas=None, **kwargs):
21
+ if timesteps is not None and sigmas is not None:
22
+ raise ValueError("Only one of `timesteps` or `sigmas` can be passed.")
23
+ if timesteps is not None:
24
+ scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
25
+ timesteps = scheduler.timesteps
26
+ num_inference_steps = len(timesteps)
27
+ elif sigmas is not None:
28
+ scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
29
+ timesteps = scheduler.timesteps
30
+ num_inference_steps = len(timesteps)
31
+ else:
32
+ scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
33
+ timesteps = scheduler.timesteps
34
+ return timesteps, num_inference_steps
35
+
36
+
37
+ class AnimaTextToImagePipeline(DiffusionPipeline):
38
+ """Pipeline for text-to-image generation using the Anima model.
39
+
40
+ Anima uses a Cosmos Predict2 backbone with a Qwen3 text encoder and an LLM adapter
41
+ that cross-attends T5 token embeddings to Qwen3 hidden states.
42
+ """
43
+
44
+ model_cpu_offload_seq = "text_encoder->llm_adapter->transformer->vae"
45
+ _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"]
46
+
47
+ def __init__(
48
+ self,
49
+ text_encoder: PreTrainedModel,
50
+ tokenizer: PreTrainedTokenizerFast,
51
+ t5_tokenizer: PreTrainedTokenizerFast,
52
+ llm_adapter,
53
+ transformer: CosmosTransformer3DModel,
54
+ vae: AutoencoderKLWan,
55
+ scheduler: FlowMatchEulerDiscreteScheduler,
56
+ ):
57
+ super().__init__()
58
+
59
+ self.register_modules(
60
+ text_encoder=text_encoder,
61
+ tokenizer=tokenizer,
62
+ t5_tokenizer=t5_tokenizer,
63
+ llm_adapter=llm_adapter,
64
+ transformer=transformer,
65
+ vae=vae,
66
+ scheduler=scheduler,
67
+ )
68
+
69
+ self.vae_scale_factor_temporal = 2 ** sum(self.vae.temperal_downsample) if getattr(self, "vae", None) else 4
70
+ self.vae_scale_factor_spatial = 2 ** len(self.vae.temperal_downsample) if getattr(self, "vae", None) else 8
71
+ self.video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial)
72
+
73
+ def _encode_prompt(
74
+ self,
75
+ prompt: Union[str, List[str]],
76
+ device: torch.device,
77
+ dtype: torch.dtype,
78
+ max_sequence_length: int = 512,
79
+ ):
80
+ """Encode prompt through Qwen3 and run LLM adapter with T5 token IDs."""
81
+ prompt = [prompt] if isinstance(prompt, str) else prompt
82
+ batch_size = len(prompt)
83
+
84
+ # Check for empty prompts - return zero embeddings directly
85
+ all_empty = all(p.strip() == "" for p in prompt)
86
+ if all_empty:
87
+ return torch.zeros(batch_size, 512, self.llm_adapter.config.target_dim, device=device, dtype=dtype)
88
+
89
+ # Tokenize with Qwen3 tokenizer
90
+ qwen_inputs = self.tokenizer(
91
+ prompt,
92
+ padding=True,
93
+ truncation=True,
94
+ max_length=max_sequence_length,
95
+ return_tensors="pt",
96
+ )
97
+ qwen_input_ids = qwen_inputs.input_ids.to(device)
98
+ qwen_attention_mask = qwen_inputs.attention_mask.to(device)
99
+
100
+ # Get Qwen3 hidden states
101
+ qwen_outputs = self.text_encoder(
102
+ input_ids=qwen_input_ids,
103
+ attention_mask=qwen_attention_mask,
104
+ )
105
+ qwen_hidden_states = qwen_outputs.last_hidden_state.to(dtype=dtype)
106
+
107
+ # Tokenize with T5 tokenizer (we only need the IDs for the adapter embedding)
108
+ t5_inputs = self.t5_tokenizer(
109
+ prompt,
110
+ padding=True,
111
+ truncation=True,
112
+ max_length=max_sequence_length,
113
+ return_tensors="pt",
114
+ )
115
+ t5_input_ids = t5_inputs.input_ids.to(device)
116
+
117
+ # Run LLM adapter: T5 token embeddings attend to Qwen3 hidden states
118
+ adapted_embeds = self.llm_adapter(
119
+ source_hidden_states=qwen_hidden_states,
120
+ target_input_ids=t5_input_ids,
121
+ )
122
+
123
+ # Pad to 512 sequence length if shorter
124
+ if adapted_embeds.shape[1] < 512:
125
+ adapted_embeds = torch.nn.functional.pad(
126
+ adapted_embeds, (0, 0, 0, 512 - adapted_embeds.shape[1])
127
+ )
128
+
129
+ return adapted_embeds
130
+
131
+ def encode_prompt(
132
+ self,
133
+ prompt: Union[str, List[str]],
134
+ negative_prompt: Optional[Union[str, List[str]]] = None,
135
+ do_classifier_free_guidance: bool = True,
136
+ num_images_per_prompt: int = 1,
137
+ prompt_embeds: Optional[torch.Tensor] = None,
138
+ negative_prompt_embeds: Optional[torch.Tensor] = None,
139
+ max_sequence_length: int = 512,
140
+ device: Optional[torch.device] = None,
141
+ dtype: Optional[torch.dtype] = None,
142
+ ):
143
+ device = device or self._execution_device
144
+ dtype = dtype or self.text_encoder.dtype
145
+ prompt = [prompt] if isinstance(prompt, str) else prompt
146
+
147
+ if prompt is not None:
148
+ batch_size = len(prompt)
149
+ else:
150
+ batch_size = prompt_embeds.shape[0]
151
+
152
+ if prompt_embeds is None:
153
+ prompt_embeds = self._encode_prompt(prompt, device, dtype, max_sequence_length)
154
+ _, seq_len, _ = prompt_embeds.shape
155
+ prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
156
+ prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
157
+
158
+ if do_classifier_free_guidance and negative_prompt_embeds is None:
159
+ negative_prompt = negative_prompt or ""
160
+ negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt
161
+ negative_prompt_embeds = self._encode_prompt(negative_prompt, device, dtype, max_sequence_length)
162
+ _, seq_len, _ = negative_prompt_embeds.shape
163
+ negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)
164
+ negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
165
+
166
+ return prompt_embeds, negative_prompt_embeds
167
+
168
+ def prepare_latents(
169
+ self,
170
+ batch_size: int,
171
+ num_channels_latents: int,
172
+ height: int,
173
+ width: int,
174
+ num_frames: int = 1,
175
+ dtype: torch.dtype = None,
176
+ device: torch.device = None,
177
+ generator=None,
178
+ latents: torch.Tensor = None,
179
+ ):
180
+ num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1
181
+ latent_height = height // self.vae_scale_factor_spatial
182
+ latent_width = width // self.vae_scale_factor_spatial
183
+
184
+ if latents is not None:
185
+ return latents.to(device=device, dtype=dtype)
186
+
187
+ shape = (batch_size, num_channels_latents, num_latent_frames, latent_height, latent_width)
188
+ latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
189
+ return latents
190
+
191
+ def check_inputs(self, prompt, height, width, prompt_embeds=None):
192
+ if height % 16 != 0 or width % 16 != 0:
193
+ raise ValueError(f"`height` and `width` have to be divisible by 16 but are {height} and {width}.")
194
+ if prompt is not None and prompt_embeds is not None:
195
+ raise ValueError("Cannot forward both `prompt` and `prompt_embeds`.")
196
+ elif prompt is None and prompt_embeds is None:
197
+ raise ValueError("Provide either `prompt` or `prompt_embeds`.")
198
+
199
+ @property
200
+ def guidance_scale(self):
201
+ return self._guidance_scale
202
+
203
+ @property
204
+ def do_classifier_free_guidance(self):
205
+ return self._guidance_scale > 1.0
206
+
207
+ @property
208
+ def num_timesteps(self):
209
+ return self._num_timesteps
210
+
211
+ @property
212
+ def interrupt(self):
213
+ return self._interrupt
214
+
215
+ @torch.no_grad()
216
+ def __call__(
217
+ self,
218
+ prompt: Union[str, List[str]] = None,
219
+ negative_prompt: Optional[Union[str, List[str]]] = None,
220
+ height: int = 768,
221
+ width: int = 1360,
222
+ num_inference_steps: int = 35,
223
+ guidance_scale: float = 7.0,
224
+ num_images_per_prompt: Optional[int] = 1,
225
+ generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
226
+ latents: Optional[torch.Tensor] = None,
227
+ prompt_embeds: Optional[torch.Tensor] = None,
228
+ negative_prompt_embeds: Optional[torch.Tensor] = None,
229
+ output_type: Optional[str] = "pil",
230
+ return_dict: bool = True,
231
+ callback_on_step_end: Optional[
232
+ Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
233
+ ] = None,
234
+ callback_on_step_end_tensor_inputs: List[str] = ["latents"],
235
+ max_sequence_length: int = 512,
236
+ ):
237
+ if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)):
238
+ callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs
239
+
240
+ num_frames = 1
241
+
242
+ self.check_inputs(prompt, height, width, prompt_embeds)
243
+ self._guidance_scale = guidance_scale
244
+ self._current_timestep = None
245
+ self._interrupt = False
246
+
247
+ device = self._execution_device
248
+
249
+ if prompt is not None and isinstance(prompt, str):
250
+ batch_size = 1
251
+ elif prompt is not None and isinstance(prompt, list):
252
+ batch_size = len(prompt)
253
+ else:
254
+ batch_size = prompt_embeds.shape[0]
255
+
256
+ # Encode prompt
257
+ prompt_embeds, negative_prompt_embeds = self.encode_prompt(
258
+ prompt=prompt,
259
+ negative_prompt=negative_prompt,
260
+ do_classifier_free_guidance=self.do_classifier_free_guidance,
261
+ num_images_per_prompt=num_images_per_prompt,
262
+ prompt_embeds=prompt_embeds,
263
+ negative_prompt_embeds=negative_prompt_embeds,
264
+ device=device,
265
+ max_sequence_length=max_sequence_length,
266
+ )
267
+
268
+ # Prepare timesteps - use default descending schedule (1→0)
269
+ timesteps, num_inference_steps = retrieve_timesteps(
270
+ self.scheduler, num_inference_steps=num_inference_steps, device=device
271
+ )
272
+
273
+ # Prepare latents
274
+ transformer_dtype = self.transformer.dtype
275
+ num_channels_latents = self.transformer.config.in_channels
276
+ latents = self.prepare_latents(
277
+ batch_size * num_images_per_prompt,
278
+ num_channels_latents,
279
+ height,
280
+ width,
281
+ num_frames,
282
+ torch.float32,
283
+ device,
284
+ generator,
285
+ latents,
286
+ )
287
+
288
+ padding_mask = latents.new_zeros(1, 1, height, width, dtype=transformer_dtype)
289
+
290
+ # Denoising loop using CONST preconditioning (flow matching velocity model):
291
+ # - c_in = 1.0 (no input scaling)
292
+ # - timestep = sigma (passed directly)
293
+ # - model output is the velocity: denoised = x - velocity * sigma
294
+ # - CFG applied to velocity (equivalent to applying to denoised for linear preconditioning)
295
+ num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
296
+ self._num_timesteps = len(timesteps)
297
+
298
+ with self.progress_bar(total=num_inference_steps) as progress_bar:
299
+ for i, t in enumerate(timesteps):
300
+ if self.interrupt:
301
+ continue
302
+
303
+ self._current_timestep = t
304
+ sigma = self.scheduler.sigmas[i]
305
+
306
+ # Pass sigma directly as timestep (CONST preconditioning)
307
+ timestep = sigma.expand(latents.shape[0]).to(transformer_dtype)
308
+ latent_model_input = latents.to(transformer_dtype)
309
+
310
+ # Model predicts velocity (raw output IS the velocity for CONST)
311
+ velocity = self.transformer(
312
+ hidden_states=latent_model_input,
313
+ timestep=timestep,
314
+ encoder_hidden_states=prompt_embeds,
315
+ padding_mask=padding_mask,
316
+ return_dict=False,
317
+ )[0].float()
318
+
319
+ if self.do_classifier_free_guidance:
320
+ velocity_uncond = self.transformer(
321
+ hidden_states=latent_model_input,
322
+ timestep=timestep,
323
+ encoder_hidden_states=negative_prompt_embeds,
324
+ padding_mask=padding_mask,
325
+ return_dict=False,
326
+ )[0].float()
327
+ velocity = velocity_uncond + self.guidance_scale * (velocity - velocity_uncond)
328
+
329
+ # Euler step: scheduler computes x_next = x + (sigma_next - sigma) * velocity
330
+ latents = self.scheduler.step(velocity, t, latents, return_dict=False)[0]
331
+
332
+ if callback_on_step_end is not None:
333
+ callback_kwargs = {}
334
+ for k in callback_on_step_end_tensor_inputs:
335
+ callback_kwargs[k] = locals()[k]
336
+ callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
337
+ latents = callback_outputs.pop("latents", latents)
338
+ prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
339
+ negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds)
340
+
341
+ if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
342
+ progress_bar.update()
343
+
344
+ self._current_timestep = None
345
+
346
+ if not output_type == "latent":
347
+ latents_mean = (
348
+ torch.tensor(self.vae.config.latents_mean)
349
+ .view(1, self.vae.config.z_dim, 1, 1, 1)
350
+ .to(latents.device, latents.dtype)
351
+ )
352
+ latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(1, self.vae.config.z_dim, 1, 1, 1).to(
353
+ latents.device, latents.dtype
354
+ )
355
+ latents = latents / latents_std + latents_mean
356
+ video = self.vae.decode(latents.to(self.vae.dtype), return_dict=False)[0]
357
+ video = self.video_processor.postprocess_video(video, output_type=output_type)
358
+ image = [batch[0] for batch in video]
359
+ if isinstance(video, torch.Tensor):
360
+ image = torch.stack(image)
361
+ elif isinstance(video, np.ndarray):
362
+ image = np.stack(image)
363
+ else:
364
+ image = latents[:, :, 0]
365
+
366
+ self.maybe_free_model_hooks()
367
+
368
+ if not return_dict:
369
+ return (image,)
370
+
371
+ return CosmosImagePipelineOutput(images=image)
scheduler/scheduler_config.json ADDED
@@ -0,0 +1,18 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "_class_name": "FlowMatchEulerDiscreteScheduler",
3
+ "_diffusers_version": "0.39.0.dev0",
4
+ "base_image_seq_len": 256,
5
+ "base_shift": 0.5,
6
+ "invert_sigmas": false,
7
+ "max_image_seq_len": 4096,
8
+ "max_shift": 1.15,
9
+ "num_train_timesteps": 1000,
10
+ "shift": 3.0,
11
+ "shift_terminal": null,
12
+ "stochastic_sampling": false,
13
+ "time_shift_type": "exponential",
14
+ "use_beta_sigmas": false,
15
+ "use_dynamic_shifting": false,
16
+ "use_exponential_sigmas": false,
17
+ "use_karras_sigmas": false
18
+ }
t5_tokenizer/tokenizer.json ADDED
The diff for this file is too large to render. See raw diff
 
t5_tokenizer/tokenizer_config.json ADDED
@@ -0,0 +1,315 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "backend": "tokenizers",
3
+ "clean_up_tokenization_spaces": true,
4
+ "eos_token": "</s>",
5
+ "extra_id_0": "<extra_id_0>",
6
+ "extra_id_1": "<extra_id_1>",
7
+ "extra_id_10": "<extra_id_10>",
8
+ "extra_id_11": "<extra_id_11>",
9
+ "extra_id_12": "<extra_id_12>",
10
+ "extra_id_13": "<extra_id_13>",
11
+ "extra_id_14": "<extra_id_14>",
12
+ "extra_id_15": "<extra_id_15>",
13
+ "extra_id_16": "<extra_id_16>",
14
+ "extra_id_17": "<extra_id_17>",
15
+ "extra_id_18": "<extra_id_18>",
16
+ "extra_id_19": "<extra_id_19>",
17
+ "extra_id_2": "<extra_id_2>",
18
+ "extra_id_20": "<extra_id_20>",
19
+ "extra_id_21": "<extra_id_21>",
20
+ "extra_id_22": "<extra_id_22>",
21
+ "extra_id_23": "<extra_id_23>",
22
+ "extra_id_24": "<extra_id_24>",
23
+ "extra_id_25": "<extra_id_25>",
24
+ "extra_id_26": "<extra_id_26>",
25
+ "extra_id_27": "<extra_id_27>",
26
+ "extra_id_28": "<extra_id_28>",
27
+ "extra_id_29": "<extra_id_29>",
28
+ "extra_id_3": "<extra_id_3>",
29
+ "extra_id_30": "<extra_id_30>",
30
+ "extra_id_31": "<extra_id_31>",
31
+ "extra_id_32": "<extra_id_32>",
32
+ "extra_id_33": "<extra_id_33>",
33
+ "extra_id_34": "<extra_id_34>",
34
+ "extra_id_35": "<extra_id_35>",
35
+ "extra_id_36": "<extra_id_36>",
36
+ "extra_id_37": "<extra_id_37>",
37
+ "extra_id_38": "<extra_id_38>",
38
+ "extra_id_39": "<extra_id_39>",
39
+ "extra_id_4": "<extra_id_4>",
40
+ "extra_id_40": "<extra_id_40>",
41
+ "extra_id_41": "<extra_id_41>",
42
+ "extra_id_42": "<extra_id_42>",
43
+ "extra_id_43": "<extra_id_43>",
44
+ "extra_id_44": "<extra_id_44>",
45
+ "extra_id_45": "<extra_id_45>",
46
+ "extra_id_46": "<extra_id_46>",
47
+ "extra_id_47": "<extra_id_47>",
48
+ "extra_id_48": "<extra_id_48>",
49
+ "extra_id_49": "<extra_id_49>",
50
+ "extra_id_5": "<extra_id_5>",
51
+ "extra_id_50": "<extra_id_50>",
52
+ "extra_id_51": "<extra_id_51>",
53
+ "extra_id_52": "<extra_id_52>",
54
+ "extra_id_53": "<extra_id_53>",
55
+ "extra_id_54": "<extra_id_54>",
56
+ "extra_id_55": "<extra_id_55>",
57
+ "extra_id_56": "<extra_id_56>",
58
+ "extra_id_57": "<extra_id_57>",
59
+ "extra_id_58": "<extra_id_58>",
60
+ "extra_id_59": "<extra_id_59>",
61
+ "extra_id_6": "<extra_id_6>",
62
+ "extra_id_60": "<extra_id_60>",
63
+ "extra_id_61": "<extra_id_61>",
64
+ "extra_id_62": "<extra_id_62>",
65
+ "extra_id_63": "<extra_id_63>",
66
+ "extra_id_64": "<extra_id_64>",
67
+ "extra_id_65": "<extra_id_65>",
68
+ "extra_id_66": "<extra_id_66>",
69
+ "extra_id_67": "<extra_id_67>",
70
+ "extra_id_68": "<extra_id_68>",
71
+ "extra_id_69": "<extra_id_69>",
72
+ "extra_id_7": "<extra_id_7>",
73
+ "extra_id_70": "<extra_id_70>",
74
+ "extra_id_71": "<extra_id_71>",
75
+ "extra_id_72": "<extra_id_72>",
76
+ "extra_id_73": "<extra_id_73>",
77
+ "extra_id_74": "<extra_id_74>",
78
+ "extra_id_75": "<extra_id_75>",
79
+ "extra_id_76": "<extra_id_76>",
80
+ "extra_id_77": "<extra_id_77>",
81
+ "extra_id_78": "<extra_id_78>",
82
+ "extra_id_79": "<extra_id_79>",
83
+ "extra_id_8": "<extra_id_8>",
84
+ "extra_id_80": "<extra_id_80>",
85
+ "extra_id_81": "<extra_id_81>",
86
+ "extra_id_82": "<extra_id_82>",
87
+ "extra_id_83": "<extra_id_83>",
88
+ "extra_id_84": "<extra_id_84>",
89
+ "extra_id_85": "<extra_id_85>",
90
+ "extra_id_86": "<extra_id_86>",
91
+ "extra_id_87": "<extra_id_87>",
92
+ "extra_id_88": "<extra_id_88>",
93
+ "extra_id_89": "<extra_id_89>",
94
+ "extra_id_9": "<extra_id_9>",
95
+ "extra_id_90": "<extra_id_90>",
96
+ "extra_id_91": "<extra_id_91>",
97
+ "extra_id_92": "<extra_id_92>",
98
+ "extra_id_93": "<extra_id_93>",
99
+ "extra_id_94": "<extra_id_94>",
100
+ "extra_id_95": "<extra_id_95>",
101
+ "extra_id_96": "<extra_id_96>",
102
+ "extra_id_97": "<extra_id_97>",
103
+ "extra_id_98": "<extra_id_98>",
104
+ "extra_id_99": "<extra_id_99>",
105
+ "extra_ids": 100,
106
+ "extra_special_tokens": [
107
+ "<extra_id_0>",
108
+ "<extra_id_1>",
109
+ "<extra_id_2>",
110
+ "<extra_id_3>",
111
+ "<extra_id_4>",
112
+ "<extra_id_5>",
113
+ "<extra_id_6>",
114
+ "<extra_id_7>",
115
+ "<extra_id_8>",
116
+ "<extra_id_9>",
117
+ "<extra_id_10>",
118
+ "<extra_id_11>",
119
+ "<extra_id_12>",
120
+ "<extra_id_13>",
121
+ "<extra_id_14>",
122
+ "<extra_id_15>",
123
+ "<extra_id_16>",
124
+ "<extra_id_17>",
125
+ "<extra_id_18>",
126
+ "<extra_id_19>",
127
+ "<extra_id_20>",
128
+ "<extra_id_21>",
129
+ "<extra_id_22>",
130
+ "<extra_id_23>",
131
+ "<extra_id_24>",
132
+ "<extra_id_25>",
133
+ "<extra_id_26>",
134
+ "<extra_id_27>",
135
+ "<extra_id_28>",
136
+ "<extra_id_29>",
137
+ "<extra_id_30>",
138
+ "<extra_id_31>",
139
+ "<extra_id_32>",
140
+ "<extra_id_33>",
141
+ "<extra_id_34>",
142
+ "<extra_id_35>",
143
+ "<extra_id_36>",
144
+ "<extra_id_37>",
145
+ "<extra_id_38>",
146
+ "<extra_id_39>",
147
+ "<extra_id_40>",
148
+ "<extra_id_41>",
149
+ "<extra_id_42>",
150
+ "<extra_id_43>",
151
+ "<extra_id_44>",
152
+ "<extra_id_45>",
153
+ "<extra_id_46>",
154
+ "<extra_id_47>",
155
+ "<extra_id_48>",
156
+ "<extra_id_49>",
157
+ "<extra_id_50>",
158
+ "<extra_id_51>",
159
+ "<extra_id_52>",
160
+ "<extra_id_53>",
161
+ "<extra_id_54>",
162
+ "<extra_id_55>",
163
+ "<extra_id_56>",
164
+ "<extra_id_57>",
165
+ "<extra_id_58>",
166
+ "<extra_id_59>",
167
+ "<extra_id_60>",
168
+ "<extra_id_61>",
169
+ "<extra_id_62>",
170
+ "<extra_id_63>",
171
+ "<extra_id_64>",
172
+ "<extra_id_65>",
173
+ "<extra_id_66>",
174
+ "<extra_id_67>",
175
+ "<extra_id_68>",
176
+ "<extra_id_69>",
177
+ "<extra_id_70>",
178
+ "<extra_id_71>",
179
+ "<extra_id_72>",
180
+ "<extra_id_73>",
181
+ "<extra_id_74>",
182
+ "<extra_id_75>",
183
+ "<extra_id_76>",
184
+ "<extra_id_77>",
185
+ "<extra_id_78>",
186
+ "<extra_id_79>",
187
+ "<extra_id_80>",
188
+ "<extra_id_81>",
189
+ "<extra_id_82>",
190
+ "<extra_id_83>",
191
+ "<extra_id_84>",
192
+ "<extra_id_85>",
193
+ "<extra_id_86>",
194
+ "<extra_id_87>",
195
+ "<extra_id_88>",
196
+ "<extra_id_89>",
197
+ "<extra_id_90>",
198
+ "<extra_id_91>",
199
+ "<extra_id_92>",
200
+ "<extra_id_93>",
201
+ "<extra_id_94>",
202
+ "<extra_id_95>",
203
+ "<extra_id_96>",
204
+ "<extra_id_97>",
205
+ "<extra_id_98>",
206
+ "<extra_id_99>"
207
+ ],
208
+ "is_local": false,
209
+ "model_max_length": 512,
210
+ "model_specific_special_tokens": {
211
+ "extra_id_0": "<extra_id_0>",
212
+ "extra_id_1": "<extra_id_1>",
213
+ "extra_id_10": "<extra_id_10>",
214
+ "extra_id_11": "<extra_id_11>",
215
+ "extra_id_12": "<extra_id_12>",
216
+ "extra_id_13": "<extra_id_13>",
217
+ "extra_id_14": "<extra_id_14>",
218
+ "extra_id_15": "<extra_id_15>",
219
+ "extra_id_16": "<extra_id_16>",
220
+ "extra_id_17": "<extra_id_17>",
221
+ "extra_id_18": "<extra_id_18>",
222
+ "extra_id_19": "<extra_id_19>",
223
+ "extra_id_2": "<extra_id_2>",
224
+ "extra_id_20": "<extra_id_20>",
225
+ "extra_id_21": "<extra_id_21>",
226
+ "extra_id_22": "<extra_id_22>",
227
+ "extra_id_23": "<extra_id_23>",
228
+ "extra_id_24": "<extra_id_24>",
229
+ "extra_id_25": "<extra_id_25>",
230
+ "extra_id_26": "<extra_id_26>",
231
+ "extra_id_27": "<extra_id_27>",
232
+ "extra_id_28": "<extra_id_28>",
233
+ "extra_id_29": "<extra_id_29>",
234
+ "extra_id_3": "<extra_id_3>",
235
+ "extra_id_30": "<extra_id_30>",
236
+ "extra_id_31": "<extra_id_31>",
237
+ "extra_id_32": "<extra_id_32>",
238
+ "extra_id_33": "<extra_id_33>",
239
+ "extra_id_34": "<extra_id_34>",
240
+ "extra_id_35": "<extra_id_35>",
241
+ "extra_id_36": "<extra_id_36>",
242
+ "extra_id_37": "<extra_id_37>",
243
+ "extra_id_38": "<extra_id_38>",
244
+ "extra_id_39": "<extra_id_39>",
245
+ "extra_id_4": "<extra_id_4>",
246
+ "extra_id_40": "<extra_id_40>",
247
+ "extra_id_41": "<extra_id_41>",
248
+ "extra_id_42": "<extra_id_42>",
249
+ "extra_id_43": "<extra_id_43>",
250
+ "extra_id_44": "<extra_id_44>",
251
+ "extra_id_45": "<extra_id_45>",
252
+ "extra_id_46": "<extra_id_46>",
253
+ "extra_id_47": "<extra_id_47>",
254
+ "extra_id_48": "<extra_id_48>",
255
+ "extra_id_49": "<extra_id_49>",
256
+ "extra_id_5": "<extra_id_5>",
257
+ "extra_id_50": "<extra_id_50>",
258
+ "extra_id_51": "<extra_id_51>",
259
+ "extra_id_52": "<extra_id_52>",
260
+ "extra_id_53": "<extra_id_53>",
261
+ "extra_id_54": "<extra_id_54>",
262
+ "extra_id_55": "<extra_id_55>",
263
+ "extra_id_56": "<extra_id_56>",
264
+ "extra_id_57": "<extra_id_57>",
265
+ "extra_id_58": "<extra_id_58>",
266
+ "extra_id_59": "<extra_id_59>",
267
+ "extra_id_6": "<extra_id_6>",
268
+ "extra_id_60": "<extra_id_60>",
269
+ "extra_id_61": "<extra_id_61>",
270
+ "extra_id_62": "<extra_id_62>",
271
+ "extra_id_63": "<extra_id_63>",
272
+ "extra_id_64": "<extra_id_64>",
273
+ "extra_id_65": "<extra_id_65>",
274
+ "extra_id_66": "<extra_id_66>",
275
+ "extra_id_67": "<extra_id_67>",
276
+ "extra_id_68": "<extra_id_68>",
277
+ "extra_id_69": "<extra_id_69>",
278
+ "extra_id_7": "<extra_id_7>",
279
+ "extra_id_70": "<extra_id_70>",
280
+ "extra_id_71": "<extra_id_71>",
281
+ "extra_id_72": "<extra_id_72>",
282
+ "extra_id_73": "<extra_id_73>",
283
+ "extra_id_74": "<extra_id_74>",
284
+ "extra_id_75": "<extra_id_75>",
285
+ "extra_id_76": "<extra_id_76>",
286
+ "extra_id_77": "<extra_id_77>",
287
+ "extra_id_78": "<extra_id_78>",
288
+ "extra_id_79": "<extra_id_79>",
289
+ "extra_id_8": "<extra_id_8>",
290
+ "extra_id_80": "<extra_id_80>",
291
+ "extra_id_81": "<extra_id_81>",
292
+ "extra_id_82": "<extra_id_82>",
293
+ "extra_id_83": "<extra_id_83>",
294
+ "extra_id_84": "<extra_id_84>",
295
+ "extra_id_85": "<extra_id_85>",
296
+ "extra_id_86": "<extra_id_86>",
297
+ "extra_id_87": "<extra_id_87>",
298
+ "extra_id_88": "<extra_id_88>",
299
+ "extra_id_89": "<extra_id_89>",
300
+ "extra_id_9": "<extra_id_9>",
301
+ "extra_id_90": "<extra_id_90>",
302
+ "extra_id_91": "<extra_id_91>",
303
+ "extra_id_92": "<extra_id_92>",
304
+ "extra_id_93": "<extra_id_93>",
305
+ "extra_id_94": "<extra_id_94>",
306
+ "extra_id_95": "<extra_id_95>",
307
+ "extra_id_96": "<extra_id_96>",
308
+ "extra_id_97": "<extra_id_97>",
309
+ "extra_id_98": "<extra_id_98>",
310
+ "extra_id_99": "<extra_id_99>"
311
+ },
312
+ "pad_token": "<pad>",
313
+ "tokenizer_class": "T5Tokenizer",
314
+ "unk_token": "<unk>"
315
+ }
text_encoder/config.json ADDED
@@ -0,0 +1,63 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "architectures": [
3
+ "Qwen3Model"
4
+ ],
5
+ "attention_bias": false,
6
+ "attention_dropout": 0.0,
7
+ "bos_token_id": null,
8
+ "dtype": "bfloat16",
9
+ "eos_token_id": null,
10
+ "head_dim": 128,
11
+ "hidden_act": "silu",
12
+ "hidden_size": 1024,
13
+ "initializer_range": 0.02,
14
+ "intermediate_size": 3072,
15
+ "layer_types": [
16
+ "full_attention",
17
+ "full_attention",
18
+ "full_attention",
19
+ "full_attention",
20
+ "full_attention",
21
+ "full_attention",
22
+ "full_attention",
23
+ "full_attention",
24
+ "full_attention",
25
+ "full_attention",
26
+ "full_attention",
27
+ "full_attention",
28
+ "full_attention",
29
+ "full_attention",
30
+ "full_attention",
31
+ "full_attention",
32
+ "full_attention",
33
+ "full_attention",
34
+ "full_attention",
35
+ "full_attention",
36
+ "full_attention",
37
+ "full_attention",
38
+ "full_attention",
39
+ "full_attention",
40
+ "full_attention",
41
+ "full_attention",
42
+ "full_attention",
43
+ "full_attention"
44
+ ],
45
+ "max_position_embeddings": 32768,
46
+ "max_window_layers": 28,
47
+ "model_type": "qwen3",
48
+ "num_attention_heads": 16,
49
+ "num_hidden_layers": 28,
50
+ "num_key_value_heads": 8,
51
+ "pad_token_id": null,
52
+ "rms_norm_eps": 1e-06,
53
+ "rope_parameters": {
54
+ "rope_theta": 1000000.0,
55
+ "rope_type": "default"
56
+ },
57
+ "sliding_window": null,
58
+ "tie_word_embeddings": false,
59
+ "transformers_version": "5.6.0.dev0",
60
+ "use_cache": false,
61
+ "use_sliding_window": false,
62
+ "vocab_size": 151936
63
+ }
text_encoder/model.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:75e8db94adb374b61baca3d898fc2c0d5086d20689a96d4e74576546433a0526
3
+ size 1192133232
tokenizer/chat_template.jinja ADDED
@@ -0,0 +1,89 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {%- if tools %}
2
+ {{- '<|im_start|>system\n' }}
3
+ {%- if messages[0].role == 'system' %}
4
+ {{- messages[0].content + '\n\n' }}
5
+ {%- endif %}
6
+ {{- "# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>" }}
7
+ {%- for tool in tools %}
8
+ {{- "\n" }}
9
+ {{- tool | tojson }}
10
+ {%- endfor %}
11
+ {{- "\n</tools>\n\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\n<tool_call>\n{\"name\": <function-name>, \"arguments\": <args-json-object>}\n</tool_call><|im_end|>\n" }}
12
+ {%- else %}
13
+ {%- if messages[0].role == 'system' %}
14
+ {{- '<|im_start|>system\n' + messages[0].content + '<|im_end|>\n' }}
15
+ {%- endif %}
16
+ {%- endif %}
17
+ {%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}
18
+ {%- for message in messages[::-1] %}
19
+ {%- set index = (messages|length - 1) - loop.index0 %}
20
+ {%- if ns.multi_step_tool and message.role == "user" and message.content is string and not(message.content.startswith('<tool_response>') and message.content.endswith('</tool_response>')) %}
21
+ {%- set ns.multi_step_tool = false %}
22
+ {%- set ns.last_query_index = index %}
23
+ {%- endif %}
24
+ {%- endfor %}
25
+ {%- for message in messages %}
26
+ {%- if message.content is string %}
27
+ {%- set content = message.content %}
28
+ {%- else %}
29
+ {%- set content = '' %}
30
+ {%- endif %}
31
+ {%- if (message.role == "user") or (message.role == "system" and not loop.first) %}
32
+ {{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }}
33
+ {%- elif message.role == "assistant" %}
34
+ {%- set reasoning_content = '' %}
35
+ {%- if message.reasoning_content is string %}
36
+ {%- set reasoning_content = message.reasoning_content %}
37
+ {%- else %}
38
+ {%- if '</think>' in content %}
39
+ {%- set reasoning_content = content.split('</think>')[0].rstrip('\n').split('<think>')[-1].lstrip('\n') %}
40
+ {%- set content = content.split('</think>')[-1].lstrip('\n') %}
41
+ {%- endif %}
42
+ {%- endif %}
43
+ {%- if loop.index0 > ns.last_query_index %}
44
+ {%- if loop.last or (not loop.last and reasoning_content) %}
45
+ {{- '<|im_start|>' + message.role + '\n<think>\n' + reasoning_content.strip('\n') + '\n</think>\n\n' + content.lstrip('\n') }}
46
+ {%- else %}
47
+ {{- '<|im_start|>' + message.role + '\n' + content }}
48
+ {%- endif %}
49
+ {%- else %}
50
+ {{- '<|im_start|>' + message.role + '\n' + content }}
51
+ {%- endif %}
52
+ {%- if message.tool_calls %}
53
+ {%- for tool_call in message.tool_calls %}
54
+ {%- if (loop.first and content) or (not loop.first) %}
55
+ {{- '\n' }}
56
+ {%- endif %}
57
+ {%- if tool_call.function %}
58
+ {%- set tool_call = tool_call.function %}
59
+ {%- endif %}
60
+ {{- '<tool_call>\n{"name": "' }}
61
+ {{- tool_call.name }}
62
+ {{- '", "arguments": ' }}
63
+ {%- if tool_call.arguments is string %}
64
+ {{- tool_call.arguments }}
65
+ {%- else %}
66
+ {{- tool_call.arguments | tojson }}
67
+ {%- endif %}
68
+ {{- '}\n</tool_call>' }}
69
+ {%- endfor %}
70
+ {%- endif %}
71
+ {{- '<|im_end|>\n' }}
72
+ {%- elif message.role == "tool" %}
73
+ {%- if loop.first or (messages[loop.index0 - 1].role != "tool") %}
74
+ {{- '<|im_start|>user' }}
75
+ {%- endif %}
76
+ {{- '\n<tool_response>\n' }}
77
+ {{- content }}
78
+ {{- '\n</tool_response>' }}
79
+ {%- if loop.last or (messages[loop.index0 + 1].role != "tool") %}
80
+ {{- '<|im_end|>\n' }}
81
+ {%- endif %}
82
+ {%- endif %}
83
+ {%- endfor %}
84
+ {%- if add_generation_prompt %}
85
+ {{- '<|im_start|>assistant\n' }}
86
+ {%- if enable_thinking is defined and enable_thinking is false %}
87
+ {{- '<think>\n\n</think>\n\n' }}
88
+ {%- endif %}
89
+ {%- endif %}
tokenizer/tokenizer.json ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:508635c756562e715417f4480d287917c332083cee37d2afca7f9b623fce55d9
3
+ size 11422916
tokenizer/tokenizer_config.json ADDED
@@ -0,0 +1,42 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "add_prefix_space": false,
3
+ "backend": "tokenizers",
4
+ "bos_token": null,
5
+ "box_end": "<|box_end|>",
6
+ "box_start": "<|box_start|>",
7
+ "clean_up_tokenization_spaces": false,
8
+ "eos_token": "<|im_end|>",
9
+ "errors": "replace",
10
+ "im_end": "<|im_end|>",
11
+ "im_start": "<|im_start|>",
12
+ "image_pad": "<|image_pad|>",
13
+ "is_local": false,
14
+ "model_max_length": 131072,
15
+ "model_specific_special_tokens": {
16
+ "box_end": "<|box_end|>",
17
+ "box_start": "<|box_start|>",
18
+ "im_end": "<|im_end|>",
19
+ "im_start": "<|im_start|>",
20
+ "image_pad": "<|image_pad|>",
21
+ "object_ref_end": "<|object_ref_end|>",
22
+ "object_ref_start": "<|object_ref_start|>",
23
+ "quad_end": "<|quad_end|>",
24
+ "quad_start": "<|quad_start|>",
25
+ "video_pad": "<|video_pad|>",
26
+ "vision_end": "<|vision_end|>",
27
+ "vision_pad": "<|vision_pad|>",
28
+ "vision_start": "<|vision_start|>"
29
+ },
30
+ "object_ref_end": "<|object_ref_end|>",
31
+ "object_ref_start": "<|object_ref_start|>",
32
+ "pad_token": "<|endoftext|>",
33
+ "quad_end": "<|quad_end|>",
34
+ "quad_start": "<|quad_start|>",
35
+ "split_special_tokens": false,
36
+ "tokenizer_class": "Qwen2Tokenizer",
37
+ "unk_token": null,
38
+ "video_pad": "<|video_pad|>",
39
+ "vision_end": "<|vision_end|>",
40
+ "vision_pad": "<|vision_pad|>",
41
+ "vision_start": "<|vision_start|>"
42
+ }
transformer/config.json ADDED
@@ -0,0 +1,36 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "_class_name": "CosmosTransformer3DModel",
3
+ "_diffusers_version": "0.39.0.dev0",
4
+ "adaln_lora_dim": 256,
5
+ "attention_head_dim": 128,
6
+ "concat_padding_mask": true,
7
+ "controlnet_block_every_n": null,
8
+ "crossattn_proj_in_channels": 1024,
9
+ "encoder_hidden_states_channels": 1024,
10
+ "extra_pos_embed_type": null,
11
+ "img_context_dim_in": null,
12
+ "img_context_dim_out": 2048,
13
+ "img_context_num_tokens": 256,
14
+ "in_channels": 16,
15
+ "max_size": [
16
+ 128,
17
+ 240,
18
+ 240
19
+ ],
20
+ "mlp_ratio": 4.0,
21
+ "num_attention_heads": 16,
22
+ "num_layers": 28,
23
+ "out_channels": 16,
24
+ "patch_size": [
25
+ 1,
26
+ 2,
27
+ 2
28
+ ],
29
+ "rope_scale": [
30
+ 1.0,
31
+ 4.0,
32
+ 4.0
33
+ ],
34
+ "text_embed_dim": 1024,
35
+ "use_crossattn_projection": false
36
+ }
transformer/diffusion_pytorch_model.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:7995c5ff0c4d649d558798e1afe5ab2edc77eb251b61d533ce6af6024ba81052
3
+ size 3912877104
vae/config.json ADDED
@@ -0,0 +1,64 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "_class_name": "AutoencoderKLWan",
3
+ "_diffusers_version": "0.39.0.dev0",
4
+ "_name_or_path": "/mnt/models/Diffusers/models--CalamitousFelicitousness--Anima-Preview-3-sdnext-diffusers/snapshots/1352a0ff5e208e66d4b1b5bf13951cf0add453ee/vae",
5
+ "attn_scales": [],
6
+ "base_dim": 96,
7
+ "decoder_base_dim": null,
8
+ "dim_mult": [
9
+ 1,
10
+ 2,
11
+ 4,
12
+ 4
13
+ ],
14
+ "dropout": 0.0,
15
+ "in_channels": 3,
16
+ "is_residual": false,
17
+ "latents_mean": [
18
+ -0.7571,
19
+ -0.7089,
20
+ -0.9113,
21
+ 0.1075,
22
+ -0.1745,
23
+ 0.9653,
24
+ -0.1517,
25
+ 1.5508,
26
+ 0.4134,
27
+ -0.0715,
28
+ 0.5517,
29
+ -0.3632,
30
+ -0.1922,
31
+ -0.9497,
32
+ 0.2503,
33
+ -0.2921
34
+ ],
35
+ "latents_std": [
36
+ 2.8184,
37
+ 1.4541,
38
+ 2.3275,
39
+ 2.6558,
40
+ 1.2196,
41
+ 1.7708,
42
+ 2.6052,
43
+ 2.0743,
44
+ 3.2687,
45
+ 2.1526,
46
+ 2.8652,
47
+ 1.5579,
48
+ 1.6382,
49
+ 1.1253,
50
+ 2.8251,
51
+ 1.916
52
+ ],
53
+ "num_res_blocks": 2,
54
+ "out_channels": 3,
55
+ "patch_size": null,
56
+ "scale_factor_spatial": 8,
57
+ "scale_factor_temporal": 4,
58
+ "temperal_downsample": [
59
+ false,
60
+ true,
61
+ true
62
+ ],
63
+ "z_dim": 16
64
+ }
vae/diffusion_pytorch_model.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:0c8bc8b758c649abef9ea407b95408389a3b2f610d0d10fcb054fe171d0a8344
3
+ size 253806966