diff --git a/docs/source/en/api/pipelines/ideogram4.md b/docs/source/en/api/pipelines/ideogram4.md index a4639ff23f94..3c8553a9fdad 100644 --- a/docs/source/en/api/pipelines/ideogram4.md +++ b/docs/source/en/api/pipelines/ideogram4.md @@ -40,12 +40,78 @@ image = pipe(prompt, height=1024, width=1024, generator=torch.Generator("cuda"). image.save("ideogram4.png") ``` +## Prompt upsampling + +Ideogram 4 is trained on a structured JSON caption rather than a free-form prompt, so a short prompt is best +expanded into that native schema before generation. There are two ways to produce the caption. + +### Remote (Ideogram API) + +For the best results, expand the prompt with Ideogram's hosted magic-prompt API and pass the returned caption +straight to the pipeline (get a key at [developer.ideogram.ai](https://developer.ideogram.ai/)): + +```python +import json +import requests +import torch +from diffusers import Ideogram4Pipeline + +pipe = Ideogram4Pipeline.from_pretrained("ideogram-ai/ideogram-4-nf4", torch_dtype=torch.bfloat16) +pipe.to("cuda") + +# Expand the prompt into a structured JSON caption with Ideogram's hosted magic-prompt API. +response = requests.post( + "https://api.ideogram.ai/v1/ideogram-v4/magic-prompt", + headers={"Api-Key": "your_ideogram_api_key"}, + json={"text_prompt": "A photo of a cat holding a sign that says hello world", "aspect_ratio": "1x1"}, +).json() +caption = json.dumps(response["json_prompt"]) + +# The caption is already upsampled, so pass it directly (no prompt_upsampling). +image = pipe(caption, height=1024, width=1024, generator=torch.Generator("cuda").manual_seed(0)).images[0] +image.save("ideogram4_upsampled.png") +``` + +### Local (on-device) + +For a fully local pipeline, load a small [`Ideogram4PromptEnhancerHead`] (the Qwen3-VL LM head) as the optional +`prompt_enhancer_head` component and pass `prompt_upsampling=True`. The head is grafted onto the shared +`text_encoder`, so no second text encoder is loaded. Install `outlines` for schema-constrained captions (the nf4 +checkpoint also needs `bitsandbytes`): + +```python +import torch +from diffusers import Ideogram4Pipeline, Ideogram4PromptEnhancerHead + +prompt_enhancer_head = Ideogram4PromptEnhancerHead.from_pretrained( + "diffusers/qwen3-vl-8b-instruct-lm-head", torch_dtype=torch.bfloat16 +) +pipe = Ideogram4Pipeline.from_pretrained( + "ideogram-ai/ideogram-4-nf4", prompt_enhancer_head=prompt_enhancer_head, torch_dtype=torch.bfloat16 +) +pipe.to("cuda") + +prompt = "A photo of a cat holding a sign that says hello world" +image = pipe( + prompt, + height=1024, + width=1024, + prompt_upsampling=True, + generator=torch.Generator("cuda").manual_seed(0), +).images[0] +image.save("ideogram4_upsampled.png") +``` + ## Ideogram4Pipeline [[autodoc]] Ideogram4Pipeline - all - __call__ +## Ideogram4PromptEnhancerHead + +[[autodoc]] Ideogram4PromptEnhancerHead + ## Ideogram4PipelineOutput [[autodoc]] pipelines.ideogram4.pipeline_output.Ideogram4PipelineOutput diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index a8957183ef99..4ba69324c4e3 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -594,6 +594,7 @@ "HunyuanVideoPipeline", "I2VGenXLPipeline", "Ideogram4Pipeline", + "Ideogram4PromptEnhancerHead", "IFImg2ImgPipeline", "IFImg2ImgSuperResolutionPipeline", "IFInpaintingPipeline", @@ -1413,6 +1414,7 @@ HunyuanVideoPipeline, I2VGenXLPipeline, Ideogram4Pipeline, + Ideogram4PromptEnhancerHead, IFImg2ImgPipeline, IFImg2ImgSuperResolutionPipeline, IFInpaintingPipeline, diff --git a/src/diffusers/modular_pipelines/ideogram4/encoders.py b/src/diffusers/modular_pipelines/ideogram4/encoders.py index e02e74856c3d..a8aac089c52d 100644 --- a/src/diffusers/modular_pipelines/ideogram4/encoders.py +++ b/src/diffusers/modular_pipelines/ideogram4/encoders.py @@ -17,7 +17,14 @@ from transformers import Qwen2Tokenizer, Qwen3VLModel from transformers.masking_utils import create_causal_mask -from ...utils import logging +from ...pipelines.ideogram4.prompt_enhancer import ( + PROMPT_UPSAMPLE_TEMPERATURE, + Ideogram4PromptEnhancerHead, + build_caption_logits_processor, + build_prompt_enhancer, + generate_captions, +) +from ...utils import is_outlines_available, logging from ..modular_pipeline import ModularPipelineBlocks, PipelineState from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam from .modular_pipeline import Ideogram4ModularPipeline @@ -31,6 +38,139 @@ QWEN3_VL_ACTIVATION_LAYERS = (0, 3, 6, 9, 12, 15, 18, 21, 24, 27, 30, 33, 35) +# auto_docstring +class Ideogram4PromptUpsampleStep(ModularPipelineBlocks): + """ + Optional step that rewrites the prompt(s) into Ideogram4's native structured JSON caption (the format the model is + trained on) when ``prompt_upsampling=True``. Requires the optional ``prompt_enhancer_head`` component, which is + grafted onto the shared ``text_encoder`` body to make it generative; install ``outlines`` for schema-constrained + captions. + + Components: + text_encoder (`Qwen3VLModel`): The Qwen3-VL text encoder. tokenizer (`Qwen2Tokenizer`): The tokenizer paired + with the text encoder. prompt_enhancer_head (`Ideogram4PromptEnhancerHead`): The LM head grafted onto the + text encoder for upsampling. + + Inputs: + prompt (`str`): + The prompt or prompts to guide image generation. + prompt_upsampling (`bool`, *optional*, defaults to False): + If True, rewrite the prompt into the native JSON caption before encoding. + prompt_upsampling_temperature (`float`, *optional*, defaults to 1.0): + Sampling temperature for prompt upsampling. + height (`int`, *optional*): + Together with width, sets the caption's target aspect ratio. + width (`int`, *optional*): + Together with height, sets the caption's target aspect ratio. + generator (`Generator`, *optional*): + Reused to make the upsampling reproducible. + + Outputs: + prompt (`str`): + The (possibly upsampled) prompt forwarded to the text encoder. + """ + + model_name = "ideogram4" + + def __init__(self): + # Built lazily on first upsample: the head-less encoder body + `prompt_enhancer_head`, combined. + self._prompt_enhancer = None + # Outlines logits processor for schema-constrained captions; built lazily on first upsample. + self._caption_logits_processor = None + super().__init__() + + @property + def description(self) -> str: + return ( + "Optional step that rewrites the prompt(s) into Ideogram4's native structured JSON caption when " + "`prompt_upsampling=True` (the format the model is trained on). Requires a generative `text_encoder` " + "(a `Qwen3VLForConditionalGeneration`); install `outlines` for schema-constrained captions." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("text_encoder", Qwen3VLModel, description="The Qwen3-VL text encoder."), + ComponentSpec("tokenizer", Qwen2Tokenizer, description="The tokenizer paired with the text encoder."), + ComponentSpec( + "prompt_enhancer_head", + Ideogram4PromptEnhancerHead, + description="LM head grafted onto the text encoder for prompt upsampling.", + ), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("prompt", required=True), + InputParam( + name="prompt_upsampling", + type_hint=bool, + default=False, + description="If True, rewrite the prompt into Ideogram4's native JSON caption before encoding.", + ), + InputParam( + name="prompt_upsampling_temperature", + type_hint=float, + default=PROMPT_UPSAMPLE_TEMPERATURE, + description="Sampling temperature for prompt upsampling.", + ), + InputParam.template("height"), + InputParam.template("width"), + InputParam.template("max_sequence_length", default=2048), + InputParam.template("generator"), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + name="prompt", + type_hint=list, + description="The (possibly upsampled) prompt forwarded to the text encoder.", + ), + ] + + @torch.no_grad() + def __call__(self, components: Ideogram4ModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + if block_state.prompt_upsampling: + if components.prompt_enhancer_head is None: + raise ValueError( + "Prompt upsampling requires the `prompt_enhancer_head` component, which is not loaded. Load an " + "`Ideogram4PromptEnhancerHead` and add it to the pipeline." + ) + if self._prompt_enhancer is None: + self._prompt_enhancer = build_prompt_enhancer(components.text_encoder, components.prompt_enhancer_head) + if self._caption_logits_processor is None and is_outlines_available(): + self._caption_logits_processor = build_caption_logits_processor( + self._prompt_enhancer, components.tokenizer + ) + if self._caption_logits_processor is None: + logger.warning_once( + "`outlines` is not installed; prompt upsampling runs unconstrained and may not return " + "schema-valid JSON. Install with `pip install outlines` for structured captions." + ) + height = block_state.height or components.default_height + width = block_state.width or components.default_width + block_state.prompt = generate_captions( + self._prompt_enhancer, + components.tokenizer, + self._caption_logits_processor, + block_state.prompt, + height, + width, + temperature=block_state.prompt_upsampling_temperature, + max_new_tokens=block_state.max_sequence_length, + generator=block_state.generator, + device=components._execution_device, + ) + + self.set_block_state(state, block_state) + return components, state + + # auto_docstring class Ideogram4TextEncoderStep(ModularPipelineBlocks): """ diff --git a/src/diffusers/modular_pipelines/ideogram4/modular_blocks_ideogram4.py b/src/diffusers/modular_pipelines/ideogram4/modular_blocks_ideogram4.py index 90cddb915ad4..5c8a726b76c4 100644 --- a/src/diffusers/modular_pipelines/ideogram4/modular_blocks_ideogram4.py +++ b/src/diffusers/modular_pipelines/ideogram4/modular_blocks_ideogram4.py @@ -24,7 +24,7 @@ ) from .decoders import Ideogram4DecodeStep from .denoise import Ideogram4AfterDenoiseStep, Ideogram4DenoiseStep -from .encoders import Ideogram4TextEncoderStep +from .encoders import Ideogram4PromptUpsampleStep, Ideogram4TextEncoderStep logger = logging.get_logger(__name__) # pylint: disable=invalid-name @@ -123,6 +123,10 @@ class Ideogram4AutoBlocks(SequentialPipelineBlocks): Inputs: prompt (`str`): The prompt or prompts to guide image generation. + prompt_upsampling (`bool`, *optional*, defaults to False): + Rewrite the prompt into Ideogram4's native structured JSON caption before encoding. + prompt_upsampling_temperature (`float`, *optional*, defaults to 1.0): + Sampling temperature for prompt upsampling. max_sequence_length (`int`, *optional*, defaults to 2048): Maximum sequence length for prompt encoding. num_images_per_prompt (`int`, *optional*, defaults to 1): @@ -154,8 +158,13 @@ class Ideogram4AutoBlocks(SequentialPipelineBlocks): """ model_name = "ideogram4" - block_classes = [Ideogram4TextEncoderStep(), Ideogram4CoreDenoiseStep(), Ideogram4DecodeStep()] - block_names = ["text_encoder", "denoise", "decode"] + block_classes = [ + Ideogram4PromptUpsampleStep(), + Ideogram4TextEncoderStep(), + Ideogram4CoreDenoiseStep(), + Ideogram4DecodeStep(), + ] + block_names = ["prompt_upsample", "text_encoder", "denoise", "decode"] # Workflow map declaring the trigger conditions for each supported workflow. # `True` means the workflow triggers when the input is not None. @@ -166,8 +175,8 @@ class Ideogram4AutoBlocks(SequentialPipelineBlocks): @property def description(self) -> str: return ( - "Auto Modular pipeline for text-to-image generation using Ideogram4: encode text -> core denoise " - "(asymmetric CFG over two transformers) -> decode." + "Auto Modular pipeline for text-to-image generation using Ideogram4: (optional) prompt upsampling -> " + "encode text -> core denoise (asymmetric CFG over two transformers) -> decode." ) @property diff --git a/src/diffusers/pipelines/__init__.py b/src/diffusers/pipelines/__init__.py index 7c38837f308d..db0e863127db 100644 --- a/src/diffusers/pipelines/__init__.py +++ b/src/diffusers/pipelines/__init__.py @@ -288,7 +288,7 @@ ] _import_structure["hunyuan_video1_5"] = ["HunyuanVideo15Pipeline", "HunyuanVideo15ImageToVideoPipeline"] _import_structure["hunyuan_image"] = ["HunyuanImagePipeline", "HunyuanImageRefinerPipeline"] - _import_structure["ideogram4"] = ["Ideogram4Pipeline"] + _import_structure["ideogram4"] = ["Ideogram4Pipeline", "Ideogram4PromptEnhancerHead"] _import_structure["kandinsky"] = [ "KandinskyCombinedPipeline", "KandinskyImg2ImgCombinedPipeline", @@ -748,7 +748,7 @@ ) from .hunyuan_video1_5 import HunyuanVideo15ImageToVideoPipeline, HunyuanVideo15Pipeline from .hunyuandit import HunyuanDiTPipeline - from .ideogram4 import Ideogram4Pipeline + from .ideogram4 import Ideogram4Pipeline, Ideogram4PromptEnhancerHead from .joyimage import JoyImageEditPipeline, JoyImageEditPipelineOutput from .kandinsky import ( KandinskyCombinedPipeline, diff --git a/src/diffusers/pipelines/ideogram4/__init__.py b/src/diffusers/pipelines/ideogram4/__init__.py index 2fd99fbc3500..28ba34775ddd 100644 --- a/src/diffusers/pipelines/ideogram4/__init__.py +++ b/src/diffusers/pipelines/ideogram4/__init__.py @@ -25,6 +25,8 @@ _import_structure["pipeline_output"] = ["Ideogram4PipelineOutput"] + _import_structure["prompt_enhancer"] = ["Ideogram4PromptEnhancerHead"] + if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: try: if not (is_transformers_available() and is_torch_available()): @@ -34,6 +36,7 @@ else: from .pipeline_ideogram4 import Ideogram4Pipeline from .pipeline_output import Ideogram4PipelineOutput + from .prompt_enhancer import Ideogram4PromptEnhancerHead else: import sys diff --git a/src/diffusers/pipelines/ideogram4/pipeline_ideogram4.py b/src/diffusers/pipelines/ideogram4/pipeline_ideogram4.py index 4a02d4e3d0f7..61ba4fa43a62 100644 --- a/src/diffusers/pipelines/ideogram4/pipeline_ideogram4.py +++ b/src/diffusers/pipelines/ideogram4/pipeline_ideogram4.py @@ -29,10 +29,17 @@ Ideogram4Transformer2DModel, ) from ...schedulers import FlowMatchEulerDiscreteScheduler -from ...utils import logging, replace_example_docstring +from ...utils import is_outlines_available, logging, replace_example_docstring from ...utils.torch_utils import randn_tensor from ..pipeline_utils import DiffusionPipeline from .pipeline_output import Ideogram4PipelineOutput +from .prompt_enhancer import ( + PROMPT_UPSAMPLE_TEMPERATURE, + Ideogram4PromptEnhancerHead, + build_caption_logits_processor, + build_prompt_enhancer, + generate_captions, +) logger = logging.get_logger(__name__) # pylint: disable=invalid-name @@ -155,8 +162,8 @@ class Ideogram4Pipeline(DiffusionPipeline): Unconditional (asymmetric-CFG) flow-matching transformer. """ - model_cpu_offload_seq = "text_encoder->transformer->unconditional_transformer->vae" - _optional_components = [] + model_cpu_offload_seq = "prompt_enhancer_head->text_encoder->transformer->unconditional_transformer->vae" + _optional_components = ["prompt_enhancer_head"] _callback_tensor_inputs = ["latents"] def __init__( @@ -167,6 +174,7 @@ def __init__( tokenizer: AutoTokenizer, transformer: Ideogram4Transformer2DModel, unconditional_transformer: Ideogram4Transformer2DModel, + prompt_enhancer_head: Ideogram4PromptEnhancerHead | None = None, ) -> None: super().__init__() @@ -177,6 +185,7 @@ def __init__( tokenizer=tokenizer, transformer=transformer, unconditional_transformer=unconditional_transformer, + prompt_enhancer_head=prompt_enhancer_head, ) self.vae_scale_factor = ( @@ -186,6 +195,58 @@ def __init__( self.patch_size = 2 self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor * self.patch_size) + # Built lazily on first upsample: the head-less encoder body + `prompt_enhancer_head`, combined. + self._prompt_enhancer = None + # Outlines logits processor for schema-constrained captions; built lazily on first upsample. + self._caption_logits_processor = None + + def upsample_prompt( + self, + prompt: str | list[str], + height: int = 2048, + width: int = 2048, + temperature: float = PROMPT_UPSAMPLE_TEMPERATURE, + max_new_tokens: int = 1024, + generator: torch.Generator | list[torch.Generator] | None = None, + device: torch.device | None = None, + ) -> list[str]: + """Rewrite each prompt into Ideogram4's native structured JSON caption. + + Requires the optional `prompt_enhancer_head` component, which is grafted onto the shared `text_encoder` body to + make it generative. Generation is schema-constrained when `outlines` is installed, otherwise it runs + unconstrained. Pass `generator` (the same one accepted by `__call__`) to make sampling reproducible. + """ + if self.prompt_enhancer_head is None: + raise ValueError( + "Prompt upsampling requires the `prompt_enhancer_head` component, which is not loaded. Load it and " + "pass it in, e.g.:\n" + " from diffusers import Ideogram4PromptEnhancerHead\n" + " head = Ideogram4PromptEnhancerHead.from_pretrained('diffusers/qwen3-vl-8b-instruct-lm-head')\n" + " pipe = Ideogram4Pipeline.from_pretrained(model_id, prompt_enhancer_head=head)" + ) + if self._prompt_enhancer is None: + self._prompt_enhancer = build_prompt_enhancer(self.text_encoder, self.prompt_enhancer_head) + if self._caption_logits_processor is None and is_outlines_available(): + self._caption_logits_processor = build_caption_logits_processor(self._prompt_enhancer, self.tokenizer) + if self._caption_logits_processor is None: + logger.warning_once( + "`outlines` is not installed; prompt upsampling runs unconstrained and may not return schema-valid " + "JSON. Install with `pip install outlines` for structured captions." + ) + + return generate_captions( + self._prompt_enhancer, + self.tokenizer, + self._caption_logits_processor, + prompt, + height, + width, + temperature=temperature, + max_new_tokens=max_new_tokens, + generator=generator, + device=device, + ) + @staticmethod def _prepare_ids( text_lengths: list[int], @@ -416,6 +477,8 @@ def __call__( guidance_schedule: list[float] | torch.Tensor | None = (7.0,) * 45 + (3.0,) * 3, mu: float = 0.0, std: float = 1.5, + prompt_upsampling: bool = False, + prompt_upsampling_temperature: float = PROMPT_UPSAMPLE_TEMPERATURE, max_sequence_length: int = 2048, num_images_per_prompt: int = 1, generator: torch.Generator | list[torch.Generator] | None = None, @@ -451,6 +514,13 @@ def __call__( the resolution ratio relative to 512x512. std (`float`, *optional*, defaults to 1.5): Standard deviation of the logit-normal flow-matching schedule. + prompt_upsampling (`bool`, *optional*, defaults to `False`): + If `True`, rewrite `prompt` into Ideogram4's native structured JSON caption via + [`~Ideogram4Pipeline.upsample_prompt`] before encoding. Requires the optional `prompt_enhancer_head` + component; install `outlines` for schema-constrained captions. `generator` is reused to make the + upsampling reproducible. + prompt_upsampling_temperature (`float`, *optional*, defaults to 1.0): + Sampling temperature for prompt upsampling when `prompt_upsampling=True`. max_sequence_length (`int`, *optional*, defaults to 2048): Maximum number of text tokens per prompt. num_images_per_prompt (`int`, *optional*, defaults to 1): @@ -492,6 +562,18 @@ def __call__( self._guidance_scale = guidance_scale self._interrupt = False + # 0. Optionally rewrite the prompt(s) into Ideogram4's native structured JSON caption. + if prompt_upsampling: + prompt = self.upsample_prompt( + prompt, + height=height, + width=width, + temperature=prompt_upsampling_temperature, + max_new_tokens=max_sequence_length, + generator=generator, + device=device, + ) + # 1. Image grid (drives both the packed layout and the latent shape). grid_h, grid_w = ( height // (self.vae_scale_factor * self.patch_size), diff --git a/src/diffusers/pipelines/ideogram4/prompt_enhancer.py b/src/diffusers/pipelines/ideogram4/prompt_enhancer.py new file mode 100644 index 000000000000..2b7867afe650 --- /dev/null +++ b/src/diffusers/pipelines/ideogram4/prompt_enhancer.py @@ -0,0 +1,218 @@ +# Copyright 2026 Ideogram AI and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Prompt-enhancement assets for Ideogram4. + +Ideogram4 is trained on a *structured JSON caption* rather than a free-form prompt. The optional prompt enhancer +rewrites a short user idea into that native caption schema by combining the head-less Qwen3-VL text encoder with the +optional `Ideogram4PromptEnhancerHead` component to form a generative model. + +This mirrors the role of Flux2's `system_messages.py`, but the target is a constrained JSON object instead of free +text, so `outlines` (an optional dependency) is used to guarantee a schema-valid result when available. + +The caption helpers here are shared by `Ideogram4Pipeline` and the modular `Ideogram4PromptUpsampleStep`. +""" + +import math + +import torch +from torch import nn + +from ...configuration_utils import ConfigMixin, register_to_config +from ...models.modeling_utils import ModelMixin +from ...utils import logging + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + +PROMPT_UPSAMPLE_TEMPERATURE = 1.0 + + +class Ideogram4PromptEnhancerHead(ModelMixin, ConfigMixin): + """LM head that makes the head-less Qwen3-VL `text_encoder` generative for prompt upsampling. + + An optional pipeline component (`prompt_enhancer_head`): its weights load via a normal `from_pretrained` (its own + small repo, or bundled in the model repo) rather than an in-pipeline download. At upsample time the pipeline + combines it with the shared `text_encoder` body to form the generative model. + """ + + @register_to_config + def __init__(self, hidden_size: int = 4096, vocab_size: int = 151936): + super().__init__() + self.lm_head = nn.Linear(hidden_size, vocab_size, bias=False) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.lm_head(hidden_states) + + +def build_prompt_enhancer(text_encoder, prompt_enhancer_head): + """Combine the head-less Qwen3-VL `text_encoder` body with `prompt_enhancer_head` into a generative model. + + The body is shared by reference (no second copy in memory); only the small head is extra. + """ + from accelerate import init_empty_weights + from transformers import Qwen3VLForConditionalGeneration + + with init_empty_weights(): + enhancer = Qwen3VLForConditionalGeneration(text_encoder.config) + enhancer.model = text_encoder + enhancer.lm_head = prompt_enhancer_head.lm_head + return enhancer.eval() + + +# System message that instructs the encoder to emit Ideogram4's native single-line JSON caption. +# docstyle-ignore +CAPTION_SYSTEM_MESSAGE = """You convert a short user idea into a structured JSON caption for an image renderer. Output ONE minified single-line JSON object and NOTHING else (no markdown, no commentary). + +SCHEMA — keys in this exact order: +{"high_level_description":"...","compositional_deconstruction":{"background":"...","elements":[ ... ]}} +- object element: {"type":"obj","desc":"..."} +- text element: {"type":"text","text":"VERBATIM CHARS","desc":"..."} + +STEP 1 — PICK THE MEDIUM. It decides what `background` and `elements` mean. Honor any medium or style the user implies; default to photograph only when nothing else fits. Render ANY subject faithfully — real, fantastical, sci-fi, surreal, abstract — in the chosen medium. + +A) DESIGNED ARTIFACT — poster, logo, album/book cover, flyer, banner, sticker, packaging, app icon, infographic, menu, card, wordmark. THE FRAME IS THE ARTIFACT — never a photo of it hanging in a room. + - high_level_description: name it as graphic design (e.g. "a minimalist jazz poster, flat graphic design..."). + - background: the design's OWN backdrop only — a flat color, gradient, or simple texture filling the frame. No room, wall, floor, easel, depth, or camera/photo language. + - elements: the design's parts as a flat 2D layout — a `text` element for every headline/label (verbatim), `obj` elements for the central graphic/illustration/shapes/badges. Place by region (top / center / bottom). + +B) SCENE — a photograph, illustration, painting, 3D render, anime frame, etc. of a real or imagined place or subject. + - high_level_description: one sentence naming the subject and the medium/style. + - background: the scene SHELL — surroundings, ground/sky/walls, atmosphere, ambient light; concrete and specific. The ground/floor/water surface lives here, never as an element. + - elements: the main subject FIRST as an `obj`, then supporting `obj` elements (props, secondary subjects) that plausibly belong. Add `text` elements only where the scene would really carry text (signs, labels, brands). + +C) ABSTRACT / CONCEPTUAL — "nostalgia", "chaos and order", "sound waves", pure pattern. Concretize the idea into a deliberate visual composition. + - background: the dominant color field, gradient, or texture of the composition. + - elements: the shapes, forms, motifs, or symbolic objects that carry the concept, as `obj` elements. Add `text` only if the idea calls for words. + +UNIVERSAL RULES (every medium): +1. The user's core subject/concept MUST appear among the elements (as an `obj`, normally first). Naming it only in high_level_description or background is NOT enough. +2. Commit to ONE concrete value each (one color, one style, one count). No hedging: ban "various", "such as", "e.g.", "or similar", "maybe", "X or Y" for one property. +3. NEVER use a transparent, empty, or plain white background UNLESS the user explicitly says "transparent", "isolated", "sticker", or "cutout". +4. A coherent subject (one animal, person, vehicle, object) is exactly ONE element; its parts go inside its `desc`. Use separate elements for genuinely separate subjects. +5. Each `desc` is 25-55 words, identity-first, standalone. Do not mention shadows, depth of field, bokeh, lens, focus, or grain. +6. high_level_description: one sentence, at most 40 words, starts with the subject, names the medium. Preserve non-ASCII characters as-is. +7. Output STRICTLY VALID JSON: double quotes around every key and string, NO trailing commas, each element object closes with "}" right after its last value. +8. Catch the "warm" impulse. Only when you are about to describe light as "warm", "golden", "amber", or "honey", stop and check: is there a specific physical source in the scene casting that colour (candle, sunset, lamp, neon, fire)? If YES, name the source and the colour it casts instead of the mood word. If NO, you are just reaching for warmth as ambience — drop it and leave the light neutral ("soft" or "even"). Don't recolour or relight anything else; this only intercepts the warm reach, every other scene and mood the user wants is untouched. +9. Describe physical reality, not impressions. Avoid mood-words — "luminous", "radiant", "vibrant", "lush", "dynamic", "gorgeous", "stunning", "breathtaking", "mesmerizing", and metaphorical "glowing" — they produce a generic AI look (the same trap as "warm"). Use observable properties: "the cheekbone catches a small highlight", not "luminous complexion". +10. Every named thing must appear as its own element. Each subject, object, sign, and quoted phrase the user names gets its own element — quoted text (single or double quotes) becomes its own verbatim `text` element. Count the named units in the prompt; the element list must hold at least that many. Don't drop or merge them. +11. Don't add what wasn't asked for. No glitch art, wireframe overlay, body fragmentation, double-exposure, "dissolving", or extra stylization unless the prompt requests it. Asked for a cinematic photo of a journalist → render that, not a glitch-art composite. +12. Name attributes concretely, anchored to landmarks. People: skin tone, hair (colour + style), each visible garment with colour, expression, pose, one distinguishing feature. Objects: shape, material, colour, a distinctive part. Place things against named references — "resting on the lower-right corner of the table", not "on the surface". +13. Name real references by name. If the user names a brand, product, character, place, or person (Nike Dunk Low, Spider-Man, the Eiffel Tower), keep that exact name in the `desc`; don't swap it for a generic look-alike unless they ask for an anonymous one. +14. "Professional photo/headshot" of a person means professional CONTEXT — neutral attire, soft even daylight, neutral backdrop, friendly expression — not dramatic studio gear; no heavy rim-light or creamy bokeh unless asked. + +EXAMPLES + +User idea: a cup of coffee on a table +Output: {"high_level_description":"A white ceramic cup of black coffee on a worn wooden cafe table, a casual overcast-daylight phone photograph with an off-center composition.","compositional_deconstruction":{"background":"Scratched oak cafe table filling the lower frame, a pale grey mortar-lined brick wall a few feet behind slightly out of focus, a tall window on the left spilling soft overcast daylight across the table, neutral white balance, muted brown and green tones.","elements":[{"type":"obj","desc":"White ceramic cup of black coffee with a thin curved handle turned to the right and a faint crema ring at the rim, resting on a matching round saucer near the center of the table, a thin wisp of steam at the surface."},{"type":"obj","desc":"Brushed-steel teaspoon lying on the saucer to the right of the cup, handle angled toward the lower-right corner, a single small water droplet on the bowl of the spoon."}]}} + +User idea: a minimalist poster for a jazz festival +Output: {"high_level_description":"A minimalist jazz festival poster, flat graphic design with bold typography and a single abstract saxophone motif on a deep teal background.","compositional_deconstruction":{"background":"Solid deep teal background filling the entire frame with a subtle fine paper-grain texture and a thin mustard-yellow keyline border just inside the edges, no scene and no depth.","elements":[{"type":"obj","desc":"A large flat geometric saxophone in mustard yellow and cream, centered in the upper two-thirds, built from simple bold shapes with no shading, angled diagonally from lower-left to upper-right."},{"type":"text","text":"JAZZ\\nFESTIVAL","desc":"Large bold condensed sans-serif headline in cream, stacked on two lines across the center of the poster, slightly overlapping the saxophone motif."},{"type":"text","text":"NOV 15 · CITY HALL","desc":"Small uppercase mustard-yellow caption centered near the bottom edge with wide letter spacing."}]}}""" + +# User turn. `{aspect_ratio}` and `{original_prompt}` are filled in by `Ideogram4Pipeline.upsample_prompt`. +# docstyle-ignore +CAPTION_USER_TEMPLATE = """TARGET IMAGE ASPECT RATIO: {aspect_ratio} (width:height). +User idea: {original_prompt}""" + + +def build_caption_logits_processor(model, tokenizer): + """Build an `outlines` logits processor that constrains generation to the Ideogram4 caption schema. + + Returns a logits processor compatible with `transformers` `generate(logits_processor=[...])`. The caller is + responsible for checking `is_outlines_available()` first; `outlines` (and its `pydantic` dependency) are imported + lazily here so they remain optional. The schema mirrors Ideogram's native caption / caption_verifier: a high-level + description plus a compositional deconstruction of background + typed elements. + """ + from typing import List, Literal, Union + + import outlines + from pydantic import BaseModel, Field + + class ObjElement(BaseModel): + type: Literal["obj"] + desc: str + + class TextElement(BaseModel): + type: Literal["text"] + text: str + desc: str + + class Composition(BaseModel): + background: str + elements: List[Union[ObjElement, TextElement]] = Field(min_length=1) + + class Caption(BaseModel): + high_level_description: str + compositional_deconstruction: Composition + + outlines_model = outlines.from_transformers(model, tokenizer) + return outlines.Generator(outlines_model, Caption).logits_processor + + +def generate_captions( + prompt_enhancer, + tokenizer, + logits_processor, + prompt: str | list[str], + height: int, + width: int, + temperature: float = PROMPT_UPSAMPLE_TEMPERATURE, + max_new_tokens: int = 1024, + generator: torch.Generator | list[torch.Generator] | None = None, + device: torch.device | None = None, +) -> list[str]: + """Rewrite each prompt into the native structured JSON caption with the grafted `prompt_enhancer`. + + Pass `generator` to make sampling reproducible (a seed is derived from it and used inside a forked RNG so the + caller's own RNG stream is untouched). + """ + device = device or prompt_enhancer.device + prompts = [prompt] if isinstance(prompt, str) else list(prompt) + divisor = math.gcd(width, height) or 1 + aspect_ratio = f"{width // divisor}:{height // divisor}" + + sampling_seed = None + if generator is not None: + gen = generator[0] if isinstance(generator, list) else generator + sampling_seed = int(torch.randint(0, 2**63 - 1, (1,), generator=gen, device=gen.device).item()) + fork_devices = [device] if getattr(device, "type", None) == "cuda" else [] + + captions = [] + for i, text_prompt in enumerate(prompts): + messages = [ + {"role": "system", "content": CAPTION_SYSTEM_MESSAGE}, + { + "role": "user", + "content": CAPTION_USER_TEMPLATE.format(aspect_ratio=aspect_ratio, original_prompt=text_prompt), + }, + ] + inputs = tokenizer.apply_chat_template( + messages, add_generation_prompt=True, tokenize=True, return_tensors="pt", return_dict=True + ).to(device) + generate_kwargs = { + "max_new_tokens": max_new_tokens, + "do_sample": temperature > 0, + "temperature": temperature, + "use_cache": True, + } + if logits_processor is not None: + logits_processor.reset() + generate_kwargs["logits_processor"] = [logits_processor] + with torch.random.fork_rng(devices=fork_devices, enabled=sampling_seed is not None): + if sampling_seed is not None: + torch.manual_seed(sampling_seed + i) + generated = prompt_enhancer.generate(**inputs, **generate_kwargs) + new_tokens = generated[:, inputs["input_ids"].shape[1] :] + captions.append(tokenizer.decode(new_tokens[0], skip_special_tokens=True).strip()) + return captions diff --git a/src/diffusers/utils/__init__.py b/src/diffusers/utils/__init__.py index 3738821a96a1..10ad75d92f17 100644 --- a/src/diffusers/utils/__init__.py +++ b/src/diffusers/utils/__init__.py @@ -101,6 +101,7 @@ is_opencv_available, is_optimum_quanto_available, is_optimum_quanto_version, + is_outlines_available, is_peft_available, is_peft_version, is_pytorch_retinaface_available, diff --git a/src/diffusers/utils/dummy_torch_and_transformers_objects.py b/src/diffusers/utils/dummy_torch_and_transformers_objects.py index 3b8004110802..56d1edeec0b3 100644 --- a/src/diffusers/utils/dummy_torch_and_transformers_objects.py +++ b/src/diffusers/utils/dummy_torch_and_transformers_objects.py @@ -2012,6 +2012,21 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch", "transformers"]) +class Ideogram4PromptEnhancerHead(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + class IFImg2ImgPipeline(metaclass=DummyObject): _backends = ["torch", "transformers"] diff --git a/src/diffusers/utils/import_utils.py b/src/diffusers/utils/import_utils.py index 937e20d6b155..ce439bfecbf2 100644 --- a/src/diffusers/utils/import_utils.py +++ b/src/diffusers/utils/import_utils.py @@ -205,6 +205,7 @@ def _is_package_available(pkg_name: str, get_dist_name: bool = False) -> tuple[b _tensorboard_available, _tensorboard_version = _is_package_available("tensorboard") _compel_available, _compel_version = _is_package_available("compel") _sentencepiece_available, _sentencepiece_version = _is_package_available("sentencepiece") +_outlines_available, _outlines_version = _is_package_available("outlines") _torchsde_available, _torchsde_version = _is_package_available("torchsde") _peft_available, _peft_version = _is_package_available("peft") _torchvision_available, _torchvision_version = _is_package_available("torchvision") @@ -375,6 +376,10 @@ def is_sentencepiece_available(): return _sentencepiece_available +def is_outlines_available(): + return _outlines_available + + def is_imageio_available(): return _imageio_available