Preserve pinned official Ideogram 4 FP8 release, components, original terms and reference code
8ebf922 verified | """Minimal example: text-to-image inference with Ideogram 4.""" | |
| from __future__ import annotations | |
| import argparse | |
| import os | |
| import sys | |
| import torch | |
| from ideogram4 import ( | |
| DEFAULT_MAGIC_PROMPT, | |
| MAGIC_PROMPTS, | |
| PRESETS, | |
| Ideogram4Pipeline, | |
| Ideogram4PipelineConfig, | |
| aspect_ratio_from_size, | |
| moderate_image, | |
| moderate_prompt, | |
| ) | |
| QUANTIZATION_REPOS = { | |
| "nf4": "ideogram-ai/ideogram-4-nf4", | |
| "fp8": "ideogram-ai/ideogram-4-fp8", | |
| } | |
| def _default_device() -> str: | |
| """Pick the best available device: cuda, then mps (Apple Silicon), then cpu.""" | |
| if torch.cuda.is_available(): | |
| return "cuda" | |
| if torch.backends.mps.is_available(): | |
| return "mps" | |
| return "cpu" | |
| def _default_quantization() -> str: | |
| """nf4 (bitsandbytes 4-bit) is CUDA-only; fall back to fp8 on mps/cpu.""" | |
| return "nf4" if torch.cuda.is_available() else "fp8" | |
| def _print_flags(label: str, flags: list[tuple[str, float]]) -> None: | |
| print(f"{label}:", file=sys.stderr) | |
| for name, score in sorted(flags, key=lambda x: -x[1]): | |
| print(f" {name}: {score:.3f}", file=sys.stderr) | |
| def main() -> None: | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument("--prompt", required=True) | |
| parser.add_argument("--output", default="output.png") | |
| parser.add_argument("--height", type=int, default=1024) | |
| parser.add_argument("--width", type=int, default=1024) | |
| parser.add_argument( | |
| "--sampler-preset", | |
| choices=sorted(PRESETS.keys()), | |
| default="V4_QUALITY_48", | |
| help=( | |
| "Named sampler preset that bundles num_steps, per-step CFG schedule, " | |
| "mu, and std. To use a different configuration, add a new entry to " | |
| "ideogram4.sampler_configs.PRESETS. Available: " | |
| + ", ".join(sorted(PRESETS.keys())) | |
| ), | |
| ) | |
| parser.add_argument("--seed", type=int, default=0) | |
| parser.add_argument("--device", default=_default_device()) | |
| parser.add_argument( | |
| "--quantization", | |
| choices=sorted(QUANTIZATION_REPOS.keys()), | |
| default=_default_quantization(), | |
| help=( | |
| "Weight precision. 'nf4' loads bitsandbytes 4-bit pre-quantized weights " | |
| f"from {QUANTIZATION_REPOS['nf4']} (CUDA only). 'fp8' loads weight-only " | |
| f"e4m3 float8 transformer weights from {QUANTIZATION_REPOS['fp8']}; " | |
| "activations stay bf16, so it runs on any device (no FP8 hardware needed). " | |
| "Defaults to nf4 when CUDA is available, otherwise fp8." | |
| ), | |
| ) | |
| parser.add_argument( | |
| "--magic-prompt", | |
| action=argparse.BooleanOptionalAction, | |
| default=True, | |
| help=( | |
| "Expand the plain --prompt into the structured JSON caption the model " | |
| "was trained on, using a magic-prompt LLM (ON by default; pass " | |
| "--no-magic-prompt to feed --prompt to the model verbatim, e.g. when you " | |
| "already have a structured caption)." | |
| ), | |
| ) | |
| parser.add_argument( | |
| "--magic-prompt-model", | |
| choices=sorted(MAGIC_PROMPTS), | |
| default=DEFAULT_MAGIC_PROMPT, | |
| help=( | |
| "Which magic-prompt configuration to use (a model + system-prompt " | |
| f"version). Default: {DEFAULT_MAGIC_PROMPT}." | |
| ), | |
| ) | |
| parser.add_argument( | |
| "--magic-prompt-key", | |
| default=os.environ.get("MAGIC_PROMPT_API_KEY") | |
| or os.environ.get("IDEOGRAM_API_KEY"), | |
| help=( | |
| "API key for the magic-prompt model (or set MAGIC_PROMPT_API_KEY). " | |
| "Required unless --no-magic-prompt. The claude-* configurations call " | |
| "OpenRouter (get a key at https://openrouter.ai); the ideogram-4-v1 " | |
| "configuration calls Ideogram's free hosted magic-prompt API and reads " | |
| "IDEOGRAM_API_KEY by default." | |
| ), | |
| ) | |
| parser.add_argument( | |
| "--warn-on-caption-issues", | |
| action="store_true", | |
| help=( | |
| "If set, caption verifier issues are emitted as warnings instead of " | |
| "aborting with an error." | |
| ), | |
| ) | |
| parser.add_argument( | |
| "--hive-text-key", | |
| default=os.environ.get("HIVE_TEXT_MODERATION_KEY"), | |
| help=( | |
| "Hive Text Moderation API key (or set HIVE_TEXT_MODERATION_KEY). " | |
| "Screens the prompt for NSFW / hate / violence / self-harm / bullying " | |
| "before generation. We strongly discourage running this inference code " | |
| "without a Hive key configured: doing so disables prompt safety " | |
| "screening entirely. Sign up at https://thehive.ai and create a Text " | |
| "Moderation project for the key." | |
| ), | |
| ) | |
| parser.add_argument( | |
| "--hive-visual-key", | |
| default=os.environ.get("HIVE_VISUAL_MODERATION_KEY"), | |
| help=( | |
| "Hive Visual Content Moderation API key (or set " | |
| "HIVE_VISUAL_MODERATION_KEY). Screens the generated image for NSFW / " | |
| "gore / drugs / weapons / hate symbols. We strongly discourage running " | |
| "this inference code without a Hive key configured: doing so disables " | |
| "output safety screening entirely. Sign up at https://thehive.ai and " | |
| "create a Visual Content Moderation project." | |
| ), | |
| ) | |
| args = parser.parse_args() | |
| if args.hive_text_key: | |
| flags = moderate_prompt(args.prompt, args.hive_text_key) | |
| if flags: | |
| _print_flags("Prompt rejected by Hive text moderation", flags) | |
| sys.exit(2) | |
| else: | |
| print( | |
| "WARNING: no Hive text moderation key configured -- prompt safety " | |
| "screening is DISABLED. We strongly discourage using this inference code " | |
| "without a Hive key. Set HIVE_TEXT_MODERATION_KEY or pass " | |
| "--hive-text-key to enable.", | |
| file=sys.stderr, | |
| ) | |
| prompt = args.prompt | |
| if args.magic_prompt: | |
| if not args.magic_prompt_key: | |
| print( | |
| "ERROR: magic prompt is enabled but no API key was found. Set " | |
| "MAGIC_PROMPT_API_KEY, pass --magic-prompt-key, or disable expansion " | |
| "with --no-magic-prompt.", | |
| file=sys.stderr, | |
| ) | |
| sys.exit(2) | |
| aspect_ratio = aspect_ratio_from_size(args.width, args.height) | |
| magic = MAGIC_PROMPTS[args.magic_prompt_model](api_key=args.magic_prompt_key) # type: ignore[call-arg] | |
| print( | |
| f"Magic prompt ({args.magic_prompt_model}): " | |
| f"expanding prompt for {aspect_ratio}...", | |
| file=sys.stderr, | |
| ) | |
| prompt = magic.expand(args.prompt, aspect_ratio=aspect_ratio) | |
| print(f"Expanded caption:\n{prompt}", file=sys.stderr) | |
| preset = PRESETS[args.sampler_preset] | |
| pipe = Ideogram4Pipeline.from_pretrained( | |
| config=Ideogram4PipelineConfig(weights_repo=QUANTIZATION_REPOS[args.quantization]), | |
| device=args.device, | |
| dtype=torch.bfloat16, | |
| ) | |
| images = pipe( | |
| prompt, | |
| height=args.height, | |
| width=args.width, | |
| num_steps=preset.num_steps, | |
| guidance_schedule=preset.guidance_schedule, | |
| mu=preset.mu, | |
| std=preset.std, | |
| seed=args.seed, | |
| raise_on_caption_issues=not args.warn_on_caption_issues, | |
| ) | |
| if args.hive_visual_key: | |
| flags = moderate_image(images[0], args.hive_visual_key) | |
| if flags: | |
| _print_flags("Generated image rejected by Hive visual moderation", flags) | |
| sys.exit(2) | |
| else: | |
| print( | |
| "WARNING: no Hive visual moderation key configured -- output safety " | |
| "screening is DISABLED. We strongly discourage using this inference code " | |
| "without a Hive key. Set HIVE_VISUAL_MODERATION_KEY or pass " | |
| "--hive-visual-key to enable.", | |
| file=sys.stderr, | |
| ) | |
| images[0].save(args.output) | |
| print(f"Saved {args.output}") | |
| if __name__ == "__main__": | |
| main() | |