feat(ml): comfyui_build_flux_workflow — builder Flux schnell+dev (camino custom-advanced)

Builder puro que arma el workflow ComfyUI de Flux en API format con el camino canonico custom-advanced (UNETLoader + DualCLIPLoader[flux] + VAELoader -> RandomNoise + KSamplerSelect + BasicScheduler -> BasicGuider -> SamplerCustomAdvanced -> VAEDecode -> SaveImage). - variant 'schnell' (~4 pasos, sin FluxGuidance) o 'dev' (~20 pasos, con FluxGuidance), con unet y steps por defecto por variante. - Parametro 'available' opcional valida los modelos contra /object_info y lanza FileNotFoundError claro (que falta + carpeta) sin romper la pureza. - width/height/seed/guidance/prefijo parametrizables. - 11 tests unitarios (class_types schnell vs dev, defaults por variante, error path, determinismo). Verificado con generaciones reales (schnell 1024 y 768, dev 768x1024) que producen PNG en disco.
2026-06-27 20:36:55 +02:00
parent cda36408d0
commit b77d223f01
3 changed files with 358 additions and 165 deletions
@@ -3,6 +3,8 @@
 import os
 import sys

+import pytest
+
 sys.path.insert(0, os.path.dirname(__file__))
 sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", ".."))

@@ -10,35 +12,54 @@ from ml.comfyui_build_flux_workflow import comfyui_build_flux_workflow
 from _comfyui_wf_assert import assert_api_format, class_types, node_by_ct


-def test_estructura_y_class_types():
-    wf = comfyui_build_flux_workflow("POS")
+_BASE_CTS = {
+    "UNETLoader",
+    "DualCLIPLoader",
+    "VAELoader",
+    "EmptyLatentImage",
+    "CLIPTextEncode",
+    "RandomNoise",
+    "KSamplerSelect",
+    "BasicScheduler",
+    "BasicGuider",
+    "SamplerCustomAdvanced",
+    "VAEDecode",
+    "SaveImage",
+}
+
+
+def test_schnell_class_types_sin_fluxguidance():
+    wf = comfyui_build_flux_workflow("POS", variant="schnell")
    assert_api_format(wf)
-    assert class_types(wf) == {
-        "UNETLoader",
-        "DualCLIPLoader",
-        "VAELoader",
-        "CLIPTextEncode",
-        "FluxGuidance",
-        "EmptySD3LatentImage",
-        "KSampler",
-        "VAEDecode",
-        "SaveImage",
-    }
+    # schnell usa el camino custom-advanced y NO incluye FluxGuidance.
+    assert class_types(wf) == _BASE_CTS
+    # BasicGuider consume el CLIPTextEncode positivo directo.
+    assert node_by_ct(wf, "BasicGuider")["inputs"]["conditioning"] == ["6", 0]
+
+
+def test_dev_class_types_con_fluxguidance():
+    wf = comfyui_build_flux_workflow("POS", variant="dev", guidance=2.5)
+    assert_api_format(wf)
+    assert class_types(wf) == _BASE_CTS | {"FluxGuidance"}
+    fg = node_by_ct(wf, "FluxGuidance")["inputs"]
+    assert fg["guidance"] == 2.5
+    assert fg["conditioning"] == ["6", 0]  # FluxGuidance aplica sobre el positivo
+    # BasicGuider consume la salida de FluxGuidance, no el CLIPTextEncode directo.
+    assert node_by_ct(wf, "BasicGuider")["inputs"]["conditioning"] == ["21", 0]


 def test_loaders_separados_de_flux():
    # Flux carga UNET + dos text encoders + VAE por separado (no checkpoint unico).
    wf = comfyui_build_flux_workflow(
        "POS",
-        unet="IMG_flux1-schnell-fp8-e4m3fn.safetensors",
-        clip_l="clip_l.safetensors",
-        t5xxl="t5xxl_fp8_e4m3fn_scaled.safetensors",
-        vae="ae.safetensors",
-        weight_dtype="fp8_e4m3fn",
+        variant="schnell",
+        clip_l_name="clip_l.safetensors",
+        t5xxl_name="t5xxl_fp8_e4m3fn_scaled.safetensors",
+        vae_name="ae.safetensors",
    )
    unet = node_by_ct(wf, "UNETLoader")["inputs"]
    assert unet["unet_name"] == "IMG_flux1-schnell-fp8-e4m3fn.safetensors"
-    assert unet["weight_dtype"] == "fp8_e4m3fn"
+    assert unet["weight_dtype"] == "default"
    dual = node_by_ct(wf, "DualCLIPLoader")["inputs"]
    assert dual["type"] == "flux"
    assert dual["clip_name1"] == "t5xxl_fp8_e4m3fn_scaled.safetensors"
@@ -46,25 +67,36 @@ def test_loaders_separados_de_flux():
    assert node_by_ct(wf, "VAELoader")["inputs"]["vae_name"] == "ae.safetensors"


-def test_guidance_y_cfg_de_flux():
-    # La guia va por FluxGuidance; el cfg del KSampler se fija a 1.0 (schnell).
-    wf = comfyui_build_flux_workflow("POS", guidance=2.5)
-    assert node_by_ct(wf, "FluxGuidance")["inputs"]["guidance"] == 2.5
-    ks = node_by_ct(wf, "KSampler")["inputs"]
-    assert ks["cfg"] == 1.0
-    # KSampler positive consume la salida de FluxGuidance, no la del CLIPTextEncode directo.
-    assert ks["positive"] == ["13", 0]
+def test_unet_default_por_variante():
+    schnell = comfyui_build_flux_workflow("POS", variant="schnell")
+    dev = comfyui_build_flux_workflow("POS", variant="dev")
+    assert (
+        node_by_ct(schnell, "UNETLoader")["inputs"]["unet_name"]
+        == "IMG_flux1-schnell-fp8-e4m3fn.safetensors"
+    )
+    assert (
+        node_by_ct(dev, "UNETLoader")["inputs"]["unet_name"]
+        == "IMG_flux1-dev-fp8-e4m3fn.safetensors"
+    )
+
+
+def test_steps_default_por_variante():
+    schnell = comfyui_build_flux_workflow("POS", variant="schnell")
+    dev = comfyui_build_flux_workflow("POS", variant="dev")
+    assert node_by_ct(schnell, "BasicScheduler")["inputs"]["steps"] == 4
+    assert node_by_ct(dev, "BasicScheduler")["inputs"]["steps"] == 20
+    # steps explicito gana al default.
+    custom = comfyui_build_flux_workflow("POS", variant="schnell", steps=6)
+    assert node_by_ct(custom, "BasicScheduler")["inputs"]["steps"] == 6


 def test_params_se_reflejan_en_los_nodos():
-    wf = comfyui_build_flux_workflow("POS", width=768, height=512, steps=8, seed=123)
-    ks = node_by_ct(wf, "KSampler")["inputs"]
-    assert ks["seed"] == 123
-    assert ks["steps"] == 8
-    lat = node_by_ct(wf, "EmptySD3LatentImage")["inputs"]
+    wf = comfyui_build_flux_workflow(
+        "POS", variant="schnell", width=768, height=512, seed=123
+    )
+    assert node_by_ct(wf, "RandomNoise")["inputs"]["noise_seed"] == 123
+    lat = node_by_ct(wf, "EmptyLatentImage")["inputs"]
    assert lat["width"] == 768 and lat["height"] == 512
-    pos = node_by_ct(wf, "FluxGuidance")["inputs"]["conditioning"]
-    assert pos == ["6", 0]  # FluxGuidance aplica sobre el CLIPTextEncode positivo


 def test_filename_prefix_en_saveimage():
@@ -72,8 +104,36 @@ def test_filename_prefix_en_saveimage():
    assert node_by_ct(wf, "SaveImage")["inputs"]["filename_prefix"] == "demo_flux"


+def test_variant_invalido_lanza_valueerror():
+    with pytest.raises(ValueError):
+        comfyui_build_flux_workflow("POS", variant="turbo")
+
+
+def test_available_valida_modelos_faltantes():
+    # Si se pasa 'available' y un modelo elegido no esta, lanza FileNotFoundError
+    # con el nombre que falta (error path: no crashea opaco).
+    available = {
+        "unet": ["otro_modelo.safetensors"],  # el schnell por defecto NO esta
+        "clip": ["clip_l.safetensors", "t5xxl_fp8_e4m3fn_scaled.safetensors"],
+        "vae": ["ae.safetensors"],
+    }
+    with pytest.raises(FileNotFoundError) as exc:
+        comfyui_build_flux_workflow("POS", variant="schnell", available=available)
+    assert "IMG_flux1-schnell-fp8-e4m3fn.safetensors" in str(exc.value)
+
+
+def test_available_ok_no_lanza():
+    available = {
+        "unet": ["IMG_flux1-schnell-fp8-e4m3fn.safetensors"],
+        "clip": ["clip_l.safetensors", "t5xxl_fp8_e4m3fn_scaled.safetensors"],
+        "vae": ["ae.safetensors"],
+    }
+    wf = comfyui_build_flux_workflow("POS", variant="schnell", available=available)
+    assert_api_format(wf)
+
+
 def test_determinista():
    # Builder puro: misma entrada -> mismo dict (sin red, seed fijo, sin estado).
-    a = comfyui_build_flux_workflow("POS", seed=123)
-    b = comfyui_build_flux_workflow("POS", seed=123)
+    a = comfyui_build_flux_workflow("POS", variant="dev", seed=123)
+    b = comfyui_build_flux_workflow("POS", variant="dev", seed=123)
    assert a == b