ideogram4/ideogram4_handler.py

import os

import torch

from shared.utils.hf import build_hf_url
from .prompt_enhancer import IDEOGRAM4_PROMPT_ENHANCER, IDEOGRAM4_PROMPT_INFOS


_PROJECT_REPO = "DeepBeepMeep/Ideogram4"
_PROJECT_FOLDER = "ideogram4"
_TEXT_ENCODER_FOLDER = "Qwen3-VL-8B-Instruct"
_VAE_REPO = "DeepBeepMeep/Flux2"
_VAE_FILENAME = "flux2_vae.safetensors"
_TURBOTIME_MODEL_TYPE = "ideogram4_turbotime"
_PRESET_CHOICES = [
    ("Quality 48", "V4_QUALITY_48"),
    ("Default 20", "V4_DEFAULT_20"),
    ("Turbo 12", "V4_TURBO_12"),
]
_DEFAULT_PRESET = "V4_DEFAULT_20"
_DEFAULT_SAMPLE_SOLVER = "euler"
_DEFAULT_FLOW_SHIFT = 1.0
_SAMPLE_SOLVERS = [
    ("Euler", "euler"),
    ("RES 2M", "res_2m"),
    ("RES 2S", "res_2s"),
]
_SAMPLE_SOLVER_IDS = {value for _, value in _SAMPLE_SOLVERS}
_CUSTOM_SETTINGS = [
    {"id": "ideogram_mu", "label": "Scheduler Mu", "name": "Scheduler Mu", "type": "float", "default": 0.0, "min": -10.0, "max": 10.0, "inc": 0.05},
    {"id": "ideogram_std", "label": "Scheduler Std", "name": "Scheduler Std", "type": "float", "default": 1.75, "min": 0.1, "max": 5.0, "inc": 0.05},
]
IDEOGRAM4_INFOS = """## Ideogram 4

Ideogram 4 is an image model specialized for prompt-following, text rendering, layout, logos, posters, typography, and graphic design.

It can take a plain text prompt, but it is designed to work especially well with structured JSON prompts. In WanGP, Magic Prompt can convert a normal text prompt into Ideogram's JSON-style caption format. For best prompt adherence, use the JSON format when possible.


## Guidance phases

Ideogram 4 uses WanGP guidance phases for CFG switching.

Phase 1 uses Guidance during the high-noise part of denoising. Phase 2 uses Guidance2 after the first switch threshold. If Three Phases is enabled, phase 3 uses Guidance3 after the second switch threshold.

The official presets use two phases: most steps at CFG 7, then final low-noise polish at CFG 3.


## Scheduler custom settings

Scheduler Mu and Scheduler Std control the Ideogram logit-normal sampling schedule.

Mu shifts where denoising samples concentrate along the noise schedule. Std controls how wide that distribution is.

The official presets set these values for you. Change them mainly when matching a Comfy workflow or experimenting with a different schedule.


## LoRA multipliers

Ideogram 4 has two parallel transformers: the conditional transformer and the unconditional transformer.

A plain LoRA multiplier applies to both branches.

To set different values per branch, use ':' as conditional:unconditional. For example, 0.9:0.4 applies 0.9 to the conditional transformer and 0.4 to the unconditional transformer.

Guidance phases still use ';'. For example, 1:1.2;0.8:1.0 means phase 1 cond/uncond = 1/1.2 and phase 2 cond/uncond = 0.8/1.0."""
IDEOGRAM4_TURBOTIME_INFOS = """## Ideogram 4 TurboTime

Ideogram 4 TurboTime uses the Ideogram 4 conditional transformer with the TurboTime LoRA applied. It is designed for a few denoising steps with no CFG and no unconditional transformer branch.

LoRA multipliers apply to the conditional transformer only."""
_PRESET_SETTINGS = {
    "V4_QUALITY_48": {"num_inference_steps": 48, "mu": 0.0, "std": 1.5, "switch_threshold": 184},
    "V4_DEFAULT_20": {"num_inference_steps": 20, "mu": 0.0, "std": 1.75, "switch_threshold": 211},
    "V4_TURBO_12": {"num_inference_steps": 12, "mu": 0.5, "std": 1.75, "switch_threshold": 302},
}
_LEGACY_CFG_OVERRIDE_KEYS = (
    "ideogram_cfg_override_enabled",
    "ideogram_cfg_override_start_percent",
    "ideogram_cfg_override_end_percent",
)
_TURBOTIME_SETTINGS = {
    "num_inference_steps": 8,
    "guidance_phases": 0,
    "guidance_scale": 0,
    "guidance2_scale": 0,
    "guidance3_scale": 0,
    "switch_threshold": 0,
    "switch_threshold2": 0,
    "sample_solver": _DEFAULT_SAMPLE_SOLVER,
    "flow_shift": _DEFAULT_FLOW_SHIFT,
}
_TURBOTIME_CUSTOM_SETTINGS = {
    "ideogram_mu": 0.5,
    "ideogram_std": 1.75,
}


def _time_snr_shift(shift, sigma):
    return shift * sigma / (1.0 + (shift - 1.0) * sigma)


def _percent_to_switch_threshold(percent, shift=1.0):
    return int(round(_time_snr_shift(float(shift), 1.0 - float(percent)) * 1000.0))


def _drop_legacy_cfg_override_settings(custom_settings):
    for key in _LEGACY_CFG_OVERRIDE_KEYS:
        custom_settings.pop(key, None)


def _apply_turbotime_settings(ui_defaults, preserve_user_values=False):
    custom_settings = ui_defaults.get("custom_settings", None)
    if not isinstance(custom_settings, dict):
        custom_settings = {}
    _drop_legacy_cfg_override_settings(custom_settings)
    for key, value in _TURBOTIME_CUSTOM_SETTINGS.items():
        if preserve_user_values:
            custom_settings.setdefault(key, value)
        else:
            custom_settings[key] = value
    for key, value in _TURBOTIME_SETTINGS.items():
        if preserve_user_values and key not in {"guidance_phases", "guidance_scale", "guidance2_scale", "guidance3_scale", "switch_threshold", "switch_threshold2"}:
            ui_defaults.setdefault(key, value)
        else:
            ui_defaults[key] = value
    ui_defaults["custom_settings"] = custom_settings
    ui_defaults["model_mode"] = None


def _apply_preset_settings(ui_defaults, preset_name):
    preset = _PRESET_SETTINGS[preset_name]
    custom_settings = ui_defaults.get("custom_settings", None)
    if not isinstance(custom_settings, dict):
        custom_settings = {}
    custom_settings.update({
        "ideogram_mu": preset["mu"],
        "ideogram_std": preset["std"],
    })
    _drop_legacy_cfg_override_settings(custom_settings)
    ui_defaults.update({
        "num_inference_steps": preset["num_inference_steps"],
        "guidance_phases": 2,
        "guidance_scale": 7.0,
        "guidance2_scale": 3.0,
        "switch_threshold": preset["switch_threshold"],
        "sample_solver": _DEFAULT_SAMPLE_SOLVER,
        "flow_shift": _DEFAULT_FLOW_SHIFT,
        "custom_settings": custom_settings,
    })


def _model_uses_nf4(model_def):
    urls = []
    if isinstance(model_def, dict):
        urls.extend(model_def.get("URLs", []) or [])
        urls.extend(model_def.get("URLs2", []) or [])
    return any("_nf4" in str(url).lower() for url in urls)


def _conditional_transformer_only(base_model_type, model_def=None):
    return base_model_type == _TURBOTIME_MODEL_TYPE or bool((model_def or {}).get("conditional_transformer_only", False))


class family_handler:
    @staticmethod
    def query_model_def(base_model_type, model_def):
        conditional_only = _conditional_transformer_only(base_model_type, model_def)
        text_encoder_filename = "Qwen3-VL-8B-Instruct_nf4.safetensors" if _model_uses_nf4(model_def) else "Qwen3-VL-8B-Instruct_fp8.safetensors"
        model_def_update = {
            "image_outputs": True,
            "flux2": True,
            "vae_upsamplers": {"flux2_vae_pid": [1]},
            "excluded_spatial_upsamplers": ["flux2_pid"],
            "guidance_max_phases": 3,
            "lora_multiplier_branches": ["cond", "uncond"],
            "inference_steps": True,
            "fit_into_canvas_image_refs": 0,
            "profiles_dir": [base_model_type],
            "preset_profiles_dir": ["ideogram4_presets"],
            "sample_solvers": _SAMPLE_SOLVERS,
            "flow_shift": True,
            "custom_settings": [one.copy() for one in _CUSTOM_SETTINGS],
            "infos": IDEOGRAM4_INFOS,
            "no_negative_prompt": True,
            "no_background_removal": True,
            "skip_prompt_template": True,
            "resolutions_categories": ["<=2k"],
            "vae_block_size": 16,
            "prompt_enhancer_button_label": "Magic Prompt",
            "prompt_infos": IDEOGRAM4_PROMPT_INFOS,
            "prompt_helper_popup_dims": [78, 100],
            "prompt_enhancer_def": {
                "selection": ["T"],
                "labels": {"T": "Ideogram JSON caption from Text Prompt"},
                "default": "",
            },
            "text_prompt_enhancer_instructions": IDEOGRAM4_PROMPT_ENHANCER,
            "image_prompt_enhancer_instructions": IDEOGRAM4_PROMPT_ENHANCER,
            "text_prompt_enhancer_max_tokens": 2048,
            "image_prompt_enhancer_max_tokens": 2048,
            "text_encoder_folder": _TEXT_ENCODER_FOLDER,
            "text_encoder_URLs": [
                build_hf_url(_PROJECT_REPO, _TEXT_ENCODER_FOLDER, text_encoder_filename),
            ],
        }
        if conditional_only:
            model_def_update.update({
                "conditional_transformer_only": True,
                "guidance_max_phases": 0,
                "lora_multiplier_phases": 1,
                "preset_profiles_dir": [],
                "infos": IDEOGRAM4_TURBOTIME_INFOS,
            })
            model_def_update.pop("lora_multiplier_branches", None)
        return model_def_update

    @staticmethod
    def query_supported_types():
        return ["ideogram4", _TURBOTIME_MODEL_TYPE]

    @staticmethod
    def query_family_maps():
        return {}, {}

    @staticmethod
    def query_model_family():
        return "ideogram4"

    @staticmethod
    def query_family_infos():
        return {"ideogram4": (1140, "Ideogram")}

    @staticmethod
    def render_prompt_helper(model_type, model_def, prompt_id, popup_id, prompt_elem_id, resolution_elem_id):
        from .prompt_helper import render_prompt_helper

        return render_prompt_helper(model_type, model_def, prompt_id, popup_id, prompt_elem_id, resolution_elem_id)

    @staticmethod
    def get_prompt_helper_css():
        from .prompt_helper import get_prompt_helper_css

        return get_prompt_helper_css()

    @staticmethod
    def get_prompt_helper_javascript():
        from .prompt_helper import get_prompt_helper_javascript

        return get_prompt_helper_javascript()

    @staticmethod
    def get_rgb_factors(base_model_type):
        from shared.RGB_factors import get_rgb_factors

        return get_rgb_factors("flux", sub_family="flux2")

    @staticmethod
    def register_lora_cli_args(parser, lora_root):
        parser.add_argument(
            "--lora-dir-ideogram4",
            type=str,
            default=None,
            help=f"Path to a directory that contains Ideogram 4 LoRAs (default: {os.path.join(lora_root, 'ideogram4')})",
        )

    @staticmethod
    def get_lora_dir(base_model_type, args, lora_root):
        return getattr(args, "lora_dir_ideogram4", None) or os.path.join(lora_root, "ideogram4")

    @staticmethod
    def query_model_files(computeList, base_model_type, model_def=None):
        transformer_configs = ["ideogram4_transformer_config.json"]
        if not _conditional_transformer_only(base_model_type, model_def):
            transformer_configs.append("ideogram4_unconditional_transformer_config.json")
        return [
            {
                "repoId": _VAE_REPO,
                "sourceFolderList": [""],
                "fileList": [[_VAE_FILENAME]],
            },
            {
                "repoId": _PROJECT_REPO,
                "sourceFolderList": [_PROJECT_FOLDER, _TEXT_ENCODER_FOLDER],
                "fileList": [
                    transformer_configs,
                    ["config.json", "tokenizer.json", "tokenizer_config.json", "chat_template.jinja"],
                ],
            }
        ]

    @staticmethod
    def load_model(
        model_filename,
        model_type=None,
        base_model_type=None,
        model_def=None,
        quantizeTransformer=False,
        text_encoder_quantization=None,
        dtype=torch.bfloat16,
        VAE_dtype=torch.float32,
        mixed_precision_transformer=False,
        save_quantized=False,
        submodel_no_list=None,
        text_encoder_filename=None,
        **kwargs,
    ):
        from .ideogram4_main import model_factory

        pipe_processor = model_factory(
            checkpoint_dir="ckpts",
            model_filename=model_filename,
            model_type=model_type,
            model_def=model_def,
            base_model_type=base_model_type,
            text_encoder_filename=text_encoder_filename,
            dtype=dtype,
            VAE_dtype=VAE_dtype,
            save_quantized=save_quantized,
        )
        pipe = {
            "transformer": pipe_processor.conditional_transformer,
            "text_encoder": pipe_processor.text_encoder,
            "vae": pipe_processor.autoencoder,
        }
        co_tenants_map = {}
        if pipe_processor.unconditional_transformer is not None:
            pipe["transformer2"] = pipe_processor.unconditional_transformer
            co_tenants_map = {
                "transformer": ["transformer2"],
                "transformer2": ["transformer"],
            }
        pipe = {
            "pipe": pipe,
            "coTenantsMap": co_tenants_map,
        }
        return pipe_processor, pipe

    @staticmethod
    def update_default_settings(base_model_type, model_def, ui_defaults):
        ui_defaults.update({"image_mode": 1, "model_mode": None, "batch_size": 1})
        if _conditional_transformer_only(base_model_type, model_def):
            _apply_turbotime_settings(ui_defaults)
            return
        _apply_preset_settings(ui_defaults, _DEFAULT_PRESET)

    @staticmethod
    def fix_settings(base_model_type, settings_version, model_def, ui_defaults):
        if _conditional_transformer_only(base_model_type, model_def):
            _apply_turbotime_settings(ui_defaults, preserve_user_values=True)
            return
        old_mode = ui_defaults.get("model_mode", None)
        if old_mode in _PRESET_SETTINGS:
            _apply_preset_settings(ui_defaults, old_mode)
        elif ui_defaults.get("sample_solver", None) in _PRESET_SETTINGS:
            _apply_preset_settings(ui_defaults, ui_defaults["sample_solver"])
        else:
            custom_settings = ui_defaults.get("custom_settings", None)
            if not isinstance(custom_settings, dict):
                custom_settings = {}
            defaults = _PRESET_SETTINGS[_DEFAULT_PRESET]
            if custom_settings.get("ideogram_cfg_override_enabled", None) not in (None, 0, "0", False) and "ideogram_cfg_override_start_percent" in custom_settings:
                ui_defaults["switch_threshold"] = _percent_to_switch_threshold(custom_settings["ideogram_cfg_override_start_percent"], ui_defaults.get("flow_shift", _DEFAULT_FLOW_SHIFT))
            _drop_legacy_cfg_override_settings(custom_settings)
            custom_settings.setdefault("ideogram_mu", defaults["mu"])
            custom_settings.setdefault("ideogram_std", defaults["std"])
            ui_defaults["custom_settings"] = custom_settings
            ui_defaults.setdefault("num_inference_steps", defaults["num_inference_steps"])
            ui_defaults.setdefault("guidance_scale", 7.0)
            ui_defaults.setdefault("guidance2_scale", 3.0)
            ui_defaults.setdefault("switch_threshold", defaults["switch_threshold"])
            if ui_defaults.get("sample_solver", "") not in _SAMPLE_SOLVER_IDS:
                ui_defaults["sample_solver"] = _DEFAULT_SAMPLE_SOLVER
            ui_defaults.setdefault("flow_shift", _DEFAULT_FLOW_SHIFT)
            ui_defaults.setdefault("guidance_phases", 2)
        ui_defaults["model_mode"] = None
        ui_defaults.setdefault("image_mode", 1)