Upload 5 files

README.md

@@ -1,3 +1,79 @@
 ---
+library_name: diffusers
+pipeline_tag: image-to-text
+base_model: google/gemma-4-E4B-it
+tags:
+  - image-to-text
+  - captioning
+  - modular-diffusers
+  - ideogram
 license: apache-2.0
 ---
+
+# Ideogram4 caption block
+
+A custom [Modular Diffusers](https://huggingface.co/docs/diffusers/main/en/modular_diffusers/overview) block that
+produces **Ideogram 4's native structured-JSON caption** from either an image or a short text idea, using a Gemma-4
+vision-language model. The caption it returns can be fed straight into the Ideogram 4 generation pipeline as `prompt`.
+
+```
+image  -> caption it -> { high_level_description, style_description, compositional_deconstruction }
+prompt -> enhance -> { high_level_description, compositional_deconstruction }
+```
+
+The mode is chosen automatically: pass `image` to caption it, or `prompt` (with no image) to enhance a short idea.
+Prompt enhancement reuses Ideogram's canonical magic-prompt system message from
+`diffusers.pipelines.ideogram4.prompt_enhancer`.
+
+## Loading & running
+
+```python
+import torch
+from diffusers import ModularPipeline
+from diffusers.utils import load_image
+
+pipe = ModularPipeline.from_pretrained("OzzyGT/ideogram4-caption-blocks", trust_remote_code=True)
+pipe.load_components(torch_dtype=torch.bfloat16)
+pipe.to("cuda")
+
+image = load_image("your_image.png")
+caption = pipe(image=image, output="caption")   # caption the image -> JSON string
+print(caption)
+```
+
+Enhance a text idea instead (no image) — Ideogram's "magic prompt":
+
+```python
+caption = pipe(
+    prompt="a cozy coffee shop on a rainy evening",
+    output="caption",
+)
+```
+
+Get the parsed dict instead (or alongside):
+
+```python
+out = pipe(image=image, output=["caption", "caption_json"])
+out["caption_json"]   # dict, or None if the model output couldn't be parsed
+```
+
+Inputs: `image` (caption it) **or** `prompt` (enhance it); `instruction` (image-mode schema prompt), `height`/`width`
+(aspect-ratio hint for enhance mode), `max_new_tokens` (2048), `temperature` (0.0 = greedy). Outputs: `caption`
+(pretty JSON string), `caption_json` (parsed dict), `caption_raw` (raw decoded text).
+
+## Caption → generate
+
+```python
+caption = pipe(image=ref, output="caption")
+# feed straight into the Ideogram 4 generation pipeline (see OzzyGT/ideogram4-modular)
+image = gen_pipe(prompt=caption, output="images")[0]
+```
+
+## Notes
+
+- **Default checkpoint:** `google/gemma-4-E4B-it` (official bf16, ~15 GB — gated, needs a license-accepted HF
+  token and a >=24 GB GPU). Requires `transformers>=5.12`. To run on a smaller GPU, point
+  `pretrained_model_name_or_path` at a quantized checkpoint of the same model — import its quantization backend
+  first, as the block no longer bundles one.
+- Ideogram 4 was trained on these JSON captions, so a caption from this block is the ideal `prompt` for
+  re-generation / auto-captioned img2img.

__init__.py

@@ -0,0 +1,12 @@
+from .ideogram4_caption import (
+    DEFAULT_INSTRUCTION,
+    Ideogram4CaptionBlocks,
+    Ideogram4CaptionStep,
+)
+
+
+__all__ = [
+    "DEFAULT_INSTRUCTION",
+    "Ideogram4CaptionBlocks",
+    "Ideogram4CaptionStep",
+]

ideogram4_caption.py

@@ -0,0 +1,252 @@
+# Copyright 2026 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.
+"""Custom modular-diffusers block: image OR text -> Ideogram 4 structured-JSON caption.
+
+Wraps a vision-language model (Gemma-4-E4B via `Gemma4ForConditionalGeneration`) with two modes, selected by input:
+  - `image` given  -> caption the image into Ideogram 4's native JSON schema (the `ideogram4-caption` Space logic).
+  - `prompt` given (no image) -> enhance the short text idea into a JSON caption (Ideogram's "magic prompt",
+    reusing the canonical `CAPTION_SYSTEM_MESSAGE` from `diffusers.pipelines.ideogram4.prompt_enhancer`).
+
+Either way the produced `caption` can be fed straight into the Ideogram 4 generation pipeline as the `prompt`
+(the model is trained on these JSON captions). This is the code-usable backend, minus the Gradio UI.
+"""
+
+import json
+import math
+import re
+
+import torch
+
+from diffusers.modular_pipelines.modular_pipeline import ModularPipelineBlocks, PipelineState, SequentialPipelineBlocks
+from diffusers.modular_pipelines.modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+from diffusers.pipelines.ideogram4.prompt_enhancer import CAPTION_SYSTEM_MESSAGE, CAPTION_USER_TEMPLATE
+from diffusers.utils import logging
+
+from transformers import AutoProcessor, Gemma4ForConditionalGeneration
+
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+# Ideogram 4's native caption schema (from the official ideogram-oss/ideogram4 docs). Key ORDER matters for
+# quality, bbox is [ymin, xmin, ymax, xmax] in 0-1000 normalized coords, colors are UPPERCASE #RRGGBB.
+DEFAULT_INSTRUCTION = (
+    "You are a captioner for the Ideogram 4 image model. Look at the image and output ONE JSON object (no prose, "
+    "no markdown fences) that reconstructs it in Ideogram 4's native caption schema. Keep keys in the exact order "
+    "shown.\n\n"
+    "SCHEMA (keys in order):\n"
+    "- high_level_description: one or two sentences summarizing the whole image (subject + medium/style).\n"
+    "- style_description: an object with EXACTLY ONE of `photo` or `art_style`:\n"
+    "    photo caption order:     aesthetics, lighting, photo, medium, color_palette\n"
+    "    non-photo caption order: aesthetics, lighting, medium, art_style, color_palette\n"
+    '    color_palette: optional, up to 16 UPPERCASE hex colors (e.g. "#1B1B2F").\n'
+    "- compositional_deconstruction: {background, elements}.\n"
+    "    background: the global setting/backdrop.\n"
+    "    elements: list; keys in order per type:\n"
+    "      obj:  type, bbox, desc, color_palette\n"
+    "      text: type, bbox, text, desc, color_palette\n\n"
+    "RULES:\n"
+    "- bbox is [ymin, xmin, ymax, xmax] in 0-1000 normalized coordinates (origin top-left), independent of "
+    "resolution. bbox and per-element color_palette are optional (<=5 colors per element).\n"
+    "- Main subject first. A coherent subject (one person/animal/object) is exactly ONE element; its parts go in "
+    "its `desc`. `desc` is identity-first and concrete (color, material, pose, distinguishing features).\n"
+    "- Commit to one concrete value each (no 'various'/'or similar'). `text` elements carry the literal in-image "
+    "characters, verbatim. Colors UPPERCASE #RRGGBB. Describe only what is visible; do not invent."
+)
+
+
+def _balance_json(frag: str) -> str:
+    """Close unbalanced strings/brackets in a truncated JSON fragment so it can be parsed."""
+    stack, in_str, esc = [], False, False
+    for ch in frag:
+        if in_str:
+            if esc:
+                esc = False
+            elif ch == "\\":
+                esc = True
+            elif ch == '"':
+                in_str = False
+        elif ch == '"':
+            in_str = True
+        elif ch in "{[":
+            stack.append(ch)
+        elif ch == "}" and stack and stack[-1] == "{":
+            stack.pop()
+        elif ch == "]" and stack and stack[-1] == "[":
+            stack.pop()
+    out = frag
+    if in_str:
+        out += '"'
+    out = out.rstrip()
+    if out.endswith(","):
+        out = out[:-1]
+    for ch in reversed(stack):
+        out += "}" if ch == "{" else "]"
+    return out
+
+
+def _parse_json(text: str):
+    """Parse the model output into (dict|None, repaired: bool). Robust to markdown fences, trailing prose, and
+    truncated output (balances brackets). `repaired=True` means the output was incomplete and was closed to parse."""
+    text = text.strip()
+    text = re.sub(r"^```(?:json)?\s*", "", text)
+    text = re.sub(r"\s*```$", "", text).strip()
+    start = text.find("{")
+    if start == -1:
+        return None, False
+    frag = text[start:]
+    try:
+        obj = json.loads(frag[: frag.rfind("}") + 1], strict=False)
+        if isinstance(obj, dict):
+            return obj, False
+    except json.JSONDecodeError:
+        pass
+    try:
+        obj, _ = json.JSONDecoder(strict=False).raw_decode(frag)
+        if isinstance(obj, dict):
+            return obj, False
+    except json.JSONDecodeError:
+        pass
+    try:
+        obj = json.loads(_balance_json(frag), strict=False)
+        return (obj, True) if isinstance(obj, dict) else (None, False)
+    except json.JSONDecodeError:
+        return None, False
+
+
+class Ideogram4CaptionStep(ModularPipelineBlocks):
+    model_name = "ideogram4_caption"
+
+    @property
+    def description(self) -> str:
+        return (
+            "Turns an image OR a short text prompt into Ideogram 4's native structured-JSON caption using a "
+            "vision-language model (Gemma-4). If `image` is given it is captioned; otherwise `prompt` is enhanced "
+            "(magic-prompt) into a caption. Outputs `caption` (JSON string), `caption_json` (parsed dict, or None if "
+            "unparseable), and `caption_raw` (raw decoded text)."
+        )
+
+    @property
+    def expected_components(self) -> list[ComponentSpec]:
+        return [
+            ComponentSpec("caption_model", Gemma4ForConditionalGeneration, description="Vision-language captioner (Gemma-4)."),
+            ComponentSpec("caption_processor", AutoProcessor, description="Processor paired with the captioner."),
+        ]
+
+    @property
+    def inputs(self) -> list[InputParam]:
+        return [
+            InputParam("image", description="PIL image to caption. If omitted, `prompt` is enhanced into a caption."),
+            InputParam("prompt", type_hint=str, description="Short text idea to enhance into a JSON caption (used when no image is given)."),
+            InputParam("instruction", default=DEFAULT_INSTRUCTION, description="Instruction for image-caption mode (ignored when enhancing a prompt)."),
+            InputParam("height", type_hint=int, description="Target height; sets the aspect ratio hint for prompt enhancement."),
+            InputParam("width", type_hint=int, description="Target width; sets the aspect ratio hint for prompt enhancement."),
+            InputParam("max_new_tokens", type_hint=int, default=2048, description="Max tokens to generate (captions can be long)."),
+            InputParam(
+                "temperature",
+                type_hint=float,
+                default=0.0,
+                description="Sampling temperature; 0 = greedy (deterministic).",
+            ),
+        ]
+
+    @property
+    def intermediate_outputs(self) -> list[OutputParam]:
+        return [
+            OutputParam("caption", type_hint=str, description="Pretty-printed JSON caption (raw text if unparseable)."),
+            OutputParam("caption_json", type_hint=dict, description="Parsed caption dict, or None if unparseable."),
+            OutputParam("caption_raw", type_hint=str, description="Raw decoded model output."),
+        ]
+
+    @torch.no_grad()
+    def __call__(self, components, state: PipelineState):
+        block_state = self.get_block_state(state)
+
+        model = components.caption_model
+        processor = components.caption_processor
+        device = model.device  # run on wherever the model physically is (offload-friendly)
+
+        if block_state.image is not None:
+            content = [
+                {"type": "image", "image": block_state.image},
+                {"type": "text", "text": block_state.instruction},
+            ]
+        elif block_state.prompt:
+            h = int(block_state.height or 1024)
+            w = int(block_state.width or 1024)
+            d = math.gcd(w, h) or 1
+            aspect_ratio = f"{w // d}:{h // d}"
+            enhance_text = (
+                CAPTION_SYSTEM_MESSAGE
+                + "\n\n"
+                + CAPTION_USER_TEMPLATE.format(aspect_ratio=aspect_ratio, original_prompt=block_state.prompt)
+            )
+            content = [{"type": "text", "text": enhance_text}]
+        else:
+            raise ValueError("Provide either `image` (to caption) or `prompt` (to enhance into a caption).")
+
+        messages = [{"role": "user", "content": content}]
+        inputs = processor.apply_chat_template(
+            messages,
+            add_generation_prompt=True,
+            tokenize=True,
+            return_dict=True,
+            return_tensors="pt",
+            enable_thinking=False,  # Gemma-4 reasoning mode off for structured-JSON output
+        ).to(device)
+
+        do_sample = float(block_state.temperature) > 0
+        generated = model.generate(
+            **inputs,
+            max_new_tokens=int(block_state.max_new_tokens),
+            do_sample=do_sample,
+            temperature=float(block_state.temperature) if do_sample else None,
+        )
+        input_len = inputs["input_ids"].shape[1]
+        text = processor.decode(generated[0][input_len:], skip_special_tokens=True)
+
+        data, _repaired = _parse_json(text)
+        block_state.caption_json = data
+        block_state.caption = json.dumps(data, indent=2, ensure_ascii=False) if data is not None else text
+        block_state.caption_raw = text
+
+        self.set_block_state(state, block_state)
+        return components, state
+
+
+# auto_docstring
+class Ideogram4CaptionBlocks(SequentialPipelineBlocks):
+    """
+    Standalone image/text -> Ideogram 4 structured-JSON caption pipeline (Gemma-4). Load with
+    `ModularPipeline.from_pretrained(repo, trust_remote_code=True)` +
+    `load_components(torch_dtype=torch.bfloat16)`, then either
+    `pipe(image=img, output="caption")` (caption an image) or `pipe(prompt="a cat...", output="caption")`
+    (enhance a text idea into a caption).
+    """
+
+    model_name = "ideogram4_caption"
+    block_classes = [Ideogram4CaptionStep]
+    block_names = ["caption"]
+
+    @property
+    def description(self) -> str:
+        return "Image or text -> Ideogram 4 structured-JSON caption (Gemma-4; captions an image, or enhances a prompt)."
+
+    @property
+    def outputs(self) -> list[OutputParam]:
+        return [
+            OutputParam("caption", type_hint=str),
+            OutputParam("caption_json", type_hint=dict),
+            OutputParam("caption_raw", type_hint=str),
+        ]

modular_config.json

@@ -0,0 +1,7 @@
+{
+  "_class_name": "Ideogram4CaptionBlocks",
+  "_diffusers_version": "0.39.0.dev0",
+  "auto_map": {
+    "ModularPipelineBlocks": "ideogram4_caption.Ideogram4CaptionBlocks"
+  }
+}

modular_model_index.json

@@ -0,0 +1,33 @@
+{
+  "_blocks_class_name": "Ideogram4CaptionBlocks",
+  "_class_name": "ModularPipeline",
+  "_diffusers_version": "0.39.0.dev0",
+  "caption_model": [
+    "transformers",
+    "Gemma4ForConditionalGeneration",
+    {
+      "pretrained_model_name_or_path": "google/gemma-4-E4B-it",
+      "revision": null,
+      "subfolder": "",
+      "type_hint": [
+        "transformers",
+        "Gemma4ForConditionalGeneration"
+      ],
+      "variant": null
+    }
+  ],
+  "caption_processor": [
+    "transformers",
+    "AutoProcessor",
+    {
+      "pretrained_model_name_or_path": "google/gemma-4-E4B-it",
+      "revision": null,
+      "subfolder": "",
+      "type_hint": [
+        "transformers",
+        "AutoProcessor"
+      ],
+      "variant": null
+    }
+  ]
+}