deepgenteam/DeepGen-1.0-diffusers
874
1"""2DeepGen Diffusers Pipeline - Standalone pipeline for DeepGen-1.0.3 4This file is self-contained and does not require the DeepGen repository.5It can be used with `trust_remote_code=True` when loading from HuggingFace Hub.6 7Usage:8 import torch9 from diffusers import DiffusionPipeline10 pipe = DiffusionPipeline.from_pretrained(11 "deepgenteam/DeepGen-1.0-diffusers",12 torch_dtype=torch.bfloat16,13 trust_remote_code=True,14 )15 pipe.to("cuda")16 17 # Text-to-Image18 image = pipe("a racoon holding a shiny red apple", height=512, width=512).images[0]19 20 # Image Edit21 from PIL import Image22 image = pipe("Place this guitar on a sandy beach.",23 image=Image.open("guitar.png"), height=512, width=512).images[0]24"""25 26import inspect27import math28import os29import json30import warnings31from functools import partial32from typing import Any, Callable, Dict, List, Optional, Tuple, Union33 34import numpy as np35import torch36import torch.nn as nn37import torch.nn.functional as F38import torch.utils.checkpoint39from torch.nn.init import _calculate_fan_in_and_fan_out40from torch.nn.utils.rnn import pad_sequence41 42from einops import rearrange43from PIL import Image44from safetensors.torch import load_file45 46from diffusers import AutoencoderKL, FlowMatchEulerDiscreteScheduler47from diffusers.configuration_utils import ConfigMixin, register_to_config48from diffusers.image_processor import PipelineImageInput, VaeImageProcessor49from diffusers.loaders import (50 FromOriginalModelMixin,51 FromSingleFileMixin,52 PeftAdapterMixin,53 SD3IPAdapterMixin,54 SD3LoraLoaderMixin,55 SD3Transformer2DLoadersMixin,56)57from diffusers.models.attention import FeedForward, JointTransformerBlock, _chunked_feed_forward58from diffusers.models.attention_processor import (59 Attention,60 AttentionProcessor,61 FusedJointAttnProcessor2_0,62 JointAttnProcessor2_0,63)64from diffusers.models.embeddings import CombinedTimestepTextProjEmbeddings, PatchEmbed65from diffusers.models.modeling_outputs import Transformer2DModelOutput66from diffusers.models.modeling_utils import ModelMixin67from diffusers.models.normalization import AdaLayerNormContinuous, AdaLayerNormZero68from diffusers.pipelines.pipeline_utils import DiffusionPipeline69from diffusers.pipelines.stable_diffusion_3.pipeline_output import StableDiffusion3PipelineOutput70from diffusers.utils import (71 USE_PEFT_BACKEND,72 is_torch_xla_available,73 logging,74 scale_lora_layers,75 unscale_lora_layers,76)77from diffusers.utils.torch_utils import maybe_allow_in_graph, randn_tensor78 79from transformers import (80 AutoTokenizer,81 CLIPTextModelWithProjection,82 CLIPTokenizer,83 Qwen2_5_VLForConditionalGeneration,84 SiglipImageProcessor,85 SiglipVisionModel,86 T5EncoderModel,87 T5TokenizerFast,88)89from transformers.activations import ACT2FN90from transformers.configuration_utils import PretrainedConfig91from transformers.utils import (92 is_flash_attn_2_available,93 is_flash_attn_greater_or_equal_2_10,94)95 96if is_flash_attn_2_available():97 from transformers.modeling_flash_attention_utils import _flash_attention_forward98 99if is_torch_xla_available():100 import torch_xla.core.xla_model as xm101 XLA_AVAILABLE = True102else:103 XLA_AVAILABLE = False104 105 106logger = logging.get_logger(__name__)107 108IMAGE_MEAN = (0.48145466, 0.4578275, 0.40821073)109IMAGE_STD = (0.26862954, 0.26130258, 0.27577711)110 111 112# =============================================================================113# Connector: Config + Attention + MLP + Encoder114# =============================================================================115 116class ConnectorConfig(PretrainedConfig):117 def __init__(118 self,119 hidden_size=768,120 intermediate_size=3072,121 num_hidden_layers=12,122 num_attention_heads=12,123 hidden_act="gelu_pytorch_tanh",124 layer_norm_eps=1e-6,125 attention_dropout=0.0,126 **kwargs,127 ):128 super().__init__(**kwargs)129 self.hidden_size = hidden_size130 self.intermediate_size = intermediate_size131 self.num_hidden_layers = num_hidden_layers132 self.num_attention_heads = num_attention_heads133 self.attention_dropout = attention_dropout134 self.layer_norm_eps = layer_norm_eps135 self.hidden_act = hidden_act136 137 138def _trunc_normal_(tensor, mean, std, a, b):139 def norm_cdf(x):140 return (1.0 + math.erf(x / math.sqrt(2.0))) / 2.0141 if (mean < a - 2 * std) or (mean > b + 2 * std):142 warnings.warn(143 "mean is more than 2 std from [a, b] in nn.init.trunc_normal_. "144 "The distribution of values may be incorrect.", stacklevel=2)145 l = norm_cdf((a - mean) / std)146 u = norm_cdf((b - mean) / std)147 tensor.uniform_(2 * l - 1, 2 * u - 1)148 tensor.erfinv_()149 tensor.mul_(std * math.sqrt(2.0))150 tensor.add_(mean)151 tensor.clamp_(min=a, max=b)152 153 154def trunc_normal_tf_(tensor, mean=0.0, std=1.0, a=-2.0, b=2.0):155 with torch.no_grad():156 _trunc_normal_(tensor, 0, 1.0, a, b)157 tensor.mul_(std).add_(mean)158 159 160def variance_scaling_(tensor, scale=1.0, mode="fan_in", distribution="normal"):161 fan_in, fan_out = _calculate_fan_in_and_fan_out(tensor)162 denom = {"fan_in": fan_in, "fan_out": fan_out, "fan_avg": (fan_in + fan_out) / 2}[mode]163 variance = scale / denom164 if distribution == "truncated_normal":165 trunc_normal_tf_(tensor, std=math.sqrt(variance) / 0.87962566103423978)166 elif distribution == "normal":167 with torch.no_grad():168 tensor.normal_(std=math.sqrt(variance))169 elif distribution == "uniform":170 bound = math.sqrt(3 * variance)171 with torch.no_grad():172 tensor.uniform_(-bound, bound)173 174 175def lecun_normal_(tensor):176 variance_scaling_(tensor, mode="fan_in", distribution="truncated_normal")177 178 179def default_flax_embed_init(tensor):180 variance_scaling_(tensor, mode="fan_in", distribution="normal")181 182 183class ConnectorAttention(nn.Module):184 def __init__(self, config):185 super().__init__()186 self.config = config187 self.embed_dim = config.hidden_size188 self.num_heads = config.num_attention_heads189 self.head_dim = self.embed_dim // self.num_heads190 if self.head_dim * self.num_heads != self.embed_dim:191 raise ValueError(192 f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} "193 f"and `num_heads`: {self.num_heads}).")194 self.scale = self.head_dim ** -0.5195 self.dropout = config.attention_dropout196 self.k_proj = nn.Linear(self.embed_dim, self.embed_dim)197 self.v_proj = nn.Linear(self.embed_dim, self.embed_dim)198 self.q_proj = nn.Linear(self.embed_dim, self.embed_dim)199 self.out_proj = nn.Linear(self.embed_dim, self.embed_dim)200 201 def forward(self, hidden_states, attention_mask=None, output_attentions=False):202 batch_size, q_len, _ = hidden_states.size()203 query_states = self.q_proj(hidden_states).view(batch_size, q_len, self.num_heads, self.head_dim).transpose(1, 2)204 key_states = self.k_proj(hidden_states).view(batch_size, q_len, self.num_heads, self.head_dim).transpose(1, 2)205 value_states = self.v_proj(hidden_states).view(batch_size, q_len, self.num_heads, self.head_dim).transpose(1, 2)206 207 k_v_seq_len = key_states.shape[-2]208 attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) * self.scale209 if attention_mask is not None:210 attn_weights = attn_weights + attention_mask211 attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)212 attn_weights = nn.functional.dropout(attn_weights, p=self.dropout, training=self.training)213 attn_output = torch.matmul(attn_weights, value_states)214 attn_output = attn_output.transpose(1, 2).contiguous().reshape(batch_size, q_len, self.embed_dim)215 attn_output = self.out_proj(attn_output)216 return attn_output, attn_weights217 218 219class ConnectorFlashAttention2(ConnectorAttention):220 is_causal = False221 222 def __init__(self, *args, **kwargs):223 super().__init__(*args, **kwargs)224 self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10()225 226 def forward(self, hidden_states, attention_mask=None, output_attentions=False):227 output_attentions = False228 batch_size, q_len, _ = hidden_states.size()229 query_states = self.q_proj(hidden_states).view(batch_size, q_len, self.num_heads, self.head_dim).transpose(1, 2)230 key_states = self.k_proj(hidden_states).view(batch_size, q_len, self.num_heads, self.head_dim).transpose(1, 2)231 value_states = self.v_proj(hidden_states).view(batch_size, q_len, self.num_heads, self.head_dim).transpose(1, 2)232 query_states = query_states.transpose(1, 2)233 key_states = key_states.transpose(1, 2)234 value_states = value_states.transpose(1, 2)235 dropout_rate = self.dropout if self.training else 0.0236 input_dtype = query_states.dtype237 if input_dtype == torch.float32:238 if torch.is_autocast_enabled():239 target_dtype = torch.get_autocast_gpu_dtype()240 elif hasattr(self.config, "_pre_quantization_dtype"):241 target_dtype = self.config._pre_quantization_dtype242 else:243 target_dtype = self.q_proj.weight.dtype244 query_states = query_states.to(target_dtype)245 key_states = key_states.to(target_dtype)246 value_states = value_states.to(target_dtype)247 attn_output = _flash_attention_forward(248 query_states, key_states, value_states, attention_mask, q_len,249 dropout=dropout_rate, is_causal=self.is_causal,250 use_top_left_mask=self._flash_attn_uses_top_left_mask)251 attn_output = attn_output.reshape(batch_size, q_len, self.embed_dim).contiguous()252 attn_output = self.out_proj(attn_output)253 return attn_output, None254 255 256class ConnectorSdpaAttention(ConnectorAttention):257 is_causal = False258 259 def forward(self, hidden_states, attention_mask=None, output_attentions=False):260 if output_attentions:261 return super().forward(hidden_states=hidden_states, attention_mask=attention_mask, output_attentions=output_attentions)262 batch_size, q_len, _ = hidden_states.size()263 query_states = self.q_proj(hidden_states).view(batch_size, q_len, self.num_heads, self.head_dim).transpose(1, 2)264 key_states = self.k_proj(hidden_states).view(batch_size, q_len, self.num_heads, self.head_dim).transpose(1, 2)265 value_states = self.v_proj(hidden_states).view(batch_size, q_len, self.num_heads, self.head_dim).transpose(1, 2)266 if query_states.device.type == "cuda" and attention_mask is not None:267 query_states = query_states.contiguous()268 key_states = key_states.contiguous()269 value_states = value_states.contiguous()270 is_causal = True if self.is_causal and q_len > 1 else False271 attn_output = torch.nn.functional.scaled_dot_product_attention(272 query_states, key_states, value_states, attn_mask=attention_mask,273 dropout_p=self.dropout if self.training else 0.0, is_causal=is_causal)274 attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, q_len, self.embed_dim)275 attn_output = self.out_proj(attn_output)276 return attn_output, None277 278 279CONNECTOR_ATTENTION_CLASSES = {280 "eager": ConnectorAttention,281 "flash_attention_2": ConnectorFlashAttention2,282 "sdpa": ConnectorSdpaAttention,283}284 285 286class ConnectorMLP(nn.Module):287 def __init__(self, config):288 super().__init__()289 self.config = config290 self.activation_fn = ACT2FN[config.hidden_act]291 self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size)292 self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size)293 294 def forward(self, hidden_states):295 hidden_states = self.fc1(hidden_states)296 hidden_states = self.activation_fn(hidden_states)297 hidden_states = self.fc2(hidden_states)298 return hidden_states299 300 301def _init_connector_weights(module):302 if isinstance(module, nn.Embedding):303 default_flax_embed_init(module.weight)304 elif isinstance(module, ConnectorAttention):305 nn.init.xavier_uniform_(module.q_proj.weight)306 nn.init.xavier_uniform_(module.k_proj.weight)307 nn.init.xavier_uniform_(module.v_proj.weight)308 nn.init.xavier_uniform_(module.out_proj.weight)309 nn.init.zeros_(module.q_proj.bias)310 nn.init.zeros_(module.k_proj.bias)311 nn.init.zeros_(module.v_proj.bias)312 nn.init.zeros_(module.out_proj.bias)313 elif isinstance(module, ConnectorMLP):314 nn.init.xavier_uniform_(module.fc1.weight)315 nn.init.xavier_uniform_(module.fc2.weight)316 nn.init.normal_(module.fc1.bias, std=1e-6)317 nn.init.normal_(module.fc2.bias, std=1e-6)318 elif isinstance(module, (nn.Linear, nn.Conv2d)):319 lecun_normal_(module.weight)320 if module.bias is not None:321 nn.init.zeros_(module.bias)322 elif isinstance(module, nn.LayerNorm):323 module.bias.data.zero_()324 module.weight.data.fill_(1.0)325 326 327class ConnectorEncoderLayer(nn.Module):328 def __init__(self, config):329 super().__init__()330 self.embed_dim = config.hidden_size331 self.self_attn = CONNECTOR_ATTENTION_CLASSES[config._attn_implementation](config=config)332 self.layer_norm1 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)333 self.mlp = ConnectorMLP(config)334 self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)335 336 def forward(self, hidden_states, attention_mask, output_attentions=False):337 residual = hidden_states338 hidden_states = self.layer_norm1(hidden_states)339 hidden_states, attn_weights = self.self_attn(340 hidden_states=hidden_states, attention_mask=attention_mask, output_attentions=output_attentions)341 hidden_states = residual + hidden_states342 residual = hidden_states343 hidden_states = self.layer_norm2(hidden_states)344 hidden_states = self.mlp(hidden_states)345 hidden_states = residual + hidden_states346 outputs = (hidden_states,)347 if output_attentions:348 outputs += (attn_weights,)349 return outputs350 351 352class ConnectorEncoder(nn.Module):353 def __init__(self, config):354 super().__init__()355 self.config = config356 self.layers = nn.ModuleList([ConnectorEncoderLayer(config) for _ in range(config.num_hidden_layers)])357 self.gradient_checkpointing = False358 self.apply(_init_connector_weights)359 360 def forward(self, inputs_embeds):361 hidden_states = inputs_embeds362 for encoder_layer in self.layers:363 if self.gradient_checkpointing and self.training:364 layer_outputs = torch.utils.checkpoint.checkpoint(365 encoder_layer.__call__, hidden_states, None, False, use_reentrant=False)366 else:367 layer_outputs = encoder_layer(hidden_states, None, output_attentions=False)368 hidden_states = layer_outputs[0]369 return hidden_states370 371 372class DeepGenConnector(nn.Module):373 """Connector module bridging VLM hidden states to DiT conditioning."""374 375 def __init__(self, connector_config, num_queries, llm_hidden_size,376 projector_1_in, projector_1_out,377 projector_2_in, projector_2_out,378 projector_3_in, projector_3_out):379 super().__init__()380 self.connector = ConnectorEncoder(ConnectorConfig(**connector_config))381 self.projector_1 = nn.Linear(projector_1_in, projector_1_out)382 self.projector_2 = nn.Linear(projector_2_in, projector_2_out)383 self.projector_3 = nn.Linear(projector_3_in, projector_3_out)384 self.meta_queries = nn.Parameter(torch.zeros(num_queries, llm_hidden_size))385 self.num_queries = num_queries386 387 def llm2dit(self, x):388 x = self.connector(self.projector_1(x))389 pooled_out = self.projector_2(x.mean(1))390 seq_out = self.projector_3(x)391 return pooled_out, seq_out392 393 394# =============================================================================395# Custom SD3 Transformer (dynamic resolution + attention mask)396# =============================================================================397 398class CustomJointAttnProcessor2_0:399 """Attention processor supporting attention masks for dynamic-resolution SD3."""400 401 def __init__(self):402 if not hasattr(F, "scaled_dot_product_attention"):403 raise ImportError("CustomJointAttnProcessor2_0 requires PyTorch 2.0+")404 405 def __call__(self, attn, hidden_states, encoder_hidden_states=None,406 attention_mask=None, *args, **kwargs):407 residual = hidden_states408 batch_size = hidden_states.shape[0]409 410 query = attn.to_q(hidden_states)411 key = attn.to_k(hidden_states)412 value = attn.to_v(hidden_states)413 414 inner_dim = key.shape[-1]415 head_dim = inner_dim // attn.heads416 417 query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)418 key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)419 value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)420 421 if attn.norm_q is not None:422 query = attn.norm_q(query)423 if attn.norm_k is not None:424 key = attn.norm_k(key)425 426 if encoder_hidden_states is not None:427 ctx_len = encoder_hidden_states.shape[1]428 encoder_hidden_states_query_proj = attn.add_q_proj(encoder_hidden_states).view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)429 encoder_hidden_states_key_proj = attn.add_k_proj(encoder_hidden_states).view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)430 encoder_hidden_states_value_proj = attn.add_v_proj(encoder_hidden_states).view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)431 432 if attn.norm_added_q is not None:433 encoder_hidden_states_query_proj = attn.norm_added_q(encoder_hidden_states_query_proj)434 if attn.norm_added_k is not None:435 encoder_hidden_states_key_proj = attn.norm_added_k(encoder_hidden_states_key_proj)436 437 query = torch.cat([query, encoder_hidden_states_query_proj], dim=2)438 key = torch.cat([key, encoder_hidden_states_key_proj], dim=2)439 value = torch.cat([value, encoder_hidden_states_value_proj], dim=2)440 441 if attention_mask is not None:442 encoder_attention_mask = torch.ones(443 batch_size, ctx_len, dtype=torch.bool, device=hidden_states.device)444 attention_mask = torch.cat([attention_mask, encoder_attention_mask], dim=1)445 446 if attention_mask is not None:447 attention_mask = attention_mask[:, None] * attention_mask[..., None]448 indices = range(attention_mask.shape[1])449 attention_mask[:, indices, indices] = True450 attention_mask = attention_mask[:, None]451 452 hidden_states = F.scaled_dot_product_attention(453 query, key, value, dropout_p=0.0, is_causal=False, attn_mask=attention_mask)454 hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)455 hidden_states = hidden_states.to(query.dtype)456 457 if encoder_hidden_states is not None:458 hidden_states, encoder_hidden_states = (459 hidden_states[:, :residual.shape[1]],460 hidden_states[:, residual.shape[1]:])461 if not attn.context_pre_only:462 encoder_hidden_states = attn.to_add_out(encoder_hidden_states)463 464 hidden_states = attn.to_out[0](hidden_states)465 hidden_states = attn.to_out[1](hidden_states)466 467 if encoder_hidden_states is not None:468 return hidden_states, encoder_hidden_states469 else:470 return hidden_states471 472 473class CustomJointTransformerBlock(JointTransformerBlock):474 def __init__(self, *args, **kwargs):475 super().__init__(*args, **kwargs)476 self.attn.set_processor(CustomJointAttnProcessor2_0())477 if self.attn2 is not None:478 self.attn2.set_processor(CustomJointAttnProcessor2_0())479 480 def forward(self, hidden_states, encoder_hidden_states, temb,481 attention_mask=None, joint_attention_kwargs=None):482 joint_attention_kwargs = joint_attention_kwargs or {}483 if self.use_dual_attention:484 norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp, norm_hidden_states2, gate_msa2 = self.norm1(hidden_states, emb=temb)485 else:486 norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(hidden_states, emb=temb)487 488 if self.context_pre_only:489 norm_encoder_hidden_states = self.norm1_context(encoder_hidden_states, temb)490 else:491 norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.norm1_context(encoder_hidden_states, emb=temb)492 493 attn_output, context_attn_output = self.attn(494 hidden_states=norm_hidden_states, attention_mask=attention_mask,495 encoder_hidden_states=norm_encoder_hidden_states, **joint_attention_kwargs)496 497 attn_output = gate_msa.unsqueeze(1) * attn_output498 hidden_states = hidden_states + attn_output499 500 if self.use_dual_attention:501 attn_output2 = self.attn2(hidden_states=norm_hidden_states2, attention_mask=attention_mask, **joint_attention_kwargs)502 attn_output2 = gate_msa2.unsqueeze(1) * attn_output2503 hidden_states = hidden_states + attn_output2504 505 norm_hidden_states = self.norm2(hidden_states)506 norm_hidden_states = norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]507 if self._chunk_size is not None:508 ff_output = _chunked_feed_forward(self.ff, norm_hidden_states, self._chunk_dim, self._chunk_size)509 else:510 ff_output = self.ff(norm_hidden_states)511 ff_output = gate_mlp.unsqueeze(1) * ff_output512 hidden_states = hidden_states + ff_output513 514 if self.context_pre_only:515 encoder_hidden_states = None516 else:517 context_attn_output = c_gate_msa.unsqueeze(1) * context_attn_output518 encoder_hidden_states = encoder_hidden_states + context_attn_output519 norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states)520 norm_encoder_hidden_states = norm_encoder_hidden_states * (1 + c_scale_mlp[:, None]) + c_shift_mlp[:, None]521 if self._chunk_size is not None:522 context_ff_output = _chunked_feed_forward(self.ff_context, norm_encoder_hidden_states, self._chunk_dim, self._chunk_size)523 else:524 context_ff_output = self.ff_context(norm_encoder_hidden_states)525 encoder_hidden_states = encoder_hidden_states + c_gate_mlp.unsqueeze(1) * context_ff_output526 527 return encoder_hidden_states, hidden_states528 529 530class SD3Transformer2DModel(531 ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin, SD3Transformer2DLoadersMixin532):533 _supports_gradient_checkpointing = True534 _no_split_modules = ["JointTransformerBlock", "CustomJointTransformerBlock"]535 _skip_layerwise_casting_patterns = ["pos_embed", "norm"]536 537 @register_to_config538 def __init__(539 self,540 sample_size: int = 128,541 patch_size: int = 2,542 in_channels: int = 16,543 num_layers: int = 18,544 attention_head_dim: int = 64,545 num_attention_heads: int = 18,546 joint_attention_dim: int = 4096,547 caption_projection_dim: int = 1152,548 pooled_projection_dim: int = 2048,549 out_channels: int = 16,550 pos_embed_max_size: int = 96,551 dual_attention_layers: Tuple[int, ...] = (),552 qk_norm: Optional[str] = None,553 ):554 super().__init__()555 self.out_channels = out_channels if out_channels is not None else in_channels556 self.inner_dim = num_attention_heads * attention_head_dim557 558 self.pos_embed = PatchEmbed(559 height=sample_size, width=sample_size, patch_size=patch_size,560 in_channels=in_channels, embed_dim=self.inner_dim,561 pos_embed_max_size=pos_embed_max_size)562 self.time_text_embed = CombinedTimestepTextProjEmbeddings(563 embedding_dim=self.inner_dim, pooled_projection_dim=pooled_projection_dim)564 self.context_embedder = nn.Linear(joint_attention_dim, caption_projection_dim)565 566 self.transformer_blocks = nn.ModuleList([567 CustomJointTransformerBlock(568 dim=self.inner_dim,569 num_attention_heads=num_attention_heads,570 attention_head_dim=attention_head_dim,571 context_pre_only=i == num_layers - 1,572 qk_norm=qk_norm,573 use_dual_attention=True if i in dual_attention_layers else False,574 ) for i in range(num_layers)575 ])576 577 self.norm_out = AdaLayerNormContinuous(self.inner_dim, self.inner_dim, elementwise_affine=False, eps=1e-6)578 self.proj_out = nn.Linear(self.inner_dim, patch_size * patch_size * self.out_channels, bias=True)579 self.gradient_checkpointing = False580 581 @property582 def attn_processors(self):583 processors = {}584 def fn_recursive_add_processors(name, module, processors):585 if hasattr(module, "get_processor"):586 processors[f"{name}.processor"] = module.get_processor()587 for sub_name, child in module.named_children():588 fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)589 return processors590 for name, module in self.named_children():591 fn_recursive_add_processors(name, module, processors)592 return processors593 594 def set_attn_processor(self, processor):595 count = len(self.attn_processors.keys())596 if isinstance(processor, dict) and len(processor) != count:597 raise ValueError(f"A dict of processors was passed, but the number of processors {len(processor)} does not match the number of attention layers: {count}.")598 def fn_recursive_attn_processor(name, module, processor):599 if hasattr(module, "set_processor"):600 if not isinstance(processor, dict):601 module.set_processor(processor)602 else:603 module.set_processor(processor.pop(f"{name}.processor"))604 for sub_name, child in module.named_children():605 fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)606 for name, module in self.named_children():607 fn_recursive_attn_processor(name, module, processor)608 609 def forward(610 self,611 hidden_states,612 encoder_hidden_states=None,613 cond_hidden_states=None,614 pooled_projections=None,615 timestep=None,616 block_controlnet_hidden_states=None,617 joint_attention_kwargs=None,618 return_dict=True,619 skip_layers=None,620 ):621 if joint_attention_kwargs is not None:622 joint_attention_kwargs = joint_attention_kwargs.copy()623 lora_scale = joint_attention_kwargs.pop("scale", 1.0)624 else:625 lora_scale = 1.0626 627 if USE_PEFT_BACKEND:628 scale_lora_layers(self, lora_scale)629 else:630 if joint_attention_kwargs is not None and joint_attention_kwargs.get("scale", None) is not None:631 logger.warning("Passing `scale` via `joint_attention_kwargs` when not using the PEFT backend is ineffective.")632 633 latent_sizes = [hs.shape[-2:] for hs in hidden_states]634 bsz = len(hidden_states)635 636 hidden_states_list = []637 for idx in range(bsz):638 hidden_states_per_sample = self.pos_embed(hidden_states[idx][None])[0]639 if cond_hidden_states is not None:640 for ref in cond_hidden_states[idx]:641 hidden_states_per_sample = torch.cat(642 [hidden_states_per_sample, self.pos_embed(ref[None])[0]])643 hidden_states_list.append(hidden_states_per_sample)644 645 max_len = max([len(hs) for hs in hidden_states_list])646 attention_mask = torch.zeros(bsz, max_len, dtype=torch.bool, device=self.device)647 for i, hs in enumerate(hidden_states_list):648 attention_mask[i, :len(hs)] = True649 650 hidden_states = pad_sequence(hidden_states_list, batch_first=True, padding_value=0.0, padding_side='right')651 652 temb = self.time_text_embed(timestep, pooled_projections)653 encoder_hidden_states = self.context_embedder(encoder_hidden_states)654 655 if joint_attention_kwargs is not None and "ip_adapter_image_embeds" in joint_attention_kwargs:656 ip_adapter_image_embeds = joint_attention_kwargs.pop("ip_adapter_image_embeds")657 ip_hidden_states, ip_temb = self.image_proj(ip_adapter_image_embeds, timestep)658 joint_attention_kwargs.update(ip_hidden_states=ip_hidden_states, temb=ip_temb)659 660 for index_block, block in enumerate(self.transformer_blocks):661 is_skip = True if skip_layers is not None and index_block in skip_layers else False662 if torch.is_grad_enabled() and self.gradient_checkpointing and not is_skip:663 encoder_hidden_states, hidden_states = self._gradient_checkpointing_func(664 block, hidden_states, encoder_hidden_states, temb, attention_mask, joint_attention_kwargs)665 elif not is_skip:666 encoder_hidden_states, hidden_states = block(667 hidden_states=hidden_states, encoder_hidden_states=encoder_hidden_states,668 temb=temb, attention_mask=attention_mask, joint_attention_kwargs=joint_attention_kwargs)669 670 if block_controlnet_hidden_states is not None and block.context_pre_only is False:671 interval_control = len(self.transformer_blocks) / len(block_controlnet_hidden_states)672 hidden_states = hidden_states + block_controlnet_hidden_states[int(index_block / interval_control)]673 674 hidden_states = self.norm_out(hidden_states, temb)675 hidden_states = self.proj_out(hidden_states)676 677 patch_size = self.config.patch_size678 latent_sizes = [(ls[0] // patch_size, ls[1] // patch_size) for ls in latent_sizes]679 680 output = [rearrange(hs[:math.prod(latent_size)], '(h w) (p q c) -> c (h p) (w q)',681 h=latent_size[0], w=latent_size[1], p=patch_size, q=patch_size)682 for hs, latent_size in zip(hidden_states, latent_sizes)]683 684 try:685 output = torch.stack(output)686 except:687 pass688 689 if USE_PEFT_BACKEND:690 unscale_lora_layers(self, lora_scale)691 692 if not return_dict:693 return (output,)694 return Transformer2DModelOutput(sample=output)695 696 697# =============================================================================698# Custom StableDiffusion3Pipeline (with cond_latents + dynamic shift)699# =============================================================================700 701def calculate_shift(image_seq_len, base_seq_len=256, max_seq_len=4096, base_shift=0.5, max_shift=1.15):702 m = (max_shift - base_shift) / (max_seq_len - base_seq_len)703 b = base_shift - m * base_seq_len704 mu = image_seq_len * m + b705 return mu706 707 708def retrieve_timesteps(scheduler, num_inference_steps=None, device=None, timesteps=None, sigmas=None, **kwargs):709 if timesteps is not None and sigmas is not None:710 raise ValueError("Only one of `timesteps` or `sigmas` can be passed.")711 if timesteps is not None:712 accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())713 if not accepts_timesteps:714 raise ValueError(f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom timestep schedules.")715 scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)716 timesteps = scheduler.timesteps717 num_inference_steps = len(timesteps)718 elif sigmas is not None:719 accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())720 if not accept_sigmas:721 raise ValueError(f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom sigmas schedules.")722 scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)723 timesteps = scheduler.timesteps724 num_inference_steps = len(timesteps)725 else:726 scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)727 timesteps = scheduler.timesteps728 return timesteps, num_inference_steps729 730 731class _SD3Pipeline(DiffusionPipeline, SD3LoraLoaderMixin, FromSingleFileMixin, SD3IPAdapterMixin):732 """Internal SD3 pipeline with cond_latents support."""733 734 model_cpu_offload_seq = "text_encoder->text_encoder_2->text_encoder_3->image_encoder->transformer->vae"735 _optional_components = ["image_encoder", "feature_extractor"]736 _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds", "negative_pooled_prompt_embeds"]737 738 def __init__(self, transformer, scheduler, vae, text_encoder, tokenizer,739 text_encoder_2, tokenizer_2, text_encoder_3, tokenizer_3,740 image_encoder=None, feature_extractor=None):741 super().__init__()742 self.register_modules(743 vae=vae, text_encoder=text_encoder, text_encoder_2=text_encoder_2,744 text_encoder_3=text_encoder_3, tokenizer=tokenizer, tokenizer_2=tokenizer_2,745 tokenizer_3=tokenizer_3, transformer=transformer, scheduler=scheduler,746 image_encoder=image_encoder, feature_extractor=feature_extractor)747 self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8748 self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)749 self.tokenizer_max_length = self.tokenizer.model_max_length if hasattr(self, "tokenizer") and self.tokenizer is not None else 77750 self.default_sample_size = self.transformer.config.sample_size if hasattr(self, "transformer") and self.transformer is not None else 128751 self.patch_size = self.transformer.config.patch_size if hasattr(self, "transformer") and self.transformer is not None else 2752 753 def check_inputs(self, prompt, prompt_2, prompt_3, height, width, negative_prompt=None,754 negative_prompt_2=None, negative_prompt_3=None, prompt_embeds=None,755 negative_prompt_embeds=None, pooled_prompt_embeds=None,756 negative_pooled_prompt_embeds=None, callback_on_step_end_tensor_inputs=None,757 max_sequence_length=None):758 if height % (self.vae_scale_factor * self.patch_size) != 0 or width % (self.vae_scale_factor * self.patch_size) != 0:759 raise ValueError(f"`height` and `width` have to be divisible by {self.vae_scale_factor * self.patch_size}.")760 if prompt_embeds is not None and pooled_prompt_embeds is None:761 raise ValueError("If `prompt_embeds` are provided, `pooled_prompt_embeds` also have to be passed.")762 if negative_prompt_embeds is not None and negative_pooled_prompt_embeds is None:763 raise ValueError("If `negative_prompt_embeds` are provided, `negative_pooled_prompt_embeds` also have to be passed.")764 765 def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype, device, generator, latents=None):766 if latents is not None:767 return latents.to(device=device, dtype=dtype)768 shape = (batch_size, num_channels_latents, int(height) // self.vae_scale_factor, int(width) // self.vae_scale_factor)769 if isinstance(generator, list) and len(generator) != batch_size:770 raise ValueError(f"You have passed a list of generators of length {len(generator)}, but requested an effective batch size of {batch_size}.")771 latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)772 return latents773 774 @property775 def guidance_scale(self):776 return self._guidance_scale777 778 @property779 def do_classifier_free_guidance(self):780 return self._guidance_scale > 1781 782 @property783 def joint_attention_kwargs(self):784 return self._joint_attention_kwargs785 786 @torch.no_grad()787 def __call__(788 self,789 prompt=None, prompt_2=None, prompt_3=None,790 height=None, width=None, num_inference_steps=28, sigmas=None,791 guidance_scale=7.0,792 negative_prompt=None, negative_prompt_2=None, negative_prompt_3=None,793 num_images_per_prompt=1, generator=None, latents=None,794 cond_latents=None,795 prompt_embeds=None, negative_prompt_embeds=None,796 pooled_prompt_embeds=None, negative_pooled_prompt_embeds=None,797 output_type="pil", return_dict=True,798 joint_attention_kwargs=None, callback_on_step_end=None,799 callback_on_step_end_tensor_inputs=["latents"],800 max_sequence_length=256, mu=None, **kwargs,801 ):802 height = height or self.default_sample_size * self.vae_scale_factor803 width = width or self.default_sample_size * self.vae_scale_factor804 805 self.check_inputs(prompt, prompt_2, prompt_3, height, width,806 negative_prompt=negative_prompt, prompt_embeds=prompt_embeds,807 negative_prompt_embeds=negative_prompt_embeds,808 pooled_prompt_embeds=pooled_prompt_embeds,809 negative_pooled_prompt_embeds=negative_pooled_prompt_embeds)810 811 self._guidance_scale = guidance_scale812 self._joint_attention_kwargs = joint_attention_kwargs813 self._interrupt = False814 815 if prompt is not None and isinstance(prompt, str):816 batch_size = 1817 elif prompt is not None and isinstance(prompt, list):818 batch_size = len(prompt)819 else:820 batch_size = prompt_embeds.shape[0]821 822 device = self._execution_device823 824 (prompt_embeds, negative_prompt_embeds, pooled_prompt_embeds, negative_pooled_prompt_embeds) = (825 prompt_embeds, negative_prompt_embeds, pooled_prompt_embeds, negative_pooled_prompt_embeds)826 827 if self.do_classifier_free_guidance:828 prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)829 pooled_prompt_embeds = torch.cat([negative_pooled_prompt_embeds, pooled_prompt_embeds], dim=0)830 831 num_channels_latents = self.transformer.config.in_channels832 latents = self.prepare_latents(833 batch_size * num_images_per_prompt, num_channels_latents, height, width,834 prompt_embeds.dtype, device, generator, latents)835 836 scheduler_kwargs = {}837 if self.scheduler.config.get("use_dynamic_shifting", None) and mu is None:838 _, _, h, w = latents.shape839 image_seq_len = (h // self.transformer.config.patch_size) * (w // self.transformer.config.patch_size)840 mu = calculate_shift(841 image_seq_len,842 self.scheduler.config.get("base_image_seq_len", 256),843 self.scheduler.config.get("max_image_seq_len", 4096),844 self.scheduler.config.get("base_shift", 0.5),845 self.scheduler.config.get("max_shift", 1.16))846 scheduler_kwargs["mu"] = mu847 elif mu is not None:848 scheduler_kwargs["mu"] = mu849 850 timesteps, num_inference_steps = retrieve_timesteps(851 self.scheduler, num_inference_steps, device, sigmas=sigmas, **scheduler_kwargs)852 num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)853 854 if cond_latents is not None and self.do_classifier_free_guidance:855 if len(cond_latents) == latents.shape[0]:856 cond_latents = cond_latents * 2857 858 with self.progress_bar(total=num_inference_steps) as progress_bar:859 for i, t in enumerate(timesteps):860 if self._interrupt:861 continue862 latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents863 timestep = t.expand(latent_model_input.shape[0])864 noise_pred = self.transformer(865 hidden_states=latent_model_input, cond_hidden_states=cond_latents,866 timestep=timestep, encoder_hidden_states=prompt_embeds,867 pooled_projections=pooled_prompt_embeds,868 joint_attention_kwargs=self.joint_attention_kwargs,869 return_dict=False)[0]870 871 if self.do_classifier_free_guidance:872 noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)873 noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond)874 875 latents_dtype = latents.dtype876 latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]877 if latents.dtype != latents_dtype:878 if torch.backends.mps.is_available():879 latents = latents.to(latents_dtype)880 881 if callback_on_step_end is not None:882 callback_kwargs = {}883 for k in callback_on_step_end_tensor_inputs:884 callback_kwargs[k] = locals()[k]885 callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)886 latents = callback_outputs.pop("latents", latents)887 888 if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):889 progress_bar.update()890 891 if XLA_AVAILABLE:892 xm.mark_step()893 894 if output_type == "latent":895 image = latents896 else:897 latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor898 image = self.vae.decode(latents, return_dict=False)[0]899 image = self.image_processor.postprocess(image, output_type=output_type)900 901 self.maybe_free_model_hooks()902 903 if not return_dict:904 return (image,)905 return StableDiffusion3PipelineOutput(images=image)906 907 908# =============================================================================909# DeepGen Pipeline (main entry point)910# =============================================================================911 912class DeepGenPipeline(DiffusionPipeline):913 """914 DeepGen 1.0 Pipeline for text-to-image generation and image editing.915 916 This pipeline integrates Qwen2.5-VL (VLM) + SCB Connector + SD3 DiT into a917 single interface. Standard diffusers components (transformer, vae, scheduler)918 are loaded by DiffusionPipeline; non-standard components (VLM, connector,919 tokenizer, prompt_template) are loaded automatically on first use.920 921 Usage:922 pipe = DiffusionPipeline.from_pretrained(923 "deepgenteam/DeepGen-1.0-diffusers",924 torch_dtype=torch.bfloat16,925 trust_remote_code=True,926 )927 pipe.to("cuda")928 result = pipe("a raccoon holding an apple", height=512, width=512)929 result.images[0].save("output.png")930 """931 932 _optional_components = []933 934 def __init__(935 self,936 transformer: SD3Transformer2DModel,937 vae: AutoencoderKL,938 scheduler: FlowMatchEulerDiscreteScheduler,939 ):940 super().__init__()941 self.register_modules(942 transformer=transformer,943 vae=vae,944 scheduler=scheduler,945 )946 self._upgrade_transformer()947 self._extras_loaded = False948 self._cpu_offload = False949 self._gpu_device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")950 self.lmm = None951 self.tokenizer = None952 self.connector_module = None953 self.prompt_template = None954 self.max_length = 1024955 self.image_token_id = None956 self.vit_mean = torch.tensor(IMAGE_MEAN)957 self.vit_std = torch.tensor(IMAGE_STD)958 959 def _upgrade_transformer(self):960 """Convert standard diffusers SD3Transformer2DModel to custom version961 with cond_latents support for image editing. No weight copying needed."""962 from diffusers.models.transformers.transformer_sd3 import SD3Transformer2DModel as _OrigSD3963 if isinstance(self.transformer, _OrigSD3) and not isinstance(self.transformer, SD3Transformer2DModel):964 self.transformer.__class__ = SD3Transformer2DModel965 for block in self.transformer.transformer_blocks:966 block.__class__ = CustomJointTransformerBlock967 block.attn.set_processor(CustomJointAttnProcessor2_0())968 if block.attn2 is not None:969 block.attn2.set_processor(CustomJointAttnProcessor2_0())970 971 def _resolve_pretrained_path(self):972 path = self.config._name_or_path973 if os.path.isdir(path):974 return path975 from huggingface_hub import snapshot_download976 return snapshot_download(repo_id=path)977 978 def _load_extras(self, vlm_model_path=None, attn_implementation="flash_attention_2"):979 """Load non-standard components (VLM, connector, tokenizer, prompt_template)."""980 if self._extras_loaded:981 return982 path = self._resolve_pretrained_path()983 dtype = next(self.transformer.parameters()).dtype984 985 model_index_path = os.path.join(path, "model_index.json")986 extra_cfg = {}987 if os.path.isfile(model_index_path):988 with open(model_index_path, "r") as f:989 extra_cfg = json.load(f)990 991 # Resolve VLM path: prefer local merged VLM (with LoRA baked in)992 vlm_path = vlm_model_path993 if vlm_path is None:994 local_merged = os.path.join(path, "vlm")995 if os.path.isdir(local_merged):996 vlm_path = local_merged997 else:998 vlm_path = extra_cfg.get("vlm", "Qwen/Qwen2.5-VL-3B-Instruct")999 if not os.path.isdir(vlm_path):1000 local_candidate = os.path.join("/data/huggingface", vlm_path.split("/")[-1])1001 if os.path.isdir(local_candidate):1002 vlm_path = local_candidate1003 print(f"Loading VLM from {vlm_path}...")1004 try:1005 self.lmm = Qwen2_5_VLForConditionalGeneration.from_pretrained(1006 vlm_path, torch_dtype=dtype, attn_implementation=attn_implementation)1007 except Exception:1008 self.lmm = Qwen2_5_VLForConditionalGeneration.from_pretrained(1009 vlm_path, torch_dtype=dtype, attn_implementation="sdpa")1010 self.lmm.requires_grad_(False)1011 1012 print("Loading tokenizer...")1013 tokenizer_path = os.path.join(path, "tokenizer")1014 if os.path.isdir(tokenizer_path):1015 self.tokenizer = AutoTokenizer.from_pretrained(1016 tokenizer_path, trust_remote_code=True, padding_side='right')1017 else:1018 self.tokenizer = AutoTokenizer.from_pretrained(1019 vlm_path, trust_remote_code=True, padding_side='right')1020 1021 print("Loading connector...")1022 connector_dir = os.path.join(path, "connector")1023 with open(os.path.join(connector_dir, "config.json"), "r") as f:1024 connector_cfg = json.load(f)1025 1026 conn_cfg = connector_cfg["connector"].copy()1027 conn_cfg["_attn_implementation"] = "sdpa"1028 1029 self.connector_module = DeepGenConnector(1030 connector_config=conn_cfg,1031 num_queries=connector_cfg["num_queries"],1032 llm_hidden_size=connector_cfg["llm_hidden_size"],1033 projector_1_in=connector_cfg["projector_1_in"],1034 projector_1_out=connector_cfg["projector_1_out"],1035 projector_2_in=connector_cfg["projector_2_in"],1036 projector_2_out=connector_cfg["projector_2_out"],1037 projector_3_in=connector_cfg["projector_3_in"],1038 projector_3_out=connector_cfg["projector_3_out"],1039 )1040 connector_state = load_file(os.path.join(connector_dir, "model.safetensors"))1041 self.connector_module.load_state_dict(connector_state, strict=True)1042 self.connector_module = self.connector_module.to(dtype=dtype)1043 1044 prompt_template_path = os.path.join(path, "prompt_template.json")1045 with open(prompt_template_path, "r") as f:1046 self.prompt_template = json.load(f)1047 1048 self.max_length = connector_cfg.get("max_length", 1024)1049 self.image_token_id = self.tokenizer.convert_tokens_to_ids(1050 self.prompt_template['IMG_CONTEXT_TOKEN'])1051 1052 if not self._cpu_offload:1053 device = self._gpu_device1054 self.lmm = self.lmm.to(device=device)1055 self.connector_module = self.connector_module.to(device=device, dtype=dtype)1056 1057 self.vit_mean = self.vit_mean.to(device=self._gpu_device)1058 self.vit_std = self.vit_std.to(device=self._gpu_device)1059 1060 self._extras_loaded = True1061 print("All components loaded.")1062 1063 @property1064 def llm(self):1065 return self.lmm.language_model1066 1067 @property1068 def num_queries(self):1069 return self.connector_module.num_queries1070 1071 def to(self, *args, **kwargs):1072 result = super().to(*args, **kwargs)1073 device = None1074 dtype = None1075 for a in args:1076 if isinstance(a, torch.device):1077 device = a1078 elif isinstance(a, str):1079 device = torch.device(a)1080 elif isinstance(a, torch.dtype):1081 dtype = a1082 device = device or kwargs.get("device")1083 dtype = dtype or kwargs.get("dtype")1084 1085 if device is not None:1086 self._gpu_device = device1087 if self._extras_loaded:1088 if device is not None:1089 self.lmm = self.lmm.to(device=device)1090 self.connector_module = self.connector_module.to(device=device)1091 self.vit_mean = self.vit_mean.to(device=device)1092 self.vit_std = self.vit_std.to(device=device)1093 if dtype is not None:1094 self.lmm = self.lmm.to(dtype=dtype)1095 self.connector_module = self.connector_module.to(dtype=dtype)1096 return result1097 1098 def enable_model_cpu_offload(self, gpu_id=None, device=None):1099 """Enable sequential CPU offload to reduce GPU memory usage (~14GB)."""1100 self._cpu_offload = True1101 if device is not None:1102 self._gpu_device = torch.device(device) if isinstance(device, str) else device1103 elif gpu_id is not None:1104 self._gpu_device = torch.device(f"cuda:{gpu_id}")1105 self.transformer = self.transformer.to("cpu")1106 self.vae = self.vae.to("cpu")1107 if self._extras_loaded:1108 self.lmm = self.lmm.to("cpu")1109 self.connector_module = self.connector_module.to("cpu")1110 self.vit_mean = self.vit_mean.to(self._gpu_device)1111 self.vit_std = self.vit_std.to(self._gpu_device)1112 torch.cuda.empty_cache()1113 1114 def _offload_to(self, module, device):1115 module.to(device)1116 if device == torch.device("cpu") or device == "cpu":1117 torch.cuda.empty_cache()1118 1119 @classmethod1120 def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):1121 """1122 Load the full pipeline. When called directly (not via DiffusionPipeline),1123 loads all components immediately including VLM and connector.1124 """1125 vlm_model_path = kwargs.pop("vlm_model_path", None)1126 attn_implementation = kwargs.pop("attn_implementation", "flash_attention_2")1127 1128 pipe = super().from_pretrained(pretrained_model_name_or_path, **kwargs)1129 1130 pipe._load_extras(vlm_model_path=vlm_model_path,1131 attn_implementation=attn_implementation)1132 return pipe1133 1134 @torch.no_grad()1135 def pixels_to_latents(self, x):1136 z = self.vae.encode(x).latent_dist.sample()1137 z = (z - self.vae.config.shift_factor) * self.vae.config.scaling_factor1138 return z1139 1140 @torch.no_grad()1141 def latents_to_pixels(self, z):1142 z = (z / self.vae.config.scaling_factor) + self.vae.config.shift_factor1143 x_rec = self.vae.decode(z).sample1144 return x_rec1145 1146 def prepare_text2image_prompts(self, texts):1147 texts = [self.prompt_template['GENERATION'].format(input=text) for text in texts]1148 texts = [self.prompt_template['INSTRUCTION'].format(input=text) for text in texts]1149 return self.tokenizer(1150 texts, add_special_tokens=True, return_tensors='pt',1151 padding=True, padding_side='left').to(self._gpu_device)1152 1153 def prepare_image2image_prompts(self, texts, num_refs, ref_lens):1154 prompts = []1155 cnt = 01156 for text, num_ref in zip(texts, num_refs):1157 image_tokens = ''1158 for _ in range(num_ref):1159 image_tokens += (self.prompt_template['IMG_START_TOKEN'] +1160 self.prompt_template['IMG_CONTEXT_TOKEN'] * ref_lens[cnt] +1161 self.prompt_template['IMG_END_TOKEN'])1162 cnt += 11163 prompts.append(self.prompt_template['INSTRUCTION'].format(1164 input=f'{image_tokens}\n{text}'))1165 return self.tokenizer(1166 prompts, add_special_tokens=True, return_tensors='pt',1167 padding=True, padding_side='left').to(self._gpu_device)1168 1169 def prepare_forward_input(self, query_embeds, input_ids=None,1170 image_embeds=None, image_grid_thw=None,1171 attention_mask=None, past_key_values=None):1172 b, l, _ = query_embeds.shape1173 attention_mask = attention_mask.to(device=self._gpu_device, dtype=torch.bool)1174 input_ids = torch.cat([input_ids, input_ids.new_zeros(b, l)], dim=1)1175 attention_mask = torch.cat([attention_mask, attention_mask.new_ones(b, l)], dim=1)1176 1177 position_ids, _ = self.lmm.model.get_rope_index(1178 input_ids=input_ids, image_grid_thw=image_grid_thw,1179 video_grid_thw=None, second_per_grid_ts=None,1180 attention_mask=attention_mask)1181 1182 if past_key_values is not None:1183 inputs_embeds = query_embeds1184 position_ids = position_ids[..., -l:]1185 else:1186 input_ids = input_ids[:, :-l]1187 if image_embeds is None:1188 inputs_embeds = self.llm.get_input_embeddings()(input_ids)1189 else:1190 inputs_embeds = torch.zeros(1191 *input_ids.shape, self.llm.config.hidden_size,1192 device=self._gpu_device, dtype=self.transformer.dtype)1193 inputs_embeds[input_ids == self.image_token_id] = \1194 image_embeds.contiguous().view(-1, self.llm.config.hidden_size)1195 inputs_embeds[input_ids != self.image_token_id] = \1196 self.llm.get_input_embeddings()(input_ids[input_ids != self.image_token_id])1197 inputs_embeds = torch.cat([inputs_embeds, query_embeds], dim=1)1198 1199 return dict(inputs_embeds=inputs_embeds, attention_mask=attention_mask,1200 position_ids=position_ids, past_key_values=past_key_values)