337 lines
10 KiB
Python
337 lines
10 KiB
Python
# SimpleSyrup - workflow-focused ComfyUI extensions for image generation
|
|
# Copyright (C) 2026 Artificial Sweetener and contributors
|
|
# SPDX-License-Identifier: AGPL-3.0-or-later
|
|
|
|
"""Tests for separate FLUX.1 and FLUX.2 loading orchestration."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass, field
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
from simple_syrup.domain.flux_profiles import (
|
|
Flux2TextEncoderProfile,
|
|
FluxGeneration,
|
|
FluxModelProfile,
|
|
)
|
|
from simple_syrup.runtime.auto_model_artifact import AutoModelArtifact
|
|
from simple_syrup.runtime.flux_artifacts import (
|
|
FLUX2_DEV_TEXT_ENCODER,
|
|
FLUX2_KLEIN_4B_TEXT_ENCODER,
|
|
FLUX2_KLEIN_9B_TEXT_ENCODER,
|
|
FLUX2_VAE,
|
|
FLUX_CLIP_L,
|
|
FLUX_T5_XXL,
|
|
FLUX_VAE,
|
|
)
|
|
from simple_syrup.runtime.model_downloads import ProgressReporter
|
|
from simple_syrup.services.flux2_loader_service import Flux2LoaderService
|
|
from simple_syrup.services.flux_loader_components import AUTO_CHOICE
|
|
from simple_syrup.services.flux_loader_service import FluxLoaderService
|
|
|
|
|
|
class RecordingDiffusionLoader:
|
|
"""Record diffusion selection and return a fixed model object."""
|
|
|
|
def __init__(self) -> None:
|
|
"""Create an empty call log."""
|
|
|
|
self.calls: list[tuple[str, str]] = []
|
|
self.model = object()
|
|
|
|
def load(self, diffusion_model: str, weight_dtype: str) -> object:
|
|
"""Record the requested standalone model."""
|
|
|
|
self.calls.append((diffusion_model, weight_dtype))
|
|
return self.model
|
|
|
|
|
|
class RecordingInspector:
|
|
"""Return a configured model profile and record inspection calls."""
|
|
|
|
def __init__(self, profile: FluxModelProfile | None) -> None:
|
|
"""Configure the profile returned by inspection."""
|
|
|
|
self.profile = profile
|
|
self.models: list[object] = []
|
|
|
|
def inspect(self, model: object) -> FluxModelProfile | None:
|
|
"""Record and return the configured profile."""
|
|
|
|
self.models.append(model)
|
|
return self.profile
|
|
|
|
|
|
@dataclass
|
|
class RecordingComponents:
|
|
"""Record automatic/manual component orchestration requests."""
|
|
|
|
encoder_requests: list[
|
|
tuple[str, AutoModelArtifact | None, ProgressReporter | None]
|
|
] = field(default_factory=list)
|
|
vae_requests: list[tuple[str, AutoModelArtifact, ProgressReporter | None]] = field(
|
|
default_factory=list
|
|
)
|
|
|
|
def resolve_text_encoder(
|
|
self,
|
|
selection: str,
|
|
automatic_artifact: AutoModelArtifact | None,
|
|
progress: ProgressReporter | None,
|
|
) -> Path:
|
|
"""Record and return a deterministic encoder path."""
|
|
|
|
self.encoder_requests.append((selection, automatic_artifact, progress))
|
|
filename = (
|
|
automatic_artifact.filename
|
|
if selection == AUTO_CHOICE and automatic_artifact is not None
|
|
else selection
|
|
)
|
|
return Path("C:/models/text_encoders") / filename
|
|
|
|
def load_vae(
|
|
self,
|
|
selection: str,
|
|
automatic_artifact: AutoModelArtifact,
|
|
progress: ProgressReporter | None,
|
|
) -> object:
|
|
"""Record and return a deterministic VAE object."""
|
|
|
|
self.vae_requests.append((selection, automatic_artifact, progress))
|
|
return "vae"
|
|
|
|
|
|
@dataclass
|
|
class RecordingTextEncoderLoader:
|
|
"""Record text encoder paths, Comfy CLIP type, and device."""
|
|
|
|
calls: list[tuple[tuple[Path, ...], str, str]] = field(default_factory=list)
|
|
|
|
def load(
|
|
self,
|
|
paths: tuple[Path, ...],
|
|
clip_type_name: str,
|
|
device: str,
|
|
) -> object:
|
|
"""Record and return a fixed CLIP object."""
|
|
|
|
self.calls.append((paths, clip_type_name, device))
|
|
return "clip"
|
|
|
|
|
|
class IdentityProgress:
|
|
"""No-op progress reporter whose identity is asserted by tests."""
|
|
|
|
def start(self, label: str, total: int | None) -> None:
|
|
"""Accept a progress start."""
|
|
|
|
def advance(self, current: int, total: int | None) -> None:
|
|
"""Accept a progress update."""
|
|
|
|
def finish(self) -> None:
|
|
"""Accept progress completion."""
|
|
|
|
|
|
def test_flux1_auto_loads_dual_encoders_and_generation_vae() -> None:
|
|
"""FLUX.1 auto resolves CLIP-L, T5-XXL, and ae through one progress path."""
|
|
|
|
diffusion = RecordingDiffusionLoader()
|
|
inspector = RecordingInspector(FluxModelProfile(FluxGeneration.FLUX))
|
|
components = RecordingComponents()
|
|
text_loader = RecordingTextEncoderLoader()
|
|
progress = IdentityProgress()
|
|
service = FluxLoaderService(diffusion, text_loader, inspector, components)
|
|
|
|
result = service.load_models(
|
|
"renamed-civit-finetune.safetensors",
|
|
"fp8_e4m3fn_fast",
|
|
AUTO_CHOICE,
|
|
AUTO_CHOICE,
|
|
"cpu",
|
|
AUTO_CHOICE,
|
|
progress,
|
|
)
|
|
|
|
assert result == (diffusion.model, "clip", "vae")
|
|
assert components.encoder_requests == [
|
|
(AUTO_CHOICE, FLUX_CLIP_L, progress),
|
|
(AUTO_CHOICE, FLUX_T5_XXL, progress),
|
|
]
|
|
assert components.vae_requests == [(AUTO_CHOICE, FLUX_VAE, progress)]
|
|
assert text_loader.calls[0][1:] == ("FLUX", "cpu")
|
|
|
|
|
|
def test_flux1_auto_rejects_unrecognized_model_with_manual_guidance() -> None:
|
|
"""Unresolved FLUX.1 auto policy fails before any support artifact resolution."""
|
|
|
|
components = RecordingComponents()
|
|
service = FluxLoaderService(
|
|
RecordingDiffusionLoader(),
|
|
RecordingTextEncoderLoader(),
|
|
RecordingInspector(None),
|
|
components,
|
|
)
|
|
|
|
with pytest.raises(ValueError, match="Select the text encoders and VAE manually"):
|
|
service.load_models(
|
|
"unknown.safetensors",
|
|
"default",
|
|
AUTO_CHOICE,
|
|
AUTO_CHOICE,
|
|
"default",
|
|
AUTO_CHOICE,
|
|
)
|
|
|
|
assert components.encoder_requests == []
|
|
|
|
|
|
def test_flux1_fully_manual_path_does_not_gatekeep_model_family() -> None:
|
|
"""Manual selections proceed without architecture detection."""
|
|
|
|
inspector = RecordingInspector(None)
|
|
components = RecordingComponents()
|
|
text_loader = RecordingTextEncoderLoader()
|
|
service = FluxLoaderService(
|
|
RecordingDiffusionLoader(),
|
|
text_loader,
|
|
inspector,
|
|
components,
|
|
)
|
|
|
|
result = service.load_models(
|
|
"possibly-not-flux.safetensors",
|
|
"default",
|
|
"manual_clip.safetensors",
|
|
"manual_t5.safetensors",
|
|
"default",
|
|
"manual_vae.safetensors",
|
|
)
|
|
|
|
assert result[1:] == ("clip", "vae")
|
|
assert inspector.models == []
|
|
assert text_loader.calls[0][1] == "FLUX"
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("encoder_profile", "expected_artifact"),
|
|
(
|
|
(Flux2TextEncoderProfile.DEV, FLUX2_DEV_TEXT_ENCODER),
|
|
(Flux2TextEncoderProfile.KLEIN_4B, FLUX2_KLEIN_4B_TEXT_ENCODER),
|
|
(Flux2TextEncoderProfile.KLEIN_9B, FLUX2_KLEIN_9B_TEXT_ENCODER),
|
|
),
|
|
)
|
|
def test_flux2_auto_selects_encoder_for_loaded_architecture(
|
|
encoder_profile: Flux2TextEncoderProfile,
|
|
expected_artifact: AutoModelArtifact,
|
|
) -> None:
|
|
"""FLUX.2 auto supports dev, Klein 4B, and Klein 9B/KV profiles."""
|
|
|
|
components = RecordingComponents()
|
|
text_loader = RecordingTextEncoderLoader()
|
|
progress = IdentityProgress()
|
|
service = Flux2LoaderService(
|
|
RecordingDiffusionLoader(),
|
|
text_loader,
|
|
RecordingInspector(FluxModelProfile(FluxGeneration.FLUX2, encoder_profile)),
|
|
components,
|
|
)
|
|
|
|
result = service.load_models(
|
|
"arbitrarily-renamed-finetune.safetensors",
|
|
"default",
|
|
AUTO_CHOICE,
|
|
"default",
|
|
AUTO_CHOICE,
|
|
progress,
|
|
)
|
|
|
|
assert result[1:] == ("clip", "vae")
|
|
assert components.encoder_requests == [(AUTO_CHOICE, expected_artifact, progress)]
|
|
assert components.vae_requests == [(AUTO_CHOICE, FLUX2_VAE, progress)]
|
|
assert text_loader.calls[0][1] == "FLUX2"
|
|
|
|
|
|
def test_flux2_unknown_width_requires_manual_text_encoder_only() -> None:
|
|
"""A recognized future FLUX.2 model can still auto-load its common VAE."""
|
|
|
|
profile = FluxModelProfile(FluxGeneration.FLUX2)
|
|
service = Flux2LoaderService(
|
|
RecordingDiffusionLoader(),
|
|
RecordingTextEncoderLoader(),
|
|
RecordingInspector(profile),
|
|
RecordingComponents(),
|
|
)
|
|
|
|
with pytest.raises(ValueError, match="Select the text encoder manually"):
|
|
service.load_models(
|
|
"future-flux2.safetensors",
|
|
"default",
|
|
AUTO_CHOICE,
|
|
"default",
|
|
AUTO_CHOICE,
|
|
)
|
|
|
|
components = RecordingComponents()
|
|
manual_service = Flux2LoaderService(
|
|
RecordingDiffusionLoader(),
|
|
RecordingTextEncoderLoader(),
|
|
RecordingInspector(profile),
|
|
components,
|
|
)
|
|
manual_service.load_models(
|
|
"future-flux2.safetensors",
|
|
"default",
|
|
"manual_encoder.safetensors",
|
|
"default",
|
|
AUTO_CHOICE,
|
|
)
|
|
|
|
assert components.encoder_requests[0][1] is None
|
|
assert components.vae_requests[0][1] is FLUX2_VAE
|
|
|
|
|
|
def test_flux2_fully_manual_path_does_not_gatekeep_model_family() -> None:
|
|
"""Manual FLUX.2 selections proceed even when Comfy cannot classify the model."""
|
|
|
|
inspector = RecordingInspector(None)
|
|
text_loader = RecordingTextEncoderLoader()
|
|
service = Flux2LoaderService(
|
|
RecordingDiffusionLoader(),
|
|
text_loader,
|
|
inspector,
|
|
RecordingComponents(),
|
|
)
|
|
|
|
service.load_models(
|
|
"non-flux-or-unknown.safetensors",
|
|
"default",
|
|
"manual_encoder.safetensors",
|
|
"default",
|
|
"manual_vae.safetensors",
|
|
)
|
|
|
|
assert inspector.models == []
|
|
assert text_loader.calls[0][1] == "FLUX2"
|
|
|
|
|
|
def test_flux2_auto_rejects_a_detected_flux1_model() -> None:
|
|
"""Cross-generation automatic selection fails with a manual escape hatch."""
|
|
|
|
service = Flux2LoaderService(
|
|
RecordingDiffusionLoader(),
|
|
RecordingTextEncoderLoader(),
|
|
RecordingInspector(FluxModelProfile(FluxGeneration.FLUX)),
|
|
RecordingComponents(),
|
|
)
|
|
|
|
with pytest.raises(ValueError, match="recognizes as FLUX.2"):
|
|
service.load_models(
|
|
"flux1.safetensors",
|
|
"default",
|
|
AUTO_CHOICE,
|
|
"default",
|
|
AUTO_CHOICE,
|
|
)
|