o Gix @sddlmZddlZddlmZddlmZmZddlmZddl m Z m Z dd l m Z dd lmZdd lmZdd lmZdd lmZeeZdededededejdejf ddZdejdejdejfddZGdddZGdddeje Z GdddejZ!Gd d!d!ejZ"Gd"d#d#ejZ#Gd$d%d%ejZ$Gd&d'd'ejZ%d(ejdedejfd)d*Z&d+ejded,ejdejfd-d.Z'Gd/d0d0eee Z(dS)1)AnyN)nn) ConfigMixinregister_to_config)logging)AttentionMixinAttentionModuleMixin)dispatch_attention_fn)get_timestep_embedding)Transformer2DModelOutput) ModelMixin)RMSNorm batch_sizeheightwidth patch_sizedevicereturncCstj||||d|d}tj|||ddddf|d<tj|||ddddf|d<|||||dd|ddS)a Generates 2D patch coordinate indices for a batch of images. Args: batch_size (`int`): Number of images in the batch. height (`int`): Height of the input images (in pixels). width (`int`): Width of the input images (in pixels). patch_size (`int`): Size of the square patches that the image is divided into. device (`torch.device`): The device on which to create the tensor. Returns: `torch.Tensor`: Tensor of shape `(batch_size, num_patches, 2)` containing the (row, col) coordinates of each patch in the image grid. r)rN.r.rr)torchzerosarangereshape unsqueezerepeat)rrrrrimg_idsr a/home/ubuntu/.local/lib/python3.10/site-packages/diffusers/models/transformers/transformer_prx.py get_image_ids!s""(r"xq freqs_ciscCsn|jg|jdddddR}|j|j|jd}|d|d|d|d}|j|j|S)a Applies rotary positional embeddings (RoPE) to a query tensor. Args: xq (`torch.Tensor`): Input tensor of shape `(..., dim)` representing the queries. freqs_cis (`torch.Tensor`): Precomputed rotary frequency components of shape `(..., dim/2, 2)` containing cosine and sine pairs. Returns: `torch.Tensor`: Tensor of the same shape as `xq` with rotary embeddings applied. Nrrrdtyperr)floatrshapetorr'type_as)r#r$xq_xq_outr r r! apply_rope=s* r.c @s`eZdZdZdZdZddZ   ddddejdejdBd ejdBd ejdBd ejf d d Z dS)PRXAttnProcessor2_0z Processor for implementing PRX-style attention with multi-source tokens and RoPE. Supports multiple attention backends (Flash Attention, Sage Attention, etc.) via dispatch_attention_fn. NcCsttjjds tddS)Nscaled_dot_product_attentionzHPRXAttnProcessor2_0 requires PyTorch 2.0, please upgrade PyTorch to 2.0.)hasattrrr functional ImportError)selfr r r!__init__[szPRXAttnProcessor2_0.__init__attn PRXAttention hidden_statesencoder_hidden_statesattention_maskimage_rotary_embrc#Ksp|durtd||}|j\}} } ||| d|j|j}|ddddd}|d|d|d} } } || } || } | |}|j\}}} |||d|j|j}|ddddd}|d|d}}| |}|durzt | |} t | |} t j || fdd}t j || fdd}d}|dur| j\}} }} |jd}|dkrtd |j|jd |krtd |jd d || j}t j||ft j|d }|j|t jd}t j ||gd d}|ddddddfd |j|d }| dd}|dd}|dd}t|||||j|jd}|j\}} }!}"||| |!|"}|jd|}t|jdkr6|jd|}|S)a Apply PRX attention using PRXAttention module. Args: attn: PRXAttention module containing projection layers hidden_states: Image tokens [B, L_img, D] encoder_hidden_states: Text tokens [B, L_txt, D] attention_mask: Boolean mask for text tokens [B, L_txt] image_rotary_emb: Rotary positional embeddings [B, 1, L_img, head_dim//2, 2, 2] NzLPRXAttnProcessor2_0 requires 'encoder_hidden_states' containing text tokens.rrrrdimz"Unsupported attention_mask shape: r%zattention_mask last dim z must equal text length r'rr&) attn_maskbackendparallel_config) ValueError img_qkv_projr)rheadshead_dimpermutenorm_qnorm_k txt_kv_proj norm_added_kr.rcatr>ronesboolr*expand transposer _attention_backend_parallel_configto_outlen)#r4r6r8r9r:r;kwargsimg_qkvBL_img_img_qimg_kimg_vtxt_kvL_txttxt_ktxt_vkvattn_mask_tensorbsl_imgl_txtrones_img joint_maskquerykeyvalue attn_outputrseq_len num_headsrFr r r!__call___sd           &    zPRXAttnProcessor2_0.__call__NNN) __name__ __module__ __qualname____doc__rQrRr5rTensorror r r r!r/Rs*r/c seZdZdZeZegZ      ddeded ed ed ed e f fd d Z   dde j de j dBde j dBde j dBde j f ddZ ZS)r7z PRX-style attention module that handles multi-source tokens and RoPE. Similar to FluxAttention but adapted for PRX's architecture. @Fư>N query_dimrEdim_headbiasout_biasepscst||_||_|||_||_tj||d|d|_t |j|dd|_ t |j|dd|_ tj||d|d|_ t |j|dd|_ tg|_|jtj|j||d|jtd|durh|}||dS)Nrr{T)r}elementwise_affiner)superr5rErF inner_dimryrLinearrDrrHrIrJrK ModuleListrSappendDropout_default_processor_cls set_processor)r4ryrErzr{r|r} processor __class__r r!r5s   zPRXAttention.__init__r8r9r:r;rcKs|j||f|||d|S)N)r9r:r;)r)r4r8r9r:r;rUr r r!forwardszPRXAttention.forward)rvrwFFrxNrp)rqrrrsrtr/r_available_processorsintrNr(r5rrur __classcell__r r rr!r7sJ$r7csfeZdZdZdededeeffdd Zdejdededejfd d Z d ejdejfd d Z Z S) PRXEmbedNDa N-dimensional rotary positional embedding. This module creates rotary embeddings (RoPE) across multiple axes, where each axis can have its own embedding dimension. The embeddings are combined and returned as a single tensor Args: dim (int): Base embedding dimension (must be even). theta (int): Scaling factor that controls the frequency spectrum of the rotary embeddings. axes_dim (list[int]): list of embedding dimensions for each axis (each must be even). r>thetaaxes_dimcs t||_||_||_dSN)rr5r>rr)r4r>rrrr r!r5 s  zPRXEmbedND.__init__posrc Cs|ddksJ|jjdk}|jjdk}|s|rtjntj}tjd|d||jd|}d||}|d|d} tjt| t |  t | t| gdd} | j g| j ddddR} | S) Nrrmpsnpur?g?r%r=) rtyperfloat32float64rrstackcossinrr)r() r4rr>ris_mpsis_npur'scaleomegaoutr r r!ropes   0"zPRXEmbedND.ropeidscs6jd}tjfddt|Ddd}|dS)Nr%cs4g|]}dddd|fj|jqSr)rrr.0irr4r r! &s4z&PRXEmbedND.forward..r=r)r)rrLranger)r4rn_axesembr rr!r#s  zPRXEmbedND.forward) rqrrrsrtrlistr5rrurrrr r rr!rs rcs@eZdZdZdedeffdd Zdejdejfdd ZZ S) MLPEmbedderan A simple 2-layer MLP used for embedding inputs. Args: in_dim (`int`): Dimensionality of the input features. hidden_dim (`int`): Dimensionality of the hidden and output embedding space. Returns: `torch.Tensor`: Tensor of shape `(..., hidden_dim)` containing the embedded representations. in_dim hidden_dimcs<ttj||dd|_t|_tj||dd|_dS)NTr~)rr5rrin_layerSiLUsilu out_layer)r4rrrr r!r5;s  zMLPEmbedder.__init__xrcCs||||Sr)rrr)r4rr r r!rAszMLPEmbedder.forward rqrrrsrtrr5rrurrr r rr!r,src sbeZdZdZdeffdd ZdejdeeejejejfeejejejfffddZ Z S) Modulationa Modulation network that generates scale, shift, and gating parameters. Given an input vector, the module projects it through a linear layer to produce six chunks, which are grouped into two tuples `(shift, scale, gate)`. Args: dim (`int`): Dimensionality of the input vector. The output will have `6 * dim` features internally. Returns: ((`torch.Tensor`, `torch.Tensor`, `torch.Tensor`), (`torch.Tensor`, `torch.Tensor`, `torch.Tensor`)): Two tuples `(shift, scale, gate)`. r>csHttj|d|dd|_tj|jjdtj|jjddS)NTr~r) rr5rrlininit constant_weightr{)r4r>rr r!r5Us zModulation.__init__vecrcCsN|tj|dddddfjddd}t|ddt|ddfS)Nrr%r=r)rrr2rchunktuple)r4rrr r r!r[s. zModulation.forward) rqrrrsrtrr5rrurrrr r rr!rEs*rcs~eZdZdZ  ddededededBffdd Z dd ejd ejd ejd ejdejdBde e e fdejfddZ Z S)PRXBlocku5 Multimodal transformer block with text–image cross-attention, modulation, and MLP. Args: hidden_size (`int`): Dimension of the hidden representations. num_heads (`int`): Number of attention heads. mlp_ratio (`float`, *optional*, defaults to 4.0): Expansion ratio for the hidden dimension inside the MLP. qk_scale (`float`, *optional*): Scale factor for queries and keys. If not provided, defaults to ``head_dim**-0.5``. Attributes: img_pre_norm (`nn.LayerNorm`): Pre-normalization applied to image tokens before attention. attention (`PRXAttention`): Multi-head attention module with built-in QKV projections and normalizations for cross-attention between image and text tokens. post_attention_layernorm (`nn.LayerNorm`): Normalization applied after attention. gate_proj / up_proj / down_proj (`nn.Linear`): Feedforward layers forming the gated MLP. mlp_act (`nn.GELU`): Nonlinear activation used in the MLP. modulation (`Modulation`): Produces scale/shift/gating parameters for modulated layers. Methods: The forward method performs cross-attention and the MLP with modulation. @N hidden_sizern mlp_ratioqk_scalec st||_||_|||_|p|jd|_t|||_||_t j |ddd|_ t |||jdddt d|_t j |ddd|_t j||jdd|_t j||jdd|_t j|j|dd|_t jdd|_t||_dS) NgFrxrr})ryrErzr{r|r}rr~tanh) approximate)rr5rrnrFrrmlp_hidden_dimrr LayerNorm img_pre_normr7r/ attentionpost_attention_layernormr gate_projup_proj down_projGELUmlp_actr modulation)r4rrnrrrr r!r5s.   zPRXBlock.__init__r8r9tembr;r:rUrc Ks||\}}|\} } } |\} } }d| ||| }|j||||d}|| |}d| ||| }||||||||}|S)a Runs modulation-gated cross-attention and MLP, with residual connections. Args: hidden_states (`torch.Tensor`): Image tokens of shape `(B, L_img, hidden_size)`. encoder_hidden_states (`torch.Tensor`): Text tokens of shape `(B, L_txt, hidden_size)`. temb (`torch.Tensor`): Conditioning vector used by `Modulation` to produce scale/shift/gates, shape `(B, hidden_size)` (or broadcastable). image_rotary_emb (`torch.Tensor`): Rotary positional embeddings applied inside attention. attention_mask (`torch.Tensor`, *optional*): Boolean mask for text tokens of shape `(B, L_txt)`, where `0` marks padding. **kwargs: Additional keyword arguments for API compatibility. Returns: `torch.Tensor`: Updated image tokens of shape `(B, L_img, hidden_size)`. r)r8r9r:r;)rrrrrrrr)r4r8r9rr;r:rUmod_attnmod_mlp attn_shift attn_scale attn_gate mlp_shift mlp_scalemlp_gatehidden_states_modattn_outrr r r!rs   (zPRXBlock.forward)rNr)rqrrrsrtrr(r5rrudictstrrrrr r rr!rbs:$. rcsJeZdZdZdededeffdd Zdejdejd ejfd d ZZ S) FinalLayera Final projection layer with adaptive LayerNorm modulation. This layer applies a normalized and modulated transformation to input tokens and projects them into patch-level outputs. Args: hidden_size (`int`): Dimensionality of the input tokens. patch_size (`int`): Size of the square image patches. out_channels (`int`): Number of output channels per pixel (e.g. RGB = 3). Forward Inputs: x (`torch.Tensor`): Input tokens of shape `(B, L, hidden_size)`, where `L` is the number of patches. vec (`torch.Tensor`): Conditioning vector of shape `(B, hidden_size)` used to generate shift and scale parameters for adaptive LayerNorm. Returns: `torch.Tensor`: Projected patch outputs of shape `(B, L, patch_size * patch_size * out_channels)`. rr out_channelscs\ttj|ddd|_tj||||dd|_tttj|d|dd|_ dS)NFrxrTr~r) rr5rr norm_finalrlinear SequentialradaLN_modulation)r4rrrrr r!r5s &zFinalLayer.__init__rrrcCs`||jddd\}}d|dddddf|||dddddf}||}|S)Nrrr=)rrrr)r4rrshiftrr r r!rs: zFinalLayer.forwardrr r rr!rs$rimgcCsT|j\}}}}|}|||||||||}td|}||d|||}|S)a Flattens an image tensor into a sequence of non-overlapping patches. Args: img (`torch.Tensor`): Input image tensor of shape `(B, C, H, W)`. patch_size (`int`): Size of each square patch. Must evenly divide both `H` and `W`. Returns: `torch.Tensor`: Flattened patch sequence of shape `(B, L, C * patch_size * patch_size)`, where `L = (H // patch_size) * (W // patch_size)` is the number of patches. znchpwq->nhwcpqr%)r)rreinsum)rrbchwpr r r!img2seqs  rseqr)c Cst|tr|dd\}}nt|tjr"t|dt|d}}n tdt|d|j\}}}|}|||} ||||||| ||}t d|}||| ||}|S)av Reconstructs an image tensor from a sequence of patches (inverse of `img2seq`). Args: seq (`torch.Tensor`): Patch sequence of shape `(B, L, C * patch_size * patch_size)`, where `L = (H // patch_size) * (W // patch_size)`. patch_size (`int`): Size of each square patch. shape (`tuple` or `torch.Tensor`): The original image spatial shape `(H, W)`. If a tensor is provided, the first two values are interpreted as height and width. Returns: `torch.Tensor`: Reconstructed image tensor of shape `(B, C, H, W)`. Nrrz shape type z not supportedznhwcpq->nchpwq) isinstancerrrurNotImplementedErrorrr)rr) rrr)rrrldrrr r r!seq2img%s     rcseZdZdZdZdZe       d'd ededededededede dededeffdd Z de j de j de j fddZ d(de j de j d e j d!e j d Bd"eeefd Bd#edee j d$feBfd%d&ZZS))PRXTransformer2DModela Transformer-based 2D model for text to image generation. Args: in_channels (`int`, *optional*, defaults to 16): Number of input channels in the latent image. patch_size (`int`, *optional*, defaults to 2): Size of the square patches used to flatten the input image. context_in_dim (`int`, *optional*, defaults to 2304): Dimensionality of the text conditioning input. hidden_size (`int`, *optional*, defaults to 1792): Dimension of the hidden representation. mlp_ratio (`float`, *optional*, defaults to 3.5): Expansion ratio for the hidden dimension inside MLP blocks. num_heads (`int`, *optional*, defaults to 28): Number of attention heads. depth (`int`, *optional*, defaults to 16): Number of transformer blocks. axes_dim (`list[int]`, *optional*): list of dimensions for each positional embedding axis. Defaults to `[32, 32]`. theta (`int`, *optional*, defaults to 10000): Frequency scaling factor for rotary embeddings. time_factor (`float`, *optional*, defaults to 1000.0): Scaling factor applied in timestep embeddings. time_max_period (`int`, *optional*, defaults to 10000): Maximum frequency period for timestep embeddings. Attributes: pe_embedder (`EmbedND`): Multi-axis rotary embedding generator for positional encodings. img_in (`nn.Linear`): Projection layer for image patch tokens. time_in (`MLPEmbedder`): Embedding layer for timestep embeddings. txt_in (`nn.Linear`): Projection layer for text conditioning. blocks (`nn.ModuleList`): Stack of transformer blocks (`PRXBlock`). final_layer (`LastLayer`): Projection layer mapping hidden tokens back to patch outputs. Methods: attn_processors: Returns a dictionary of all attention processors in the model. set_attn_processor(processor): Replaces attention processors across all attention layers. process_inputs(image_latent, txt): Converts inputs into patch tokens, encodes text, and produces positional encodings. compute_timestep_embedding(timestep, dtype): Creates a timestep embedding of dimension 256, scaled and projected. forward_transformers(image_latent, cross_attn_conditioning, timestep, time_embedding, attention_mask, **block_kwargs): Runs the sequence of transformer blocks over image and text tokens. forward(image_latent, timestep, cross_attn_conditioning, micro_conditioning, cross_attn_mask=None, attention_kwargs=None, return_dict=True): Full forward pass from latent input to reconstructed output image. Returns: `Transformer2DModelOutput` if `return_dict=True` (default), otherwise a tuple containing: - `sample` (`torch.Tensor`): Reconstructed image of shape `(B, C, H, W)`. z config.jsonTr  @N'@@ in_channelsrcontext_in_dimrrrndepthrr time_factortime_max_periodc s$t|dur ddg}|_|_jjd_| _| _||dkr2td|d|||} t|| krFtd|d| |_ |_ t | | |d_ t jjjdj d d _td j d _t |j _t fd dt|D_tj dj_d_dS)N rrz Hidden size z must be divisible by num_heads zGot z but expected positional dim )r>rrTr~)rrcsg|] }tjjdqS))r)rrrnrrr4r r!rsz2PRXTransformer2DModel.__init__..rF)rr5rrrrrrCsumrrnr pe_embedderrrimg_inrtime_intxt_inrrblocksr final_layergradient_checkpointing) r4rrrrrrnrrrrrpe_dimrrr!r5s4      zPRXTransformer2DModel.__init__timestepr'rc Cs$|t|d|j|jddd|S)NrTr) timesteps embedding_dim max_periodrflip_sin_to_cosdownscale_freq_shift)r r rrr*)r4rr'r r r!_compute_timestep_embeddingsz1PRXTransformer2DModel._compute_timestep_embeddingr8r9r:attention_kwargs return_dict.c Cs||}t||j}||}|j\} } } } t| | | |j|jd} || }|j||j d}|j D]}t rG|j rG||j|||||}q2||||||d}q2|||}t||j|j}|sd|fSt|dS)a Forward pass of the PRXTransformer2DModel. The latent image is split into patch tokens, combined with text conditioning, and processed through a stack of transformer blocks modulated by the timestep. The output is reconstructed into the latent image space. Args: hidden_states (`torch.Tensor`): Input latent image tensor of shape `(B, C, H, W)`. timestep (`torch.Tensor`): Timestep tensor of shape `(B,)` or `(1,)`, used for temporal conditioning. encoder_hidden_states (`torch.Tensor`): Text conditioning tensor of shape `(B, L_txt, context_in_dim)`. attention_mask (`torch.Tensor`, *optional*): Boolean mask of shape `(B, L_txt)`, where `0` marks padding in the text sequence. attention_kwargs (`dict`, *optional*): Additional arguments passed to attention layers. return_dict (`bool`, *optional*, defaults to `True`): Whether to return a `Transformer2DModelOutput` or a tuple. Returns: `Transformer2DModelOutput` if `return_dict=True`, otherwise a tuple: - `sample` (`torch.Tensor`): Output latent image of shape `(B, C, H, W)`. )rr)r')r8r9rr;r:)sample)r rrrr)r"rrrr'r ris_grad_enabledr _gradient_checkpointing_funcror rr )r4r8rr9r:rrtxtrrdrYrrrperblockoutputr r r!rs: #       zPRXTransformer2DModel.forward) rrrrrrrNrrr)NNT)rqrrrsrt config_name _supports_gradient_checkpointingrrr(rr5rrur'rrrrrNrr rrr r rr!rNst>    9r))typingrrrconfiguration_utilsrrutilsrrr r attention_dispatchr embeddingsr modeling_outputsr modeling_utilsr normalizationr get_loggerrqloggerrrrur"r.r/Moduler7rrrrrrrrr r r r!s0         &n=/}( )