diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 3c22b25d32..364df883ec 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -48,6 +48,7 @@ from .diffusion_device import ( resolve_diffusion_device_target, ) from .diffusion_ideogram4 import ideogram4_repo_is_fp8, load_ideogram4_pipeline +from .diffusion_hidream import HIDREAM_FAMILY_NAME, hidream_te4_kwargs from .diffusion_krea2 import KREA2_FAMILY_NAME, load_krea2_pipeline from .diffusion_memory import ( MEMORY_MODE_BALANCED, @@ -302,6 +303,12 @@ _TRUSTED_NON_GGUF_REPOS = frozenset( # HunyuanImage 2.1: the community diffusers mirror (open, tencent-hunyuan-community # license), safetensors-only, including the diffusers-native guider components. "hunyuanvideo-community/hunyuanimage-2.1-diffusers", + # HiDream-I1: open MIT-weights repos, all three variants one family. The Llama TE the + # model_index names comes from the unsloth mirror (diffusion_hidream.py), which the + # unsloth/ org prefix already trusts. + "hidream-ai/hidream-i1-full", + "hidream-ai/hidream-i1-dev", + "hidream-ai/hidream-i1-fast", # Ideogram 4: no bf16 ships. -fp8 stores the two DiTs as raw float8 (the family base); # the two nf4 repos are identical bnb-4bit exports (both listed so either id loads). "ideogram-ai/ideogram-4-fp8", @@ -1388,6 +1395,10 @@ class DiffusionBackend: pipe_kwargs: dict[str, Any] = {"torch_dtype": dtype} if hf_token: pipe_kwargs["token"] = hf_token + if fam.name == HIDREAM_FAMILY_NAME: + # The repo names a Llama text_encoder_4 it does not ship; + # supply it from the open mirror (diffusion_hidream.py). + pipe_kwargs.update(hidream_te4_kwargs(dtype, hf_token)) # The prefetched snapshot dir keeps from_pretrained off the hub (its # sweep re-pulls files the scoped prefetch skipped: 24 GB per FLUX.1). pipe = pipeline_cls.from_pretrained( @@ -1428,6 +1439,9 @@ class DiffusionBackend: pipe_kwargs = {"torch_dtype": dtype, "transformer": transformer} if hf_token: pipe_kwargs["token"] = hf_token + if fam.name == HIDREAM_FAMILY_NAME: + # Same Llama TE4 assembly as the full-pipeline branch above. + pipe_kwargs.update(hidream_te4_kwargs(dtype, hf_token)) pipe = pipeline_cls.from_pretrained( _base_local_dir or base, **pipe_kwargs ) @@ -1864,6 +1878,10 @@ class DiffusionBackend: pipe_kwargs: dict[str, Any] = {"torch_dtype": dtype, "transformer": transformer} if hf_token: pipe_kwargs["token"] = hf_token + if getattr(fam, "name", None) == HIDREAM_FAMILY_NAME: + # The repo ships no Llama text_encoder_4; assemble it from the open mirror + # (diffusion_hidream.py) exactly like the full-pipeline load branch. + pipe_kwargs.update(hidream_te4_kwargs(dtype, hf_token)) pipe = pipeline_cls.from_pretrained(base_local_dir or base, **pipe_kwargs) pipe.to(device) return pipe diff --git a/studio/backend/core/inference/diffusion_auto_policy.py b/studio/backend/core/inference/diffusion_auto_policy.py index 851ea65115..32e3a3ee02 100644 --- a/studio/backend/core/inference/diffusion_auto_policy.py +++ b/studio/backend/core/inference/diffusion_auto_policy.py @@ -54,6 +54,10 @@ _FAMILY_BF16_GB: dict[str, tuple[float, float, float]] = { "lumina-2": (5.2, 5.2, 0.2), # 17B dual-stream DiT (32.5 GB bf16 on disk) + Qwen2.5-VL 15.5 GB + ByT5 0.8 GB. "hunyuanimage-2.1": (32.5, 16.3, 0.8), + # 17B MoE DiT (34.2 GB bf16) + FOUR text encoders: CLIP-L 0.5 + CLIP-G 2.8 + T5-XXL 9.5 + # from the repo, plus the Llama-3.1-8B text_encoder_4 (~16 GB bf16) assembled from the + # open mirror at load time (diffusion_hidream.py). + "hidream-i1": (34.2, 28.8, 0.2), # Two ~9.3B DiTs (conditional + unconditional_transformer for Ideogram's dual-branch CFG), # both resident, plus a Qwen3-VL encoder. The vendor stores them as raw float8; these are the # bf16-resident sizes after the dtype cast, so each doubles (37.2 = 2 x 18.6, encoder 16.3). diff --git a/studio/backend/core/inference/diffusion_families.py b/studio/backend/core/inference/diffusion_families.py index 87b897d88f..6cdb6c6685 100644 --- a/studio/backend/core/inference/diffusion_families.py +++ b/studio/backend/core/inference/diffusion_families.py @@ -347,6 +347,21 @@ _FAMILIES: tuple[DiffusionFamily, ...] = ( # Exported bf16-only; keep the fp16 fallback off like z-image / krea-2. fp16_incompatible = True, ), + # HiDream-I1: a 17B MoE DiT (16 double + 32 single layers, 4 routed experts) with FOUR text + # encoders. The repos ship CLIP-L/CLIP-G/T5-XXL but NOT the Llama-3.1-8B text_encoder_4 their + # model_index names: the loader assembles it from the open unsloth mirror + # (diffusion_hidream.py). Full / Dev / Fast share the arch, so one family covers all three + # (per-variant step/guidance defaults below). city96 publishes a GGUF but the GGUF path would + # need the same TE4 assembly for tiny demand, so no GGUF artifact is wired yet. + DiffusionFamily( + name = "hidream-i1", + pipeline_class = "HiDreamImagePipeline", + transformer_class = "HiDreamImageTransformer2DModel", + base_repo = "HiDream-ai/HiDream-I1-Full", + aliases = ("hidream", "hidream-i1-full", "hidream-i1-dev", "hidream-i1-fast"), + # Exported bf16-only; keep the fp16 fallback off like the other modern DiTs. + fp16_incompatible = True, + ), # Ideogram 4 (diffusers >= 0.39): a 34-layer DiT PAIR (conditional + unconditional_transformer # for dual-branch CFG, both ~9B, so memory planning counts two DiTs) with a Qwen3-VL encoder. # No bf16 checkpoint: ideogram-4-fp8 (raw float8, upcast on load) is the highest-precision @@ -532,6 +547,11 @@ _GENERATION_DEFAULTS: tuple[tuple[str, int, float], ...] = ( # HunyuanImage 2.1 model-card: 50 steps; the guidance value feeds the call's # distilled_guidance_scale (default 3.25), while real CFG runs inside the repo guiders. ("hunyuanimage", 50, 3.25), + # HiDream-I1 upstream inference.py: Full 50 steps / guidance 5; the distilled Dev (28) and + # Fast (16) run guidance-free. Specific keys precede the generic "hidream" (Full + fallback). + ("hidream-i1-dev", 28, 0.0), + ("hidream-i1-fast", 16, 0.0), + ("hidream", 50, 5.0), # Ideogram 4 model-card: 48 steps, guidance 7 (its schedule tapers the last 3 steps to 3.0; # the loader keeps that taper when the request matches these defaults exactly). ("ideogram", 48, 7.0), diff --git a/studio/backend/core/inference/diffusion_hidream.py b/studio/backend/core/inference/diffusion_hidream.py new file mode 100644 index 0000000000..82ac27f7d5 --- /dev/null +++ b/studio/backend/core/inference/diffusion_hidream.py @@ -0,0 +1,47 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""HiDream-I1 Llama text-encoder assembly. + +The HiDream-ai/HiDream-I1-* repos name ``text_encoder_4`` (LlamaForCausalLM) and +``tokenizer_4`` in their model_index but do NOT ship the weights: the official example +loads meta-llama/Meta-Llama-3.1-8B-Instruct separately and passes both components into +``HiDreamImagePipeline.from_pretrained``. That upstream repo is Hub-gated (manual +approval), so Studio loads the open unsloth mirror instead -- byte-identical weights, +no license wall at load time, and the unsloth org is already inside the loader's +non-GGUF trust gate. ``output_hidden_states=True`` matches the official example: the +pipeline's prompt encoder consumes the Llama hidden states, not the logits. +""" + +from __future__ import annotations + +from typing import Any, Optional + +from loggers import get_logger + +logger = get_logger(__name__) + +HIDREAM_FAMILY_NAME = "hidream-i1" + +# Open mirror of the gated meta-llama/Meta-Llama-3.1-8B-Instruct the pipeline expects. +HIDREAM_LLAMA_REPO = "unsloth/Meta-Llama-3.1-8B-Instruct" + + +def hidream_te4_kwargs(dtype: Any, hf_token: Optional[str] = None) -> dict[str, Any]: + """``{text_encoder_4, tokenizer_4}`` kwargs for a HiDream pipeline ``from_pretrained``. + + Loaded eagerly (~16 GB bf16) before the pipeline call so a failure surfaces as a + clear error instead of a half-built pipeline.""" + import torch # noqa: F401 -- dtype values are torch dtypes; import keeps parity with callers + from transformers import AutoTokenizer, LlamaForCausalLM + + logger.info("diffusion.hidream: loading Llama TE4 from %s", HIDREAM_LLAMA_REPO) + tokenizer_4 = AutoTokenizer.from_pretrained(HIDREAM_LLAMA_REPO, token = hf_token) + text_encoder_4 = LlamaForCausalLM.from_pretrained( + HIDREAM_LLAMA_REPO, + output_hidden_states = True, + output_attentions = True, + torch_dtype = dtype, + token = hf_token, + ) + return {"text_encoder_4": text_encoder_4, "tokenizer_4": tokenizer_4} diff --git a/studio/backend/tests/test_diffusion_more_families.py b/studio/backend/tests/test_diffusion_more_families.py index f59988b7ac..afd0f3df7c 100644 --- a/studio/backend/tests/test_diffusion_more_families.py +++ b/studio/backend/tests/test_diffusion_more_families.py @@ -202,6 +202,62 @@ def test_hunyuanimage21_bf16_component_table_present(): assert vae_gb <= 1.0 +# ── hidream-i1 family ──────────────────────────────────────────────────────── +@pytest.mark.parametrize( + "repo_id", + [ + "HiDream-ai/HiDream-I1-Full", + "HiDream-ai/HiDream-I1-Dev", + "HiDream-ai/HiDream-I1-Fast", + ], +) +def test_detect_family_hidream_repos(repo_id): + # One family covers all three variants (same 17B MoE arch + 4-TE stack). + fam = detect_family(repo_id) + assert fam is not None and fam.name == "hidream-i1" + assert fam.pipeline_class == "HiDreamImagePipeline" + assert fam.transformer_class == "HiDreamImageTransformer2DModel" + assert fam.base_repo == "HiDream-ai/HiDream-I1-Full" + # Published bf16-only upstream; the fp16 fallback stays off. + assert fam.fp16_incompatible is True + + +def test_hidream_override_and_trust(): + assert detect_family("x", override = "hidream-i1").name == "hidream-i1" + assert detect_family("x", override = "hidream").name == "hidream-i1" + # The three official repos load via from_pretrained -> allowlisted; the Llama TE4 + # comes from the unsloth mirror, which the org prefix already trusts. + for rid in ( + "HiDream-ai/HiDream-I1-Full", + "HiDream-ai/HiDream-I1-Dev", + "HiDream-ai/HiDream-I1-Fast", + ): + assert _is_trusted_diffusion_repo(rid) + assert not _is_trusted_diffusion_repo("HiDream-ai/some-future-repo") + assert _is_trusted_diffusion_repo("unsloth/Meta-Llama-3.1-8B-Instruct") + + +def test_hidream_generation_defaults(): + # Upstream inference.py: Full 50 steps / guidance 5; Dev and Fast are distilled and + # run guidance-free at 28 / 16 steps. The specific keys must beat the generic + # "hidream" (which also appears in the owner segment of every variant id). + assert default_generation_params("HiDream-ai/HiDream-I1-Full") == (50, 5.0) + assert default_generation_params("HiDream-ai/HiDream-I1-Dev") == (28, 0.0) + assert default_generation_params("HiDream-ai/HiDream-I1-Fast") == (16, 0.0) + + +def test_hidream_bf16_component_table_present(): + fam = detect_family("HiDream-ai/HiDream-I1-Full") + sizes = family_bf16_components_gb(fam) + assert sizes is not None + transformer_gb, encoders_gb, vae_gb = sizes + # 17B MoE DiT 34.2 GB; TEs = CLIP-L 0.5 + CLIP-G 2.8 + T5-XXL 9.5 from the repo plus + # the ~16 GB Llama TE4 assembled from the mirror -> ~28.8 GB. + assert 32.0 <= transformer_gb <= 37.0 + assert 26.0 <= encoders_gb <= 32.0 + assert vae_gb <= 0.5 + + def test_ideogram4_generation_defaults(): # Model-card settings: 48 steps, guidance 7 (the backend keeps the pipeline's # recommended tapered schedule when the request matches exactly).