diff --git a/src/maxdiffusion/checkpointing/ideogram_checkpointer.py b/src/maxdiffusion/checkpointing/ideogram_checkpointer.py new file mode 100644 index 000000000..cf14cd5e5 --- /dev/null +++ b/src/maxdiffusion/checkpointing/ideogram_checkpointer.py @@ -0,0 +1,113 @@ +""" +Copyright 2025 Google LLC + +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 + + https://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. +""" + +import json +import jax +import numpy as np +from typing import Optional, Tuple +from maxdiffusion.pipelines.ideogram.ideogram_pipeline import IdeogramPipeline +from maxdiffusion import max_logging +from maxdiffusion.checkpointing.checkpointing_utils import create_orbax_checkpoint_manager +import orbax.checkpoint as ocp +from etils import epath + +IDEOGRAM_CHECKPOINT = "IDEOGRAM_CHECKPOINT" + + +class IdeogramCheckpointer: + + def __init__(self, config, checkpoint_type: str = IDEOGRAM_CHECKPOINT): + self.config = config + self.checkpoint_type = checkpoint_type + self.opt_state = None + + self.checkpoint_manager: ocp.CheckpointManager = create_orbax_checkpoint_manager( + getattr(self.config, "checkpoint_dir", ""), + enable_checkpointing=True, + save_interval_steps=1, + checkpoint_type=checkpoint_type, + dataset_type=getattr(config, "dataset_type", None), + ) + + def load_ideogram_configs_from_orbax(self, step: Optional[int]) -> Tuple[Optional[dict], Optional[int]]: + if self.checkpoint_manager is None: + max_logging.log("No checkpoint manager configured, skipping Orbax load.") + return None, None + + if step is None: + step = self.checkpoint_manager.latest_step() + max_logging.log(f"Latest Ideogram checkpoint step: {step}") + if step is None: + max_logging.log("No Ideogram checkpoint found.") + return None, None + max_logging.log(f"Loading Ideogram checkpoint from step {step}") + metadatas = self.checkpoint_manager.item_metadata(step) + transformer_metadata = metadatas.ideogram_state + abstract_tree_structure_params = jax.tree_util.tree_map(ocp.utils.to_shape_dtype_struct, transformer_metadata) + params_restore = ocp.args.PyTreeRestore( + restore_args=jax.tree.map( + lambda _: ocp.RestoreArgs(restore_type=np.ndarray), + abstract_tree_structure_params, + ) + ) + + max_logging.log("Restoring Ideogram checkpoint") + restored_checkpoint = self.checkpoint_manager.restore( + directory=epath.Path(self.config.checkpoint_dir), + step=step, + args=ocp.args.Composite( + ideogram_state=params_restore, + ideogram_config=ocp.args.JsonRestore(), + ), + ) + max_logging.log(f"restored checkpoint {restored_checkpoint.keys()}") + max_logging.log(f"restored checkpoint ideogram_state {restored_checkpoint.ideogram_state.keys()}") + max_logging.log(f"optimizer found in checkpoint {'opt_state' in restored_checkpoint.ideogram_state.keys()}") + return restored_checkpoint, step + + def load_checkpoint( + self, step=None, vae_only=False, load_transformer=True + ) -> Tuple[IdeogramPipeline, Optional[dict], Optional[int]]: + restored_checkpoint, step = self.load_ideogram_configs_from_orbax(step) + opt_state = None + + if restored_checkpoint: + max_logging.log("Loading Ideogram pipeline from checkpoint") + pipeline = IdeogramPipeline.from_checkpoint(self.config, restored_checkpoint, vae_only, load_transformer) + if "opt_state" in restored_checkpoint.ideogram_state.keys(): + opt_state = restored_checkpoint.ideogram_state["opt_state"] + else: + max_logging.log("No checkpoint found, loading pipeline from pretrained hub") + pipeline = IdeogramPipeline.from_pretrained(self.config, vae_only, load_transformer) + + return pipeline, opt_state, step + + def save_checkpoint(self, train_step, pipeline: IdeogramPipeline, train_states: dict): + """Saves the training state and model configurations.""" + + def config_to_json(model_or_config): + return json.loads(model_or_config.to_json_string()) + + max_logging.log(f"Saving checkpoint for step {train_step}") + items = { + "ideogram_config": ocp.args.JsonSave(config_to_json(pipeline.transformer)), + } + + items["ideogram_state"] = ocp.args.PyTreeSave(train_states) + + # Save the checkpoint + self.checkpoint_manager.save(train_step, args=ocp.args.Composite(**items)) + max_logging.log(f"Checkpoint for step {train_step} saved.") diff --git a/src/maxdiffusion/common_types.py b/src/maxdiffusion/common_types.py index 77da6900b..2840d3dfd 100644 --- a/src/maxdiffusion/common_types.py +++ b/src/maxdiffusion/common_types.py @@ -53,6 +53,7 @@ WAN2_2 = "wan2.2" LTX2_VIDEO = "ltx2_video" LTX2_3 = "ltx2.3" +IDEOGRAM4 = "ideogram4" WAN_MODEL = WAN2_1 diff --git a/src/maxdiffusion/configs/base_ideogram.yml b/src/maxdiffusion/configs/base_ideogram.yml new file mode 100644 index 000000000..ac42604b6 --- /dev/null +++ b/src/maxdiffusion/configs/base_ideogram.yml @@ -0,0 +1,79 @@ +# Copyright 2025 Google LLC +# +# 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. + +run_name: "" +metrics_dir: "" +tensorboard_dir: "" +output_dir: "" + +model_name: "ideogram4" +pretrained_model_name_or_path: "ideogram-ai/ideogram-4-fp8" + +prompt: > + {"high_level_description": "A magical, cinematic shot of a couple holding hands on a flying carpet soaring high above a historic Arabian desert town at sunset.", "style_description": {"aesthetics": "historic, magical realism, epic fantasy, cinematic", "lighting": "golden hour, dramatic sunset lighting, long shadows, warm glowing atmosphere", "photo": "shot on 35mm lens, wide aperture, photorealistic, hyper-detailed", "medium": "digital photography, cinematic film still", "color_palette": ["#FF9B54", "#7A3B69", "#D35B40", "#2C3E50", "#F4D03F"]}, "compositional_deconstruction": {"background": "A vast, sweeping desert landscape with historic Arabian architecture, sand dunes, and a vibrant sunset sky.", "elements": [{"type": "obj", "bbox": [200, 300, 700, 700], "desc": "A richly detailed, ornate flying carpet suspended in mid-air."}, {"type": "obj", "bbox": [100, 350, 600, 650], "desc": "A romantic couple sitting on the flying carpet, holding hands and looking out over the city."}, {"type": "obj", "bbox": [600, 0, 1000, 1000], "desc": "An ancient Arabian desert town with sandstone buildings, domes, and bustling streets far below."}]}} +negative_prompt: "" +height: 1024 +width: 1024 +num_inference_steps: 50 +guidance_scale: 7.0 +seed: 42 + +global_batch_size_to_train_on: 1 +per_device_batch_size: 1 + +enable_profiler: False +enable_ml_diagnostics: False +profiler_steps: 5 +skip_jax_distributed_system: False +skip_jax_compilation: False +hardware: "tpu" +jax_cache_dir: "" +weights_dtype: "bfloat16" +activations_dtype: "bfloat16" +save_config_to_gcs: False +attention: "flash" +attention_sharding_uniform: True +mesh_axes: ['data', 'fsdp', 'context', 'tensor'] +logical_axis_rules: [ + ['batch', 'data'], + ['activation_heads', 'fsdp'], + ['activation_batch', 'data'], + ['activation_kv', 'tensor'], + ['mlp','tensor'], + ['embed','fsdp'], + ['heads', 'tensor'], + ['norm', 'fsdp'], + ['conv_batch', ['data','fsdp']], + ['out_channels', 'tensor'], + ['conv_out', 'fsdp'], + ['conv_in', 'fsdp'] + ] +data_sharding: [['data', 'fsdp', 'context', 'tensor']] +dcn_data_parallelism: 1 +dcn_fsdp_parallelism: -1 +dcn_context_parallelism: 1 +dcn_tensor_parallelism: 1 +ici_data_parallelism: 1 +ici_fsdp_parallelism: -1 +ici_context_parallelism: 1 +ici_tensor_parallelism: 1 +allow_split_physical_axes: False +learning_rate_schedule_steps: -1 +max_train_steps: 500 +unet_checkpoint: '' +dataset_name: '' +dataset_save_location: '' +compile_topology_num_slices: -1 +quantization_local_shard_count: -1 +use_qwix_quantization: False diff --git a/src/maxdiffusion/generate_ideogram.py b/src/maxdiffusion/generate_ideogram.py new file mode 100644 index 000000000..bc7193ff2 --- /dev/null +++ b/src/maxdiffusion/generate_ideogram.py @@ -0,0 +1,254 @@ +# Copyright 2025 Google LLC +# +# 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 Sequence +import jax +from jax.sharding import Mesh + +import time + +import subprocess +import numpy as np +from PIL import Image +from flax import nnx +from absl import app + +from maxdiffusion import pyconfig, max_logging, max_utils +from maxdiffusion.checkpointing.ideogram_checkpointer import IdeogramCheckpointer + + +def _add_sharding_rule(vs: nnx.Variable, logical_axis_rules) -> nnx.Variable: + vs.set_metadata(sharding_rules=logical_axis_rules) + return vs + + +def create_sharded_logical_model(model, logical_axis_rules, mesh): + if model is None: + return None + graphdef, state, rest_of_state = nnx.split(model, nnx.Param, ...) + + def map_leaf(path, leaf): + if not isinstance(leaf, nnx.Variable): + return jax.sharding.PartitionSpec() + path_str = ".".join([str(p.key) if hasattr(p, "key") else str(p) for p in path]) + # Manually implement FSDP by matching layer names since nnx.Linear lacks axis_names + if "qkv.kernel" in path_str: + return jax.sharding.PartitionSpec("fsdp", None) + elif "o.kernel" in path_str: + return jax.sharding.PartitionSpec(None, "fsdp") + elif "w1.kernel" in path_str or "w3.kernel" in path_str: + return jax.sharding.PartitionSpec("fsdp", None) + elif "w2.kernel" in path_str: + return jax.sharding.PartitionSpec(None, "fsdp") + elif "adaln_modulation.kernel" in path_str: + return jax.sharding.PartitionSpec(None, "fsdp") + elif "final_layer.linear.kernel" in path_str: + return jax.sharding.PartitionSpec("fsdp", None) + elif "input_proj.kernel" in path_str: + return jax.sharding.PartitionSpec(None, "fsdp") + elif "llm_cond_proj.kernel" in path_str: + return jax.sharding.PartitionSpec(None, "fsdp") + elif "embed_image_indicator.embedding" in path_str: + return jax.sharding.PartitionSpec(None, "fsdp") + + # Fallback to replicated for small arrays like biases or norms + leaf_shape = leaf.shape if hasattr(leaf, "shape") else leaf.value.shape + if len(leaf_shape) == 1: + return jax.sharding.PartitionSpec(None) + elif len(leaf_shape) == 2: + return jax.sharding.PartitionSpec(None, None) + else: + return jax.sharding.PartitionSpec(*([None] * len(leaf_shape))) + + pspecs = jax.tree_util.tree_map_with_path(map_leaf, state, is_leaf=lambda x: isinstance(x, nnx.Variable)) + + sharded_state = jax.tree.map( + lambda x, p: x.replace( + value=jax.device_put(x.get_value() if hasattr(x, "get_value") else x.value, jax.sharding.NamedSharding(mesh, p)) + ), + state, + pspecs, + is_leaf=lambda x: isinstance(x, nnx.Variable), + ) + model = nnx.merge(graphdef, sharded_state, rest_of_state) + return model + + +def get_git_commit_hash(): + try: + commit_hash = subprocess.check_output(["git", "rev-parse", "HEAD"]).strip().decode("utf-8") + return commit_hash + except subprocess.CalledProcessError: + max_logging.log("Warning: 'git rev-parse HEAD' failed.") + return None + except FileNotFoundError: + max_logging.log("Warning: 'git' command not found.") + return None + + +jax.config.update("jax_use_shardy_partitioner", True) + + +def call_pipeline(config, pipeline, prompt, negative_prompt=None): + seed = getattr(config, "seed", 42) + height = getattr(config, "height", 256) + width = getattr(config, "width", 256) + num_inference_steps = getattr(config, "num_inference_steps", 50) + guidance_scale = getattr(config, "guidance_scale", 7.0) + + # Convert single prompt to list of prompts to match pipeline batch dimension + if isinstance(prompt, str): + data_parallelism = getattr(config, "dcn_data_parallelism", 1) * getattr(config, "ici_data_parallelism", 1) + num_prompts = getattr(config, "per_device_batch_size", 1) * data_parallelism + prompts = [prompt] * num_prompts + if negative_prompt is None: + negative_prompts = [""] * num_prompts + elif isinstance(negative_prompt, str): + negative_prompts = [negative_prompt] * num_prompts + else: + negative_prompts = negative_prompt + else: + prompts = prompt + negative_prompts = negative_prompt + + images = pipeline.generate( + prompts=prompts, + negative_prompts=negative_prompts, + height=height, + width=width, + num_steps=num_inference_steps, + guidance_scale=guidance_scale, + seed=seed, + ) + return images + + +def run(config, filename_prefix="", commit_hash=None): + writer = max_utils.initialize_summary_writer(config) + if jax.process_index() == 0 and writer: + max_logging.log(f"TensorBoard logs will be written to: {config.tensorboard_dir}") + if commit_hash: + writer.add_text("inference/git_commit_hash", commit_hash, global_step=0) + max_logging.log(f"Git Commit Hash: {commit_hash}") + + t0_load = time.perf_counter() + max_logging.log("Loading pipeline weights for Ideogram via checkpointer...") + + checkpointer = IdeogramCheckpointer(config) + pipeline, _, _ = checkpointer.load_checkpoint(load_transformer=True) + + load_time = time.perf_counter() - t0_load + max_logging.log(f"Model loaded: {load_time:.1f}s") + + # Apply sharding over the device mesh + max_logging.log("Applying sharding constraints to models...") + devices_array = max_utils.create_device_mesh(config) + mesh = Mesh(devices_array, config.mesh_axes) + logical_axis_rules = tuple(tuple(rule) for rule in config.logical_axis_rules) + with mesh: + pipeline.conditional_transformer = create_sharded_logical_model( + pipeline.conditional_transformer, logical_axis_rules, mesh + ) + pipeline.unconditional_transformer = create_sharded_logical_model( + pipeline.unconditional_transformer, logical_axis_rules, mesh + ) + pipeline.autoencoder = create_sharded_logical_model(pipeline.autoencoder, logical_axis_rules, mesh) + + s0 = time.perf_counter() + prompt = getattr(config, "prompt", "A cute dog") + negative_prompt = getattr(config, "negative_prompt", "") + + max_logging.log(f"Num steps: {config.num_inference_steps}, height: {config.height}, width: {config.width}") + max_logging.log("===================== Model details =======================") + max_logging.log(f"hardware: {jax.devices()[0].platform}") + max_logging.log(f"number of devices: {jax.device_count()}") + max_logging.log("============================================================") + + original_enable_profiler = config.get_keys().get("enable_profiler", False) + original_enable_mld = config.get_keys().get("enable_ml_diagnostics", False) + original_num_steps = config.get_keys().get("num_inference_steps", 40) + + # 1. Warmup Compilation + config.get_keys()["enable_profiler"] = False + config.get_keys()["enable_ml_diagnostics"] = False + config.get_keys()["num_inference_steps"] = 2 # lower for warmup + + max_logging.log("🚀 Starting warmup compilation pass (2 steps)...") + with mesh: + _ = call_pipeline(config, pipeline, prompt, negative_prompt) + + compile_time = time.perf_counter() - s0 + max_logging.log(f"compile_time: {compile_time}") + + # 2. Actual Generation + config.get_keys()["num_inference_steps"] = original_num_steps + s0 = time.perf_counter() + max_logging.log(f"🚀 Starting actual full-length generation pass ({original_num_steps} steps)...") + with mesh: + out_images = call_pipeline(config, pipeline, prompt, negative_prompt) + generation_time = time.perf_counter() - s0 + max_logging.log(f"generation_time: {generation_time}") + + # Save images + saved_image_paths = [] + actual_prefix = filename_prefix + if not actual_prefix and getattr(config, "run_name", None): + actual_prefix = getattr(config, "run_name") + "_" + + for i in range(len(out_images)): + image_path = f"{actual_prefix}ideogram_output_{getattr(config, 'seed', 42)}_{i}.png" + image_np = np.array(out_images[i]) + image_np = (image_np * 255).astype(np.uint8) + img = Image.fromarray(image_np) + img.save(image_path) + saved_image_paths.append(image_path) + max_logging.log(f"Saved image to {image_path}") + + timing_str = ( + f"\n{'=' * 50}\n" + f" TIMING SUMMARY\n" + f"{'=' * 50}\n" + f" Load (checkpoint): {load_time:>7.1f}s\n" + f" Compile: {compile_time:>7.1f}s\n" + f" {'─' * 40}\n" + f" Inference: {generation_time:>7.1f}s\n" + f"{'=' * 50}" + ) + max_logging.log(timing_str) + + # 3. Profiling Run + if original_enable_profiler or original_enable_mld: + profiling_steps = config.get_keys().get("profiler_steps", 5) + config.get_keys()["enable_profiler"] = original_enable_profiler + config.get_keys()["enable_ml_diagnostics"] = original_enable_mld + config.get_keys()["num_inference_steps"] = profiling_steps + + max_logging.log(f"🚀 Starting Profiling run ({profiling_steps} steps)...") + profiler = max_utils.Profiler(config, session_name=f"denoise_profile_{profiling_steps}_steps") + profiler.start() + _ = call_pipeline(config, pipeline, prompt, negative_prompt) + profiler.stop() + + return saved_image_paths + + +def main(argv: Sequence[str]) -> None: + commit_hash = get_git_commit_hash() + pyconfig.initialize(argv) + max_utils.ensure_machinelearning_job_runs(pyconfig.config) + run(pyconfig.config, commit_hash=commit_hash) + + +if __name__ == "__main__": + app.run(main) diff --git a/src/maxdiffusion/models/__init__.py b/src/maxdiffusion/models/__init__.py index 7ff8fd8fb..7240a2c81 100644 --- a/src/maxdiffusion/models/__init__.py +++ b/src/maxdiffusion/models/__init__.py @@ -30,6 +30,8 @@ from .lora import * from .flux.transformers.transformer_flux_flax import FluxTransformer2DModel from .ltx_video.transformers.transformer3d import Transformer3DModel + from .ideogram.transformer_ideogram import Ideogram4Transformer + from .ideogram.autoencoder_ideogram import AutoEncoder else: import sys diff --git a/src/maxdiffusion/models/ideogram/__init__.py b/src/maxdiffusion/models/ideogram/__init__.py new file mode 100644 index 000000000..642d03fcd --- /dev/null +++ b/src/maxdiffusion/models/ideogram/__init__.py @@ -0,0 +1,4 @@ +from .transformer_ideogram import Ideogram4Transformer, Ideogram4Config, Ideogram4TransformerBlock +from .autoencoder_ideogram import AutoEncoder, AutoEncoderParams +from .constants import * +from .ideogram_utils import * diff --git a/src/maxdiffusion/models/ideogram/autoencoder_ideogram.py b/src/maxdiffusion/models/ideogram/autoencoder_ideogram.py new file mode 100644 index 000000000..4ec6cd12b --- /dev/null +++ b/src/maxdiffusion/models/ideogram/autoencoder_ideogram.py @@ -0,0 +1,308 @@ +from typing import List +import jax +import jax.numpy as jnp +from flax import nnx + +from dataclasses import dataclass + + +@dataclass +class AutoEncoderParams: + resolution: int = 256 + in_channels: int = 3 + ch: int = 128 + out_ch: int = 3 + ch_mult: List[int] = (1, 2, 4, 4) + num_res_blocks: int = 2 + z_channels: int = 32 + + +def swish(x: jax.Array) -> jax.Array: + return x * jax.nn.sigmoid(x) + + +class AttnBlock(nnx.Module): + + def __init__(self, rngs: nnx.Rngs, in_channels: int): + self.in_channels = in_channels + self.norm = nnx.GroupNorm(in_channels, num_groups=32, epsilon=1e-6, use_bias=True, use_scale=True, rngs=rngs) + self.q = nnx.Conv(in_channels, in_channels, kernel_size=(1, 1), rngs=rngs) + self.k = nnx.Conv(in_channels, in_channels, kernel_size=(1, 1), rngs=rngs) + self.v = nnx.Conv(in_channels, in_channels, kernel_size=(1, 1), rngs=rngs) + self.proj_out = nnx.Conv(in_channels, in_channels, kernel_size=(1, 1), rngs=rngs) + + def attention(self, h_: jax.Array) -> jax.Array: + h_ = self.norm(h_) + q = self.q(h_) + k = self.k(h_) + v = self.v(h_) + + b, h_dim, w_dim, c = q.shape + q = q.reshape((b, h_dim * w_dim, c)) + k = k.reshape((b, h_dim * w_dim, c)) + v = v.reshape((b, h_dim * w_dim, c)) + + attn_weights = jnp.einsum("bqc,bkc->bqk", q, k) * (c**-0.5) + attn_weights = jax.nn.softmax(attn_weights, axis=-1) + h_ = jnp.einsum("bqk,bkc->bqc", attn_weights, v) + + return h_.reshape((b, h_dim, w_dim, c)) + + def __call__(self, x: jax.Array) -> jax.Array: + return x + self.proj_out(self.attention(x)) + + +class ResnetBlock(nnx.Module): + + def __init__(self, rngs: nnx.Rngs, in_channels: int, out_channels: int): + self.in_channels = in_channels + self.out_channels = out_channels + + self.norm1 = nnx.GroupNorm(in_channels, num_groups=32, epsilon=1e-6, rngs=rngs) + self.conv1 = nnx.Conv(in_channels, out_channels, kernel_size=(3, 3), strides=(1, 1), padding=((1, 1), (1, 1)), rngs=rngs) + + self.norm2 = nnx.GroupNorm(out_channels, num_groups=32, epsilon=1e-6, rngs=rngs) + self.conv2 = nnx.Conv( + out_channels, out_channels, kernel_size=(3, 3), strides=(1, 1), padding=((1, 1), (1, 1)), rngs=rngs + ) + + if self.in_channels != self.out_channels: + self.nin_shortcut = nnx.Conv(in_channels, out_channels, kernel_size=(1, 1), strides=(1, 1), padding=0, rngs=rngs) + + def __call__(self, x: jax.Array) -> jax.Array: + h = x + h = self.norm1(h) + h = swish(h) + h = self.conv1(h) + + h = self.norm2(h) + h = swish(h) + h = self.conv2(h) + + if self.in_channels != self.out_channels: + x = self.nin_shortcut(x) + + return x + h + + +class Downsample(nnx.Module): + + def __init__(self, rngs: nnx.Rngs, in_channels: int): + self.conv = nnx.Conv(in_channels, in_channels, kernel_size=(3, 3), strides=(2, 2), padding=0, rngs=rngs) + + def __call__(self, x: jax.Array) -> jax.Array: + x = jnp.pad(x, ((0, 0), (0, 1), (0, 1), (0, 0)), mode="constant") + return self.conv(x) + + +class Upsample(nnx.Module): + + def __init__(self, rngs: nnx.Rngs, in_channels: int): + self.conv = nnx.Conv(in_channels, in_channels, kernel_size=(3, 3), strides=(1, 1), padding=((1, 1), (1, 1)), rngs=rngs) + + def __call__(self, x: jax.Array) -> jax.Array: + b, h, w, c = x.shape + x = jax.image.resize(x, (b, h * 2, w * 2, c), method="nearest") + return self.conv(x) + + +class Encoder(nnx.Module): + + def __init__( + self, + rngs: nnx.Rngs, + resolution: int, + in_channels: int, + ch: int, + ch_mult: List[int], + num_res_blocks: int, + z_channels: int, + ): + self.quant_conv = nnx.Conv(2 * z_channels, 2 * z_channels, kernel_size=(1, 1), rngs=rngs) + self.ch = ch + self.num_resolutions = len(ch_mult) + self.num_res_blocks = num_res_blocks + self.resolution = resolution + self.in_channels = in_channels + self.conv_in = nnx.Conv(in_channels, self.ch, kernel_size=(3, 3), strides=(1, 1), padding=((1, 1), (1, 1)), rngs=rngs) + + curr_res = resolution + in_ch_mult = (1,) + tuple(ch_mult) + + down_blocks = [] + down_attns = [] + downsamples = [] + + block_in = self.ch + for i_level in range(self.num_resolutions): + block = [] + attn = [] + block_in = ch * in_ch_mult[i_level] + block_out = ch * ch_mult[i_level] + for _ in range(self.num_res_blocks): + block.append(ResnetBlock(rngs, in_channels=block_in, out_channels=block_out)) + block_in = block_out + + downsample = None + if i_level != self.num_resolutions - 1: + downsample = Downsample(rngs, block_in) + curr_res = curr_res // 2 + + down_blocks.append(nnx.List(block)) + down_attns.append(nnx.List(attn)) + downsamples.append(downsample) + + self.down_blocks = nnx.List(down_blocks) + self.down_attns = nnx.List(down_attns) + self.downsamples = nnx.List(downsamples) + + self.mid_block_1 = ResnetBlock(rngs, in_channels=block_in, out_channels=block_in) + self.mid_attn_1 = AttnBlock(rngs, block_in) + self.mid_block_2 = ResnetBlock(rngs, in_channels=block_in, out_channels=block_in) + + self.norm_out = nnx.GroupNorm(block_in, num_groups=32, epsilon=1e-6, rngs=rngs) + self.conv_out = nnx.Conv( + block_in, 2 * z_channels, kernel_size=(3, 3), strides=(1, 1), padding=((1, 1), (1, 1)), rngs=rngs + ) + + def __call__(self, x: jax.Array) -> jax.Array: + hs = [self.conv_in(x)] + for i_level in range(self.num_resolutions): + block = self.down_blocks[i_level] + attn = self.down_attns[i_level] + downsample = self.downsamples[i_level] + for i_block in range(self.num_res_blocks): + h = block[i_block](hs[-1]) + if len(attn) > 0: + h = attn[i_block](h) + hs.append(h) + if downsample is not None: + hs.append(downsample(hs[-1])) + + h = hs[-1] + h = self.mid_block_1(h) + h = self.mid_attn_1(h) + h = self.mid_block_2(h) + + h = self.norm_out(h) + h = swish(h) + h = self.conv_out(h) + h = self.quant_conv(h) + return h + + +class Decoder(nnx.Module): + + def __init__( + self, + rngs: nnx.Rngs, + ch: int, + out_ch: int, + ch_mult: List[int], + num_res_blocks: int, + in_channels: int, + resolution: int, + z_channels: int, + ): + self.post_quant_conv = nnx.Conv(z_channels, z_channels, kernel_size=(1, 1), rngs=rngs) + self.ch = ch + self.num_resolutions = len(ch_mult) + self.num_res_blocks = num_res_blocks + self.resolution = resolution + self.in_channels = in_channels + self.ffactor = 2 ** (self.num_resolutions - 1) + + block_in = ch * ch_mult[self.num_resolutions - 1] + curr_res = resolution // 2 ** (self.num_resolutions - 1) + self.z_shape = (1, curr_res, curr_res, z_channels) + + self.conv_in = nnx.Conv(z_channels, block_in, kernel_size=(3, 3), strides=(1, 1), padding=((1, 1), (1, 1)), rngs=rngs) + + self.mid_block_1 = ResnetBlock(rngs, in_channels=block_in, out_channels=block_in) + self.mid_attn_1 = AttnBlock(rngs, block_in) + self.mid_block_2 = ResnetBlock(rngs, in_channels=block_in, out_channels=block_in) + + up_blocks = [] + up_attns = [] + upsamples = [] + + for i_level in reversed(range(self.num_resolutions)): + block = [] + attn = [] + block_out = ch * ch_mult[i_level] + for _ in range(self.num_res_blocks + 1): + block.append(ResnetBlock(rngs, in_channels=block_in, out_channels=block_out)) + block_in = block_out + + upsample = None + if i_level != 0: + upsample = Upsample(rngs, block_in) + curr_res = curr_res * 2 + + up_blocks.insert(0, nnx.List(block)) + up_attns.insert(0, nnx.List(attn)) + upsamples.insert(0, upsample) + + self.up_blocks = nnx.List(up_blocks) + self.up_attns = nnx.List(up_attns) + self.upsamples = nnx.List(upsamples) + + self.norm_out = nnx.GroupNorm(block_in, num_groups=32, epsilon=1e-6, rngs=rngs) + self.conv_out = nnx.Conv(block_in, out_ch, kernel_size=(3, 3), strides=(1, 1), padding=((1, 1), (1, 1)), rngs=rngs) + + def __call__(self, z: jax.Array) -> jax.Array: + z = self.post_quant_conv(z) + h = self.conv_in(z) + + h = self.mid_block_1(h) + h = self.mid_attn_1(h) + h = self.mid_block_2(h) + + for i_level in reversed(range(self.num_resolutions)): + block = self.up_blocks[i_level] + attn = self.up_attns[i_level] + upsample = self.upsamples[i_level] + for i_block in range(self.num_res_blocks + 1): + h = block[i_block](h) + if len(attn) > 0: + h = attn[i_block](h) + if upsample is not None: + h = upsample(h) + + h = self.norm_out(h) + h = swish(h) + h = self.conv_out(h) + return h + + +class AutoEncoder(nnx.Module): + + def __init__(self, rngs: nnx.Rngs, params: AutoEncoderParams): + self.params = params + self.encoder = Encoder( + rngs, + resolution=params.resolution, + in_channels=params.in_channels, + ch=params.ch, + ch_mult=params.ch_mult, + num_res_blocks=params.num_res_blocks, + z_channels=params.z_channels, + ) + self.decoder = Decoder( + rngs, + resolution=params.resolution, + in_channels=params.in_channels, + ch=params.ch, + out_ch=params.out_ch, + ch_mult=params.ch_mult, + num_res_blocks=params.num_res_blocks, + z_channels=params.z_channels, + ) + + def encode(self, x: jax.Array) -> jax.Array: + # Note: input x should be in NHWC for flax. + return self.encoder(x) + + def decode(self, z: jax.Array) -> jax.Array: + # Note: z should be in NHWC for flax. + return self.decoder(z) diff --git a/src/maxdiffusion/models/ideogram/constants.py b/src/maxdiffusion/models/ideogram/constants.py new file mode 100644 index 000000000..442929f26 --- /dev/null +++ b/src/maxdiffusion/models/ideogram/constants.py @@ -0,0 +1,11 @@ +SEQUENCE_PADDING_INDICATOR = -1 + +OUTPUT_IMAGE_INDICATOR = 2 +LLM_TOKEN_INDICATOR = 3 + +# Image grid coordinates start at this offset so they never collide with text token indices +# (text positions start at 0 and never exceed max_text_tokens, which is well below this). +IMAGE_POSITION_OFFSET = 65536 + +# Layers of Qwen3-VL whose hidden states are concatenated and fed to the transformer. +QWEN3_VL_ACTIVATION_LAYERS = (0, 3, 6, 9, 12, 15, 18, 21, 24, 27, 30, 33, 35) diff --git a/src/maxdiffusion/models/ideogram/ideogram_utils.py b/src/maxdiffusion/models/ideogram/ideogram_utils.py new file mode 100644 index 000000000..ba152132e --- /dev/null +++ b/src/maxdiffusion/models/ideogram/ideogram_utils.py @@ -0,0 +1,291 @@ +# pylint: disable=missing-module-docstring, missing-function-docstring, too-many-positional-arguments, consider-using-dict-items, unused-argument, unspecified-encoding +import json +import re +from typing import Optional + +import jax +import jax.numpy as jnp + +from flax.traverse_util import flatten_dict, unflatten_dict +from safetensors import safe_open +import torch + +from maxdiffusion import max_logging +from ..modeling_flax_pytorch_utils import ( + rename_key_and_reshape_tensor, + torch2jax, + validate_flax_state_dict, +) +from huggingface_hub import hf_hub_download +from huggingface_hub.utils import EntryNotFoundError + +_NUM_RESOLUTIONS = 4 + + +def _rewrite_diffusers_key(key: str) -> Optional[str]: + if key.startswith("bn."): + return key + + if key.startswith("quant_conv."): + return key.replace("quant_conv.", "encoder.quant_conv.", 1) + if key.startswith("post_quant_conv."): + return key.replace("post_quant_conv.", "decoder.post_quant_conv.", 1) + + if key == "encoder.conv_norm_out.weight": + return "encoder.norm_out.weight" + if key == "encoder.conv_norm_out.bias": + return "encoder.norm_out.bias" + if key == "decoder.conv_norm_out.weight": + return "decoder.norm_out.weight" + if key == "decoder.conv_norm_out.bias": + return "decoder.norm_out.bias" + + m = re.match(r"^(encoder|decoder)\.mid_block\.resnets\.(\d+)\.(.+)$", key) + if m: + side, idx, rest = m.group(1), int(m.group(2)), m.group(3) + rest = rest.replace("conv_shortcut", "nin_shortcut") + return f"{side}.mid_block_{idx + 1}.{rest}" + + m = re.match(r"^(encoder|decoder)\.mid_block\.attentions\.0\.(.+)$", key) + if m: + side, rest = m.group(1), m.group(2) + rest = ( + rest.replace("group_norm.", "norm.") + .replace("to_q.", "q.") + .replace("to_k.", "k.") + .replace("to_v.", "v.") + .replace("to_out.0.", "proj_out.") + ) + return f"{side}.mid_attn_1.{rest}" + + m = re.match(r"^encoder\.down_blocks\.(\d+)\.resnets\.(\d+)\.(.+)$", key) + if m: + level, res_idx, rest = m.group(1), m.group(2), m.group(3) + rest = rest.replace("conv_shortcut", "nin_shortcut") + return f"encoder.down_blocks.{level}.{res_idx}.{rest}" + + m = re.match(r"^encoder\.down_blocks\.(\d+)\.downsamplers\.0\.conv\.(.+)$", key) + if m: + return f"encoder.downsamples.{m.group(1)}.conv.{m.group(2)}" + + m = re.match(r"^decoder\.up_blocks\.(\d+)\.resnets\.(\d+)\.(.+)$", key) + if m: + diffusers_idx = int(m.group(1)) + res_idx = m.group(2) + rest = m.group(3).replace("conv_shortcut", "nin_shortcut") + return f"decoder.up_blocks.{_NUM_RESOLUTIONS - 1 - diffusers_idx}.{res_idx}.{rest}" + + m = re.match(r"^decoder\.up_blocks\.(\d+)\.upsamplers\.0\.conv\.(.+)$", key) + if m: + diffusers_idx = int(m.group(1)) + return f"decoder.upsamples.{_NUM_RESOLUTIONS - 1 - diffusers_idx}.conv.{m.group(2)}" + + if key.startswith(( + "encoder.conv_in.", + "encoder.conv_out.", + "decoder.conv_in.", + "decoder.conv_out.", + )): + return key + + return None + + +def convert_diffusers_state_dict(src: dict) -> dict: + out = {} + attn_substrings = (".mid_attn_1.",) + for src_key, tensor in src.items(): + dst_key = _rewrite_diffusers_key(src_key) + if dst_key is None: + continue + if dst_key.startswith("bn."): + continue + if any(s in dst_key for s in attn_substrings) and dst_key.endswith(".weight") and tensor.ndim == 2: + tensor = jnp.expand_dims(tensor, axis=(-1, -2)) + out[dst_key] = tensor + return out + + +def _tuple_str_to_int(tuple_key): + return tuple(int(item) if item.isdigit() else item for item in tuple_key) + + +def load_sharded_checkpoint(pretrained_model_name_or_path, subfolder, device, filename=None): + tensors = {} + + if filename is not None: + try: + ckpt_path = hf_hub_download(pretrained_model_name_or_path, subfolder=subfolder, filename=filename) + if filename.endswith(".safetensors"): + with safe_open(ckpt_path, framework="pt") as f: + for k in f.keys(): + tensors[k] = torch2jax(f.get_tensor(k)) + else: + loaded_state_dict = torch.load(ckpt_path, map_location="cpu") + for k, v in loaded_state_dict.items(): + tensors[k] = torch2jax(v) + return tensors + except EntryNotFoundError: + max_logging.log(f"Warning: Specific file {filename} not found. Falling back to default logic.") + + index_file = "model.safetensors.index.json" + try: + index_path = hf_hub_download(pretrained_model_name_or_path, subfolder=subfolder, filename=index_file) + with open(index_path, "r") as f: + index_data = json.load(f) + weight_map = index_data["weight_map"] + shards = set(weight_map.values()) + + for shard in shards: + shard_path = hf_hub_download(pretrained_model_name_or_path, subfolder=subfolder, filename=shard) + with safe_open(shard_path, framework="pt") as f: + for k in f.keys(): + tensors[k] = torch2jax(f.get_tensor(k)) + except EntryNotFoundError: + # Fallback for non-sharded model + try: + filename = "model.safetensors" + ckpt_path = hf_hub_download(pretrained_model_name_or_path, subfolder=subfolder, filename=filename) + except EntryNotFoundError: + filename = "diffusion_pytorch_model.safetensors" + ckpt_path = hf_hub_download(pretrained_model_name_or_path, subfolder=subfolder, filename=filename) + + if filename.endswith(".safetensors"): + with safe_open(ckpt_path, framework="pt") as f: + for k in f.keys(): + tensors[k] = torch2jax(f.get_tensor(k)) + else: + loaded_state_dict = torch.load(ckpt_path, map_location="cpu") + for k, v in loaded_state_dict.items(): + tensors[k] = torch2jax(v) + + return tensors + + +def rename_for_ideogram_transformer(pt_key: str) -> str: + # Rename rules specifically for ideogram. + # The pytorch model uses `transformer_blocks.0.attn.qkv.weight` + # and we map to nnx arrays. `rename_key` handles weight->kernel. + return pt_key + + +def get_key_and_value(pt_tuple_key, tensor, flax_state_dict, random_flax_state_dict, scan_layers, num_layers): + block_index = None + m = re.match(r"^transformer_blocks\.(\d+)", ".".join(pt_tuple_key)) + if m: + block_index = int(m.group(1)) + if scan_layers: + # Map transformer_blocks.N -> layers + pt_tuple_key = ("layers",) + pt_tuple_key[2:] + else: + # Map transformer_blocks.N -> layers.index + pt_tuple_key = ("layers", str(block_index)) + pt_tuple_key[2:] + + flax_key, flax_tensor = rename_key_and_reshape_tensor(pt_tuple_key, tensor, random_flax_state_dict, scan_layers) + flax_key_final = [] + for k in flax_key: + if isinstance(k, str) and k.isdigit(): + flax_key_final.append(int(k)) + else: + flax_key_final.append(k) + flax_key = tuple(flax_key_final) + + if scan_layers and block_index is not None: + if "layers" in flax_key: + if flax_key in flax_state_dict: + new_tensor = flax_state_dict[flax_key] + else: + new_tensor = jnp.zeros((num_layers,) + flax_tensor.shape, dtype=flax_tensor.dtype) + + new_tensor = new_tensor.at[block_index].set(flax_tensor) + flax_tensor = new_tensor + + return flax_key, flax_tensor + + +def load_transformer_weights( + pretrained_model_name_or_path: str, + eval_shapes: dict, + device: str, + hf_download: bool = True, + num_layers: int = 42, + scan_layers: bool = True, + subfolder: str = "transformer", +): + device = jax.local_devices(backend=device)[0] + max_logging.log(f"Load and port {pretrained_model_name_or_path} {subfolder} on {device}") + + with jax.default_device(device): + tensors = load_sharded_checkpoint(pretrained_model_name_or_path, subfolder, device) + + flax_state_dict = {} + cpu = jax.local_devices(backend="cpu")[0] + flattened_dict = flatten_dict(eval_shapes) + + random_flax_state_dict = {} + for key in flattened_dict: + random_flax_state_dict[tuple(str(item) for item in key)] = flattened_dict[key] + + for pt_key, tensor in tensors.items(): + if pt_key.endswith(".weight") and pt_key + "_scale" in tensors: + scale = tensors[pt_key + "_scale"] + if scale.ndim == 1 and tensor.ndim == 2 and scale.shape[0] == tensor.shape[0]: + scale = jnp.expand_dims(scale, axis=-1) + tensor = tensor * scale + + renamed_pt_key = pt_key + renamed_pt_key = rename_for_ideogram_transformer(renamed_pt_key) + + pt_tuple_key = tuple(renamed_pt_key.split(".")) + + flax_key, flax_tensor = get_key_and_value( + pt_tuple_key, tensor, flax_state_dict, random_flax_state_dict, scan_layers, num_layers + ) + + flax_state_dict[flax_key] = jax.device_put(jnp.asarray(flax_tensor), device=cpu) + + validate_flax_state_dict(eval_shapes, flax_state_dict) + flax_state_dict = unflatten_dict(flax_state_dict) + del tensors + jax.clear_caches() + return flax_state_dict + + +def load_vae_weights( + pretrained_model_name_or_path: str, eval_shapes: dict, device: str, hf_download: bool = True, subfolder: str = "vae" +): + device = jax.local_devices(backend=device)[0] + + max_logging.log(f"Load and port {pretrained_model_name_or_path} VAE on {device}") + + with jax.default_device(device): + tensors = load_sharded_checkpoint(pretrained_model_name_or_path, subfolder, device) + tensors = convert_diffusers_state_dict(tensors) + + flax_state_dict = {} + cpu = jax.local_devices(backend="cpu")[0] + flattened_eval = flatten_dict(eval_shapes) + + random_flax_state_dict = {} + for key in flattened_eval: + random_flax_state_dict[tuple(str(item) for item in key)] = flattened_eval[key] + + for pt_key, tensor in tensors.items(): + if pt_key.endswith(".weight") and pt_key + "_scale" in tensors: + scale = tensors[pt_key + "_scale"] + if scale.ndim == 1 and tensor.ndim == 2 and scale.shape[0] == tensor.shape[0]: + scale = jnp.expand_dims(scale, axis=-1) + tensor = tensor * scale + renamed_pt_key = pt_key + pt_tuple_key = tuple(renamed_pt_key.split(".")) + + flax_key, flax_tensor = rename_key_and_reshape_tensor(pt_tuple_key, tensor, random_flax_state_dict) + + flax_key = _tuple_str_to_int(flax_key) + flax_state_dict[flax_key] = jax.device_put(jnp.asarray(flax_tensor), device=cpu) + + validate_flax_state_dict(eval_shapes, flax_state_dict) + flax_state_dict = unflatten_dict(flax_state_dict) + del tensors + jax.clear_caches() + return flax_state_dict diff --git a/src/maxdiffusion/models/ideogram/latent_norm.py b/src/maxdiffusion/models/ideogram/latent_norm.py new file mode 100644 index 000000000..d61533ce3 --- /dev/null +++ b/src/maxdiffusion/models/ideogram/latent_norm.py @@ -0,0 +1,273 @@ +from __future__ import annotations + +import jax.numpy as jnp +from typing import Any + +LATENT_SHIFT: tuple[float, ...] = ( + 0.01984364, + 0.10149707, + 0.29689495, + 0.27188619, + -0.21445648, + -0.15979549, + 0.05021099, + -0.15083604, + -0.15360136, + -0.20131799, + 0.01922352, + 0.0622626, + 0.10140969, + -0.06739428, + 0.3758261, + -0.233712, + 0.35164491, + -0.02590912, + -0.0271935, + -0.10833897, + -0.1476848, + -0.01130957, + -0.2298372, + 0.23526423, + -0.10893522, + 0.11957631, + 0.04047799, + 0.3134589, + -0.17225064, + -0.18646109, + -0.34691978, + -0.03571246, + 0.02583857, + 0.10190072, + 0.28402294, + 0.26952152, + -0.21634675, + -0.17938656, + 0.04358909, + -0.15007621, + -0.1548502, + -0.18971131, + 0.02710861, + 0.05609494, + 0.10697846, + -0.06854968, + 0.38167698, + -0.24269937, + 0.35705471, + -0.03063305, + -0.02946109, + -0.11244286, + -0.14336038, + -0.01362137, + -0.21863696, + 0.23228983, + -0.11739769, + 0.11693044, + 0.02563311, + 0.31356594, + -0.17420591, + -0.19006285, + -0.34905377, + -0.04025005, + 0.01924137, + 0.07652984, + 0.2995608, + 0.2628057, + -0.22011674, + -0.12715361, + 0.04879879, + -0.14075719, + -0.15935895, + -0.2123584, + 0.01974813, + 0.05523547, + 0.10011992, + -0.06428964, + 0.37781868, + -0.21491644, + 0.34254215, + -0.03153528, + -0.0310082, + -0.10761415, + -0.14730405, + -0.02475182, + -0.2285588, + 0.2515081, + -0.10445128, + 0.12446, + 0.07062869, + 0.30880162, + -0.18016875, + -0.18869164, + -0.34533499, + -0.0129177, + 0.02578168, + 0.07993659, + 0.28642181, + 0.26038408, + -0.22459419, + -0.14820155, + 0.04059549, + -0.14043529, + -0.16111187, + -0.2020305, + 0.02602069, + 0.04852717, + 0.10432153, + -0.06309942, + 0.38402443, + -0.22397003, + 0.34814481, + -0.03774432, + -0.03381438, + -0.11245691, + -0.14128767, + -0.02853208, + -0.21752016, + 0.24872463, + -0.11399775, + 0.1222687, + 0.05620835, + 0.309178, + -0.18065738, + -0.19401479, + -0.34495114, + -0.01760592, +) + +LATENT_SCALE: tuple[float, ...] = ( + 1.63933691, + 1.70204478, + 1.73642566, + 1.90004803, + 1.6675316, + 1.69059584, + 1.56853198, + 1.62314944, + 1.89106626, + 1.58086668, + 1.60822129, + 1.60962993, + 1.63322129, + 1.56074359, + 1.73419528, + 1.7919265, + 1.64040632, + 1.66802808, + 1.60390303, + 1.75480492, + 1.63187587, + 1.64334594, + 1.61722884, + 1.60146046, + 1.63459219, + 1.55291476, + 1.68771497, + 1.68415657, + 1.78966054, + 1.66631641, + 1.65626686, + 1.65976433, + 1.63487607, + 1.69513249, + 1.72933756, + 1.91310663, + 1.67035057, + 1.72286863, + 1.56719251, + 1.61934825, + 1.88628859, + 1.56911539, + 1.59455129, + 1.60829869, + 1.62470611, + 1.56052853, + 1.73677003, + 1.77563606, + 1.63732541, + 1.66370527, + 1.59508952, + 1.75153949, + 1.63029275, + 1.64517667, + 1.61659342, + 1.59722044, + 1.64103121, + 1.5408531, + 1.68610394, + 1.67772755, + 1.78998563, + 1.66621713, + 1.65458955, + 1.66041308, + 1.64710857, + 1.68163503, + 1.74000294, + 1.92784786, + 1.67411194, + 1.67395548, + 1.57406532, + 1.62199356, + 1.87618195, + 1.5584375, + 1.57438785, + 1.61711053, + 1.63094305, + 1.55644029, + 1.73124302, + 1.80666627, + 1.6463621, + 1.65932006, + 1.60816188, + 1.75682671, + 1.64695873, + 1.63121722, + 1.61380832, + 1.60478651, + 1.63396035, + 1.53505068, + 1.65534289, + 1.67132281, + 1.80317197, + 1.6767314, + 1.65700938, + 1.68426259, + 1.65339716, + 1.67540638, + 1.73298504, + 1.94067348, + 1.67893609, + 1.70635117, + 1.5730906, + 1.61928553, + 1.87148809, + 1.56244866, + 1.56697152, + 1.61584394, + 1.62759496, + 1.55480378, + 1.73484107, + 1.79055143, + 1.64688773, + 1.66121492, + 1.60135887, + 1.75254572, + 1.64798332, + 1.62989921, + 1.61381592, + 1.60792883, + 1.63939668, + 1.53075757, + 1.65371318, + 1.66801185, + 1.80029087, + 1.67591476, + 1.65655173, + 1.68533454, +) + + +def get_latent_norm() -> tuple[jnp.ndarray, "Any"]: + shift = jnp.array(LATENT_SHIFT, dtype=jnp.float32) + scale = jnp.array(LATENT_SCALE, dtype=jnp.float32) + assert shift.shape == (128,) and scale.shape == (128,) + return shift, scale diff --git a/src/maxdiffusion/models/ideogram/quantized_loading.py b/src/maxdiffusion/models/ideogram/quantized_loading.py new file mode 100644 index 000000000..a00412282 --- /dev/null +++ b/src/maxdiffusion/models/ideogram/quantized_loading.py @@ -0,0 +1,283 @@ +from __future__ import annotations + +import warnings + +try: + import bitsandbytes as bnb +except ImportError: + bnb = None +import torch +import torch.nn as nn +import torch.nn.functional as F + + +_BNB_SIBLING_SUFFIXES = ( + ".absmax", + ".quant_map", + ".nested_absmax", + ".nested_quant_map", +) + +# Largest magnitude representable by the e4m3 float8 format. Per-row weight +# scales map each row's max abs value onto this so we use the full range. +FP8_E4M3_MAX = 448.0 +FP8_WEIGHT_DTYPE = torch.float8_e4m3fn +FP8_SCALE_SUFFIX = ".weight_scale" +# Marker written into the text encoder's config.json so the loader knows to take +# the custom weight-only FP8 path instead of transformers' from_pretrained. +FP8_TEXT_ENCODER_CONFIG_FLAG = "ideogram_fp8_weight_only" + + +def is_bnb4bit_state_dict(state_dict: dict[str, torch.Tensor]) -> bool: + """True if any key looks like a bnb 4-bit quant_state sibling.""" + return any(".quant_state.bitsandbytes__" in k for k in state_dict) + + +def swap_linears_to_bnb4bit( + module: nn.Module, + compute_dtype: torch.dtype, + *, + quant_type: str = "nf4", + compress_statistics: bool = False, +) -> None: + if bnb is None: + raise ImportError("bitsandbytes is not installed. 4-bit quantization is not available on this device.") + for name, child in list(module.named_children()): + if isinstance(child, nn.Linear): + new_linear = bnb.nn.Linear4bit( + child.in_features, + child.out_features, + bias=child.bias is not None, + compute_dtype=compute_dtype, + compress_statistics=compress_statistics, + quant_type=quant_type, + ) + setattr(module, name, new_linear) + else: + swap_linears_to_bnb4bit( + child, + compute_dtype, + quant_type=quant_type, + compress_statistics=compress_statistics, + ) + + +def load_bnb4bit_state_dict( + model: nn.Module, + state_dict: dict[str, torch.Tensor], + device: torch.device, + dtype: torch.dtype, +) -> None: + if bnb is None: + raise ImportError("bitsandbytes is not installed. 4-bit quantization is not available on this device.") + consumed: set[str] = set() + for full_name, tensor in state_dict.items(): + if ".quant_state." in full_name or full_name.endswith(_BNB_SIBLING_SUFFIXES): + continue + parent_path, _, param_name = full_name.rpartition(".") + parent = model.get_submodule(parent_path) if parent_path else model + current = parent._parameters.get(param_name) + if not isinstance(current, bnb.nn.Params4bit): + continue + prefix = full_name + "." + quantized_stats = {k: v for k, v in state_dict.items() if k.startswith(prefix)} + # bnb's from_prequantized pops keys it consumes from the dict, so snapshot + # the names first. + consumed.add(full_name) + consumed.update(quantized_stats.keys()) + parent._parameters[param_name] = bnb.nn.Params4bit.from_prequantized( + data=tensor, + quantized_stats=quantized_stats, + requires_grad=False, + device=device, + ) + + remaining = {k: v for k, v in state_dict.items() if k not in consumed} + for k in list(remaining): + if remaining[k].is_floating_point(): + remaining[k] = remaining[k].to(device=device, dtype=dtype) + else: + remaining[k] = remaining[k].to(device=device) + + missing, unexpected = model.load_state_dict(remaining, strict=False) + # Quantized weights are loaded via from_prequantized above, so they appear in + # `missing` from load_state_dict's perspective — filter those out. + real_missing = [m for m in missing if m not in consumed] + if real_missing: + raise RuntimeError(f"missing keys after quantized load: {real_missing[:10]}") + if unexpected: + raise RuntimeError(f"unexpected keys after quantized load: {unexpected[:10]}") + + for p in model.parameters(): + if isinstance(p, bnb.nn.Params4bit): + continue + if p.is_floating_point() and p.dtype != dtype: + p.data = p.data.to(dtype=dtype) + if p.device != device: + p.data = p.data.to(device=device) + for name, b in list(model.named_buffers()): + if b.is_floating_point() and b.dtype != dtype: + parent_path, _, leaf = name.rpartition(".") + parent = model.get_submodule(parent_path) if parent_path else model + parent.register_buffer( + leaf, + b.to(device=device, dtype=dtype), + persistent=leaf not in parent._non_persistent_buffers_set, + ) + elif b.device != device: + parent_path, _, leaf = name.rpartition(".") + parent = model.get_submodule(parent_path) if parent_path else model + parent.register_buffer( + leaf, + b.to(device=device), + persistent=leaf not in parent._non_persistent_buffers_set, + ) + + +# --------------------------------------------------------------------------- +# Weight-only FP8 (e4m3) +# +# Activations stay in the compute dtype (e.g. bfloat16); only Linear weights are +# stored as float8 with a per-output-channel (per-row) float32 scale. At forward +# time the weight is dequantized back to the compute dtype and a normal bf16 +# matmul runs, so this needs no FP8 tensor-core hardware and works on any device +# that can store float8 (CPU included). The win is ~2x smaller Linear weights. +# --------------------------------------------------------------------------- + + +def quantize_weight_to_fp8( + weight: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Quantize a 2-D Linear weight to e4m3 float8 with per-row scales. + + Returns ``(weight_fp8, scale)`` where ``weight_fp8`` has shape ``(out, in)`` + in ``float8_e4m3fn`` and ``scale`` has shape ``(out,)`` in float32 such that + ``weight ≈ weight_fp8.to(dtype) * scale[:, None]``. + """ + w = weight.detach().to(torch.float32) + amax = w.abs().amax(dim=1, keepdim=True).clamp(min=1e-12) + scale = amax / FP8_E4M3_MAX + q = (w / scale).clamp(-FP8_E4M3_MAX, FP8_E4M3_MAX).to(FP8_WEIGHT_DTYPE) + return q, scale.squeeze(1).to(torch.float32) + + +def is_fp8_state_dict(state_dict: dict[str, torch.Tensor]) -> bool: + """True if the checkpoint carries weight-only FP8 Linear weights.""" + return any(k.endswith(FP8_SCALE_SUFFIX) for k in state_dict) or any( + v.dtype == FP8_WEIGHT_DTYPE for v in state_dict.values() + ) + + +class Fp8Linear(nn.Module): + """Linear layer holding an e4m3 float8 weight + per-row float32 scale. + + The weight and scale are registered as buffers (not parameters) so they load + via ``load_state_dict`` and are excluded from optimizer/grad machinery. The + dequantized matmul runs in ``compute_dtype``. + """ + + weight: torch.Tensor + weight_scale: torch.Tensor + bias: torch.Tensor | None + + def __init__( + self, + in_features: int, + out_features: int, + bias: bool, + compute_dtype: torch.dtype, + ) -> None: + super().__init__() + self.in_features = in_features + self.out_features = out_features + self.compute_dtype = compute_dtype + self.register_buffer( + "weight", + torch.empty(out_features, in_features, dtype=FP8_WEIGHT_DTYPE), + ) + self.register_buffer("weight_scale", torch.empty(out_features, dtype=torch.float32)) + if bias: + self.register_buffer("bias", torch.empty(out_features, dtype=compute_dtype)) + else: + self.bias = None + + def forward(self, x: torch.Tensor) -> torch.Tensor: + w = self.weight.to(x.dtype) * self.weight_scale.to(x.dtype).unsqueeze(1) + bias = self.bias.to(x.dtype) if self.bias is not None else None + return F.linear(x, w, bias) + + +def swap_linears_to_fp8( + module: nn.Module, + state_dict: dict[str, torch.Tensor], + compute_dtype: torch.dtype, + *, + prefix: str = "", +) -> None: + """Replace each ``nn.Linear`` that has a saved FP8 scale with an ``Fp8Linear``. + + Gating on the presence of ``.weight_scale`` means only layers that were + actually quantized at save time are swapped; everything else loads normally in + the compute dtype. + """ + for name, child in list(module.named_children()): + child_prefix = f"{prefix}{name}" + if isinstance(child, nn.Linear) and f"{child_prefix}{FP8_SCALE_SUFFIX}" in state_dict: + setattr( + module, + name, + Fp8Linear( + child.in_features, + child.out_features, + bias=child.bias is not None, + compute_dtype=compute_dtype, + ), + ) + else: + swap_linears_to_fp8(child, state_dict, compute_dtype, prefix=f"{child_prefix}.") + + +def load_fp8_state_dict( + model: nn.Module, + state_dict: dict[str, torch.Tensor], + device: torch.device, + dtype: torch.dtype, + *, + assign: bool = False, + strict: bool = True, +) -> None: + """Load a weight-only FP8 checkpoint into ``model``. + + ``model`` must already have its FP8 Linear layers swapped in (see + ``swap_linears_to_fp8``). FP8 weights are kept as float8, scales stay float32, + and every other floating tensor is cast to ``dtype``. + + ``assign=True`` replaces the module's tensors with the prepared ones rather than + copying into them. Use it when the model was built with ``from_config`` so the + non-quantized params take the loaded dtype directly and computed non-persistent + buffers (e.g. rotary caches) are left untouched. With ``assign=False`` (default), + the caller must have already put the unquantized params in ``dtype``. + + ``strict=False`` downgrades missing keys to a warning (e.g. tied weights that a + ``transformers`` model resolves itself); unexpected keys always raise. + """ + prepared: dict[str, torch.Tensor] = {} + for k, v in state_dict.items(): + if v.dtype == FP8_WEIGHT_DTYPE: + prepared[k] = v.to(device=device) + elif k.endswith(FP8_SCALE_SUFFIX): + prepared[k] = v.to(device=device, dtype=torch.float32) + elif v.is_floating_point(): + prepared[k] = v.to(device=device, dtype=dtype) + else: + prepared[k] = v.to(device=device) + + missing, unexpected = model.load_state_dict(prepared, strict=False, assign=assign) + if unexpected: + raise RuntimeError(f"unexpected keys after fp8 load: {unexpected[:10]}") + if missing: + if strict: + raise RuntimeError(f"missing keys after fp8 load: {missing[:10]}") + warnings.warn(f"missing keys after fp8 load: {missing[:10]}", stacklevel=2) + + model.to(device) diff --git a/src/maxdiffusion/models/ideogram/scheduler.py b/src/maxdiffusion/models/ideogram/scheduler.py new file mode 100644 index 000000000..2b377eff1 --- /dev/null +++ b/src/maxdiffusion/models/ideogram/scheduler.py @@ -0,0 +1,70 @@ +"""Logit-normal schedule and Euler flow-matching sampler.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass + +import jax +import jax.numpy as jnp + + +@dataclass(frozen=True) +class LogitNormalSchedule: + mean: float + std: float = 1.0 + logsnr_min: float = -15.0 + logsnr_max: float = 18.0 + + def __call__(self, t: jax.Array) -> jax.Array: + t = t.astype(jnp.float64) + z = jax.scipy.special.ndtri(t) + y = self.mean + self.std * z + t_ = jax.scipy.special.expit(y) + t_ = 1 - t_ + t_min = 1.0 / (1 + math.exp(0.5 * self.logsnr_max)) + t_max = 1.0 / (1 + math.exp(0.5 * self.logsnr_min)) + return jnp.clip(t_, t_min, t_max).astype(jnp.float32) + + +def get_schedule_for_resolution( + image_resolution: tuple[int, int], + known_resolution: tuple[int, int] = (512, 512), + known_mean: float = 1.0, + std: float = 1.0, +) -> LogitNormalSchedule: + """Resolution-aware schedule used at eval time.""" + num_pixels = image_resolution[0] * image_resolution[1] + known_pixels = known_resolution[0] * known_resolution[1] + mean = known_mean + 0.5 * math.log(num_pixels / known_pixels) + return LogitNormalSchedule(mean=mean, std=std) + + +def make_step_intervals(num_steps: int) -> jax.Array: + """Default linear step schedule used by the v4 eval config.""" + return jnp.linspace(0.0, 1.0, num_steps + 1, dtype=jnp.float32) + + +@dataclass(frozen=True, kw_only=True) +class SamplerParameters: + """Bundle of sampling hyperparameters for a named preset. + + ``guidance_schedule`` is in LOOP-INDEX order: index 0 is the LAST sampling + step (final polish), index ``num_steps - 1`` is the FIRST sampling step. + ``mu`` and ``std`` are the mean and stddev of the logit-normal noise + schedule passed to ``get_schedule_for_resolution`` (as ``known_mean`` and + ``std`` respectively). + + See ``ideogram4.sampler_configs.PRESETS`` for the named preset registry. + """ + + num_steps: int + guidance_schedule: tuple[float, ...] + mu: float + std: float = 1.0 + + def __post_init__(self) -> None: + if len(self.guidance_schedule) != self.num_steps: + raise ValueError( + f"guidance_schedule has length {len(self.guidance_schedule)}, " f"expected num_steps={self.num_steps}" + ) diff --git a/src/maxdiffusion/models/ideogram/torchax_text_encoder.py b/src/maxdiffusion/models/ideogram/torchax_text_encoder.py new file mode 100644 index 000000000..033be3bc7 --- /dev/null +++ b/src/maxdiffusion/models/ideogram/torchax_text_encoder.py @@ -0,0 +1,138 @@ +import jax + + +class TorchaxQwen3VLTextEncoder: + + def __init__(self, model): + self.model = model + + @classmethod + def from_pretrained(cls, pretrained_model_name_or_path: str, subfolder: str = "text_encoder"): + from transformers import AutoModel, AutoConfig + from huggingface_hub import hf_hub_download + from safetensors.torch import load_file + import json + import os + import torch + from .quantized_loading import swap_linears_to_fp8, load_fp8_state_dict + + kwargs = {"trust_remote_code": True} + if subfolder: + kwargs["subfolder"] = subfolder + + config = AutoConfig.from_pretrained(pretrained_model_name_or_path, **kwargs) + if getattr(config, "text_config", None) is not None: + if getattr(config.text_config, "rope_scaling", None) is None: + config.text_config.rope_scaling = {} + + # Instantiate from config (random weights, but creates all non-persistent buffers) + model = AutoModel.from_config(config, trust_remote_code=True) + + # Download and merge sharded state dict + index_filename = f"{subfolder}/model.safetensors.index.json" if subfolder else "model.safetensors.index.json" + index_path = hf_hub_download(repo_id=pretrained_model_name_or_path, filename=index_filename) + with open(index_path) as f: + index = json.load(f) + + shard_dir = os.path.dirname(index_filename) + shard_filenames = sorted(set(index["weight_map"].values())) + + state_dict = {} + for shard in shard_filenames: + shard_repo_path = os.path.join(shard_dir, shard) if shard_dir else shard + shard_path = hf_hub_download(repo_id=pretrained_model_name_or_path, filename=shard_repo_path) + state_dict.update(load_file(shard_path)) + + # Swap to Fp8Linear for quantized weights, and load state dict. + compute_dtype = torch.bfloat16 + swap_linears_to_fp8(model, state_dict, compute_dtype=compute_dtype) + # We don't move the entire model to bfloat16 explicitly before loading, + # because load_fp8_state_dict(assign=True) takes care of floating tensors. + load_fp8_state_dict(model, state_dict, device=torch.device("cpu"), dtype=compute_dtype, assign=True, strict=False) + model.eval() + return cls(model) + + def __call__( + self, + input_ids: jax.Array, + attention_mask: jax.Array, + pos_2d: jax.Array, + ) -> jax.Array: + import torch + import numpy as np + + # Run natively in PyTorch + device = next(self.model.parameters()).device + pt_input_ids = torch.from_numpy(np.array(input_ids)).to(device) + pt_attention_mask = torch.from_numpy(np.array(attention_mask)).to(device) + pt_pos_2d = torch.from_numpy(np.array(pos_2d)).to(device) + + with torch.no_grad(): + output = self._forward_inner( + self.model, + input_ids=pt_input_ids, + attention_mask=pt_attention_mask, + pos_2d=pt_pos_2d, + ) + + return jax.numpy.array(output.to(torch.float32).cpu().numpy()) + + @staticmethod + def _forward_inner(model, input_ids, attention_mask, pos_2d): + from transformers.masking_utils import create_causal_mask + import torch + + language_model = model.language_model + inputs_embeds = language_model.embed_tokens(input_ids) + + position_ids_4d = pos_2d[None, ...].expand(4, pos_2d.shape[0], -1) + text_position_ids = position_ids_4d[0] + mrope_position_ids = position_ids_4d[1:] + + import inspect + + sig = inspect.signature(create_causal_mask) + mask_kwargs = { + "config": language_model.config, + "attention_mask": attention_mask, + "past_key_values": None, + "position_ids": text_position_ids, + } + if "input_embeds" in sig.parameters: + mask_kwargs["input_embeds"] = inputs_embeds + else: + mask_kwargs["inputs_embeds"] = inputs_embeds + if "cache_position" in sig.parameters: + mask_kwargs["cache_position"] = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + + causal_mask = create_causal_mask(**mask_kwargs) + if language_model.rotary_emb.inv_freq.device.type != "jax": + language_model.rotary_emb.inv_freq = language_model.rotary_emb.inv_freq.to(inputs_embeds.device) + position_embeddings = language_model.rotary_emb(inputs_embeds, mrope_position_ids) + + from .constants import QWEN3_VL_ACTIVATION_LAYERS + + tap_set = set(QWEN3_VL_ACTIVATION_LAYERS) + captured = {} + hidden_states = inputs_embeds + for layer_idx, decoder_layer in enumerate(language_model.layers): + layer_out = decoder_layer( + hidden_states, + attention_mask=causal_mask, + position_ids=text_position_ids, + past_key_values=None, + position_embeddings=position_embeddings, + ) + hidden_states = layer_out[0] if isinstance(layer_out, (tuple, list)) else layer_out + if layer_idx in tap_set: + captured[layer_idx] = hidden_states + + selected = [captured[i] for i in QWEN3_VL_ACTIVATION_LAYERS] + + # Interleave features per token by stacking and permuting + batch_size, seq_len, _ = selected[0].shape + stacked = torch.stack(selected, dim=0) # (num_taps, B, L, H) + stacked = torch.permute(stacked, (1, 2, 3, 0)) + stacked = stacked.reshape(batch_size, seq_len, -1) + + return stacked diff --git a/src/maxdiffusion/models/ideogram/transformer_ideogram.py b/src/maxdiffusion/models/ideogram/transformer_ideogram.py new file mode 100644 index 000000000..a1bae9c9e --- /dev/null +++ b/src/maxdiffusion/models/ideogram/transformer_ideogram.py @@ -0,0 +1,328 @@ +import math +from typing import Tuple, Any +import jax +import jax.numpy as jnp +from flax import nnx +from dataclasses import dataclass + + +@dataclass +class Ideogram4Config: + emb_dim: int = 4608 + num_heads: int = 18 + in_channels: int = 128 + llm_features_dim: int = 53248 + adanln_dim: int = 512 + rope_theta: int = 5000000 + mrope_section: Tuple[int, int, int] = (24, 20, 20) + intermediate_size: int = 12288 + norm_eps: float = 1e-5 + num_layers: int = 34 + patch_size: int = 2 + + +class Ideogram4MRoPE(nnx.Module): + + def __init__(self, head_dim: int, base: int, mrope_section: Tuple[int, int, int]): + self.head_dim = head_dim + self.mrope_section = mrope_section + inv_freq = 1.0 / (base ** (jnp.arange(0, head_dim, 2, dtype=jnp.float32) / head_dim)) + self.inv_freq = nnx.Variable(inv_freq) + + def __call__(self, position_ids: jax.Array) -> Tuple[jax.Array, jax.Array]: + # position_ids: (B, L, 3) + inv_freq = self.inv_freq.value + + freqs_axes = [] + for i in range(3): + pos_axis = position_ids[..., i].astype(jnp.float32) + f = jnp.einsum("i, bl -> bli", inv_freq, pos_axis) + freqs_axes.append(f) + + freqs_t = freqs_axes[0] + + # Interleave logic + # In PyTorch: + # for axis, offset in ((1, 1), (2, 2)): + # length = self.mrope_section[axis] * 3 + # idx = torch.arange(offset, length, 3) + # freqs_t[..., idx] = freqs_axes[axis][..., idx] + + # In JAX, we can create an array of indices. + inv_freq_size = freqs_t.shape[-1] + indices = jnp.arange(inv_freq_size) + + cond_h = (indices % 3 == 1) & (indices < self.mrope_section[1] * 3) + cond_w = (indices % 3 == 2) & (indices < self.mrope_section[2] * 3) + + freqs_t = jnp.where(cond_h, freqs_axes[1], freqs_t) + freqs_t = jnp.where(cond_w, freqs_axes[2], freqs_t) + + emb = jnp.concatenate([freqs_t, freqs_t], axis=-1) + return jnp.cos(emb), jnp.sin(emb) + + +def _rotate_half(x: jax.Array) -> jax.Array: + half = x.shape[-1] // 2 + x1 = x[..., :half] + x2 = x[..., half:] + return jnp.concatenate([-x2, x1], axis=-1) + + +def _apply_rotary_pos_emb(q: jax.Array, k: jax.Array, cos: jax.Array, sin: jax.Array) -> Tuple[jax.Array, jax.Array]: + cos = jnp.expand_dims(cos, axis=1) + sin = jnp.expand_dims(sin, axis=1) + q_embed = (q * cos) + (_rotate_half(q) * sin) + k_embed = (k * cos) + (_rotate_half(k) * sin) + return q_embed, k_embed + + +class Ideogram4Attention(nnx.Module): + + def __init__(self, rngs: nnx.Rngs, hidden_size: int, num_heads: int, eps: float = 1e-5, dtype=jnp.float32): + self.hidden_size = hidden_size + self.num_heads = num_heads + self.head_dim = hidden_size // num_heads + self.dtype = dtype + + self.qkv = nnx.Linear(hidden_size, hidden_size * 3, use_bias=False, rngs=rngs, dtype=dtype) + self.norm_q = nnx.RMSNorm(self.head_dim, epsilon=eps, dtype=dtype, rngs=rngs) + self.norm_k = nnx.RMSNorm(self.head_dim, epsilon=eps, dtype=dtype, rngs=rngs) + self.o = nnx.Linear(hidden_size, hidden_size, use_bias=False, rngs=rngs, dtype=dtype) + + def __call__(self, x: jax.Array, segment_ids: jax.Array, cos: jax.Array, sin: jax.Array) -> jax.Array: + batch_size, seq_len, _ = x.shape + + qkv = self.qkv(x) + qkv = qkv.reshape((batch_size, seq_len, 3, self.num_heads, self.head_dim)) + + q = qkv[:, :, 0] + k = qkv[:, :, 1] + v = qkv[:, :, 2] + + q = self.norm_q(q) + k = self.norm_k(k) + + # Transpose to (B, num_heads, L, head_dim) + q = jnp.transpose(q, (0, 2, 1, 3)) + k = jnp.transpose(k, (0, 2, 1, 3)) + v = jnp.transpose(v, (0, 2, 1, 3)) + + q, k = _apply_rotary_pos_emb(q, k, cos, sin) + + # Block-diagonal mask from segment ids + attn_mask = jnp.expand_dims(segment_ids, axis=2) == jnp.expand_dims(segment_ids, axis=1) + attn_mask = jnp.expand_dims(attn_mask, axis=1) # (B, 1, L, L) + + # JAX scaled dot product attention + scale = 1.0 / math.sqrt(self.head_dim) + attn_weights = jnp.einsum("bhqd,bhkd->bhqk", q, k) * scale + + # apply mask + attn_weights = jnp.where(attn_mask, attn_weights, -1e10) + + attn_weights = jax.nn.softmax(attn_weights, axis=-1) + out = jnp.einsum("bhqk,bhkd->bhqd", attn_weights, v) + + out = jnp.transpose(out, (0, 2, 1, 3)).reshape((batch_size, seq_len, self.hidden_size)) + return self.o(out) + + +class Ideogram4MLP(nnx.Module): + + def __init__(self, rngs: nnx.Rngs, dim: int, hidden_dim: int, dtype=jnp.float32): + self.w1 = nnx.Linear(dim, hidden_dim, use_bias=False, rngs=rngs, dtype=dtype) + self.w2 = nnx.Linear(hidden_dim, dim, use_bias=False, rngs=rngs, dtype=dtype) + self.w3 = nnx.Linear(dim, hidden_dim, use_bias=False, rngs=rngs, dtype=dtype) + + def __call__(self, x: jax.Array) -> jax.Array: + return self.w2(jax.nn.silu(self.w1(x)) * self.w3(x)) + + +class Ideogram4TransformerBlock(nnx.Module): + + def __init__( + self, + rngs: nnx.Rngs, + hidden_size: int, + intermediate_size: int, + num_heads: int, + norm_eps: float, + adanln_dim: int, + dtype=jnp.float32, + ): + self.attention = Ideogram4Attention(rngs, hidden_size, num_heads, eps=1e-5, dtype=dtype) + self.feed_forward = Ideogram4MLP(rngs, hidden_size, intermediate_size, dtype=dtype) + + self.attention_norm1 = nnx.RMSNorm(hidden_size, epsilon=norm_eps, dtype=dtype, rngs=rngs) + self.ffn_norm1 = nnx.RMSNorm(hidden_size, epsilon=norm_eps, dtype=dtype, rngs=rngs) + self.attention_norm2 = nnx.RMSNorm(hidden_size, epsilon=norm_eps, dtype=dtype, rngs=rngs) + self.ffn_norm2 = nnx.RMSNorm(hidden_size, epsilon=norm_eps, dtype=dtype, rngs=rngs) + + self.adaln_modulation = nnx.Linear(adanln_dim, 4 * hidden_size, use_bias=True, rngs=rngs, dtype=dtype) + + def __call__( + self, x: jax.Array, segment_ids: jax.Array, cos: jax.Array, sin: jax.Array, adaln_input: jax.Array + ) -> jax.Array: + mod = self.adaln_modulation(adaln_input) + + # mod is split into 4 parts + hidden_size = x.shape[-1] + scale_msa = mod[..., 0 * hidden_size : 1 * hidden_size] + gate_msa = mod[..., 1 * hidden_size : 2 * hidden_size] + scale_mlp = mod[..., 2 * hidden_size : 3 * hidden_size] + gate_mlp = mod[..., 3 * hidden_size : 4 * hidden_size] + + gate_msa = jnp.tanh(gate_msa) + gate_mlp = jnp.tanh(gate_mlp) + scale_msa = 1.0 + scale_msa + scale_mlp = 1.0 + scale_mlp + + attn_out = self.attention( + self.attention_norm1(x) * scale_msa, + segment_ids=segment_ids, + cos=cos, + sin=sin, + ) + x = x + gate_msa * self.attention_norm2(attn_out) + x = x + gate_mlp * self.ffn_norm2(self.feed_forward(self.ffn_norm1(x) * scale_mlp)) + return x + + +def _sinusoidal_embedding(t: jax.Array, dim: int, scale: float = 1e4) -> jax.Array: + t = t.astype(jnp.float32) + half = dim // 2 + freq = math.log(scale) / (half - 1) + freq = jnp.exp(jnp.arange(half, dtype=jnp.float32) * -freq) + emb = jnp.expand_dims(t, -1) * freq + emb = jnp.concatenate([jnp.sin(emb), jnp.cos(emb)], axis=-1) + if dim % 2 == 1: + emb = jnp.pad(emb, ((0, 0), (0, 1))) + return emb + + +class Ideogram4EmbedScalar(nnx.Module): + + def __init__(self, rngs: nnx.Rngs, dim: int, input_range: Tuple[float, float], dtype=jnp.float32): + self.dim = dim + self.range_min, self.range_max = input_range + self.mlp_in = nnx.Linear(dim, dim, use_bias=True, rngs=rngs, dtype=dtype) + self.mlp_out = nnx.Linear(dim, dim, use_bias=True, rngs=rngs, dtype=dtype) + + def __call__(self, x: jax.Array) -> jax.Array: + x = x.astype(jnp.float32) + scaled = 1e4 * (x - self.range_min) / (self.range_max - self.range_min) + emb = _sinusoidal_embedding(scaled, self.dim) + emb = emb.astype(self.mlp_in.dtype) + emb = jax.nn.silu(self.mlp_in(emb)) + return self.mlp_out(emb) + + +class Ideogram4FinalLayer(nnx.Module): + + def __init__(self, rngs: nnx.Rngs, hidden_size: int, out_channels: int, adanln_dim: int, dtype=jnp.float32): + self.norm_final = nnx.LayerNorm(hidden_size, epsilon=1e-6, use_bias=False, use_scale=False, dtype=dtype, rngs=rngs) + self.linear = nnx.Linear(hidden_size, out_channels, use_bias=True, rngs=rngs, dtype=dtype) + self.adaln_modulation = nnx.Linear(adanln_dim, hidden_size, use_bias=True, rngs=rngs, dtype=dtype) + + def __call__(self, x: jax.Array, c: jax.Array) -> jax.Array: + scale = 1.0 + self.adaln_modulation(jax.nn.silu(c)) + return self.linear(self.norm_final(x) * scale) + + +class Ideogram4Transformer(nnx.Module): + + def __init__(self, rngs: nnx.Rngs, config: Any, dtype=jnp.float32): + self.config = config + self.dtype = dtype + + head_dim = config.emb_dim // config.num_heads + + self.input_proj = nnx.Linear(config.in_channels, config.emb_dim, use_bias=True, rngs=rngs, dtype=dtype) + self.llm_cond_norm = nnx.RMSNorm(config.llm_features_dim, epsilon=1e-6, dtype=dtype, rngs=rngs) + self.llm_cond_proj = nnx.Linear(config.llm_features_dim, config.emb_dim, use_bias=True, rngs=rngs, dtype=dtype) + + self.t_embedding = Ideogram4EmbedScalar(rngs, config.emb_dim, input_range=(0.0, 1.0), dtype=dtype) + self.adaln_proj = nnx.Linear(config.emb_dim, config.adanln_dim, use_bias=True, rngs=rngs, dtype=dtype) + + self.embed_image_indicator = nnx.Embed(2, config.emb_dim, rngs=rngs, dtype=dtype) + + self.rotary_emb = Ideogram4MRoPE( + head_dim=head_dim, + base=config.rope_theta, + mrope_section=config.mrope_section, + ) + + self.layers = nnx.List( + [ + Ideogram4TransformerBlock( + rngs, + hidden_size=config.emb_dim, + intermediate_size=config.intermediate_size, + num_heads=config.num_heads, + norm_eps=config.norm_eps, + adanln_dim=config.adanln_dim, + dtype=dtype, + ) + for _ in range(config.num_layers) + ] + ) + + self.final_layer = Ideogram4FinalLayer( + rngs, + hidden_size=config.emb_dim, + out_channels=config.in_channels, + adanln_dim=config.adanln_dim, + dtype=dtype, + ) + + def __call__( + self, + llm_features: jax.Array, + x: jax.Array, + t: jax.Array, + position_ids: jax.Array, + segment_ids: jax.Array, + indicator: jax.Array, + ) -> jax.Array: + batch_size, seq_len, in_channels = x.shape + + x = x.astype(self.dtype) + t = t.astype(self.dtype) + llm_features = llm_features.astype(self.dtype) + + indicator = indicator.astype(jnp.int32) + + from .constants import LLM_TOKEN_INDICATOR, OUTPUT_IMAGE_INDICATOR + + llm_token_mask = jnp.expand_dims((indicator == LLM_TOKEN_INDICATOR).astype(self.dtype), -1) + output_image_mask = jnp.expand_dims((indicator == OUTPUT_IMAGE_INDICATOR).astype(self.dtype), -1) + + llm_features = llm_features * llm_token_mask + x = x * output_image_mask + + x = self.input_proj(x) * output_image_mask + + t_cond = self.t_embedding(t) + if t.ndim == 1: + t_cond = jnp.expand_dims(t_cond, 1) + + adaln_input = jax.nn.silu(self.adaln_proj(t_cond)) + + llm_features = self.llm_cond_norm(llm_features) + llm_features = self.llm_cond_proj(llm_features) * llm_token_mask + + h = x + llm_features + + image_indicator_embedding = self.embed_image_indicator((indicator == OUTPUT_IMAGE_INDICATOR).astype(jnp.int32)) + h = h + image_indicator_embedding + + cos, sin = self.rotary_emb(position_ids) + cos = cos.astype(self.dtype) + sin = sin.astype(self.dtype) + + for layer in self.layers: + h = layer(h, segment_ids=segment_ids, cos=cos, sin=sin, adaln_input=adaln_input) + + out = self.final_layer(h, c=adaln_input) + return out.astype(jnp.float32) diff --git a/src/maxdiffusion/models/modeling_flax_pytorch_utils.py b/src/maxdiffusion/models/modeling_flax_pytorch_utils.py index 1239ddbc1..be0dee154 100644 --- a/src/maxdiffusion/models/modeling_flax_pytorch_utils.py +++ b/src/maxdiffusion/models/modeling_flax_pytorch_utils.py @@ -58,9 +58,12 @@ def validate_flax_state_dict(expected_pytree: dict, new_pytree: dict): def torch2jax(torch_tensor: torch.Tensor) -> Array: is_bfloat16 = torch_tensor.dtype == torch.bfloat16 - if is_bfloat16: + is_float8 = torch_tensor.dtype in (getattr(torch, "float8_e4m3fn", None), getattr(torch, "float8_e5m2", None)) + if is_bfloat16 or is_float8: # upcast the tensor to fp32 torch_tensor = torch_tensor.float() + if is_float8: + is_bfloat16 = True if torch.device.type != "cpu": torch_tensor = torch_tensor.to("cpu") @@ -137,7 +140,7 @@ def rename_key_and_reshape_tensor(pt_tuple_key, pt_tensor, random_flax_state_dic # embedding if pt_tuple_key[-1] == "weight" and pt_tuple_key[:-1] + ("embedding",) in random_flax_state_dict: pt_tuple_key = pt_tuple_key[:-1] + ("embedding",) - return renamed_pt_tuple_key, pt_tensor + return pt_tuple_key, pt_tensor # conv layer renamed_pt_tuple_key = pt_tuple_key[:-1] + ("kernel",) @@ -151,6 +154,10 @@ def rename_key_and_reshape_tensor(pt_tuple_key, pt_tensor, random_flax_state_dic pt_tensor = pt_tensor.transpose(2, 3, 4, 1, 0) return renamed_pt_tuple_key, pt_tensor + # direct match + if pt_tuple_key in random_flax_state_dict: + return pt_tuple_key, pt_tensor + # linear layer renamed_pt_tuple_key = pt_tuple_key[:-1] + ("kernel",) if pt_tuple_key[-1] == "weight": diff --git a/src/maxdiffusion/pipelines/__init__.py b/src/maxdiffusion/pipelines/__init__.py index 019c79a84..9df619fe1 100644 --- a/src/maxdiffusion/pipelines/__init__.py +++ b/src/maxdiffusion/pipelines/__init__.py @@ -133,6 +133,7 @@ FlaxStableDiffusionPipeline, ) from .stable_diffusion_xl import FlaxStableDiffusionXLPipeline + from .ideogram.ideogram_pipeline import IdeogramPipeline try: if not (is_torch_available() and is_note_seq_available()): diff --git a/src/maxdiffusion/pipelines/ideogram/__init__.py b/src/maxdiffusion/pipelines/ideogram/__init__.py new file mode 100644 index 000000000..c6d234858 --- /dev/null +++ b/src/maxdiffusion/pipelines/ideogram/__init__.py @@ -0,0 +1 @@ +from .ideogram_pipeline import IdeogramPipeline diff --git a/src/maxdiffusion/pipelines/ideogram/ideogram_pipeline.py b/src/maxdiffusion/pipelines/ideogram/ideogram_pipeline.py new file mode 100644 index 000000000..7bde84ec3 --- /dev/null +++ b/src/maxdiffusion/pipelines/ideogram/ideogram_pipeline.py @@ -0,0 +1,342 @@ +# pylint: disable=missing-module-docstring, missing-class-docstring, missing-function-docstring, too-many-positional-arguments, import-outside-toplevel, redefined-outer-name +from typing import Optional, Any, List + +import numpy as np + +import jax +import jax.numpy as jnp + +from flax import nnx + +from ...models.ideogram.transformer_ideogram import Ideogram4Transformer, Ideogram4Config +from ...models.ideogram.autoencoder_ideogram import AutoEncoder, AutoEncoderParams +from ...models.ideogram.constants import LLM_TOKEN_INDICATOR +from ...models.ideogram.ideogram_utils import load_transformer_weights, load_vae_weights +from ...models.ideogram.torchax_text_encoder import TorchaxQwen3VLTextEncoder +from ...models.ideogram.latent_norm import get_latent_norm +from ...models.ideogram.scheduler import get_schedule_for_resolution, make_step_intervals +from maxdiffusion import max_logging + + +class IdeogramPipeline: + + def __init__( + self, + conditional_transformer: Ideogram4Transformer, + unconditional_transformer: Ideogram4Transformer, + autoencoder: AutoEncoder, + text_encoder: Any, + tokenizer: Any, + ): + self.conditional_transformer = conditional_transformer + self.unconditional_transformer = unconditional_transformer + self.autoencoder = autoencoder + self.text_encoder = text_encoder + self.tokenizer = tokenizer + + @classmethod + def from_pretrained(cls, config, vae_only=False, load_transformer=True): + return cls._load_and_init(config, None, vae_only, load_transformer) + + @classmethod + def from_checkpoint(cls, config, restored_checkpoint, vae_only=False, load_transformer=True): + return cls._load_and_init(config, restored_checkpoint, vae_only, load_transformer) + + @classmethod + def _load_and_init(cls, config, restored_checkpoint, vae_only=False, load_transformer=True): + max_logging.log("Loading Ideogram pipeline components...") + ae_config = AutoEncoderParams() + rngs = nnx.Rngs(0) + + autoencoder = nnx.eval_shape(lambda rngs: AutoEncoder(rngs, ae_config), rngs) + ae_state = nnx.state(autoencoder).to_pure_dict() + + ae_params = load_vae_weights(config.pretrained_model_name_or_path, ae_state, "cpu") + autoencoder = AutoEncoder(rngs, ae_config) + nnx.update(autoencoder, ae_params) + + if vae_only: + return cls(None, None, autoencoder, None, None) + + conditional_transformer = None + unconditional_transformer = None + if load_transformer: + transformer_config = Ideogram4Config() + + # Load Conditional Transformer + conditional_transformer = nnx.eval_shape(lambda rngs: Ideogram4Transformer(rngs, transformer_config), rngs) + transformer_state = nnx.state(conditional_transformer).to_pure_dict() + + if restored_checkpoint: + cond_params = restored_checkpoint["ideogram_state"] + else: + cond_params = load_transformer_weights( + config.pretrained_model_name_or_path, + transformer_state, + "cpu", + num_layers=34, + scan_layers=False, + subfolder="transformer", + ) + + conditional_transformer = Ideogram4Transformer(rngs, transformer_config) + nnx.update(conditional_transformer, cond_params) + + # Load Unconditional Transformer + unconditional_transformer = nnx.eval_shape(lambda rngs: Ideogram4Transformer(rngs, transformer_config), rngs) + if restored_checkpoint: + uncond_params = restored_checkpoint["unconditional_ideogram_state"] + else: + uncond_params = load_transformer_weights( + config.pretrained_model_name_or_path, + transformer_state, + "cpu", + num_layers=34, + scan_layers=False, + subfolder="unconditional_transformer", + ) + + unconditional_transformer = Ideogram4Transformer(rngs, transformer_config) + nnx.update(unconditional_transformer, uncond_params) + + # Skip text encoder for pure test for now unless requested + max_logging.log("Initializing Torchax Text Encoder...") + text_encoder_repo = config.pretrained_model_name_or_path + subfolder = "text_encoder" + text_encoder = TorchaxQwen3VLTextEncoder.from_pretrained(text_encoder_repo, subfolder=subfolder) + + from transformers import AutoTokenizer + + tokenizer = AutoTokenizer.from_pretrained( + config.pretrained_model_name_or_path, subfolder="tokenizer", extra_special_tokens={} + ) + + return cls(conditional_transformer, unconditional_transformer, autoencoder, text_encoder, tokenizer) + + def _reorder_caption_keys(self, parsed: dict) -> dict: + canonical_keys = ["high_level_description", "style_description", "compositional_deconstruction"] + reordered = {} + for key in canonical_keys: + if key in parsed: + if key == "style_description" and isinstance(parsed[key], dict): + sd = parsed[key] + if "art_style" in sd and "photo" not in sd: + sd_keys = ["aesthetics", "lighting", "medium", "art_style", "color_palette"] + else: + sd_keys = ["aesthetics", "lighting", "photo", "medium", "color_palette"] + reordered_sd = {} + for sk in sd_keys: + if sk in sd: + reordered_sd[sk] = sd[sk] + for sk in sd: + if sk not in reordered_sd: + reordered_sd[sk] = sd[sk] + reordered[key] = reordered_sd + else: + reordered[key] = parsed[key] + for key in parsed: + if key not in reordered: + reordered[key] = parsed[key] + return reordered + + def _build_inputs_cpu(self, prompts, height, width): + batch_size = len(prompts) + + max_text_tokens = 0 + grid_h, grid_w = None, None + num_image_tokens = None + + tokenized = [] + for prompt in prompts: + import json + + try: + parsed = json.loads(prompt) + parsed = self._reorder_caption_keys(parsed) + prompt = json.dumps(parsed, ensure_ascii=False, separators=(",", ":")) + except json.JSONDecodeError: + pass + + if prompt != "": + encoded = self.tokenizer(prompt, return_tensors="np", add_special_tokens=True) + token_ids = encoded["input_ids"][0] + num_text_tokens = int(token_ids.shape[0]) + else: + num_text_tokens = 256 + token_ids = np.zeros((num_text_tokens,), dtype=np.int32) + tokenized.append((token_ids, num_text_tokens)) + + max_text_tokens = max(num_text for _, num_text in tokenized) + + patch_size = 2 + ae_scale_factor = 8 + patch = patch_size * ae_scale_factor + grid_h = height // patch + grid_w = width // patch + num_image_tokens = grid_h * grid_w + + total_seq_len = max_text_tokens + num_image_tokens + + h_idx = np.broadcast_to(np.arange(grid_h).reshape(-1, 1), (grid_h, grid_w)).reshape(-1) + w_idx = np.broadcast_to(np.arange(grid_w).reshape(1, -1), (grid_h, grid_w)).reshape(-1) + t_idx = np.zeros_like(h_idx) + IMAGE_POSITION_OFFSET = 65536 + image_pos = np.stack([t_idx, h_idx, w_idx], axis=1) + IMAGE_POSITION_OFFSET + + token_ids_out = np.zeros((batch_size, total_seq_len), dtype=np.int32) + text_position_ids = np.zeros((batch_size, total_seq_len, 3), dtype=np.int32) + position_ids = np.zeros((batch_size, total_seq_len, 3), dtype=np.int32) + + from ...models.ideogram.constants import LLM_TOKEN_INDICATOR, OUTPUT_IMAGE_INDICATOR, SEQUENCE_PADDING_INDICATOR + + segment_ids = np.full((batch_size, total_seq_len), SEQUENCE_PADDING_INDICATOR, dtype=np.int32) + indicator = np.zeros((batch_size, total_seq_len), dtype=np.int32) + + for b in range(batch_size): + toks, num_text = tokenized[b] + pad_len = max_text_tokens - num_text + total_unpadded = num_text + num_image_tokens + offset = pad_len + + token_ids_out[b, offset : offset + num_text] = toks + + text_pos = np.arange(num_text) + text_pos_3d = np.stack([text_pos, text_pos, text_pos], axis=1) + text_position_ids[b, offset : offset + num_text] = text_pos_3d + position_ids[b, offset : offset + num_text] = text_pos_3d + position_ids[b, offset + num_text :] = image_pos + + indicator[b, offset : offset + num_text] = LLM_TOKEN_INDICATOR + indicator[b, offset + num_text :] = OUTPUT_IMAGE_INDICATOR + + segment_ids[b, offset : offset + total_unpadded] = 1 + + return { + "token_ids": token_ids_out, + "text_position_ids": text_position_ids, + "position_ids": position_ids, + "segment_ids": segment_ids, + "indicator": indicator, + "num_image_tokens": num_image_tokens, + "max_text_tokens": max_text_tokens, + "grid_h": grid_h, + "grid_w": grid_w, + } + + def generate( + self, + prompts: List[str], + negative_prompts: Optional[List[str]] = None, + height: int = 1024, + width: int = 1024, + num_steps: int = 50, + guidance_scale: float = 7.0, + seed: int = 42, + ): + if negative_prompts is None: + negative_prompts = [""] * len(prompts) + + all_prompts = prompts + negative_prompts + inputs = self._build_inputs_cpu(all_prompts, height, width) + + batch_size = len(prompts) + max_text_tokens = inputs["max_text_tokens"] + num_image_tokens = inputs["num_image_tokens"] + + # 1. Text Encoding (using TorchAX text encoder) + llm_attention_mask = (inputs["indicator"][:, :max_text_tokens] == LLM_TOKEN_INDICATOR).astype(jnp.int32) + llm_features = self.text_encoder( + inputs["token_ids"][:, :max_text_tokens], + llm_attention_mask, + inputs["text_position_ids"][:, :max_text_tokens, 0], # Extract the 1D positional index for Qwen + ) + # Zero out non-LLM positions (left padding) + llm_features = llm_features * jnp.expand_dims(llm_attention_mask.astype(jnp.float32), -1) + + # Build padded LLM features: text features for the positive branch, zeros for image positions. + image_llm_padding = jnp.zeros((batch_size, num_image_tokens, llm_features.shape[-1]), dtype=jnp.float32) + # llm_features covers only the text portion; pad with zeros for the image token positions. + pos_llm_features = jnp.concatenate([llm_features[:batch_size], image_llm_padding], axis=1) + + # Initialize z + key = jax.random.PRNGKey(seed) + latent_dim = 128 # 32 * patch_size * patch_size + z = jax.random.normal(key, (batch_size, num_image_tokens, latent_dim), dtype=jnp.float32) + + # Precompute the scheduler timesteps array using the Ideogram 4 scheduler + schedule_fn = get_schedule_for_resolution((height, width), known_mean=0.5) + step_intervals = make_step_intervals(num_steps) + # Compute schedule evaluated at all intervals. schedule_fn returns a continuous scalar mapping. + sigmas = jnp.array( + [float(schedule_fn(jnp.array([step_intervals[i]]))[0]) for i in range(num_steps + 1)], dtype=jnp.float32 + ) + + # Padding for text latents + text_z_padding = jnp.zeros((batch_size, max_text_tokens, latent_dim), dtype=jnp.float32) + + # 2. Denoising loop in JAX + # Define the loop step + def denoise_step(i_fori, val): + z_curr, llm_pos, llm_neg = val + + # Euler step in model-time (mt) convention where mt increases from 0 (noisy) to 1 (clean). + i = (num_steps - 1) - i_fori + mt_curr = sigmas[i + 1] + mt_next = sigmas[i] + + t = jnp.full((batch_size,), mt_curr, dtype=jnp.float32) + + # Positive branch (conditional, text + image) + pos_z = jnp.concatenate([text_z_padding, z_curr], axis=1) + pos_v = self.conditional_transformer(llm_pos, pos_z, t, pos_position_ids, pos_segment_ids, pos_indicator)[ + :, max_text_tokens: + ] + + # Negative branch (unconditional, image only, asymmetric CFG) + neg_z = z_curr + neg_v = self.unconditional_transformer(llm_neg, neg_z, t, neg_position_ids, neg_segment_ids, neg_indicator) + + # CFG: standard formula v = uncond + guidance * (cond - uncond) + v = neg_v + guidance_scale * (pos_v - neg_v) + + # Euler step: z_next = z + v * delta_mt + delta_mt = mt_next - mt_curr + z_next = z_curr + v * delta_mt + + return z_next, llm_pos, llm_neg + + # Setup negative branch inputs for asymmetric CFG (image tokens only) + neg_llm_features = jnp.zeros((batch_size, num_image_tokens, llm_features.shape[-1]), dtype=jnp.float32) + neg_position_ids = inputs["position_ids"][:batch_size, max_text_tokens:] + neg_segment_ids = inputs["segment_ids"][:batch_size, max_text_tokens:] + neg_indicator = inputs["indicator"][:batch_size, max_text_tokens:] + + # Setup positive branch inputs (text + image tokens) + pos_position_ids = inputs["position_ids"][:batch_size] + pos_segment_ids = inputs["segment_ids"][:batch_size] + pos_indicator = inputs["indicator"][:batch_size] + + # Loop + init_val = (z, pos_llm_features, neg_llm_features) + z, _, _ = jax.lax.fori_loop(0, num_steps, denoise_step, init_val) + + # 3. Decode + # Unpatching logic + patch = 2 + ae_channels = z.shape[-1] // (patch * patch) + + # Apply latent scale and shift + shift, scale = get_latent_norm() + z = z * scale + shift + + z = z.reshape((batch_size, inputs["grid_h"], inputs["grid_w"], patch, patch, ae_channels)) + z = jnp.transpose(z, (0, 5, 1, 3, 2, 4)) + z = z.reshape((batch_size, ae_channels, inputs["grid_h"] * patch, inputs["grid_w"] * patch)) + + # Convert to NHWC for our Flax Autoencoder and cast to BF16 + z = jnp.transpose(z, (0, 2, 3, 1)).astype(jnp.bfloat16) + + images = self.autoencoder.decode(z) + images = jnp.clip((images + 1.0) / 2.0, 0.0, 1.0) + + return images diff --git a/src/maxdiffusion/pyconfig.py b/src/maxdiffusion/pyconfig.py index 3b2e9dd1a..f811bad6f 100644 --- a/src/maxdiffusion/pyconfig.py +++ b/src/maxdiffusion/pyconfig.py @@ -39,9 +39,10 @@ SEQUENCE_PARALLEL_AXIS_RULES, ULYSSES_ATTENTION_AXIS_RULES, ULYSSES_RING_ATTENTION_AXIS_RULES, + IDEOGRAM4, ) -_ALLOWED_MODEL_NAMES = {WAN2_1, WAN2_2, LTX2_VIDEO, LTX2_3} +_ALLOWED_MODEL_NAMES = {WAN2_1, WAN2_2, LTX2_VIDEO, LTX2_3, IDEOGRAM4} _ALLOWED_TRAINING_MODEL_NAMES = {WAN2_1} diff --git a/src/maxdiffusion/tests/generate_ideogram_end_to_end.py b/src/maxdiffusion/tests/generate_ideogram_end_to_end.py new file mode 100644 index 000000000..7351fa7ff --- /dev/null +++ b/src/maxdiffusion/tests/generate_ideogram_end_to_end.py @@ -0,0 +1,51 @@ +import jax.numpy as jnp +from flax import nnx +import numpy as np +from PIL import Image + +from maxdiffusion.models.ideogram import Ideogram4Transformer, Ideogram4Config, AutoEncoder, AutoEncoderParams +from maxdiffusion.pipelines.ideogram.ideogram_pipeline import IdeogramPipeline + + +def test_end_to_end(): + print("Initializing components...") + rngs = nnx.Rngs(0) + + # Use small dimensions for fast testing + config = Ideogram4Config(emb_dim=128, num_heads=2, in_channels=64, llm_features_dim=128, adanln_dim=128, num_layers=2) + transformer = Ideogram4Transformer(rngs, config) + + params = AutoEncoderParams(resolution=256, in_channels=3, ch=32, out_ch=3, ch_mult=(1, 2), num_res_blocks=1, z_channels=16) + autoencoder = AutoEncoder(rngs, params) + + # Dummy text encoder returning zero features + def mock_text_encoder(token_ids, _attention_mask, _pos_2d): + batch_size, seq_len = token_ids.shape + return jnp.zeros((batch_size, seq_len, config.llm_features_dim), dtype=jnp.float32) + + pipeline = IdeogramPipeline( + conditional_transformer=transformer, + unconditional_transformer=transformer, + autoencoder=autoencoder, + text_encoder=mock_text_encoder, + tokenizer=None, + ) + + print("Running generate...") + # Generate image with dummy prompt + # Use 256x256 image and 2 steps to make it fast + images = pipeline.generate(prompts=["a cute dog"], height=256, width=256, num_steps=2, guidance_scale=7.0, seed=42) + + print("Generation complete! Output shape:", images.shape) + + # Save the output image + image_np = np.array(images[0]) + image_np = (image_np * 255).astype(np.uint8) + img = Image.fromarray(image_np) + out_path = "ideogram_end_to_end_test.png" + img.save(out_path) + print(f"Saved test image to {out_path}") + + +if __name__ == "__main__": + test_end_to_end() diff --git a/src/maxdiffusion/tests/generate_ideogram_smoke_test.py b/src/maxdiffusion/tests/generate_ideogram_smoke_test.py new file mode 100644 index 000000000..baf5bb2cd --- /dev/null +++ b/src/maxdiffusion/tests/generate_ideogram_smoke_test.py @@ -0,0 +1,26 @@ +import unittest + +from flax import nnx + +from maxdiffusion.models.ideogram.transformer_ideogram import Ideogram4Transformer, Ideogram4Config +from maxdiffusion.models.ideogram.autoencoder_ideogram import AutoEncoder, AutoEncoderParams + + +class TestIdeogram(unittest.TestCase): + + def test_instantiate_transformer(self): + rngs = nnx.Rngs(0) + config = Ideogram4Config(emb_dim=128, num_heads=2, in_channels=64, llm_features_dim=128, adanln_dim=128, num_layers=2) + model = Ideogram4Transformer(rngs, config) + + self.assertIsNotNone(model) + + def test_instantiate_autoencoder(self): + rngs = nnx.Rngs(0) + params = AutoEncoderParams(resolution=32, in_channels=3, ch=32, out_ch=3, ch_mult=(1, 2), num_res_blocks=1, z_channels=8) + ae = AutoEncoder(rngs, params) + self.assertIsNotNone(ae) + + +if __name__ == "__main__": + unittest.main()