diff --git a/comfy_extras/nodes_ideogram4.py b/comfy_extras/nodes_ideogram4.py index d5827db4f..3ac5f759e 100644 --- a/comfy_extras/nodes_ideogram4.py +++ b/comfy_extras/nodes_ideogram4.py @@ -1,6 +1,7 @@ """Ideogram 4 sampling helper """ +import enum import math import torch @@ -10,6 +11,31 @@ from comfy_api.latest import ComfyExtension, io _LOGSNR_MIN = -15.0 _LOGSNR_MAX = 18.0 +class Ideogram4Enum(enum.Enum): + QUALITY = "Quality" + DEFAULT = "Default" + TURBO = "Turbo" + +IDEOGRAM4_PRESET_CONFIGS = { + "Quality": { + "num_steps": 48, + "mu": 0.0, + "std": 1.5, + "preset_id": "V4_QUALITY_48" + }, + "Default": { + "num_steps": 20, + "mu": 0.0, + "std": 1.75, + "preset_id": "V4_DEFAULT_20" + }, + "Turbo": { + "num_steps": 12, + "mu": 0.5, + "std": 1.75, + "preset_id": "V4_TURBO_12" + } +} def _logit_normal_schedule(u, mean, std): # Reference time (0=noise..1=clean) via the probit/ndtri quantile. @@ -54,10 +80,40 @@ class Ideogram4Scheduler(io.ComfyNode): return io.NodeOutput(ideogram4_sigmas(steps, width, height, mu, std)) +class Ideogram4SchedulerPreset(Ideogram4Scheduler): + @classmethod + def define_schema(cls) -> io.Schema: + return io.Schema( + node_id="Ideogram4SchedulerPreset", + display_name="Ideogram 4 Scheduler (Presets)", + category="sampling/custom_sampling/schedulers", + inputs=[ + io.Combo.Input("preset", options=[e.value for e in Ideogram4Enum], default=Ideogram4Enum.DEFAULT.value), + io.Int.Input("width", default=1024, min=256, max=8192, step=16), + io.Int.Input("height", default=1024, min=256, max=8192, step=16), + ], + outputs=[io.Sigmas.Output()], + ) + + @classmethod + def execute(cls, preset, width, height) -> io.NodeOutput: + config = IDEOGRAM4_PRESET_CONFIGS.get(preset) + if not config: + raise ValueError(f"Invalid preset: {preset}") + + return super().execute( + steps=config["num_steps"], + width=width, + height=height, + mu=config["mu"], + std=config["std"] + ) + + class Ideogram4Extension(ComfyExtension): @override async def get_node_list(self) -> list[type[io.ComfyNode]]: - return [Ideogram4Scheduler] + return [Ideogram4Scheduler, Ideogram4SchedulerPreset] async def comfy_entrypoint() -> Ideogram4Extension: