diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 305d143540..5f3a75fea4 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -33,6 +33,7 @@ from .diffusion_families import ( DIFFUSION_CANCELLED_MSG, DIFFUSION_NOT_LOADED_MSG, IDEOGRAM4_FAMILY_NAME, + LUMINA2_FAMILY_NAME, DiffusionFamily, default_generation_params, detect_family_for_pick, @@ -295,6 +296,9 @@ _TRUSTED_NON_GGUF_REPOS = frozenset( # undistilled base to train LoRAs on (train on Raw, run adapters on Turbo). "krea/krea-2-turbo", "krea/krea-2-raw", + # Lumina Image 2.0: standard diffusers layout (Gemma2-2B encoder), safetensors-only, + # loads through the generic from_pretrained pipeline path. + "alpha-vllm/lumina-image-2.0", # 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", @@ -2644,6 +2648,11 @@ class DiffusionBackend: kwargs.pop(state.family.cfg_kwarg, None) else: kwargs["guidance_schedule"] = None + if state.family.name == LUMINA2_FAMILY_NAME and "cfg_trunc_ratio" in call_params: + # Lumina 2's card recipe runs the CFG double-forward only over the FIRST + # quarter of the trajectory (cfg_trunc_ratio=0.25); the pipeline default (1.0) + # applies it everywhere, visibly oversaturating output. Constant card value. + kwargs["cfg_trunc_ratio"] = 0.25 if init_pil is not None: # Reference passes the whole list (FLUX.2 combines); others take the single image. kwargs["image"] = [init_pil, *ref_extra] if ref_extra else init_pil diff --git a/studio/backend/core/inference/diffusion_auto_policy.py b/studio/backend/core/inference/diffusion_auto_policy.py index 74f629c257..00dd981953 100644 --- a/studio/backend/core/inference/diffusion_auto_policy.py +++ b/studio/backend/core/inference/diffusion_auto_policy.py @@ -50,6 +50,8 @@ _FAMILY_BF16_GB: dict[str, tuple[float, float, float]] = { "qwen-image-edit": (40.9, 16.6, 0.3), "z-image": (12.3, 8.0, 0.2), "krea-2": (26.3, 8.9, 0.5), + # Ships fp32 (10.4 + 10.5 + 0.3 GB of shards); bf16-resident is half. + "lumina-2": (5.2, 5.2, 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 3865e5849b..4522215670 100644 --- a/studio/backend/core/inference/diffusion_families.py +++ b/studio/backend/core/inference/diffusion_families.py @@ -309,6 +309,20 @@ _FAMILIES: tuple[DiffusionFamily, ...] = ( # Exported bf16-only; fp16 unvalidated upstream, so keep the fp16 fallback off like z-image. fp16_incompatible = True, ), + # Lumina Image 2.0: a 2.6B single-stream DiT with a Gemma2-2B encoder and a standard + # 16-channel AutoencoderKL, all transformers-4.x-compatible, so the generic + # from_pretrained pipeline path loads it. No GGUF/sd.cpp mapping exists upstream. + # NOT aliased to bare "lumina": Lumina-Next checkpoints are a different arch + # (LuminaText2ImgPipeline) and must stay unknown rather than crash mid-load. + DiffusionFamily( + name = "lumina-2", + pipeline_class = "Lumina2Pipeline", + transformer_class = "Lumina2Transformer2DModel", + base_repo = "Alpha-VLLM/Lumina-Image-2.0", + aliases = ("lumina-image-2.0", "lumina-image-2", "lumina2"), + # Published and validated bf16-only upstream; keep the fp16 fallback off like z-image. + 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 @@ -359,6 +373,10 @@ def trainable_family_names() -> tuple[str, ...]: # call). Named here so the two modules can't drift. IDEOGRAM4_FAMILY_NAME = "ideogram-4" +# The family whose generate call carries the card's CFG-truncation ratio (the loader +# special-cases the call). Named here so the two modules can't drift. +LUMINA2_FAMILY_NAME = "lumina-2" + # Models Studio deliberately does NOT support, reason surfaced verbatim in the load error (vs the # generic unknown-family message). Keyed by a lowercase repo-id substring. The bar is a diffusers @@ -484,6 +502,9 @@ _GENERATION_DEFAULTS: tuple[tuple[str, int, float], ...] = ( ("flux.2-dev", 28, 4.0), # full (non-distilled) ("qwen-image", 20, 4.0), ("z-image", 20, 4.0), + # Lumina Image 2.0 model-card: 50 steps, guidance 4 (plus cfg_trunc_ratio 0.25, which the + # loader passes itself; see LUMINA2_FAMILY_NAME). + ("lumina", 50, 4.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/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index 4d6891f5b4..e2f49d2c6f 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -257,6 +257,7 @@ class _FakePipe: callback_on_step_end = None, guidance_scale = None, true_cfg_scale = None, + cfg_trunc_ratio = None, **kwargs, ): self.last_kwargs = { @@ -265,6 +266,7 @@ class _FakePipe: "callback_on_step_end": callback_on_step_end, "guidance_scale": guidance_scale, "true_cfg_scale": true_cfg_scale, + "cfg_trunc_ratio": cfg_trunc_ratio, **kwargs, } n = kwargs.get("num_images_per_prompt", 1) @@ -406,6 +408,10 @@ def fake_runtime(monkeypatch): # that to a fake pipe so the guidance path is reachable without real weights. diffusers.Ideogram4Pipeline = _FakePipeline diffusers.Ideogram4Transformer2DModel = _FakeTransformer + # Lumina 2, so the cfg_trunc_ratio special case is exercisable (the fake pipe's + # signature carries the kwarg, mirroring the real Lumina2Pipeline). + diffusers.Lumina2Pipeline = _FakePipeline + diffusers.Lumina2Transformer2DModel = _FakeTransformer # SDXL: a U-Net family. Its single-file checkpoint is the whole pipeline, so the pipeline # class carries from_single_file; UNet2DConditionModel is the denoiser class (fetched but # unused on the pipeline/single-file-pipeline paths). @@ -1755,6 +1761,41 @@ def test_generate_ideogram_custom_guidance_nulls_schedule(fake_runtime, tmp_path assert "guidance_schedule" in call and call["guidance_schedule"] is None +def _load_lumina(backend, tmp_path): + # Lumina 2 loads through the GENERIC pipeline path (standard diffusers layout); + # a local pipeline dir is enough here. + (tmp_path / "model_index.json").write_text("{}") + backend.load_pipeline(str(tmp_path), family_override = "lumina-2") + + +def test_generate_lumina2_passes_cfg_trunc_ratio(fake_runtime, tmp_path): + # The card recipe truncates the CFG double-forward to the first quarter of the + # trajectory; the pipeline default (1.0) applies it everywhere. The backend passes + # the constant card value on every lumina-2 generate. + backend = DiffusionBackend() + _load_lumina(backend, tmp_path) + backend.generate(prompt = "a sloth", steps = 50, guidance = 4.0) + call = backend._state.pipe.last_kwargs + assert call["cfg_trunc_ratio"] == 0.25 + assert call["guidance_scale"] == 4.0 + + +def test_generate_other_family_never_passes_cfg_trunc_ratio(fake_runtime, tmp_path): + # The kwarg is family-gated, not just signature-gated: another family whose pipeline + # happens to accept cfg_trunc_ratio must not inherit Lumina's recipe constant. + backend = DiffusionBackend() + (tmp_path / "model.gguf").write_bytes(b"weights") + backend.load_pipeline( + str(tmp_path), + gguf_filename = "model.gguf", + base_repo = "base/repo", + family_override = "z-image", + ) + backend.generate(prompt = "a sloth", steps = 9, guidance = 0.0) + call = backend._state.pipe.last_kwargs + assert call["cfg_trunc_ratio"] is None + + def test_begin_load_rejects_concurrent(monkeypatch): backend = DiffusionBackend() # The worker resolves the base + downloads, both over the network; stub them diff --git a/studio/backend/tests/test_diffusion_more_families.py b/studio/backend/tests/test_diffusion_more_families.py index 9bda972c5e..f3f2e9c291 100644 --- a/studio/backend/tests/test_diffusion_more_families.py +++ b/studio/backend/tests/test_diffusion_more_families.py @@ -85,6 +85,56 @@ def test_flux1_krea_dev_generation_defaults(): assert default_generation_params("krea/Krea-2-Raw") == (52, 3.5) +# ── lumina-2 family ────────────────────────────────────────────────────────── +@pytest.mark.parametrize( + "repo_id", + [ + "Alpha-VLLM/Lumina-Image-2.0", + # A same-arch finetune must group here via the lumina-image-2.0 token. + "neta-art/NetaYume-Lumina-Image-2.0", + ], +) +def test_detect_family_lumina2_repos(repo_id): + fam = detect_family(repo_id) + assert fam is not None and fam.name == "lumina-2" + assert fam.pipeline_class == "Lumina2Pipeline" + assert fam.transformer_class == "Lumina2Transformer2DModel" + assert fam.base_repo == "Alpha-VLLM/Lumina-Image-2.0" + # Published bf16-only upstream; the fp16 fallback stays off. + assert fam.fp16_incompatible is True + + +def test_detect_family_lumina2_override_and_next_rejected(): + assert detect_family("x", override = "lumina-2").name == "lumina-2" + assert detect_family("x", override = "lumina2").name == "lumina-2" + # Lumina-Next is a DIFFERENT arch (LuminaText2ImgPipeline): it must stay unknown + # instead of resolving here and crashing mid-load. + assert detect_family("Alpha-VLLM/Lumina-Next-SFT-diffusers") is None + + +def test_lumina2_is_trusted_non_gguf(): + # The official pipeline loads via from_pretrained -> needs the allowlist. + assert _is_trusted_diffusion_repo("Alpha-VLLM/Lumina-Image-2.0") + assert not _is_trusted_diffusion_repo("Alpha-VLLM/some-future-repo") + + +def test_lumina2_generation_defaults(): + # Model-card recipe: 50 steps at guidance 4.0 (cfg_trunc_ratio is added by the + # backend generate call itself, not the defaults table). + assert default_generation_params("Alpha-VLLM/Lumina-Image-2.0") == (50, 4.0) + + +def test_lumina2_bf16_component_table_present(): + fam = detect_family("Alpha-VLLM/Lumina-Image-2.0") + sizes = family_bf16_components_gb(fam) + assert sizes is not None + transformer_gb, encoders_gb, vae_gb = sizes + # 2.6B DiT + Gemma2-2B, both fp32 on disk -> ~5.2 GB each bf16-resident. + assert 4.0 <= transformer_gb <= 7.0 + assert 4.0 <= encoders_gb <= 7.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). diff --git a/studio/frontend/src/components/assistant-ui/model-selector/model-catalog.check.ts b/studio/frontend/src/components/assistant-ui/model-selector/model-catalog.check.ts index e1f095feda..95b046ab24 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/model-catalog.check.ts +++ b/studio/frontend/src/components/assistant-ui/model-selector/model-catalog.check.ts @@ -361,6 +361,16 @@ assert.equal( groupForRepoId("QuantStack/FLUX.1-Krea-dev-GGUF", IMAGE_CATALOG), kreaDevRoute, ); +// Lumina Image 2.0: a single ungated bf16 pipeline artifact (11 GB) -- auto-routed on a +// 24 GB GPU (11 <= 0.7 * 24) and resolvable through its canonical id. +const lumina = groupForRepoId("Alpha-VLLM/Lumina-Image-2.0", IMAGE_CATALOG); +assert.ok(lumina); +assert.equal( + pickDefaultArtifact(lumina, { gpuGb: 24, systemRamGb: 64, isDownloaded: notDownloaded }) + .repoId, + "Alpha-VLLM/Lumina-Image-2.0", +); +assert.equal(loadSpecFor("Alpha-VLLM/Lumina-Image-2.0", IMAGE_CATALOG)?.kind, "pipeline"); // FLUX.1-schnell is Apache-2.0 (not gated): its BF16 IS auto-routed on a GPU that fits it. const fluxSchnellRoute = groupForRepoId("unsloth/FLUX.1-schnell", IMAGE_CATALOG); assert.ok(fluxSchnellRoute); diff --git a/studio/frontend/src/components/assistant-ui/model-selector/model-catalog.ts b/studio/frontend/src/components/assistant-ui/model-selector/model-catalog.ts index ca8dac38c0..43a30fa5b2 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/model-catalog.ts +++ b/studio/frontend/src/components/assistant-ui/model-selector/model-catalog.ts @@ -253,6 +253,16 @@ export const IMAGE_CATALOG: CatalogGroup[] = [ scope: "image", artifacts: [bf16Pipeline("krea/Krea-2-Turbo", 18)], }, + { + // 2.6B DiT + Gemma2-2B encoder, ~11 GB bf16-resident (ships fp32, cast on + // load). Apache-2.0, not gated. No GGUF quants exist upstream, so the + // official pipeline is the only artifact. + canonicalId: "Alpha-VLLM/Lumina-Image-2.0", + displayName: "Lumina Image 2.0", + description: "Text-to-image", + scope: "image", + artifacts: [bf16Pipeline("Alpha-VLLM/Lumina-Image-2.0", 11)], + }, { // No bf16 repo exists for Ideogram 4: -fp8 stores its two DiTs as raw // float8 (~46 GB resident after the bf16 cast); -nf4-diffusers is the diff --git a/studio/frontend/src/features/images/images-page.tsx b/studio/frontend/src/features/images/images-page.tsx index 55e2bfc1e7..9e31d140e0 100644 --- a/studio/frontend/src/features/images/images-page.tsx +++ b/studio/frontend/src/features/images/images-page.tsx @@ -170,6 +170,8 @@ const MODEL_DEFAULTS: Array<{ match: string; steps: number; guidance: number }> // defaults the backend keeps the pipeline's recommended tapered guidance schedule // instead of a flat constant. { match: "ideogram", steps: 48, guidance: 7 }, + // Lumina Image 2.0 model-card recipe (the backend adds cfg_trunc_ratio itself). + { match: "lumina", steps: 50, guidance: 4 }, // SDXL: Turbo is distilled (few steps, no CFG); base/full SDXL wants ~30 steps and // real CFG (~7). "sdxl-turbo" must precede the generic "sdxl" substring match. { match: "sdxl-turbo", steps: 3, guidance: 0 },