diff --git a/README.md b/README.md index f0f775c..b971922 100644 --- a/README.md +++ b/README.md @@ -27,6 +27,7 @@ AI Toolkit is an easy to use all in one training suite for diffusion models. I t - [baidu/ERNIE-Image](https://huggingface.co/baidu/ERNIE-Image) (ERNIE-Image) - [NucleusAI/Nucleus-Image](https://huggingface.co/NucleusAI/Nucleus-Image) (Nucleus-Image) - [HiDream-ai/HiDream-O1-Image](https://huggingface.co/HiDream-ai/HiDream-O1-Image) (HiDream O1) +- [Photoroom/prxpixel-t2i](https://huggingface.co/Photoroom/prxpixel-t2i) (PRXPixel) ### Instruction / Edit - [black-forest-labs/FLUX.1-Kontext-dev](https://huggingface.co/black-forest-labs/FLUX.1-Kontext-dev) (FLUX.1-Kontext-dev) diff --git a/extensions_built_in/diffusion_models/__init__.py b/extensions_built_in/diffusion_models/__init__.py index c3a0c9d..a410b1e 100644 --- a/extensions_built_in/diffusion_models/__init__.py +++ b/extensions_built_in/diffusion_models/__init__.py @@ -14,6 +14,7 @@ from .nucleus_image import NucleusImageModel from .hidream.hidream_o1_model import HidreamO1Model from .z_image.z_image_l2p_model import ZImageL2PModel from .ideogram4 import Ideogram4Model +from .prx_pixel_t2i import PRXPixelT2IModel AI_TOOLKIT_MODELS = [ # put a list of models here @@ -42,4 +43,5 @@ AI_TOOLKIT_MODELS = [ HidreamO1Model, ZImageL2PModel, Ideogram4Model, + PRXPixelT2IModel, ] diff --git a/extensions_built_in/diffusion_models/prx_pixel_t2i/README.md b/extensions_built_in/diffusion_models/prx_pixel_t2i/README.md new file mode 100644 index 0000000..19f33b1 --- /dev/null +++ b/extensions_built_in/diffusion_models/prx_pixel_t2i/README.md @@ -0,0 +1,53 @@ +# PRXPixel (Photoroom PRX-7B, pixel-space text-to-image) + +ai-toolkit integration for [`Photoroom/prxpixel-t2i`](https://huggingface.co/Photoroom/prxpixel-t2i), +a ~7B pixel-space diffusion transformer. + +It is implemented **from scratch** so ai-toolkit does not depend on the (still +unmerged) diffusers PR [huggingface/diffusers#13928](https://github.com/huggingface/diffusers/pull/13928): +the transformer is vendored in [src/transformer_prx.py](src/transformer_prx.py) +and a minimal preview sampler lives in [src/pipeline.py](src/pipeline.py). + +## What makes this model unusual + +PRXPixel differs from a typical latent flow-matching model in three ways, each +handled in [prx_pixel_t2i.py](prx_pixel_t2i.py): + +| Property | What it means | How it's handled | +|---|---|---| +| **Pixel space** | No VAE — the transformer denoises raw RGB (`in_channels=3`, `patch_size=16`) | A `FakeVAE` (identity, scaling 1) so encode/decode are no-ops; "latents" are the image in `[-1, 1]` | +| **x-prediction** | The model predicts the clean image `x0`, not the flow velocity | `get_noise_prediction` returns `x0`; `get_loss_target` is the clean latents. The `x0 → velocity` conversion only happens at sampling time | +| **noise_scale = 2.0** | Trains/samples from `randn * 2.0`, not unit noise | `get_latent_noise_from_latents` scales the training noise; the pipeline scales the starting noise | + +Text is encoded by the Qwen3-VL text tower (`Qwen3VLTextModel`, hidden size +2048 → the transformer's `context_in_dim`), padded to 256 tokens. + +The x-prediction objective follows *"Back to Basics: Let Denoising Generative +Models Denoise"* (https://arxiv.org/abs/2511.13720). + +## Architecture (released checkpoint) + +`depth=24`, `hidden_size=3584`, `num_heads=28`, `mlp_ratio=3.5`, +`in_channels=3`, `patch_size=16`, `context_in_dim=2048`, `bottleneck_size=768`, +`axes_dim=[64, 64]`, `resolution_embeds=True`, flow-matching scheduler with +`shift=3.0`. + +## Train it + +```yaml +model: + arch: "prx_pixel" + name_or_path: "/path/to/prxpixel-t2i" # diffusers folder: transformer/, + # text_encoder/, tokenizer/, scheduler/ + quantize: true # optional: qfloat8 the transformer + quantize_te: true # optional: qfloat8 the Qwen3-VL text encoder +train: + gradient_checkpointing: true +sample: + guidance_scale: 5.0 + sample_steps: 28 +``` + +Datasets bucket to multiples of 16px (`vae_scale_factor * patch_size`). +See [../example_model/README.md](../example_model/README.md) for the generic +lifecycle, registration and LoRA conventions. diff --git a/extensions_built_in/diffusion_models/prx_pixel_t2i/__init__.py b/extensions_built_in/diffusion_models/prx_pixel_t2i/__init__.py new file mode 100644 index 0000000..23cf447 --- /dev/null +++ b/extensions_built_in/diffusion_models/prx_pixel_t2i/__init__.py @@ -0,0 +1,3 @@ +from .prx_pixel_t2i import PRXPixelT2IModel + +__all__ = ["PRXPixelT2IModel"] diff --git a/extensions_built_in/diffusion_models/prx_pixel_t2i/prx_pixel_t2i.py b/extensions_built_in/diffusion_models/prx_pixel_t2i/prx_pixel_t2i.py new file mode 100644 index 0000000..fb44ded --- /dev/null +++ b/extensions_built_in/diffusion_models/prx_pixel_t2i/prx_pixel_t2i.py @@ -0,0 +1,345 @@ +"""PRXPixelT2IModel -- the Photoroom pixel-space PRX-7B text-to-image model +(https://huggingface.co/Photoroom/prxpixel-t2i) wired into ai-toolkit. + +This is implemented from scratch (the diffusers support is an unmerged PR, +https://github.com/huggingface/diffusers/pull/13928) so ai-toolkit does not +depend on it: the transformer architecture is vendored in ``src/transformer_prx.py`` +and a minimal sampler lives in ``src/pipeline.py``. + +What makes PRXPixel unusual (and what each override below is doing about it): + + - **Pixel space, no VAE.** The transformer denoises raw RGB directly + (``in_channels=3``, ``patch_size=16``). We use a ``FakeVAE`` (identity, + scaling_factor=1) so BaseModel's encode_images/decode_latents become no-ops + and the "latents" everywhere in the toolkit are just the image in [-1, 1]. + Same trick as ``../chroma/chroma_radiance_model.py`` and + ``extensions/z_image_pixel``. + + - **x-prediction.** The model predicts the CLEAN image x0, not the + flow-matching velocity. ai-toolkit's MSE compares ``get_noise_prediction`` + against ``get_loss_target``; we set BOTH to the x0 space (prediction = the + model's x0 output, target = the clean latents), which is PRXPixel's native + training objective ("Back to Basics: Let Denoising Generative Models + Denoise", https://arxiv.org/abs/2511.13720). The x0->velocity conversion + only happens at sampling time, inside the pipeline. + + - **noise_scale = 2.0.** PRXPixel trains with a non-unit initial-noise std, + so the noise mixed into the latents (training) and the starting noise + (sampling) are ``randn * noise_scale``. We override + ``get_latent_noise_from_latents`` for the training side; the pipeline + handles the sampling side. + + - **Qwen3-VL text tower.** Prompts are encoded by ``Qwen3VLTextModel`` + (hidden size 2048 -> the transformer's ``context_in_dim``). We keep the + per-token ``last_hidden_state`` plus an attention mask. + +See ../example_model/README.md for the generic lifecycle/registration guide. +""" + +import os +from typing import List, Optional + +import torch +import yaml + +from transformers import AutoTokenizer, Qwen3VLTextModel +from optimum.quanto import freeze + +from toolkit.accelerator import unwrap_model +from toolkit.basic import flush +from toolkit.config_modules import GenerateImageConfig, ModelConfig +from toolkit.models.base_model import BaseModel +from toolkit.models.FakeVAE import FakeVAE +from toolkit.prompt_utils import PromptEmbeds +from toolkit.samplers.custom_flowmatch_sampler import ( + CustomFlowMatchEulerDiscreteScheduler, +) +from toolkit.train_tools import apply_noise_offset +from toolkit.util.quantize import quantize, get_qtype, quantize_model + +from .src.transformer_prx import PRXTransformer2DModel +from .src.pipeline import PRXPixelPipeline + + +# Flow-matching scheduler config, matching the released model's +# scheduler/scheduler_config.json (shift 3.0, 1000 train timesteps). +scheduler_config = { + "num_train_timesteps": 1000, + "use_dynamic_shifting": False, + "shift": 3.0, +} + +# Number of text tokens PRXPixel was trained with (the Qwen tokenizer's own +# model_max_length is far larger). Matches PRX_PIXEL_DEFAULT_MAX_TOKENS. +PROMPT_MAX_TOKENS = 256 +# Initial-noise std PRXPixel trains/samples with. +NOISE_SCALE = 2.0 + + +class PRXPixelT2IModel(BaseModel): + # ``model.arch: "prx_pixel"`` in the training config YAML selects this class. + arch = "prx_pixel" + + def __init__( + self, + device, + model_config: ModelConfig, + dtype="bf16", + custom_pipeline=None, + noise_scheduler=None, + **kwargs, + ): + super().__init__( + device, model_config, dtype, custom_pipeline, noise_scheduler, **kwargs + ) + self.is_flow_matching = True + self.is_transformer = True + # Matched against type(module).__name__ to place LoRA layers. + self.target_lora_modules = ["PRXTransformer2DModel"] + + # used by our overrides below + self.patch_size = 16 + self.vae_scale_factor = 1 # pixel space: no VAE downsampling + self.max_text_length = PROMPT_MAX_TOKENS + self.noise_scale = NOISE_SCALE + + @staticmethod + def get_train_scheduler(): + return CustomFlowMatchEulerDiscreteScheduler(**scheduler_config) + + def get_bucket_divisibility(self): + # pixels must be divisible by patch_size (no VAE downsample): 1 * 16 = 16 + return self.vae_scale_factor * self.patch_size + + # ------------------------------------------------------------------ + # Loading + # ------------------------------------------------------------------ + def load_model(self): + dtype = self.torch_dtype + self.print_and_status_update("Loading PRXPixel model") + # Expected diffusers-style layout under name_or_path: + # transformer/ (PRXTransformer2DModel) text_encoder/ (Qwen3VLTextModel) + # tokenizer/ scheduler/ + model_path = self.model_config.name_or_path + + # --- transformer (vendored PRXTransformer2DModel) --- + self.print_and_status_update("Loading transformer") + # from_pretrained reads config.json (bottleneck_size, resolution_embeds, + # in_channels=3, ...) and the safetensors in one shot. + transformer = PRXTransformer2DModel.from_pretrained( + model_path, subfolder="transformer", torch_dtype=dtype + ) + transformer.to(dtype=dtype) + flush() + + if self.model_config.quantize: + self.print_and_status_update("Quantizing transformer") + quantize_model(self, transformer) + flush() + + if self.model_config.low_vram: + transformer.to("cpu") + else: + transformer.to(self.device_torch, dtype=dtype) + flush() + + # --- text encoder + tokenizer (Qwen3-VL text tower) --- + self.print_and_status_update("Loading text encoder") + tokenizer = AutoTokenizer.from_pretrained(model_path, subfolder="tokenizer") + text_encoder = Qwen3VLTextModel.from_pretrained( + model_path, subfolder="text_encoder", torch_dtype=dtype + ) + text_encoder.to(self.te_device_torch) + text_encoder.eval() + text_encoder.requires_grad_(False) + flush() + + if self.model_config.quantize_te: + self.print_and_status_update("Quantizing text encoder") + quantize(text_encoder, weights=get_qtype(self.model_config.qtype_te)) + freeze(text_encoder) + flush() + + # --- "VAE": identity, since PRXPixel is pixel space --- + self.print_and_status_update("Preparing pixel-space VAE (identity)") + vae = FakeVAE(scaling_factor=1.0) + vae.to(self.vae_device_torch, dtype=self.vae_torch_dtype) + + # --- scheduler + store everything --- + self.noise_scheduler = PRXPixelT2IModel.get_train_scheduler() + self.vae = vae + self.text_encoder = text_encoder + self.tokenizer = tokenizer + self.model = transformer # aliased as self.transformer / self.unet + self.pipeline = PRXPixelPipeline(self) + self.print_and_status_update("Model Loaded") + + # ------------------------------------------------------------------ + # Sampling (training previews) + # ------------------------------------------------------------------ + def get_generation_pipeline(self): + return PRXPixelPipeline(self) + + def generate_single_image( + self, + pipeline: PRXPixelPipeline, + gen_config: GenerateImageConfig, + conditional_embeds: PromptEmbeds, + unconditional_embeds: PromptEmbeds, + generator: torch.Generator, + extra: dict, + ): + if self.model.device == torch.device("cpu"): + self.model.to(self.device_torch) + + # snap requested size to the model's divisibility (16px) + sc = self.get_bucket_divisibility() + gen_config.width = int(gen_config.width // sc * sc) + gen_config.height = int(gen_config.height // sc * sc) + + img = pipeline( + conditional_embeds=conditional_embeds, + unconditional_embeds=unconditional_embeds, + height=gen_config.height, + width=gen_config.width, + num_inference_steps=gen_config.num_inference_steps, + guidance_scale=gen_config.guidance_scale, + latents=gen_config.latents, + generator=generator, + )[0] + return img + + # ------------------------------------------------------------------ + # Training hooks + # ------------------------------------------------------------------ + def get_noise_prediction( + self, + latent_model_input: torch.Tensor, + timestep: torch.Tensor, # 0..1000 scale, 1000 = pure noise + text_embeddings: PromptEmbeds, + **kwargs, + ): + """Forward pass. Returns the model's predicted CLEAN image x0. + + PRXPixel is an x-prediction model, so the raw transformer output is the + prediction we compare against ``get_loss_target`` (the clean latents). + No velocity conversion happens here -- that is a sampling-time concern. + """ + if self.model.device == torch.device("cpu"): + self.model.to(self.device_torch) + + # toolkit timestep (0..1000) -> PRX flow time in [0, 1]. + t01 = timestep.to(self.device_torch, dtype=torch.float32) / 1000.0 + + feats = text_embeddings.text_embeds.to(self.device_torch, self.torch_dtype) + mask = getattr(text_embeddings, "attention_mask", None) + if mask is not None: + mask = mask.to(self.device_torch) + + x0_pred = self.model( + hidden_states=latent_model_input.to(self.device_torch, self.torch_dtype), + timestep=t01, + encoder_hidden_states=feats, + attention_mask=mask, + return_dict=False, + )[0] + return x0_pred + + def get_prompt_embeds(self, prompt) -> PromptEmbeds: + """Encode prompt text with the Qwen3-VL text tower. + + Returns a PromptEmbeds whose ``text_embeds`` is (B, L, 2048) and whose + ``attention_mask`` is (B, L). Prompts are padded to a fixed length + (PROMPT_MAX_TOKENS) so cached embeds are concatenatable for CFG. + """ + if isinstance(prompt, str): + prompt = [prompt] + + if self.text_encoder.device == torch.device("cpu"): + self.text_encoder.to(self.device_torch) + + text_inputs = self.tokenizer( + prompt, + padding="max_length", + max_length=self.max_text_length, + truncation=True, + return_attention_mask=True, + return_tensors="pt", + ) + input_ids = text_inputs.input_ids.to(self.text_encoder.device) + attention_mask = text_inputs.attention_mask.to(self.text_encoder.device) + + with torch.no_grad(): + output = self.text_encoder( + input_ids=input_ids, + attention_mask=attention_mask, + output_hidden_states=True, + ) + embeds = output["last_hidden_state"].to(self.torch_dtype) + + pe = PromptEmbeds(embeds) + # keep the mask as bool/long; it must not be dtype-cast with the embeds + pe.attention_mask = attention_mask.bool() + return pe + + def get_loss_target(self, *args, **kwargs): + """x-prediction target: the clean image (x0). + + PRXPixel predicts x0 directly, so the MSE target is simply the clean + latents (the pixel image in [-1, 1]), not the flow velocity. + """ + batch = kwargs.get("batch") + return batch.latents.detach() + + def get_latent_noise_from_latents(self, latents: torch.Tensor, noise_offset=0.0): + """Noise for the forward flow, scaled by the model's noise_scale. + + PRXPixel trains with a non-unit initial-noise std: the noise mixed into + the latents is ``randn * noise_scale``. The scheduler's add_noise then + forms ``x_t = (1 - t) * clean + t * noise``. + """ + noise = torch.randn_like(latents) * self.noise_scale + if noise_offset is not None and noise_offset != 0.0: + noise = apply_noise_offset(noise, noise_offset) + return noise + + def condition_noisy_latents(self, latents: torch.Tensor, batch): + # plain text-to-image: nothing to inject + return latents + + # ------------------------------------------------------------------ + # Saving / bookkeeping + # ------------------------------------------------------------------ + def get_model_has_grad(self): + return False + + def get_te_has_grad(self): + return False + + def save_model(self, output_path, meta, save_dtype): + """Full fine-tune save: write the transformer back in diffusers layout.""" + transformer: PRXTransformer2DModel = unwrap_model(self.model) + transformer.save_pretrained( + os.path.join(output_path, "transformer"), + safe_serialization=True, + ) + with open(os.path.join(output_path, "aitk_meta.yaml"), "w") as f: + yaml.dump(meta, f) + + def get_base_model_version(self): + return "prx_pixel" + + def get_transformer_block_names(self) -> Optional[List[str]]: + return ["blocks"] + + def convert_lora_weights_before_save(self, state_dict): + return { + k.replace("transformer.", "diffusion_model."): v + for k, v in state_dict.items() + } + + def convert_lora_weights_before_load(self, state_dict): + return { + k.replace("diffusion_model.", "transformer."): v + for k, v in state_dict.items() + } diff --git a/extensions_built_in/diffusion_models/prx_pixel_t2i/src/__init__.py b/extensions_built_in/diffusion_models/prx_pixel_t2i/src/__init__.py new file mode 100644 index 0000000..fa89452 --- /dev/null +++ b/extensions_built_in/diffusion_models/prx_pixel_t2i/src/__init__.py @@ -0,0 +1,4 @@ +# Everything diffusers does NOT ship (until PR #13928 lands) lives in src/: +# the vendored PRX transformer architecture and a minimal pixel-space sampler. +from .transformer_prx import PRXTransformer2DModel +from .pipeline import PRXPixelPipeline diff --git a/extensions_built_in/diffusion_models/prx_pixel_t2i/src/pipeline.py b/extensions_built_in/diffusion_models/prx_pixel_t2i/src/pipeline.py new file mode 100644 index 0000000..a940468 --- /dev/null +++ b/extensions_built_in/diffusion_models/prx_pixel_t2i/src/pipeline.py @@ -0,0 +1,154 @@ +"""Minimal preview sampler for the pixel-space PRX (PRXPixel) model. + +ai-toolkit only uses this pipeline to render preview/sample images during +training (BaseModel.generate_images -> PRXPixelT2IModel.generate_single_image). +It does not need to be a diffusers DiffusionPipeline, and ai-toolkit always +encodes the prompts itself, so the pipeline only ever receives already-encoded +``PromptEmbeds`` (text features + attention mask), never raw text. + +PRXPixel specifics this sampler bakes in (see ../prx_pixel_t2i.py for the why): + - **pixel space**: there is no VAE. The model's "latents" are the RGB image + itself in [-1, 1]; the final "decode" is just a clamp + uint8 cast. + - **x-prediction**: the transformer predicts the clean image x0, not the + flow-matching velocity. Each step converts it to velocity + ``v = (x_t - x0) / t`` (t clamped for stability) before the scheduler step, + exactly like the diffusers PRXPixelPipeline. + - **noise_scale**: PRXPixel trains with a non-unit initial-noise std, so the + starting noise is ``randn * noise_scale`` (2.0 for the released model). + - **CFG** is applied on the x0 prediction (before the velocity conversion). +""" + +from typing import List, Optional + +import torch +from PIL import Image +from diffusers.utils.torch_utils import randn_tensor + + +# Minimum normalized timestep used when converting an x0 prediction to a flow +# velocity ``v = (x_t - x0) / t``. Mirrors the 0.05 clamp in the diffusers +# PRXPixelPipeline; without it the division blows up as t -> 0. +X_PRED_T_MIN = 0.05 + + +class PRXPixelPipeline: + """Lightweight pixel-space flow-matching sampler used for training previews.""" + + def __init__(self, model): + # ``model`` is the PRXPixelT2IModel (a BaseModel subclass), giving us + # access to model.transformer, model.decode_latents, device/dtype, the + # scheduler factory and the noise scale. + self.model = model + + @property + def device(self): + return self.model.device_torch + + def to(self, *args, **kwargs): + # BaseModel.generate_images may call pipeline.to(device); we manage + # devices through the model itself, so this is a no-op. + return self + + def set_progress_bar_config(self, **kwargs): + # called by the sampler harness (inside a try/except, so optional) + pass + + def _embeds_and_mask(self, embeds, device, dtype): + """Pull (features, attention_mask) out of a PromptEmbeds onto device/dtype. + + ``text_embeds`` is (B, L, D). The mask is (B, L) bool (1 = real token); + it is kept as long for the transformer's boolean masking and never cast + to the model dtype. + """ + feats = embeds.text_embeds.to(device, dtype=dtype) + mask = getattr(embeds, "attention_mask", None) + if mask is not None: + mask = mask.to(device) + return feats, mask + + @torch.no_grad() + def __call__( + self, + conditional_embeds, # PromptEmbeds: .text_embeds (B,L,D) + .attention_mask + unconditional_embeds, # PromptEmbeds or None (negative prompt) + height: int = 1024, + width: int = 1024, + num_inference_steps: int = 28, + guidance_scale: float = 5.0, + latents: Optional[torch.Tensor] = None, # pre-made noise, usually None + generator: Optional[ + torch.Generator + ] = None, # seeded RNG for reproducible samples + **kwargs, + ) -> List[Image.Image]: + model = self.model + device = model.device_torch + dtype = model.torch_dtype + transformer = model.transformer + + # Always sample with a FRESH scheduler -- the training scheduler is + # stateful and mutating it mid-training would corrupt the train step. + scheduler = model.get_train_scheduler() + scheduler.set_timesteps(num_inference_steps, device=device) + timesteps = scheduler.timesteps # 1000 -> 0 scale + + do_cfg = unconditional_embeds is not None and guidance_scale != 1.0 + + # 1. starting noise -- pixel space, so channels = transformer.in_channels (3), + # spatial size = the requested pixels (no VAE downsample). Scaled by the + # model's noise_scale to match the learned flow-matching trajectory. + if latents is None: + shape = (1, transformer.in_channels, height, width) + latents = randn_tensor( + shape, generator=generator, device=device, dtype=torch.float32 + ) + latents = latents * model.noise_scale + latents = latents.to(device, dtype=torch.float32) + + # 2. text features + masks + cond_feats, cond_mask = self._embeds_and_mask(conditional_embeds, device, dtype) + if do_cfg: + uncond_feats, uncond_mask = self._embeds_and_mask( + unconditional_embeds, device, dtype + ) + + # 3. denoising loop + for t in timesteps: + # scheduler timesteps are 0-1000; the transformer wants [0, 1] + t01 = (t / 1000.0).to(device).float().expand(latents.shape[0]) + + x0_cond = transformer( + hidden_states=latents.to(dtype), + timestep=t01, + encoder_hidden_states=cond_feats, + attention_mask=cond_mask, + return_dict=False, + )[0] + if do_cfg: + x0_uncond = transformer( + hidden_states=latents.to(dtype), + timestep=t01, + encoder_hidden_states=uncond_feats, + attention_mask=uncond_mask, + return_dict=False, + )[0] + # classifier-free guidance is applied on the x0 prediction + x0 = x0_uncond + guidance_scale * (x0_cond - x0_uncond) + else: + x0 = x0_cond + + # convert the x0 (clean-image) prediction to the flow velocity the + # scheduler consumes: v = (x_t - x0) / t, t clamped for stability. + t_x = torch.clamp(t01.to(torch.float32), min=X_PRED_T_MIN).view(-1, 1, 1, 1) + v = (latents - x0.to(torch.float32)) / t_x + + latents = scheduler.step(v, t, latents, return_dict=False)[0] + + # 4. pixel space: the denoised latents ARE the image in [-1, 1]. + # decode_latents is an identity (FakeVAE) but we route through it so + # any future latent normalization stays in one place. + images = model.decode_latents(latents, device=device, dtype=torch.float32) + images = images.float().clamp(-1.0, 1.0) + images = ((images + 1.0) * 127.5).round().to(torch.uint8) + images = images.permute(0, 2, 3, 1).cpu().numpy() + return [Image.fromarray(arr) for arr in images] diff --git a/extensions_built_in/diffusion_models/prx_pixel_t2i/src/transformer_prx.py b/extensions_built_in/diffusion_models/prx_pixel_t2i/src/transformer_prx.py new file mode 100644 index 0000000..e83d8a1 --- /dev/null +++ b/extensions_built_in/diffusion_models/prx_pixel_t2i/src/transformer_prx.py @@ -0,0 +1,881 @@ +# PRXTransformer2DModel -- vendored, self-contained copy of the Photoroom PRX +# transformer so ai-toolkit does not depend on an unmerged diffusers PR +# (https://github.com/huggingface/diffusers/pull/13928). +# +# This is the architecture as it appears on the ``prx-pixel`` branch, with two +# small edits so it runs against a stock installed diffusers: +# - the relative ``from ...`` / ``from ..`` imports are made absolute +# (``from diffusers...``) -- every building block used here ships in stock +# diffusers already, only the assembled PRX model is missing. +# - ``maybe_adjust_dtype_for_device`` (added in a newer diffusers) is inlined +# below so we don't depend on its presence. +# +# The pixel-space PRX-7B variant (Photoroom/prxpixel-t2i) sets +# ``in_channels=3`` (raw RGB, no VAE), ``bottleneck_size`` (two-layer patch +# projection) and ``resolution_embeds=True`` (resolution-aware modulation). +# +# Copyright 2025 The Photoroom and The HuggingFace Teams. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Any, Optional + +import torch +from torch import nn + +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.utils import logging +from diffusers.models.attention import AttentionMixin, AttentionModuleMixin +from diffusers.models.attention_dispatch import dispatch_attention_fn +from diffusers.models.embeddings import get_timestep_embedding +from diffusers.models.modeling_outputs import Transformer2DModelOutput +from diffusers.models.modeling_utils import ModelMixin +from diffusers.models.normalization import RMSNorm + + +logger = logging.get_logger(__name__) + + +def maybe_adjust_dtype_for_device( + dtype: torch.dtype, device: torch.device +) -> torch.dtype: + r""" + MPS does not implement float64; fall back to float32 there. On every other device the requested dtype is returned + unchanged. Inlined from newer diffusers so this file works on older installs. + """ + if ( + device is not None + and getattr(device, "type", None) == "mps" + and dtype == torch.float64 + ): + return torch.float32 + return dtype + + +def get_image_ids( + batch_size: int, height: int, width: int, patch_size: int, device: torch.device +) -> torch.Tensor: + r""" + 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. + """ + + img_ids = torch.zeros(height // patch_size, width // patch_size, 2, device=device) + img_ids[..., 0] = torch.arange(height // patch_size, device=device)[:, None] + img_ids[..., 1] = torch.arange(width // patch_size, device=device)[None, :] + return ( + img_ids.reshape((height // patch_size) * (width // patch_size), 2) + .unsqueeze(0) + .repeat(batch_size, 1, 1) + ) + + +def apply_rope(xq: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor: + r""" + 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. + """ + xq_ = xq.float().reshape(*xq.shape[:-1], -1, 1, 2) + # Ensure freqs_cis is on the same device as queries to avoid device mismatches with offloading + freqs_cis = freqs_cis.to(device=xq.device, dtype=xq_.dtype) + xq_out = freqs_cis[..., 0] * xq_[..., 0] + freqs_cis[..., 1] * xq_[..., 1] + return xq_out.reshape(*xq.shape).type_as(xq) + + +class PRXAttnProcessor2_0: + r""" + 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. + """ + + _attention_backend = None + _parallel_config = None + + def __init__(self): + if not hasattr(torch.nn.functional, "scaled_dot_product_attention"): + raise ImportError( + "PRXAttnProcessor2_0 requires PyTorch 2.0, please upgrade PyTorch to 2.0." + ) + + def __call__( + self, + attn: "PRXAttention", + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + image_rotary_emb: torch.Tensor | None = None, + **kwargs, + ) -> torch.Tensor: + """ + 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] + """ + + if encoder_hidden_states is None: + 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) + B, L_img, _ = img_qkv.shape + img_qkv = img_qkv.reshape(B, L_img, 3, attn.heads, attn.head_dim) + img_qkv = img_qkv.permute(2, 0, 3, 1, 4) # [3, B, H, L_img, D] + img_q, img_k, img_v = img_qkv[0], img_qkv[1], img_qkv[2] + + # Apply QK normalization to image tokens + img_q = attn.norm_q(img_q) + img_k = attn.norm_k(img_k) + + # Project text tokens to K, V + txt_kv = attn.txt_kv_proj(encoder_hidden_states) + B, L_txt, _ = txt_kv.shape + txt_kv = txt_kv.reshape(B, L_txt, 2, attn.heads, attn.head_dim) + txt_kv = txt_kv.permute(2, 0, 3, 1, 4) # [2, B, H, L_txt, D] + txt_k, txt_v = txt_kv[0], txt_kv[1] + + # Apply K normalization to text tokens + txt_k = attn.norm_added_k(txt_k) + + # Apply RoPE to image queries and keys + if image_rotary_emb is not None: + img_q = apply_rope(img_q, image_rotary_emb) + img_k = apply_rope(img_k, image_rotary_emb) + + # Concatenate text and image keys/values + k = torch.cat((txt_k, img_k), dim=2) # [B, H, L_txt + L_img, D] + v = torch.cat((txt_v, img_v), dim=2) # [B, H, L_txt + L_img, D] + + # Build attention mask if provided + attn_mask_tensor = None + if attention_mask is not None: + bs, _, l_img, _ = img_q.shape + l_txt = txt_k.shape[2] + + if attention_mask.dim() != 2: + raise ValueError( + f"Unsupported attention_mask shape: {attention_mask.shape}" + ) + if attention_mask.shape[-1] != l_txt: + raise ValueError( + f"attention_mask last dim {attention_mask.shape[-1]} must equal text length {l_txt}" + ) + + device = img_q.device + ones_img = torch.ones((bs, l_img), dtype=torch.bool, device=device) + attention_mask = attention_mask.to(device=device, dtype=torch.bool) + joint_mask = torch.cat([attention_mask, ones_img], dim=-1) + attn_mask_tensor = joint_mask[:, None, None, :].expand( + -1, attn.heads, l_img, -1 + ) + + # Apply attention using dispatch_attention_fn for backend support + # Reshape to match dispatch_attention_fn expectations: [B, L, H, D] + query = img_q.transpose(1, 2) # [B, L_img, H, D] + key = k.transpose(1, 2) # [B, L_txt + L_img, H, D] + value = v.transpose(1, 2) # [B, L_txt + L_img, H, D] + + attn_output = dispatch_attention_fn( + query, + key, + value, + attn_mask=attn_mask_tensor, + backend=self._attention_backend, + parallel_config=self._parallel_config, + ) + + # Reshape from [B, L_img, H, D] to [B, L_img, H*D] + batch_size, seq_len, num_heads, head_dim = attn_output.shape + attn_output = attn_output.reshape(batch_size, seq_len, num_heads * head_dim) + + # Apply output projection + attn_output = attn.to_out[0](attn_output) + if len(attn.to_out) > 1: + attn_output = attn.to_out[1](attn_output) # dropout if present + + return attn_output + + +class PRXAttention(nn.Module, AttentionModuleMixin): + r""" + PRX-style attention module that handles multi-source tokens and RoPE. Similar to FluxAttention but adapted for + PRX's architecture. + """ + + _default_processor_cls = PRXAttnProcessor2_0 + _available_processors = [PRXAttnProcessor2_0] + + def __init__( + self, + query_dim: int, + heads: int = 8, + dim_head: int = 64, + bias: bool = False, + out_bias: bool = False, + eps: float = 1e-6, + processor=None, + ): + super().__init__() + + self.heads = heads + self.head_dim = dim_head + self.inner_dim = dim_head * heads + self.query_dim = query_dim + + self.img_qkv_proj = nn.Linear(query_dim, query_dim * 3, bias=bias) + + self.norm_q = RMSNorm(self.head_dim, eps=eps, elementwise_affine=True) + self.norm_k = RMSNorm(self.head_dim, eps=eps, elementwise_affine=True) + + self.txt_kv_proj = nn.Linear(query_dim, query_dim * 2, bias=bias) + self.norm_added_k = RMSNorm(self.head_dim, eps=eps, elementwise_affine=True) + + self.to_out = nn.ModuleList([]) + self.to_out.append(nn.Linear(self.inner_dim, query_dim, bias=out_bias)) + self.to_out.append(nn.Dropout(0.0)) + + if processor is None: + processor = self._default_processor_cls() + self.set_processor(processor) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + image_rotary_emb: torch.Tensor | None = None, + **kwargs, + ) -> torch.Tensor: + return self.processor( + self, + hidden_states, + encoder_hidden_states=encoder_hidden_states, + attention_mask=attention_mask, + image_rotary_emb=image_rotary_emb, + **kwargs, + ) + + +# inspired from https://github.com/black-forest-labs/flux/blob/main/src/flux/modules/layers.py +class PRXEmbedND(nn.Module): + r""" + 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). + """ + + def __init__(self, dim: int, theta: int, axes_dim: list[int]): + super().__init__() + self.dim = dim + self.theta = theta + self.axes_dim = axes_dim + + def rope(self, pos: torch.Tensor, dim: int, theta: int) -> torch.Tensor: + assert dim % 2 == 0 + + dtype = maybe_adjust_dtype_for_device(torch.float64, pos.device) + + scale = torch.arange(0, dim, 2, dtype=dtype, device=pos.device) / dim + omega = 1.0 / (theta**scale) + out = pos.unsqueeze(-1) * omega.unsqueeze(0) + out = torch.stack( + [torch.cos(out), -torch.sin(out), torch.sin(out), torch.cos(out)], dim=-1 + ) + # Native PyTorch equivalent of: Rearrange("b n d (i j) -> b n d i j", i=2, j=2) + # out shape: (b, n, d, 4) -> reshape to (b, n, d, 2, 2) + out = out.reshape(*out.shape[:-1], 2, 2) + return out.float() + + def forward(self, ids: torch.Tensor) -> torch.Tensor: + n_axes = ids.shape[-1] + emb = torch.cat( + [ + self.rope(ids[:, :, i], self.axes_dim[i], self.theta) + for i in range(n_axes) + ], + dim=-3, + ) + return emb.unsqueeze(1) + + +class MLPEmbedder(nn.Module): + r""" + 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. + """ + + def __init__(self, in_dim: int, hidden_dim: int): + super().__init__() + self.in_layer = nn.Linear(in_dim, hidden_dim, bias=True) + self.silu = nn.SiLU() + self.out_layer = nn.Linear(hidden_dim, hidden_dim, bias=True) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.out_layer(self.silu(self.in_layer(x))) + + +class PRXResolutionEmbedder(nn.Module): + r""" + Embeds the spatial resolution `(height, width)` of the latent into a vector that is added to the timestep + embedding, so the model can condition its modulation on the generation resolution. + + A sinusoidal embedding of dimension 128 is built for the height and the width separately and concatenated into a + 256-dim vector, which is then projected to `hidden_size` by a 2-layer MLP. This matches the `"vec"` mode of the + resolution-aware conditioning used during PRX-7B training. + + Args: + hidden_size (`int`): + Dimension of the output embedding (must match the timestep embedding dimension). + max_period (`int`, *optional*, defaults to 10000): + Maximum frequency period for the sinusoidal resolution embedding. + """ + + def __init__(self, hidden_size: int, max_period: int = 10000): + super().__init__() + self.max_period = max_period + self.mlp = MLPEmbedder(in_dim=256, hidden_dim=hidden_size) + + def forward( + self, height: torch.Tensor, width: torch.Tensor, dtype: torch.dtype + ) -> torch.Tensor: + h_emb = get_timestep_embedding( + timesteps=height, + embedding_dim=128, + max_period=self.max_period, + scale=1.0, + flip_sin_to_cos=True, + downscale_freq_shift=0.0, + ) + w_emb = get_timestep_embedding( + timesteps=width, + embedding_dim=128, + max_period=self.max_period, + scale=1.0, + flip_sin_to_cos=True, + downscale_freq_shift=0.0, + ) + hw_emb = torch.cat([h_emb, w_emb], dim=-1).to(dtype) + return self.mlp(hw_emb) + + +class Modulation(nn.Module): + r""" + 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)`. + """ + + def __init__(self, dim: int): + super().__init__() + self.lin = nn.Linear(dim, 6 * dim, bias=True) + nn.init.constant_(self.lin.weight, 0) + nn.init.constant_(self.lin.bias, 0) + + def forward( + self, vec: torch.Tensor + ) -> tuple[ + tuple[torch.Tensor, torch.Tensor, torch.Tensor], + tuple[torch.Tensor, torch.Tensor, torch.Tensor], + ]: + out = self.lin(nn.functional.silu(vec))[:, None, :].chunk(6, dim=-1) + return tuple(out[:3]), tuple(out[3:]) + + +class PRXBlock(nn.Module): + r""" + 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``. + """ + + def __init__( + self, + hidden_size: int, + num_heads: int, + mlp_ratio: float = 4.0, + qk_scale: float | None = None, + ): + super().__init__() + + self.hidden_dim = hidden_size + self.num_heads = num_heads + self.head_dim = hidden_size // num_heads + self.scale = qk_scale or self.head_dim**-0.5 + + self.mlp_hidden_dim = int(hidden_size * mlp_ratio) + self.hidden_size = hidden_size + + # Pre-attention normalization for image tokens + self.img_pre_norm = nn.LayerNorm( + hidden_size, elementwise_affine=False, eps=1e-6 + ) + + # 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=PRXAttnProcessor2_0(), + ) + + # mlp + self.post_attention_layernorm = nn.LayerNorm( + hidden_size, elementwise_affine=False, eps=1e-6 + ) + self.gate_proj = nn.Linear(hidden_size, self.mlp_hidden_dim, bias=False) + self.up_proj = nn.Linear(hidden_size, self.mlp_hidden_dim, bias=False) + self.down_proj = nn.Linear(self.mlp_hidden_dim, hidden_size, bias=False) + self.mlp_act = nn.GELU(approximate="tanh") + + self.modulation = Modulation(hidden_size) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + temb: torch.Tensor, + image_rotary_emb: torch.Tensor, + attention_mask: torch.Tensor | None = None, + **kwargs: dict[str, Any], + ) -> torch.Tensor: + r""" + 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. + + Returns: + `torch.Tensor`: + Updated image tokens of shape `(B, L_img, hidden_size)`. + """ + + mod_attn, mod_mlp = self.modulation(temb) + attn_shift, attn_scale, attn_gate = mod_attn + mlp_shift, mlp_scale, mlp_gate = mod_mlp + + hidden_states_mod = (1 + attn_scale) * self.img_pre_norm( + hidden_states + ) + attn_shift + + attn_out = self.attention( + hidden_states=hidden_states_mod, + encoder_hidden_states=encoder_hidden_states, + attention_mask=attention_mask, + image_rotary_emb=image_rotary_emb, + ) + + hidden_states = hidden_states + attn_gate * attn_out + + x = (1 + mlp_scale) * self.post_attention_layernorm(hidden_states) + mlp_shift + hidden_states = hidden_states + mlp_gate * ( + self.down_proj(self.mlp_act(self.gate_proj(x)) * self.up_proj(x)) + ) + return hidden_states + + +class FinalLayer(nn.Module): + r""" + 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). + """ + + def __init__(self, hidden_size: int, patch_size: int, out_channels: int): + super().__init__() + self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.linear = nn.Linear( + hidden_size, patch_size * patch_size * out_channels, bias=True + ) + self.adaLN_modulation = nn.Sequential( + nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size, bias=True) + ) + + def forward(self, x: torch.Tensor, vec: torch.Tensor) -> torch.Tensor: + shift, scale = self.adaLN_modulation(vec).chunk(2, dim=1) + x = (1 + scale[:, None, :]) * self.norm_final(x) + shift[:, None, :] + x = self.linear(x) + return x + + +def img2seq(img: torch.Tensor, patch_size: int) -> torch.Tensor: + r""" + 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. + """ + b, c, h, w = img.shape + p = patch_size + + # Reshape to (B, C, H//p, p, W//p, p) separating grid and patch dimensions + img = img.reshape(b, c, h // p, p, w // p, p) + + # Permute to (B, H//p, W//p, C, p, p) using einsum + # n=batch, c=channels, h=grid_height, p=patch_height, w=grid_width, q=patch_width + img = torch.einsum("nchpwq->nhwcpq", img) + + # Flatten to (B, L, C * p * p) + img = img.reshape(b, -1, c * p * p) + return img + + +def seq2img(seq: torch.Tensor, patch_size: int, shape: torch.Tensor) -> torch.Tensor: + r""" + 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)`. + """ + if isinstance(shape, tuple): + h, w = shape[-2:] + elif isinstance(shape, torch.Tensor): + h, w = (int(shape[0]), int(shape[1])) + else: + raise NotImplementedError(f"shape type {type(shape)} not supported") + + b, l, d = seq.shape + p = patch_size + c = d // (p * p) + + # Reshape back to grid structure: (B, H//p, W//p, C, p, p) + seq = seq.reshape(b, h // p, w // p, c, p, p) + + # Permute back to image layout: (B, C, H//p, p, W//p, p) + # n=batch, h=grid_height, w=grid_width, c=channels, p=patch_height, q=patch_width + seq = torch.einsum("nhwcpq->nchpwq", seq) + + # Final reshape to (B, C, H, W) + seq = seq.reshape(b, c, h, w) + return seq + + +class PRXTransformer2DModel(ModelMixin, ConfigMixin, AttentionMixin): + r""" + 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. The pixel-space PRX uses `3` (raw RGB, no VAE). + 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. + bottleneck_size (`int`, *optional*): + If set, the image patch projection (`img_in`) uses a two-layer bottleneck (`patch_dim -> bottleneck_size -> + hidden_size`) instead of a single linear layer. Used by the pixel-space PRX-7B variant where the patch + dimension is large. + resolution_embeds (`bool`, *optional*, defaults to `False`): + Whether to condition the timestep modulation on the latent resolution `(H, W)` via a + `PRXResolutionEmbedder`. Used by the PRX-7B variant. + """ + + config_name = "config.json" + _supports_gradient_checkpointing = True + + @register_to_config + def __init__( + self, + in_channels: int = 16, + patch_size: int = 2, + context_in_dim: int = 2304, + hidden_size: int = 1792, + mlp_ratio: float = 3.5, + num_heads: int = 28, + depth: int = 16, + axes_dim: list = None, + theta: int = 10000, + time_factor: float = 1000.0, + time_max_period: int = 10000, + bottleneck_size: Optional[int] = None, + resolution_embeds: bool = False, + ): + super().__init__() + + if axes_dim is None: + axes_dim = [32, 32] + + # Store parameters directly + self.in_channels = in_channels + self.patch_size = patch_size + self.out_channels = self.in_channels * self.patch_size**2 + + self.time_factor = time_factor + self.time_max_period = time_max_period + + if hidden_size % num_heads != 0: + raise ValueError( + f"Hidden size {hidden_size} must be divisible by num_heads {num_heads}" + ) + + pe_dim = hidden_size // num_heads + + if sum(axes_dim) != pe_dim: + raise ValueError(f"Got {axes_dim} but expected positional dim {pe_dim}") + + self.hidden_size = hidden_size + self.num_heads = num_heads + self.pe_embedder = PRXEmbedND(dim=pe_dim, theta=theta, axes_dim=axes_dim) + patch_dim = self.in_channels * self.patch_size**2 + if bottleneck_size is not None: + # Two-layer bottleneck projection (used by pixel-space PRX where the patch dimension is large). + self.img_in = nn.Sequential( + nn.Linear(patch_dim, bottleneck_size, bias=True), + nn.Linear(bottleneck_size, self.hidden_size, bias=True), + ) + else: + self.img_in = nn.Linear(patch_dim, 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.resolution_embedder = ( + PRXResolutionEmbedder(self.hidden_size, max_period=time_max_period) + if resolution_embeds + else None + ) + + self.blocks = nn.ModuleList( + [ + PRXBlock( + self.hidden_size, + self.num_heads, + mlp_ratio=mlp_ratio, + ) + for i in range(depth) + ] + ) + + self.final_layer = FinalLayer(self.hidden_size, 1, self.out_channels) + + self.gradient_checkpointing = False + + def _compute_timestep_embedding( + self, timestep: torch.Tensor, dtype: torch.dtype + ) -> torch.Tensor: + return self.time_in( + get_timestep_embedding( + timesteps=timestep, + embedding_dim=256, + max_period=self.time_max_period, + scale=self.time_factor, + flip_sin_to_cos=True, # Match original cos, sin order + downscale_freq_shift=0.0, + ).to(dtype) + ) + + def forward( + self, + hidden_states: torch.Tensor, + timestep: torch.Tensor, + encoder_hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + attention_kwargs: dict[str, Any] | None = None, + return_dict: bool = True, + ) -> tuple[torch.Tensor, ...] | Transformer2DModelOutput: + r""" + 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,)` on a `[0, 1]` scale, 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. + + Returns: + `Transformer2DModelOutput` if `return_dict=True`, otherwise a tuple `(sample,)` with the predicted clean + image of shape `(B, C, H, W)`. + """ + # Process text conditioning + txt = self.txt_in(encoder_hidden_states) + + # Convert image to sequence and embed + img = img2seq(hidden_states, self.patch_size) + img = self.img_in(img) + + # Generate positional embeddings + bs, _, h, w = hidden_states.shape + img_ids = get_image_ids( + bs, h, w, patch_size=self.patch_size, device=hidden_states.device + ) + pe = self.pe_embedder(img_ids) + + # Compute time embedding + vec = self._compute_timestep_embedding(timestep, dtype=img.dtype) + + # Add resolution conditioning (PRX-7B "vec" mode): embed the latent (H, W) and add it to the timestep vector + # so every block's modulation is resolution-aware. + if self.resolution_embedder is not None: + height = torch.full( + (bs,), h, device=hidden_states.device, dtype=torch.float32 + ) + width = torch.full( + (bs,), w, device=hidden_states.device, dtype=torch.float32 + ) + vec = vec + self.resolution_embedder(height, width, dtype=vec.dtype) + + # Apply transformer blocks + for block in self.blocks: + if torch.is_grad_enabled() and self.gradient_checkpointing: + img = self._gradient_checkpointing_func( + block.__call__, + img, + txt, + vec, + pe, + attention_mask, + ) + else: + img = block( + hidden_states=img, + encoder_hidden_states=txt, + temb=vec, + image_rotary_emb=pe, + attention_mask=attention_mask, + ) + + # Final layer and convert back to image + img = self.final_layer(img, vec) + output = seq2img(img, self.patch_size, hidden_states.shape) + + if not return_dict: + return (output,) + return Transformer2DModelOutput(sample=output) diff --git a/ui/src/app/jobs/new/options.ts b/ui/src/app/jobs/new/options.ts index 745f285..aed063d 100644 --- a/ui/src/app/jobs/new/options.ts +++ b/ui/src/app/jobs/new/options.ts @@ -1004,6 +1004,27 @@ export const modelArchs: ModelArch[] = [ ], hasMultiLinePrompts: true, }, + { + name: 'prx_pixel', + label: 'PRXPixel (pixel space)', + group: 'image', + defaults: { + 'config.process[0].model.name_or_path': ['Photoroom/prxpixel-t2i', defaultNameOrPath], + 'config.process[0].model.quantize': [true, false], + 'config.process[0].model.quantize_te': [true, false], + 'config.process[0].train.timestep_type': ['linear', 'sigmoid'], + 'config.process[0].network.conv': [undefined, 16], + 'config.process[0].network.conv_alpha': [undefined, 16], + 'config.process[0].model.low_vram': [true, false], + }, + disableSections: [ + 'network.conv', + ], + additionalSections: [ + 'model.low_vram', + 'model.layer_offloading', + ], + }, ].sort((a, b) => { // Sort by label, case-insensitive return a.label.localeCompare(b.label, undefined, { sensitivity: 'base' });