Add the HiDream-I1 family to the image backend
A 17B MoE DiT (16 double + 32 single layers, 4 routed experts) with four text encoders, on HiDreamImagePipeline (diffusers 0.39). One family covers the open Full / Dev / Fast repos (same arch); per-variant generation defaults follow the upstream inference recipes (Full 50 steps at guidance 5, the distilled Dev 28 and Fast 16 guidance-free). The repos name a Llama-3.1-8B text_encoder_4 in their model_index but do not ship its weights; the official example passes the gated meta-llama repo in by hand. The loader instead assembles the component from the open unsloth mirror (byte-identical weights, already inside the non-GGUF trust gate), injected at the three pipeline from_pretrained sites, with output_hidden_states matching the official example. Memory planning counts the assembled TE4: 34.2 GB DiT + 28.8 GB encoders, ~63 GB bf16-resident.
This commit is contained in:
Daniel Han
parent
5887d2489c
commit
e098ee52ea
5 changed files with 145 additions and 0 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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).
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
|
|
|
|||
47
studio/backend/core/inference/diffusion_hidream.py
Normal file
47
studio/backend/core/inference/diffusion_hidream.py
Normal file
|
|
@ -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}
|
||||
|
|
@ -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).
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue