Preserve pinned official Ideogram 4 FP8 release, components, original terms and reference code

8ebf922 verified 9 days ago

22.5 kB

from __future__ import annotations

import json
import warnings
from dataclasses import dataclass
from posixpath import dirname as _posix_dirname, join as _posix_join
from typing import Optional, Sequence

import torch
from huggingface_hub import hf_hub_download
from huggingface_hub.errors import EntryNotFoundError
from PIL import Image
from safetensors.torch import load_file
from transformers import AutoConfig, AutoModel, AutoTokenizer
from transformers.masking_utils import create_causal_mask

from ideogram4.autoencoder import (
  AutoEncoder,
  AutoEncoderParams,
  convert_diffusers_state_dict,
)
from ideogram4.caption_verifier import CaptionVerifier
from ideogram4.constants import (
  IMAGE_POSITION_OFFSET,
  LLM_TOKEN_INDICATOR,
  OUTPUT_IMAGE_INDICATOR,
  SEQUENCE_PADDING_INDICATOR,
  QWEN3_VL_ACTIVATION_LAYERS,
)
from ideogram4.latent_norm import get_latent_norm
from ideogram4.modeling_ideogram4 import Ideogram4Config, Ideogram4Transformer
from ideogram4.quantized_loading import (
  FP8_TEXT_ENCODER_CONFIG_FLAG,
  is_bnb4bit_state_dict,
  is_fp8_state_dict,
  load_bnb4bit_state_dict,
  load_fp8_state_dict,
  swap_linears_to_bnb4bit,
  swap_linears_to_fp8,
)
from ideogram4.scheduler import (
  LogitNormalSchedule,
  get_schedule_for_resolution,
  make_step_intervals,
)


def _load_subfolder_state_dict(
  repo_id: str, subfolder: str, basename: str
) -> dict[str, torch.Tensor]:
  """Download a component's weights, whether sharded (index) or a single file.
  ``basename`` is the safetensors stem (``model`` for transformers components,
  ``diffusion_pytorch_model`` for diffusers ones).
  """
  prefix = f"{subfolder}/" if subfolder else ""
  index_filename = f"{prefix}{basename}.safetensors.index.json"
  try:
    return _load_sharded_state_dict(repo_id, index_filename)
  except EntryNotFoundError:
    single_path = hf_hub_download(
      repo_id=repo_id, filename=f"{prefix}{basename}.safetensors"
    )
    return load_file(single_path)


def _load_fp8_text_encoder(
  repo_id: str,
  device: torch.device,
  dtype: torch.dtype,
  *,
  text_encoder_subfolder: str,
):
  """Rebuild the text encoder from its config and load weight-only FP8 weights.
  transformers' ``from_pretrained`` can't read our float8 layout, so we
  instantiate the architecture with ``from_config`` (which also computes the
  non-persistent buffers such as rotary caches), swap the quantized Linears, and
  load the FP8 state dict with ``assign=True``.
  """
  config = AutoConfig.from_pretrained(
    repo_id, subfolder=text_encoder_subfolder, trust_remote_code=True
  )
  model = AutoModel.from_config(config, trust_remote_code=True)
  state_dict = _load_subfolder_state_dict(repo_id, text_encoder_subfolder, "model")
  swap_linears_to_fp8(model, state_dict, compute_dtype=dtype)
  # assign=True so unquantized params take the loaded dtype and the computed
  # rotary buffers (absent from the checkpoint) survive; tied weights, if any,
  # surface as benign missing keys.
  load_fp8_state_dict(
    model, state_dict, device=device, dtype=dtype, assign=True, strict=False
  )
  model.eval()
  return model


def _load_qwen3_vl(
  repo_id: str,
  device: torch.device,
  dtype: torch.dtype,
  *,
  tokenizer_subfolder: str | None = None,
  text_encoder_subfolder: str | None = None,
):
  """Load the Qwen3-VL tokenizer + model, optionally from named subfolders of ``repo_id``.
  When the weights are published in diffusers layout the tokenizer lives at ``tokenizer/``
  and the model at ``text_encoder/`` within the same repo as the transformer weights, so
  there is no need to fetch them from a separate upstream repo.
  If the saved ``text_encoder/config.json`` carries a ``quantization_config`` (e.g.
  a bitsandbytes 4-bit checkpoint), transformers handles the bnb placement via
  ``device_map`` and we skip the explicit ``.to(device)`` move afterwards.
  """
  tokenizer_kwargs = {"subfolder": tokenizer_subfolder} if tokenizer_subfolder else {}
  model_kwargs = {"subfolder": text_encoder_subfolder} if text_encoder_subfolder else {}
  tokenizer = AutoTokenizer.from_pretrained(repo_id, **tokenizer_kwargs)

  cfg_path = hf_hub_download(
    repo_id=repo_id,
    filename=f"{text_encoder_subfolder}/config.json"
    if text_encoder_subfolder
    else "config.json",
  )
  with open(cfg_path) as f:
    cfg_data = json.load(f)
  is_quantized = "quantization_config" in cfg_data
  is_fp8 = bool(cfg_data.get(FP8_TEXT_ENCODER_CONFIG_FLAG, False))

  if is_fp8:
    model = _load_fp8_text_encoder(
      repo_id,
      device,
      dtype,
      text_encoder_subfolder=text_encoder_subfolder or "",
    )
  elif is_quantized:
    model = AutoModel.from_pretrained(
      repo_id,
      torch_dtype=dtype,
      trust_remote_code=True,
      device_map={"": device},
      **model_kwargs,
    )
    model.eval()
  else:
    model = AutoModel.from_pretrained(
      repo_id, torch_dtype=dtype, trust_remote_code=True, **model_kwargs
    )
    model.to(device)
    model.eval()
  return tokenizer, model


def _build_transformer(
  transformer_config: "Ideogram4Config",
  state_dict: dict[str, torch.Tensor],
  device: torch.device,
  dtype: torch.dtype,
) -> "Ideogram4Transformer":
  model = Ideogram4Transformer(transformer_config)
  if is_bnb4bit_state_dict(state_dict):
    if device.type != "cuda":
      raise ValueError(f"bnb 4-bit weights require a CUDA device, got device={device}")
    swap_linears_to_bnb4bit(model, compute_dtype=dtype)
    load_bnb4bit_state_dict(model, state_dict, device=device, dtype=dtype)
  elif is_fp8_state_dict(state_dict):
    # Weight-only FP8: cast the unquantized params to the compute dtype first,
    # then swap in Fp8Linear layers (which keep their weights as float8).
    model.to(dtype)
    swap_linears_to_fp8(model, state_dict, compute_dtype=dtype)
    load_fp8_state_dict(model, state_dict, device=device, dtype=dtype)
  else:
    model.load_state_dict(state_dict)
    model.to(device=device, dtype=dtype)
  model.eval()
  return model


def _load_autoencoder(weights_path: str, device: torch.device, dtype: torch.dtype):
  ae = AutoEncoder(AutoEncoderParams())
  state_dict = convert_diffusers_state_dict(load_file(weights_path))
  ae.load_state_dict(state_dict)
  ae.to(device=device, dtype=dtype)
  ae.eval()
  return ae


def _load_sharded_state_dict(
  repo_id: str, index_filename: str
) -> dict[str, torch.Tensor]:
  """Download a sharded safetensors checkpoint and merge it into one state dict.
  ``index_filename`` is the path of the safetensors index file inside the repo
  (e.g. ``conditional_model/model.safetensors.index.json``). Shard filenames in
  the index are interpreted relative to that index's directory, matching the
  layout written by ``huggingface_hub.save_torch_state_dict``.
  """
  index_path = hf_hub_download(repo_id=repo_id, filename=index_filename)
  with open(index_path) as f:
    index = json.load(f)
  weight_map: dict[str, str] = index["weight_map"]
  shard_dir = _posix_dirname(index_filename)
  shard_filenames = sorted(set(weight_map.values()))

  state_dict: dict[str, torch.Tensor] = {}
  for shard in shard_filenames:
    shard_repo_path = _posix_join(shard_dir, shard) if shard_dir else shard
    shard_path = hf_hub_download(repo_id=repo_id, filename=shard_repo_path)
    state_dict.update(load_file(shard_path))
  return state_dict


def _load_indexed_or_single_state_dict(
  repo_id: str, index_filename: str
) -> dict[str, torch.Tensor]:
  """Load a component whether published as a sharded index or a single file.
  Some repos publish each component as a single ``.safetensors`` file rather
  than a sharded checkpoint with an ``.index.json``. Try the index first and
  fall back to the single file (the index filename with ``.index.json``
  dropped) when it isn't present.
  """
  try:
    return _load_sharded_state_dict(repo_id, index_filename)
  except EntryNotFoundError:
    single_filename = index_filename.removesuffix(".index.json")
    single_path = hf_hub_download(repo_id=repo_id, filename=single_filename)
    return load_file(single_path)


@dataclass
class Ideogram4PipelineConfig:
  weights_repo: str = "ideogram-ai/ideogram-4-nf4"
  conditional_index_filename: str = (
    "transformer/diffusion_pytorch_model.safetensors.index.json"
  )
  unconditional_index_filename: str = (
    "unconditional_transformer/diffusion_pytorch_model.safetensors.index.json"
  )
  autoencoder_filename: str = "vae/diffusion_pytorch_model.safetensors"
  text_encoder_subfolder: str = "text_encoder"
  tokenizer_subfolder: str = "tokenizer"
  patch_size: int = 2
  ae_scale_factor: int = 8
  max_text_tokens: int = 2048


class Ideogram4Pipeline:
  """Ideogram 4 text-to-image pipeline."""

  def __init__(
    self,
    conditional_transformer: Ideogram4Transformer,
    unconditional_transformer: Ideogram4Transformer,
    text_encoder,
    text_tokenizer,
    autoencoder,
    config: Ideogram4PipelineConfig,
    device: torch.device,
    dtype: torch.dtype,
  ) -> None:
    self.conditional_transformer = conditional_transformer
    self.unconditional_transformer = unconditional_transformer
    self.text_encoder = text_encoder
    self.text_tokenizer = text_tokenizer
    self.autoencoder = autoencoder
    self.config = config
    self.device = device
    self.dtype = dtype
    self.caption_verifier = CaptionVerifier()

    shift, scale = get_latent_norm()
    self.latent_shift = shift.to(device)
    self.latent_scale = scale.to(device)

  @classmethod
  def from_pretrained(
    cls,
    *,
    config: Optional[Ideogram4PipelineConfig] = None,
    device: str | torch.device = "cuda",
    dtype: torch.dtype = torch.bfloat16,
    transformer_config: Optional[Ideogram4Config] = None,
  ) -> "Ideogram4Pipeline":
    config = config or Ideogram4PipelineConfig()
    transformer_config = transformer_config or Ideogram4Config()
    device = torch.device(device)

    conditional_state_dict = _load_indexed_or_single_state_dict(
      config.weights_repo, config.conditional_index_filename
    )
    unconditional_state_dict = _load_indexed_or_single_state_dict(
      config.weights_repo, config.unconditional_index_filename
    )
    autoencoder_weights = hf_hub_download(
      repo_id=config.weights_repo, filename=config.autoencoder_filename
    )

    conditional_transformer = _build_transformer(
      transformer_config, conditional_state_dict, device, dtype
    )
    del conditional_state_dict
    unconditional_transformer = _build_transformer(
      transformer_config, unconditional_state_dict, device, dtype
    )
    del unconditional_state_dict

    text_tokenizer, text_encoder = _load_qwen3_vl(
      config.weights_repo,
      device,
      dtype,
      tokenizer_subfolder=config.tokenizer_subfolder,
      text_encoder_subfolder=config.text_encoder_subfolder,
    )
    autoencoder = _load_autoencoder(autoencoder_weights, device, dtype)

    return cls(
      conditional_transformer=conditional_transformer,
      unconditional_transformer=unconditional_transformer,
      text_encoder=text_encoder,
      text_tokenizer=text_tokenizer,
      autoencoder=autoencoder,
      config=config,
      device=device,
      dtype=dtype,
    )

  def _tokenize(self, prompt: str) -> tuple[torch.Tensor, int]:
    """Build chat-formatted token ids for a single prompt."""
    messages = [{"role": "user", "content": [{"type": "text", "text": prompt}]}]
    text = self.text_tokenizer.apply_chat_template(
      messages, add_generation_prompt=True, tokenize=False
    )
    encoded = self.text_tokenizer(text, return_tensors="pt", add_special_tokens=False)
    token_ids = encoded["input_ids"][0]
    num_text_tokens = int(token_ids.shape[0])
    if num_text_tokens > self.config.max_text_tokens:
      raise ValueError(
        f"prompt has {num_text_tokens} tokens, exceeds max_text_tokens={self.config.max_text_tokens}"
      )
    return token_ids, num_text_tokens

  def _build_inputs(
    self,
    prompts: list[str],
    height: int,
    width: int,
  ) -> dict[str, torch.Tensor]:
    """Build the packed sequence (text tokens + image tokens) for one batch."""
    tokenized = [self._tokenize(p) for p in prompts]
    batch_size = len(prompts)

    patch = self.config.patch_size * self.config.ae_scale_factor
    if height % patch != 0 or width % patch != 0:
      raise ValueError(
        f"height/width must be divisible by patch_size*ae_scale_factor={patch}"
      )
    grid_h = height // patch
    grid_w = width // patch
    num_image_tokens = grid_h * grid_w

    max_text_tokens = max(num_text for _, num_text in tokenized)
    total_seq_len = max_text_tokens + num_image_tokens

    # Image position ids (t=0, h, w) offset to keep them disjoint from text positions.
    h_idx = torch.arange(grid_h).view(-1, 1).expand(grid_h, grid_w).reshape(-1)
    w_idx = torch.arange(grid_w).view(1, -1).expand(grid_h, grid_w).reshape(-1)
    t_idx = torch.zeros_like(h_idx)
    image_pos = torch.stack([t_idx, h_idx, w_idx], dim=1) + IMAGE_POSITION_OFFSET

    token_ids = torch.zeros(batch_size, total_seq_len, dtype=torch.long)
    text_position_ids = torch.zeros(batch_size, total_seq_len, 3, dtype=torch.long)
    position_ids = torch.zeros(batch_size, total_seq_len, 3, dtype=torch.long)

    segment_ids = torch.full(
      (batch_size, total_seq_len), SEQUENCE_PADDING_INDICATOR, dtype=torch.long
    )
    indicator = torch.zeros(batch_size, total_seq_len, dtype=torch.long)

    for b, (toks, num_text) in enumerate(tokenized):
      pad_len = max_text_tokens - num_text
      total_unpadded = num_text + num_image_tokens

      # Layout: [pad_len zeros] [text tokens] [image tokens]
      offset = pad_len
      token_ids[b, offset : offset + num_text] = toks
      # Image token slots stay at 0.

      text_pos = torch.arange(num_text)
      text_pos_3d = torch.stack([text_pos, text_pos, text_pos], dim=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 id 1 for the (text+image) sample, padding stays at 0.
      segment_ids[b, offset : offset + total_unpadded] = 1

    return {
      "token_ids": token_ids.to(self.device),
      "text_position_ids": text_position_ids.to(self.device),
      "position_ids": position_ids.to(self.device),
      "segment_ids": segment_ids.to(self.device),
      "indicator": indicator.to(self.device),
      "num_image_tokens": num_image_tokens,  # type: ignore[dict-item]
      "grid_h": grid_h,  # type: ignore[dict-item]
      "grid_w": grid_w,  # type: ignore[dict-item]
      "max_text_tokens": max_text_tokens,  # type: ignore[dict-item]
    }

  def _get_qwen3_vl_embeddings(
    self,
    token_ids: torch.Tensor,
    attention_mask: torch.Tensor,
    pos_2d: torch.Tensor,
  ) -> list[torch.Tensor]:
    language_model = self.text_encoder.language_model

    inputs_embeds = language_model.embed_tokens(token_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:]

    causal_mask = create_causal_mask(
      config=language_model.config,
      inputs_embeds=inputs_embeds,
      attention_mask=attention_mask,
      past_key_values=None,
      position_ids=text_position_ids,
    )
    position_embeddings = language_model.rotary_emb(inputs_embeds, mrope_position_ids)

    tap_set = set(QWEN3_VL_ACTIVATION_LAYERS)
    captured: dict[int, torch.Tensor] = {}
    hidden_states = inputs_embeds
    for layer_idx, decoder_layer in enumerate(language_model.layers):
      hidden_states = decoder_layer(
        hidden_states,
        attention_mask=causal_mask,
        position_ids=text_position_ids,
        past_key_values=None,
        position_embeddings=position_embeddings,
      )
      if layer_idx in tap_set:
        captured[layer_idx] = hidden_states

    return [captured[i] for i in QWEN3_VL_ACTIVATION_LAYERS]

  def _encode_text(
    self,
    token_ids: torch.Tensor,
    text_position_ids: torch.Tensor,
    indicator: torch.Tensor,
  ) -> torch.Tensor:
    """Run Qwen3-VL and stack hidden states from the activation layers.
    Returns a (B, L, hidden_size * num_layers) float32 tensor.
    """
    batch_size, seq_len = token_ids.shape

    # Real text positions are exactly the LLM_TOKEN_INDICATOR positions.
    attention_mask = (indicator == LLM_TOKEN_INDICATOR).to(torch.long)

    pos_2d = text_position_ids[..., 0].contiguous()

    with torch.no_grad():
      selected = self._get_qwen3_vl_embeddings(token_ids, attention_mask, pos_2d)
    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)

    # Zero out non-LLM positions (left padding) so the transformer only sees real
    # text features at LLM_TOKEN_INDICATOR positions.
    text_mask = attention_mask.to(stacked.dtype).unsqueeze(-1)
    stacked = stacked * text_mask
    return stacked.to(torch.float32)

  def _verify_prompts(
    self, prompts: list[str], *, raise_on_issues: bool = True
  ) -> None:
    """Run each prompt through the caption verifier.
    Raises ``ValueError`` if any prompt has issues. When ``raise_on_issues``
    is False, issues are emitted as warnings instead.
    """
    messages: list[str] = []
    for i, prompt in enumerate(prompts):
      issues = self.caption_verifier.verify_raw(prompt)
      if not issues:
        continue
      messages.append(f"caption verifier flagged prompt[{i}]:\n" + "\n".join(issues))
    if not messages:
      return
    combined = "\n".join(messages)
    if raise_on_issues:
      raise ValueError(combined)
    warnings.warn(combined, stacklevel=2)

  @torch.no_grad()
  def __call__(
    self,
    prompts: str | list[str],
    *,
    height: int = 1024,
    width: int = 1024,
    num_steps: int = 128,
    guidance_scale: float = 7.0,
    guidance_schedule: Optional[Sequence[float] | torch.Tensor] = None,
    mu: float = 0.5,
    std: float = 1.0,
    seed: Optional[int] = None,
    schedule: Optional[LogitNormalSchedule] = None,
    raise_on_caption_issues: bool = True,
  ) -> list[Image.Image]:
    """Generate images for the given prompts."""
    if isinstance(prompts, str):
      prompts = [prompts]

    self._verify_prompts(prompts, raise_on_issues=raise_on_caption_issues)

    schedule = schedule or get_schedule_for_resolution(
      (height, width), known_mean=mu, std=std
    )
    step_intervals = make_step_intervals(num_steps).to(self.device)

    if guidance_schedule is not None:
      gw_per_step = torch.as_tensor(
        guidance_schedule, dtype=torch.float32, device=self.device
      )
      if gw_per_step.shape != (num_steps,):
        raise ValueError(
          f"guidance_schedule must have shape ({num_steps},), "
          f"got {tuple(gw_per_step.shape)}"
        )
    else:
      gw_per_step = torch.full(
        (num_steps,), float(guidance_scale), dtype=torch.float32, device=self.device
      )

    inputs = self._build_inputs(prompts, height=height, width=width)
    batch_size = len(prompts)
    num_image_tokens = inputs["num_image_tokens"]
    grid_h, grid_w = inputs["grid_h"], inputs["grid_w"]
    max_text_tokens = inputs["max_text_tokens"]
    latent_dim = self.conditional_transformer.config.in_channels

    llm_features = self._encode_text(
      inputs["token_ids"], inputs["text_position_ids"], inputs["indicator"]
    )

    # Negative branch is image-only (asymmetric CFG) with zeroed conditioning.
    neg_position_ids = inputs["position_ids"][:, max_text_tokens:]
    neg_segment_ids = inputs["segment_ids"][:, max_text_tokens:]
    neg_indicator = inputs["indicator"][:, max_text_tokens:]
    neg_llm_features = torch.zeros(  # type: ignore[call-overload]
      batch_size,
      num_image_tokens,
      llm_features.shape[-1],
      dtype=llm_features.dtype,
      device=self.device,
    )

    generator = torch.Generator(device=self.device)
    if seed is not None:
      generator.manual_seed(seed)
    z = torch.randn(  # type: ignore[call-overload]
      batch_size,
      num_image_tokens,
      latent_dim,
      dtype=torch.float32,
      device=self.device,
      generator=generator,
    )

    text_z_padding = torch.zeros(  # type: ignore[call-overload]
      batch_size,
      max_text_tokens,
      latent_dim,
      dtype=torch.float32,
      device=self.device,
    )

    for i in range(num_steps - 1, -1, -1):
      t_val = float(schedule(step_intervals[i + 1].unsqueeze(0)).item())
      s_val = float(schedule(step_intervals[i].unsqueeze(0)).item())
      t = torch.full((batch_size,), t_val, dtype=torch.float32, device=self.device)

      pos_z = torch.cat([text_z_padding, z], dim=1)
      pos_out = self.conditional_transformer(
        llm_features=llm_features,
        x=pos_z,
        t=t,
        position_ids=inputs["position_ids"],
        segment_ids=inputs["segment_ids"],
        indicator=inputs["indicator"],
      )
      pos_v = pos_out[:, max_text_tokens:]

      neg_v = self.unconditional_transformer(
        llm_features=neg_llm_features,
        x=z,
        t=t,
        position_ids=neg_position_ids,
        segment_ids=neg_segment_ids,
        indicator=neg_indicator,
      )

      gw_i = gw_per_step[i]
      v = gw_i * pos_v + (1.0 - gw_i) * neg_v
      delta = s_val - t_val
      z = z + v * delta

    return self._decode(z, grid_h=grid_h, grid_w=grid_w)  # type: ignore[arg-type]

  def _decode(self, z: torch.Tensor, *, grid_h: int, grid_w: int) -> list[Image.Image]:
    """Unpatch and run the autoencoder decoder."""
    batch_size = z.shape[0]
    patch = self.config.patch_size

    z = z * self.latent_scale + self.latent_shift

    ae_channels = z.shape[-1] // (patch * patch)
    z = z.view(batch_size, grid_h, grid_w, patch, patch, ae_channels)
    z = z.permute(0, 5, 1, 3, 2, 4).contiguous()
    z = z.view(batch_size, ae_channels, grid_h * patch, grid_w * patch)

    z = z.to(self.dtype)
    decoded = self.autoencoder.decoder(z)

    decoded = decoded.float().clamp(-1.0, 1.0)
    decoded = ((decoded + 1.0) * 127.5).round().to(torch.uint8)
    decoded = decoded.permute(0, 2, 3, 1).cpu().numpy()
    return [Image.fromarray(arr) for arr in decoded]

	from __future__ import annotations

	import json
	import warnings
	from dataclasses import dataclass
	from posixpath import dirname as _posix_dirname, join as _posix_join
	from typing import Optional, Sequence

	import torch
	from huggingface_hub import hf_hub_download
	from huggingface_hub.errors import EntryNotFoundError
	from PIL import Image
	from safetensors.torch import load_file
	from transformers import AutoConfig, AutoModel, AutoTokenizer
	from transformers.masking_utils import create_causal_mask

	from ideogram4.autoencoder import (
	AutoEncoder,
	AutoEncoderParams,
	convert_diffusers_state_dict,
	)
	from ideogram4.caption_verifier import CaptionVerifier
	from ideogram4.constants import (
	IMAGE_POSITION_OFFSET,
	LLM_TOKEN_INDICATOR,
	OUTPUT_IMAGE_INDICATOR,
	SEQUENCE_PADDING_INDICATOR,
	QWEN3_VL_ACTIVATION_LAYERS,
	)
	from ideogram4.latent_norm import get_latent_norm
	from ideogram4.modeling_ideogram4 import Ideogram4Config, Ideogram4Transformer
	from ideogram4.quantized_loading import (
	FP8_TEXT_ENCODER_CONFIG_FLAG,
	is_bnb4bit_state_dict,
	is_fp8_state_dict,
	load_bnb4bit_state_dict,
	load_fp8_state_dict,
	swap_linears_to_bnb4bit,
	swap_linears_to_fp8,
	)
	from ideogram4.scheduler import (
	LogitNormalSchedule,
	get_schedule_for_resolution,
	make_step_intervals,
	)


	def _load_subfolder_state_dict(
	repo_id: str, subfolder: str, basename: str
	) -> dict[str, torch.Tensor]:
	"""Download a component's weights, whether sharded (index) or a single file.

	``basename`` is the safetensors stem (``model`` for transformers components,
	``diffusion_pytorch_model`` for diffusers ones).
	"""
	prefix = f"{subfolder}/" if subfolder else ""
	index_filename = f"{prefix}{basename}.safetensors.index.json"
	try:
	return _load_sharded_state_dict(repo_id, index_filename)
	except EntryNotFoundError:
	single_path = hf_hub_download(
	repo_id=repo_id, filename=f"{prefix}{basename}.safetensors"
	)
	return load_file(single_path)


	def _load_fp8_text_encoder(
	repo_id: str,
	device: torch.device,
	dtype: torch.dtype,
	*,
	text_encoder_subfolder: str,
	):
	"""Rebuild the text encoder from its config and load weight-only FP8 weights.

	transformers' ``from_pretrained`` can't read our float8 layout, so we
	instantiate the architecture with ``from_config`` (which also computes the
	non-persistent buffers such as rotary caches), swap the quantized Linears, and
	load the FP8 state dict with ``assign=True``.
	"""
	config = AutoConfig.from_pretrained(
	repo_id, subfolder=text_encoder_subfolder, trust_remote_code=True
	)
	model = AutoModel.from_config(config, trust_remote_code=True)
	state_dict = _load_subfolder_state_dict(repo_id, text_encoder_subfolder, "model")
	swap_linears_to_fp8(model, state_dict, compute_dtype=dtype)
	# assign=True so unquantized params take the loaded dtype and the computed
	# rotary buffers (absent from the checkpoint) survive; tied weights, if any,
	# surface as benign missing keys.
	load_fp8_state_dict(
	model, state_dict, device=device, dtype=dtype, assign=True, strict=False
	)
	model.eval()
	return model


	def _load_qwen3_vl(
	repo_id: str,
	device: torch.device,
	dtype: torch.dtype,
	*,
	tokenizer_subfolder: str \| None = None,
	text_encoder_subfolder: str \| None = None,
	):
	"""Load the Qwen3-VL tokenizer + model, optionally from named subfolders of ``repo_id``.

	When the weights are published in diffusers layout the tokenizer lives at ``tokenizer/``
	and the model at ``text_encoder/`` within the same repo as the transformer weights, so
	there is no need to fetch them from a separate upstream repo.

	If the saved ``text_encoder/config.json`` carries a ``quantization_config`` (e.g.
	a bitsandbytes 4-bit checkpoint), transformers handles the bnb placement via
	``device_map`` and we skip the explicit ``.to(device)`` move afterwards.
	"""
	tokenizer_kwargs = {"subfolder": tokenizer_subfolder} if tokenizer_subfolder else {}
	model_kwargs = {"subfolder": text_encoder_subfolder} if text_encoder_subfolder else {}
	tokenizer = AutoTokenizer.from_pretrained(repo_id, **tokenizer_kwargs)

	cfg_path = hf_hub_download(
	repo_id=repo_id,
	filename=f"{text_encoder_subfolder}/config.json"
	if text_encoder_subfolder
	else "config.json",
	)
	with open(cfg_path) as f:
	cfg_data = json.load(f)
	is_quantized = "quantization_config" in cfg_data
	is_fp8 = bool(cfg_data.get(FP8_TEXT_ENCODER_CONFIG_FLAG, False))

	if is_fp8:
	model = _load_fp8_text_encoder(
	repo_id,
	device,
	dtype,
	text_encoder_subfolder=text_encoder_subfolder or "",
	)
	elif is_quantized:
	model = AutoModel.from_pretrained(
	repo_id,
	torch_dtype=dtype,
	trust_remote_code=True,
	device_map={"": device},
	**model_kwargs,
	)
	model.eval()
	else:
	model = AutoModel.from_pretrained(
	repo_id, torch_dtype=dtype, trust_remote_code=True, **model_kwargs
	)
	model.to(device)
	model.eval()
	return tokenizer, model


	def _build_transformer(
	transformer_config: "Ideogram4Config",
	state_dict: dict[str, torch.Tensor],
	device: torch.device,
	dtype: torch.dtype,
	) -> "Ideogram4Transformer":
	model = Ideogram4Transformer(transformer_config)
	if is_bnb4bit_state_dict(state_dict):
	if device.type != "cuda":
	raise ValueError(f"bnb 4-bit weights require a CUDA device, got device={device}")
	swap_linears_to_bnb4bit(model, compute_dtype=dtype)
	load_bnb4bit_state_dict(model, state_dict, device=device, dtype=dtype)
	elif is_fp8_state_dict(state_dict):
	# Weight-only FP8: cast the unquantized params to the compute dtype first,
	# then swap in Fp8Linear layers (which keep their weights as float8).
	model.to(dtype)
	swap_linears_to_fp8(model, state_dict, compute_dtype=dtype)
	load_fp8_state_dict(model, state_dict, device=device, dtype=dtype)
	else:
	model.load_state_dict(state_dict)
	model.to(device=device, dtype=dtype)
	model.eval()
	return model


	def _load_autoencoder(weights_path: str, device: torch.device, dtype: torch.dtype):
	ae = AutoEncoder(AutoEncoderParams())
	state_dict = convert_diffusers_state_dict(load_file(weights_path))
	ae.load_state_dict(state_dict)
	ae.to(device=device, dtype=dtype)
	ae.eval()
	return ae


	def _load_sharded_state_dict(
	repo_id: str, index_filename: str
	) -> dict[str, torch.Tensor]:
	"""Download a sharded safetensors checkpoint and merge it into one state dict.

	``index_filename`` is the path of the safetensors index file inside the repo
	(e.g. ``conditional_model/model.safetensors.index.json``). Shard filenames in
	the index are interpreted relative to that index's directory, matching the
	layout written by ``huggingface_hub.save_torch_state_dict``.
	"""
	index_path = hf_hub_download(repo_id=repo_id, filename=index_filename)
	with open(index_path) as f:
	index = json.load(f)
	weight_map: dict[str, str] = index["weight_map"]
	shard_dir = _posix_dirname(index_filename)
	shard_filenames = sorted(set(weight_map.values()))

	state_dict: dict[str, torch.Tensor] = {}
	for shard in shard_filenames:
	shard_repo_path = _posix_join(shard_dir, shard) if shard_dir else shard
	shard_path = hf_hub_download(repo_id=repo_id, filename=shard_repo_path)
	state_dict.update(load_file(shard_path))
	return state_dict


	def _load_indexed_or_single_state_dict(
	repo_id: str, index_filename: str
	) -> dict[str, torch.Tensor]:
	"""Load a component whether published as a sharded index or a single file.

	Some repos publish each component as a single ``.safetensors`` file rather
	than a sharded checkpoint with an ``.index.json``. Try the index first and
	fall back to the single file (the index filename with ``.index.json``
	dropped) when it isn't present.
	"""
	try:
	return _load_sharded_state_dict(repo_id, index_filename)
	except EntryNotFoundError:
	single_filename = index_filename.removesuffix(".index.json")
	single_path = hf_hub_download(repo_id=repo_id, filename=single_filename)
	return load_file(single_path)


	@dataclass
	class Ideogram4PipelineConfig:
	weights_repo: str = "ideogram-ai/ideogram-4-nf4"
	conditional_index_filename: str = (
	"transformer/diffusion_pytorch_model.safetensors.index.json"
	)
	unconditional_index_filename: str = (
	"unconditional_transformer/diffusion_pytorch_model.safetensors.index.json"
	)
	autoencoder_filename: str = "vae/diffusion_pytorch_model.safetensors"
	text_encoder_subfolder: str = "text_encoder"
	tokenizer_subfolder: str = "tokenizer"
	patch_size: int = 2
	ae_scale_factor: int = 8
	max_text_tokens: int = 2048


	class Ideogram4Pipeline:
	"""Ideogram 4 text-to-image pipeline."""

	def __init__(
	self,
	conditional_transformer: Ideogram4Transformer,
	unconditional_transformer: Ideogram4Transformer,
	text_encoder,
	text_tokenizer,
	autoencoder,
	config: Ideogram4PipelineConfig,
	device: torch.device,
	dtype: torch.dtype,
	) -> None:
	self.conditional_transformer = conditional_transformer
	self.unconditional_transformer = unconditional_transformer
	self.text_encoder = text_encoder
	self.text_tokenizer = text_tokenizer
	self.autoencoder = autoencoder
	self.config = config
	self.device = device
	self.dtype = dtype
	self.caption_verifier = CaptionVerifier()

	shift, scale = get_latent_norm()
	self.latent_shift = shift.to(device)
	self.latent_scale = scale.to(device)

	@classmethod
	def from_pretrained(
	cls,
	*,
	config: Optional[Ideogram4PipelineConfig] = None,
	device: str \| torch.device = "cuda",
	dtype: torch.dtype = torch.bfloat16,
	transformer_config: Optional[Ideogram4Config] = None,
	) -> "Ideogram4Pipeline":
	config = config or Ideogram4PipelineConfig()
	transformer_config = transformer_config or Ideogram4Config()
	device = torch.device(device)

	conditional_state_dict = _load_indexed_or_single_state_dict(
	config.weights_repo, config.conditional_index_filename
	)
	unconditional_state_dict = _load_indexed_or_single_state_dict(
	config.weights_repo, config.unconditional_index_filename
	)
	autoencoder_weights = hf_hub_download(
	repo_id=config.weights_repo, filename=config.autoencoder_filename
	)

	conditional_transformer = _build_transformer(
	transformer_config, conditional_state_dict, device, dtype
	)
	del conditional_state_dict
	unconditional_transformer = _build_transformer(
	transformer_config, unconditional_state_dict, device, dtype
	)
	del unconditional_state_dict

	text_tokenizer, text_encoder = _load_qwen3_vl(
	config.weights_repo,
	device,
	dtype,
	tokenizer_subfolder=config.tokenizer_subfolder,
	text_encoder_subfolder=config.text_encoder_subfolder,
	)
	autoencoder = _load_autoencoder(autoencoder_weights, device, dtype)

	return cls(
	conditional_transformer=conditional_transformer,
	unconditional_transformer=unconditional_transformer,
	text_encoder=text_encoder,
	text_tokenizer=text_tokenizer,
	autoencoder=autoencoder,
	config=config,
	device=device,
	dtype=dtype,
	)

	def _tokenize(self, prompt: str) -> tuple[torch.Tensor, int]:
	"""Build chat-formatted token ids for a single prompt."""
	messages = [{"role": "user", "content": [{"type": "text", "text": prompt}]}]
	text = self.text_tokenizer.apply_chat_template(
	messages, add_generation_prompt=True, tokenize=False
	)
	encoded = self.text_tokenizer(text, return_tensors="pt", add_special_tokens=False)
	token_ids = encoded["input_ids"][0]
	num_text_tokens = int(token_ids.shape[0])
	if num_text_tokens > self.config.max_text_tokens:
	raise ValueError(
	f"prompt has {num_text_tokens} tokens, exceeds max_text_tokens={self.config.max_text_tokens}"
	)
	return token_ids, num_text_tokens

	def _build_inputs(
	self,
	prompts: list[str],
	height: int,
	width: int,
	) -> dict[str, torch.Tensor]:
	"""Build the packed sequence (text tokens + image tokens) for one batch."""
	tokenized = [self._tokenize(p) for p in prompts]
	batch_size = len(prompts)

	patch = self.config.patch_size * self.config.ae_scale_factor
	if height % patch != 0 or width % patch != 0:
	raise ValueError(
	f"height/width must be divisible by patch_size*ae_scale_factor={patch}"
	)
	grid_h = height // patch
	grid_w = width // patch
	num_image_tokens = grid_h * grid_w

	max_text_tokens = max(num_text for _, num_text in tokenized)
	total_seq_len = max_text_tokens + num_image_tokens

	# Image position ids (t=0, h, w) offset to keep them disjoint from text positions.
	h_idx = torch.arange(grid_h).view(-1, 1).expand(grid_h, grid_w).reshape(-1)
	w_idx = torch.arange(grid_w).view(1, -1).expand(grid_h, grid_w).reshape(-1)
	t_idx = torch.zeros_like(h_idx)
	image_pos = torch.stack([t_idx, h_idx, w_idx], dim=1) + IMAGE_POSITION_OFFSET

	token_ids = torch.zeros(batch_size, total_seq_len, dtype=torch.long)
	text_position_ids = torch.zeros(batch_size, total_seq_len, 3, dtype=torch.long)
	position_ids = torch.zeros(batch_size, total_seq_len, 3, dtype=torch.long)

	segment_ids = torch.full(
	(batch_size, total_seq_len), SEQUENCE_PADDING_INDICATOR, dtype=torch.long
	)
	indicator = torch.zeros(batch_size, total_seq_len, dtype=torch.long)

	for b, (toks, num_text) in enumerate(tokenized):
	pad_len = max_text_tokens - num_text
	total_unpadded = num_text + num_image_tokens

	# Layout: [pad_len zeros] [text tokens] [image tokens]
	offset = pad_len
	token_ids[b, offset : offset + num_text] = toks
	# Image token slots stay at 0.

	text_pos = torch.arange(num_text)
	text_pos_3d = torch.stack([text_pos, text_pos, text_pos], dim=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 id 1 for the (text+image) sample, padding stays at 0.
	segment_ids[b, offset : offset + total_unpadded] = 1

	return {
	"token_ids": token_ids.to(self.device),
	"text_position_ids": text_position_ids.to(self.device),
	"position_ids": position_ids.to(self.device),
	"segment_ids": segment_ids.to(self.device),
	"indicator": indicator.to(self.device),
	"num_image_tokens": num_image_tokens, # type: ignore[dict-item]
	"grid_h": grid_h, # type: ignore[dict-item]
	"grid_w": grid_w, # type: ignore[dict-item]
	"max_text_tokens": max_text_tokens, # type: ignore[dict-item]
	}

	def _get_qwen3_vl_embeddings(
	self,
	token_ids: torch.Tensor,
	attention_mask: torch.Tensor,
	pos_2d: torch.Tensor,
	) -> list[torch.Tensor]:
	language_model = self.text_encoder.language_model

	inputs_embeds = language_model.embed_tokens(token_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:]

	causal_mask = create_causal_mask(
	config=language_model.config,
	inputs_embeds=inputs_embeds,
	attention_mask=attention_mask,
	past_key_values=None,
	position_ids=text_position_ids,
	)
	position_embeddings = language_model.rotary_emb(inputs_embeds, mrope_position_ids)

	tap_set = set(QWEN3_VL_ACTIVATION_LAYERS)
	captured: dict[int, torch.Tensor] = {}
	hidden_states = inputs_embeds
	for layer_idx, decoder_layer in enumerate(language_model.layers):
	hidden_states = decoder_layer(
	hidden_states,
	attention_mask=causal_mask,
	position_ids=text_position_ids,
	past_key_values=None,
	position_embeddings=position_embeddings,
	)
	if layer_idx in tap_set:
	captured[layer_idx] = hidden_states

	return [captured[i] for i in QWEN3_VL_ACTIVATION_LAYERS]

	def _encode_text(
	self,
	token_ids: torch.Tensor,
	text_position_ids: torch.Tensor,
	indicator: torch.Tensor,
	) -> torch.Tensor:
	"""Run Qwen3-VL and stack hidden states from the activation layers.

	Returns a (B, L, hidden_size * num_layers) float32 tensor.
	"""
	batch_size, seq_len = token_ids.shape

	# Real text positions are exactly the LLM_TOKEN_INDICATOR positions.
	attention_mask = (indicator == LLM_TOKEN_INDICATOR).to(torch.long)

	pos_2d = text_position_ids[..., 0].contiguous()

	with torch.no_grad():
	selected = self._get_qwen3_vl_embeddings(token_ids, attention_mask, pos_2d)
	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)

	# Zero out non-LLM positions (left padding) so the transformer only sees real
	# text features at LLM_TOKEN_INDICATOR positions.
	text_mask = attention_mask.to(stacked.dtype).unsqueeze(-1)
	stacked = stacked * text_mask
	return stacked.to(torch.float32)

	def _verify_prompts(
	self, prompts: list[str], *, raise_on_issues: bool = True
	) -> None:
	"""Run each prompt through the caption verifier.

	Raises ``ValueError`` if any prompt has issues. When ``raise_on_issues``
	is False, issues are emitted as warnings instead.
	"""
	messages: list[str] = []
	for i, prompt in enumerate(prompts):
	issues = self.caption_verifier.verify_raw(prompt)
	if not issues:
	continue
	messages.append(f"caption verifier flagged prompt[{i}]:\n" + "\n".join(issues))
	if not messages:
	return
	combined = "\n".join(messages)
	if raise_on_issues:
	raise ValueError(combined)
	warnings.warn(combined, stacklevel=2)

	@torch.no_grad()
	def __call__(
	self,
	prompts: str \| list[str],
	*,
	height: int = 1024,
	width: int = 1024,
	num_steps: int = 128,
	guidance_scale: float = 7.0,
	guidance_schedule: Optional[Sequence[float] \| torch.Tensor] = None,
	mu: float = 0.5,
	std: float = 1.0,
	seed: Optional[int] = None,
	schedule: Optional[LogitNormalSchedule] = None,
	raise_on_caption_issues: bool = True,
	) -> list[Image.Image]:
	"""Generate images for the given prompts."""
	if isinstance(prompts, str):
	prompts = [prompts]

	self._verify_prompts(prompts, raise_on_issues=raise_on_caption_issues)

	schedule = schedule or get_schedule_for_resolution(
	(height, width), known_mean=mu, std=std
	)
	step_intervals = make_step_intervals(num_steps).to(self.device)

	if guidance_schedule is not None:
	gw_per_step = torch.as_tensor(
	guidance_schedule, dtype=torch.float32, device=self.device
	)
	if gw_per_step.shape != (num_steps,):
	raise ValueError(
	f"guidance_schedule must have shape ({num_steps},), "
	f"got {tuple(gw_per_step.shape)}"
	)
	else:
	gw_per_step = torch.full(
	(num_steps,), float(guidance_scale), dtype=torch.float32, device=self.device
	)

	inputs = self._build_inputs(prompts, height=height, width=width)
	batch_size = len(prompts)
	num_image_tokens = inputs["num_image_tokens"]
	grid_h, grid_w = inputs["grid_h"], inputs["grid_w"]
	max_text_tokens = inputs["max_text_tokens"]
	latent_dim = self.conditional_transformer.config.in_channels

	llm_features = self._encode_text(
	inputs["token_ids"], inputs["text_position_ids"], inputs["indicator"]
	)

	# Negative branch is image-only (asymmetric CFG) with zeroed conditioning.
	neg_position_ids = inputs["position_ids"][:, max_text_tokens:]
	neg_segment_ids = inputs["segment_ids"][:, max_text_tokens:]
	neg_indicator = inputs["indicator"][:, max_text_tokens:]
	neg_llm_features = torch.zeros( # type: ignore[call-overload]
	batch_size,
	num_image_tokens,
	llm_features.shape[-1],
	dtype=llm_features.dtype,
	device=self.device,
	)

	generator = torch.Generator(device=self.device)
	if seed is not None:
	generator.manual_seed(seed)
	z = torch.randn( # type: ignore[call-overload]
	batch_size,
	num_image_tokens,
	latent_dim,
	dtype=torch.float32,
	device=self.device,
	generator=generator,
	)

	text_z_padding = torch.zeros( # type: ignore[call-overload]
	batch_size,
	max_text_tokens,
	latent_dim,
	dtype=torch.float32,
	device=self.device,
	)

	for i in range(num_steps - 1, -1, -1):
	t_val = float(schedule(step_intervals[i + 1].unsqueeze(0)).item())
	s_val = float(schedule(step_intervals[i].unsqueeze(0)).item())
	t = torch.full((batch_size,), t_val, dtype=torch.float32, device=self.device)

	pos_z = torch.cat([text_z_padding, z], dim=1)
	pos_out = self.conditional_transformer(
	llm_features=llm_features,
	x=pos_z,
	t=t,
	position_ids=inputs["position_ids"],
	segment_ids=inputs["segment_ids"],
	indicator=inputs["indicator"],
	)
	pos_v = pos_out[:, max_text_tokens:]

	neg_v = self.unconditional_transformer(
	llm_features=neg_llm_features,
	x=z,
	t=t,
	position_ids=neg_position_ids,
	segment_ids=neg_segment_ids,
	indicator=neg_indicator,
	)

	gw_i = gw_per_step[i]
	v = gw_i * pos_v + (1.0 - gw_i) * neg_v
	delta = s_val - t_val
	z = z + v * delta

	return self._decode(z, grid_h=grid_h, grid_w=grid_w) # type: ignore[arg-type]

	def _decode(self, z: torch.Tensor, *, grid_h: int, grid_w: int) -> list[Image.Image]:
	"""Unpatch and run the autoencoder decoder."""
	batch_size = z.shape[0]
	patch = self.config.patch_size

	z = z * self.latent_scale + self.latent_shift

	ae_channels = z.shape[-1] // (patch * patch)
	z = z.view(batch_size, grid_h, grid_w, patch, patch, ae_channels)
	z = z.permute(0, 5, 1, 3, 2, 4).contiguous()
	z = z.view(batch_size, ae_channels, grid_h * patch, grid_w * patch)

	z = z.to(self.dtype)
	decoded = self.autoencoder.decoder(z)

	decoded = decoded.float().clamp(-1.0, 1.0)
	decoded = ((decoded + 1.0) * 127.5).round().to(torch.uint8)
	decoded = decoded.permute(0, 2, 3, 1).cpu().numpy()
	return [Image.fromarray(arr) for arr in decoded]