diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 91a2acb699..25dc36bee2 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -1714,9 +1714,17 @@ class DiffusionBackend: is_narrow_base = ideogram4_repo_is_fp8(repo_id) if is_narrow_base: table = family_bf16_components_gb(fam, fam.base_repo) - if table is not None and model_dense_mib is not None: + if table is not None: + # family_bf16_components_gb is a network-free constant, so reserve the bf16 + # footprint even when the cache-derived estimate is absent (empty blob cache, + # or a best-effort download probe that swallowed a transient HF error and + # returned nothing). Otherwise model_dense_mib stays None and the planner + # reads "size unknown -> stay resident", so the ~54 GB fp8 pipeline plans a + # resident placement and OOMs a card that offload would have fit. table_mib = int(sum(table) * (1000.0**3) / (1024.0 * 1024.0)) - model_dense_mib = max(model_dense_mib, table_mib) + model_dense_mib = ( + table_mib if model_dense_mib is None else max(model_dense_mib, table_mib) + ) companion_mib = None else: if transformer_resident_override_mib is not None: diff --git a/studio/backend/tests/test_diffusion_more_families.py b/studio/backend/tests/test_diffusion_more_families.py index 9d89eb5ab3..b2ffcd093e 100644 --- a/studio/backend/tests/test_diffusion_more_families.py +++ b/studio/backend/tests/test_diffusion_more_families.py @@ -59,6 +59,18 @@ def test_ideogram4_generation_defaults(): assert default_generation_params("ideogram-ai/ideogram-4-fp8") == (48, 7.0) +def test_ideogram4_bf16_reservation_table_present(): + # The memory planner reserves this bf16 footprint for a narrow (fp8) ideogram-4 base even + # when the blob-cache estimate is absent (empty cache / a best-effort download probe that + # swallowed a transient HF error), so the ~54 GB pipeline never plans a resident placement + # it cannot fit. If this constant table ever went None, that fp8 OOM safeguard would + # silently disable, so pin that it is present and sums to the expected ~54 GB. + fam = detect_family("ideogram-ai/ideogram-4-fp8") + table = family_bf16_components_gb(fam, fam.base_repo) + assert table is not None + assert sum(table) > 50.0 # transformer (37.2) + bf16 text encoder (16.3) + VAE (0.2) + + def test_ideogram4_memory_table_counts_both_dits(): fam = detect_family("ideogram-ai/ideogram-4-fp8") components = family_bf16_components_gb(fam)