rename photon into prx

Author: DavidBert

docs/source/en/_toctree.yml

@@ -541,12 +541,12 @@
         title: PAG
       - local: api/pipelines/paint_by_example
         title: Paint by Example
-      - local: api/pipelines/prx
-        title: PRX
       - local: api/pipelines/pixart
         title: PixArt-α
       - local: api/pipelines/pixart_sigma
         title: PixArt-Σ
+      - local: api/pipelines/prx
+        title: PRX
       - local: api/pipelines/qwenimage
         title: QwenImage
       - local: api/pipelines/sana

docs/source/en/api/pipelines/photon.md docs/source/en/api/pipelines/prx.mddocs/source/en/api/pipelines/photon.md renamed to docs/source/en/api/pipelines/prx.md

@@ -232,9 +232,9 @@
             "MultiControlNetModel",
             "OmniGenTransformer2DModel",
             "ParallelConfig",
-            "PRXTransformer2DModel",
             "PixArtTransformer2DModel",
             "PriorTransformer",
+            "PRXTransformer2DModel",
             "QwenImageControlNetModel",
             "QwenImageMultiControlNetModel",
             "QwenImageTransformer2DModel",
@@ -516,11 +516,11 @@
             "MusicLDMPipeline",
             "OmniGenPipeline",
             "PaintByExamplePipeline",
-            "PRXPipeline",
             "PIAPipeline",
             "PixArtAlphaPipeline",
             "PixArtSigmaPAGPipeline",
             "PixArtSigmaPipeline",
+            "PRXPipeline",
             "QwenImageControlNetInpaintPipeline",
             "QwenImageControlNetPipeline",
             "QwenImageEditInpaintPipeline",
@@ -928,9 +928,9 @@
             MultiControlNetModel,
             OmniGenTransformer2DModel,
             ParallelConfig,
-            PRXTransformer2DModel,
             PixArtTransformer2DModel,
             PriorTransformer,
+            PRXTransformer2DModel,
             QwenImageControlNetModel,
             QwenImageMultiControlNetModel,
             QwenImageTransformer2DModel,
@@ -1182,11 +1182,11 @@
             MusicLDMPipeline,
             OmniGenPipeline,
             PaintByExamplePipeline,
-            PRXPipeline,
             PIAPipeline,
             PixArtAlphaPipeline,
             PixArtSigmaPAGPipeline,
             PixArtSigmaPipeline,
+            PRXPipeline,
             QwenImageControlNetInpaintPipeline,
             QwenImageControlNetPipeline,
             QwenImageEditInpaintPipeline,

scripts/convert_photon_to_diffusers.py scripts/convert_prx_to_diffusers.pyscripts/convert_photon_to_diffusers.py renamed to scripts/convert_prx_to_diffusers.py

@@ -191,9 +191,9 @@
             LuminaNextDiT2DModel,
             MochiTransformer3DModel,
             OmniGenTransformer2DModel,
-            PRXTransformer2DModel,
             PixArtTransformer2DModel,
             PriorTransformer,
+            PRXTransformer2DModel,
             QwenImageTransformer2DModel,
             SanaTransformer2DModel,
             SD3Transformer2DModel,

src/diffusers/__init__.py

@@ -80,9 +80,9 @@ def apply_rope(xq: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:
     return xq_out.reshape(*xq.shape).type_as(xq)
 
 
-class PhotonAttnProcessor2_0:
+class PRXAttnProcessor2_0:
     r"""
-    Processor for implementing Photon-style attention with multi-source tokens and RoPE. Supports multiple attention
+    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.
     """
 
@@ -91,30 +91,30 @@ class PhotonAttnProcessor2_0:
 
     def __init__(self):
         if not hasattr(torch.nn.functional, "scaled_dot_product_attention"):
-            raise ImportError("PhotonAttnProcessor2_0 requires PyTorch 2.0, please upgrade PyTorch to 2.0.")
+            raise ImportError("PRXAttnProcessor2_0 requires PyTorch 2.0, please upgrade PyTorch to 2.0.")
 
     def __call__(
         self,
-        attn: "PhotonAttention",
+        attn: "PRXAttention",
         hidden_states: torch.Tensor,
         encoder_hidden_states: Optional[torch.Tensor] = None,
         attention_mask: Optional[torch.Tensor] = None,
         image_rotary_emb: Optional[torch.Tensor] = None,
         **kwargs,
     ) -> torch.Tensor:
         """
-        Apply Photon attention using PhotonAttention module.
+        Apply PRX attention using PRXAttention module.
 
         Args:
-            attn: PhotonAttention module containing projection layers
+            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]
         """
 
         if encoder_hidden_states is None:
-            raise ValueError("PhotonAttnProcessor2_0 requires 'encoder_hidden_states' containing text tokens.")
+            raise ValueError("PRXAttnProcessor2_0 requires 'encoder_hidden_states' containing text tokens.")
 
         # Project image tokens to Q, K, V
         img_qkv = attn.img_qkv_proj(hidden_states)
@@ -190,14 +190,14 @@ def __call__(
         return attn_output
 
 
-class PhotonAttention(nn.Module, AttentionModuleMixin):
+class PRXAttention(nn.Module, AttentionModuleMixin):
     r"""
-    Photon-style attention module that handles multi-source tokens and RoPE. Similar to FluxAttention but adapted for
-    Photon's architecture.
+    PRX-style attention module that handles multi-source tokens and RoPE. Similar to FluxAttention but adapted for
+    PRX's architecture.
     """
 
-    _default_processor_cls = PhotonAttnProcessor2_0
-    _available_processors = [PhotonAttnProcessor2_0]
+    _default_processor_cls = PRXAttnProcessor2_0
+    _available_processors = [PRXAttnProcessor2_0]
 
     def __init__(
         self,
@@ -251,7 +251,7 @@ def forward(
 
 
 # inspired from https://github.com/black-forest-labs/flux/blob/main/src/flux/modules/layers.py
-class PhotonEmbedND(nn.Module):
+class PRXEmbedND(nn.Module):
     r"""
     N-dimensional rotary positional embedding.
 
@@ -347,7 +347,7 @@ def forward(
         return tuple(out[:3]), tuple(out[3:])
 
 
-class PhotonBlock(nn.Module):
+class PRXBlock(nn.Module):
     r"""
     Multimodal transformer block with text–image cross-attention, modulation, and MLP.
 
@@ -364,7 +364,7 @@ class PhotonBlock(nn.Module):
     Attributes:
         img_pre_norm (`nn.LayerNorm`):
             Pre-normalization applied to image tokens before attention.
-        attention (`PhotonAttention`):
+        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`):
@@ -400,15 +400,15 @@ def __init__(
         # Pre-attention normalization for image tokens
         self.img_pre_norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
 
-        # PhotonAttention module with built-in projections and norms
-        self.attention = PhotonAttention(
+        # PRXAttention module with built-in projections and norms
+        self.attention = PRXAttention(
             query_dim=hidden_size,
             heads=num_heads,
             dim_head=self.head_dim,
             bias=False,
             out_bias=False,
             eps=1e-6,
-            processor=PhotonAttnProcessor2_0(),
+            processor=PRXAttnProcessor2_0(),
         )
 
         # mlp
@@ -557,7 +557,7 @@ def seq2img(seq: torch.Tensor, patch_size: int, shape: torch.Tensor) -> torch.Te
     return fold(seq.transpose(1, 2), shape, kernel_size=patch_size, stride=patch_size)
 
 
-class PhotonTransformer2DModel(ModelMixin, ConfigMixin, AttentionMixin):
+class PRXTransformer2DModel(ModelMixin, ConfigMixin, AttentionMixin):
     r"""
     Transformer-based 2D model for text to image generation.
 
@@ -595,7 +595,7 @@ class PhotonTransformer2DModel(ModelMixin, ConfigMixin, AttentionMixin):
         txt_in (`nn.Linear`):
             Projection layer for text conditioning.
         blocks (`nn.ModuleList`):
-            Stack of transformer blocks (`PhotonBlock`).
+            Stack of transformer blocks (`PRXBlock`).
         final_layer (`LastLayer`):
             Projection layer mapping hidden tokens back to patch outputs.
 
@@ -661,14 +661,14 @@ def __init__(
 
         self.hidden_size = hidden_size
         self.num_heads = num_heads
-        self.pe_embedder = PhotonEmbedND(dim=pe_dim, theta=theta, axes_dim=axes_dim)
+        self.pe_embedder = PRXEmbedND(dim=pe_dim, theta=theta, axes_dim=axes_dim)
         self.img_in = nn.Linear(self.in_channels * self.patch_size**2, self.hidden_size, bias=True)
         self.time_in = MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size)
         self.txt_in = nn.Linear(context_in_dim, self.hidden_size)
 
         self.blocks = nn.ModuleList(
             [
-                PhotonBlock(
+                PRXBlock(
                     self.hidden_size,
                     self.num_heads,
                     mlp_ratio=mlp_ratio,
@@ -702,7 +702,7 @@ def forward(
         return_dict: bool = True,
     ) -> Union[Tuple[torch.Tensor, ...], Transformer2DModelOutput]:
         r"""
-        Forward pass of the PhotonTransformer2DModel.
+        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.

src/diffusers/models/__init__.py

@@ -718,9 +718,9 @@
             StableDiffusionXLPAGPipeline,
         )
         from .paint_by_example import PaintByExamplePipeline
-        from .prx import PRXPipeline
         from .pia import PIAPipeline
         from .pixart_alpha import PixArtAlphaPipeline, PixArtSigmaPipeline
+        from .prx import PRXPipeline
         from .qwenimage import (
             QwenImageControlNetInpaintPipeline,
             QwenImageControlNetPipeline,

…odels/transformers/transformer_photon.py …s/models/transformers/transformer_prx.pysrc/diffusers/models/transformers/transformer_photon.py renamed to src/diffusers/models/transformers/transformer_prx.py src/diffusers/models/transformers/transformer_photon.py renamed to src/diffusers/models/transformers/transformer_prx.py

@@ -30,9 +30,9 @@
 from diffusers.image_processor import PixArtImageProcessor
 from diffusers.loaders import FromSingleFileMixin, LoraLoaderMixin, TextualInversionLoaderMixin
 from diffusers.models import AutoencoderDC, AutoencoderKL
-from diffusers.models.transformers.transformer_photon import PhotonTransformer2DModel
-from diffusers.pipelines.photon.pipeline_output import PhotonPipelineOutput
+from diffusers.models.transformers.transformer_prx import PRXTransformer2DModel
 from diffusers.pipelines.pipeline_utils import DiffusionPipeline
+from diffusers.pipelines.prx.pipeline_output import PRXPipelineOutput
 from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
 from diffusers.utils import (
     logging,
@@ -73,7 +73,7 @@
 
 
 class TextPreprocessor:
-    """Text preprocessing utility for PhotonPipeline."""
+    """Text preprocessing utility for PRXPipeline."""
 
     def __init__(self):
         """Initialize text preprocessor."""
@@ -203,34 +203,34 @@ def clean_text(self, text: str) -> str:
     Examples:
         ```py
         >>> import torch
-        >>> from diffusers import PhotonPipeline
+        >>> from diffusers import PRXPipeline
 
         >>> # Load pipeline with from_pretrained
-        >>> pipe = PhotonPipeline.from_pretrained("Photoroom/photon-512-t2i-sft")
+        >>> pipe = PRXPipeline.from_pretrained("Photoroom/prx-512-t2i-sft")
         >>> pipe.to("cuda")
 
         >>> prompt = "A digital painting of a rusty, vintage tram on a sandy beach"
         >>> image = pipe(prompt, num_inference_steps=28, guidance_scale=5.0).images[0]
-        >>> image.save("photon_output.png")
+        >>> image.save("prx_output.png")
         ```
 """
 
 
-class PhotonPipeline(
+class PRXPipeline(
     DiffusionPipeline,
     LoraLoaderMixin,
     FromSingleFileMixin,
     TextualInversionLoaderMixin,
 ):
     r"""
-    Pipeline for text-to-image generation using Photon Transformer.
+    Pipeline for text-to-image generation using PRX Transformer.
 
     This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the
     library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)
 
     Args:
-        transformer ([`PhotonTransformer2DModel`]):
-            The Photon transformer model to denoise the encoded image latents.
+        transformer ([`PRXTransformer2DModel`]):
+            The PRX transformer model to denoise the encoded image latents.
         scheduler ([`FlowMatchEulerDiscreteScheduler`]):
             A scheduler to be used in combination with `transformer` to denoise the encoded image latents.
         text_encoder ([`T5GemmaEncoder`]):
@@ -248,7 +248,7 @@ class PhotonPipeline(
 
     def __init__(
         self,
-        transformer: PhotonTransformer2DModel,
+        transformer: PRXTransformer2DModel,
         scheduler: FlowMatchEulerDiscreteScheduler,
         text_encoder: T5GemmaEncoder,
         tokenizer: Union[T5TokenizerFast, GemmaTokenizerFast, AutoTokenizer],
@@ -257,9 +257,9 @@ def __init__(
     ):
         super().__init__()
 
-        if PhotonTransformer2DModel is None:
+        if PRXTransformer2DModel is None:
             raise ImportError(
-                "PhotonTransformer2DModel is not available. Please ensure the transformer_photon module is properly installed."
+                "PRXTransformer2DModel is not available. Please ensure the transformer_prx module is properly installed."
             )
 
         self.text_preprocessor = TextPreprocessor()
@@ -567,7 +567,7 @@ def __call__(
                 The output format of the generate image. Choose between
                 [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
             return_dict (`bool`, *optional*, defaults to `True`):
-                Whether or not to return a [`~pipelines.photon.PhotonPipelineOutput`] instead of a plain tuple.
+                Whether or not to return a [`~pipelines.prx.PRXPipelineOutput`] instead of a plain tuple.
             use_resolution_binning (`bool`, *optional*, defaults to `True`):
                 If set to `True`, the requested height and width are first mapped to the closest resolutions using
                 predefined aspect ratio bins. After the produced latents are decoded into images, they are resized back
@@ -585,9 +585,8 @@ def __call__(
         Examples:
 
         Returns:
-            [`~pipelines.photon.PhotonPipelineOutput`] or `tuple`: [`~pipelines.photon.PhotonPipelineOutput`] if
-            `return_dict` is True, otherwise a `tuple. When returning a tuple, the first element is a list with the
-            generated images.
+            [`~pipelines.prx.PRXPipelineOutput`] or `tuple`: [`~pipelines.prx.PRXPipelineOutput`] if `return_dict` is
+            True, otherwise a `tuple. When returning a tuple, the first element is a list with the generated images.
         """
 
         # 0. Set height and width
@@ -765,4 +764,4 @@ def __call__(
         if not return_dict:
             return (image,)
 
-        return PhotonPipelineOutput(images=image)
+        return PRXPipelineOutput(images=image)