from typing import Optional, Sequence, Union import torch from PIL import Image from tqdm import tqdm from ..core.device.npu_compatible_device import get_device_type from ..diffusion import FlowMatchScheduler from ..diffusion.base_pipeline import BasePipeline, PipelineUnit from ..core import ModelConfig from ..models.ideogram4_dit import Ideogram4DiT, LLM_TOKEN_INDICATOR, OUTPUT_IMAGE_INDICATOR, IMAGE_POSITION_OFFSET from ..models.ideogram4_text_encoder import Ideogram4TextEncoder from ..models.flux2_vae import Flux2VAE from ..models.ideogram4_vae import encode, decode from transformers import AutoTokenizer class Ideogram4Pipeline(BasePipeline): def __init__(self, device=get_device_type(), torch_dtype=torch.bfloat16): super().__init__( device=device, torch_dtype=torch_dtype, height_division_factor=16, width_division_factor=16, ) self.scheduler = FlowMatchScheduler("Ideogram4") self.text_encoder: Ideogram4TextEncoder = None self.dit: Ideogram4DiT = None self.dit_uncond: Ideogram4DiT = None self.vae: Flux2VAE = None self.tokenizer: AutoTokenizer = None self.in_iteration_models = ("dit", "dit_uncond") self.units = [ Ideogram4Unit_ShapeChecker(), Ideogram4Unit_PromptEmbedder(), Ideogram4Unit_NoiseInitializer(), Ideogram4Unit_InputImageEmbedder(), ] self.model_fn = model_fn_ideogram4 @staticmethod def from_pretrained( torch_dtype: torch.dtype = torch.bfloat16, device: Union[str, torch.device] = get_device_type(), model_configs: list[ModelConfig] = [], tokenizer_config: ModelConfig = None, vram_limit: float = None, ): pipe = Ideogram4Pipeline(device=device, torch_dtype=torch_dtype) model_pool = pipe.download_and_load_models(model_configs, vram_limit) transformers = model_pool.fetch_model("ideogram4_dit", index=2) if isinstance(transformers, list): pipe.dit = transformers[0] pipe.dit_uncond = transformers[1] else: pipe.dit = transformers pipe.text_encoder = model_pool.fetch_model("ideogram4_text_encoder") pipe.vae = model_pool.fetch_model("flux2_vae") if tokenizer_config is not None: tokenizer_config.download_if_necessary() pipe.tokenizer = AutoTokenizer.from_pretrained(tokenizer_config.path) pipe.vram_management_enabled = pipe.check_vram_management_state() return pipe @torch.no_grad() def __call__( self, # Prompt prompt: str = "", negative_prompt: str = "", cfg_scale: float = 7.0, # Input image input_image: Image.Image = None, denoising_strength: float = 1.0, # Shape height: int = 1024, width: int = 1024, # Randomness seed: int = None, rand_device: str = "cpu", # Steps num_inference_steps: int = 50, # Progress bar progress_bar_cmd=tqdm, ): self.scheduler.set_timesteps(num_inference_steps, denoising_strength=denoising_strength, image_resolution=(height, width)) inputs_posi = { "prompt": prompt, } inputs_nega = { "prompt": negative_prompt, } inputs_shared = { "cfg_scale": cfg_scale, "input_image": input_image, "denoising_strength": denoising_strength, "height": height, "width": width, "seed": seed, "rand_device": rand_device, "num_inference_steps": num_inference_steps, } for unit in self.units: inputs_shared, inputs_posi, inputs_nega = self.unit_runner(unit, self, inputs_shared, inputs_posi, inputs_nega) self.load_models_to_device(self.in_iteration_models) for progress_id, timestep in enumerate(progress_bar_cmd(self.scheduler.timesteps)): timestep = timestep.unsqueeze(0).to(dtype=torch.float32, device=self.device) models = {"dit": self.dit} noise_pred_posi = self.model_fn(timestep=timestep, **models, **inputs_shared, **inputs_posi) if cfg_scale != 1: models = {"dit": self.dit_uncond if self.dit_uncond is not None else self.dit} noise_pred_nega = self.model_fn(timestep=timestep, **models, **inputs_shared, **inputs_nega) noise_pred = noise_pred_nega + cfg_scale * (noise_pred_posi - noise_pred_nega) else: noise_pred = noise_pred_posi inputs_shared["latents"] = self.step(self.scheduler, progress_id=progress_id, noise_pred=noise_pred, **inputs_shared) # Decode self.load_models_to_device(["vae"]) image = decode(self.vae, inputs_shared["latents"], height, width, self.torch_dtype) image = self.vae_output_to_image(image) self.load_models_to_device([]) return image class Ideogram4Unit_ShapeChecker(PipelineUnit): def __init__(self): super().__init__( input_params=("height", "width"), output_params=("height", "width"), ) def process(self, pipe: "Ideogram4Pipeline", height, width): height, width = pipe.check_resize_height_width(height, width) return {"height": height, "width": width} class Ideogram4Unit_PromptEmbedder(PipelineUnit): def __init__(self): super().__init__( take_over=True, output_params=("llm_features", "position_ids", "segment_ids", "indicator", "max_text_tokens"), onload_model_names=("text_encoder",) ) def process(self, pipe: "Ideogram4Pipeline", inputs_shared, inputs_posi, inputs_nega): prompt = inputs_posi.get("prompt", "") height = inputs_shared.get("height") width = inputs_shared.get("width") max_text_tokens = 2048 pipe.load_models_to_device(self.onload_model_names) messages = [{"role": "user", "content": [{"type": "text", "text": prompt}]}] text = pipe.tokenizer.apply_chat_template( messages, add_generation_prompt=True, tokenize=False ) encoded = pipe.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 > max_text_tokens: raise ValueError( f"prompt has {num_text_tokens} tokens, exceeds max_text_tokens={max_text_tokens}" ) patch = pipe.dit.patch_size * 8 grid_h = height // patch grid_w = width // patch num_image_tokens = grid_h * grid_w max_text_tokens = num_text_tokens total_seq_len = max_text_tokens + num_image_tokens 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_padded = torch.zeros(1, total_seq_len, dtype=torch.long) text_position_ids = torch.zeros(1, total_seq_len, 3, dtype=torch.long) position_ids = torch.zeros(1, total_seq_len, 3, dtype=torch.long) segment_ids = torch.zeros(1, total_seq_len, dtype=torch.long) indicator = torch.zeros(1, total_seq_len, dtype=torch.long) token_ids_padded[0, :num_text_tokens] = token_ids text_pos = torch.arange(num_text_tokens) text_pos_3d = torch.stack([text_pos, text_pos, text_pos], dim=1) text_position_ids[0, :num_text_tokens] = text_pos_3d position_ids[0, :num_text_tokens] = text_pos_3d position_ids[0, num_text_tokens:] = image_pos indicator[0, :num_text_tokens] = LLM_TOKEN_INDICATOR indicator[0, num_text_tokens:] = OUTPUT_IMAGE_INDICATOR segment_ids[0, :total_seq_len] = 1 token_ids_padded = token_ids_padded.to(pipe.device) text_position_ids = text_position_ids.to(pipe.device) position_ids = position_ids.to(pipe.device) segment_ids = segment_ids.to(pipe.device) indicator = indicator.to(pipe.device) attention_mask = (indicator == LLM_TOKEN_INDICATOR).to(torch.long) pos_2d = text_position_ids[..., 0].contiguous() with torch.no_grad(): llm_features = pipe.text_encoder(token_ids_padded, attention_mask, pos_2d) text_mask = attention_mask.to(llm_features.dtype).unsqueeze(-1) llm_features = llm_features * text_mask llm_features = llm_features.to(torch.float32) inputs_posi.update({ "llm_features": llm_features, "position_ids": position_ids, "segment_ids": segment_ids, "indicator": indicator, "max_text_tokens": max_text_tokens, }) inputs_nega.update({ "llm_features": torch.zeros(1, num_image_tokens, llm_features.shape[-1], dtype=llm_features.dtype, device=llm_features.device), "position_ids": position_ids[:, max_text_tokens:], "segment_ids": segment_ids[:, max_text_tokens:], "indicator": indicator[:, max_text_tokens:], "max_text_tokens": 0, }) return inputs_shared, inputs_posi, inputs_nega class Ideogram4Unit_NoiseInitializer(PipelineUnit): def __init__(self): super().__init__( input_params=("height", "width", "seed", "rand_device"), output_params=("noise",), ) def process(self, pipe: "Ideogram4Pipeline", height, width, seed, rand_device): patch = pipe.dit.patch_size * 8 grid_h = height // patch grid_w = width // patch num_image_tokens = grid_h * grid_w latent_dim = pipe.dit.config.in_channels noise = pipe.generate_noise((1, num_image_tokens, latent_dim), seed=seed, rand_device=rand_device, rand_torch_dtype=torch.float32) return {"noise": noise, "grid_h": grid_h, "grid_w": grid_w} class Ideogram4Unit_InputImageEmbedder(PipelineUnit): def __init__(self): super().__init__( input_params=("input_image", "noise", "height", "width"), output_params=("latents", "input_latents"), onload_model_names=("vae",) ) def process(self, pipe: "Ideogram4Pipeline", input_image, noise, height, width): if input_image is None: return {"latents": noise, "input_latents": None} pipe.load_models_to_device(["vae"]) image = pipe.preprocess_image(input_image) input_latents = encode(pipe.vae, image, height, width, torch.bfloat16) if pipe.scheduler.training: return {"latents": noise, "input_latents": input_latents} else: latents = pipe.scheduler.add_noise(input_latents, noise, timestep=pipe.scheduler.timesteps[0]) return {"latents": latents, "input_latents": input_latents} def model_fn_ideogram4( dit: Ideogram4DiT = None, latents=None, timestep=None, llm_features=None, position_ids=None, segment_ids=None, indicator=None, max_text_tokens=0, use_gradient_checkpointing=False, use_gradient_checkpointing_offload=False, **kwargs, ): t_ideogram4 = timestep.to(torch.float32) text_z_padding = torch.zeros( 1, max_text_tokens, latents.shape[-1], dtype=torch.float32, device=latents.device, ) z = torch.cat([text_z_padding, latents], dim=1) out = dit( llm_features=llm_features, x=z, t=t_ideogram4, position_ids=position_ids, segment_ids=segment_ids, indicator=indicator, use_gradient_checkpointing=use_gradient_checkpointing, use_gradient_checkpointing_offload=use_gradient_checkpointing_offload, ) return -out[:, max_text_tokens:]