Files
pi_mcps/plans/heretic-flux2-klein-integration.md
T
Patrick Plate ea0c5d39c4 fix(mcp-image-gen): fix Heretic/FLUX2 integration bugs
- Fix syntax error in server.py (dangling docstring lines)
- Correct model filename: flux-2-klein-4b.safetensors (without -fp8)
- Fix _WORKFLOW_REGISTRY key to match actual downloaded filename
- Update get_models() to always include registry models as fallback
- Fix test expectations to match corrected model names
- All 37 tests passing
2026-04-10 19:21:51 +02:00

8.9 KiB

Plan: FLUX.2 Klein 4B + Heretic Abliterated Text Encoder in mcp-image-gen

Datum: 2026-04-10
Autor: Lumen / Patrick Plate
Status: Ready for Implementation


Ziel

Das bestehende mcp-image-gen ComfyUI-Backend um ein zweites Modell erweitern: FLUX.2 Klein 4B mit dem abliterierten Qwen3-4B-Heretic als Text-Encoder.

Ergebnis: generate_image kann via model-Parameter zwischen zwei Workflows wählen:

  • flux1-schnell.safetensors → bestehender Workflow (unverändert)
  • flux-2-klein-4b-fp8.safetensors → neuer Heretic-Workflow (keine Prompt-Refusals)

Technischer Hintergrund

Warum Heretic + FLUX.2 Klein?

FLUX.2 Klein 4B verwendet Qwen3-4B als LLM Text-Encoder (statt CLIP+T5 wie bei FLUX.1). Dieser LLM-Encoder hat Safety-Alignment → verweigert bestimmte Prompts → abliterieren.

DreamFast/qwen3-4b-heretic (HuggingFace):

  • KL Divergenz: 0.0000 — null messbarer Modell-Schaden
  • Nur 3/100 Refusals nach Heretic v1.2.0 (200 Trials)
  • Drop-in Replacement für qwen_3_4b.safetensors

Modell-Architektur Unterschied

FLUX.1-schnell FLUX.2 Klein 4B
Diffusion Model flux1-schnell.safetensors (UNet) flux-2-klein-4b-fp8.safetensors
Text Encoder DualCLIPLoader (T5+CLIP) CLIPLoader (Qwen3-4B)
VAE ae.safetensors flux2-vae.safetensors
Steps 4 4 (distilled)
VRAM ~8GB ~8.4GB
Refusals keine (kein LLM-Encoder) keine (abliteriert)

Dateien & Ordner

Neue Modell-Dateien (herunterzuladen)

~/ComfyUI/models/
├── diffusion_models/
│   └── flux-2-klein-4b-fp8.safetensors    ← FLUX.2 Klein distilled 4B
├── text_encoders/
│   └── qwen_3_4b_heretic.safetensors      ← Heretic abliteriert (von DreamFast/qwen3-4b-heretic)
└── vae/
    └── flux2-vae.safetensors              ← VAE für FLUX.2

Neue/geänderte Projekt-Dateien

mcp/mcp-image-gen/
├── src/
│   ├── server.py                          ← Workflow-Registry ergänzen
│   └── workflows/
│       ├── flux_schnell.json              ← unverändert
│       └── flux2_klein_heretic.json       ← NEU
├── tests/
│   └── test_server.py                     ← neue Tests für Registry + Workflow
└── USAGE.md                               ← Download-Anleitung ergänzen

Phase 1: Modelle herunterladen

1a. FLUX.2 Klein 4B (Diffusion Model)

# Von Black Forest Labs HuggingFace
huggingface-cli download black-forest-labs/FLUX.2-klein-4B \
  flux-2-klein-4b-fp8.safetensors \
  --local-dir ~/ComfyUI/models/diffusion_models/

1b. FLUX.2 VAE

huggingface-cli download black-forest-labs/FLUX.2-klein-4B \
  flux2-vae.safetensors \
  --local-dir ~/ComfyUI/models/vae/

1c. Qwen3-4B-Heretic (abliterierter Text-Encoder)

# Von DreamFast — bereits abliteriert, kein Heretic-Run nötig
huggingface-cli download DreamFast/qwen3-4b-heretic \
  --local-dir /tmp/qwen3-4b-heretic/

# Safetensors-Datei in ComfyUI text_encoders ablegen
cp /tmp/qwen3-4b-heretic/model.safetensors \
   ~/ComfyUI/models/text_encoders/qwen_3_4b_heretic.safetensors

Hinweis: DreamFast/qwen3-4b-heretic ist ein GGUF-/SafeTensors-Mix. Wir brauchen die .safetensors Variante für ComfyUI. Falls nur GGUF verfügbar: huggingface-cli download Lockout/qwen3-4b-heretic-zimage qwen-4b-zimage-hereticV2-q8.gguf


Phase 2: Neues Workflow-JSON

Datei: mcp/mcp-image-gen/src/workflows/flux2_klein_heretic.json

FLUX.2 Klein verwendet andere ComfyUI-Nodes als FLUX.1-schnell:

  • DualCLIPLoaderCLIPLoader (einzelner Qwen-Encoder)
  • UNETLoader mit diffusion_models/ Pfad statt checkpoints/
  • EmptySD3LatentImage → gleich (kompatibel)
  • KSampler → gleich aber sampler_name: "euler", scheduler: "beta", steps: 4
{
  "6": {
    "class_type": "CLIPTextEncode",
    "inputs": {
      "clip": ["30", 0],
      "text": "PROMPT_PLACEHOLDER"
    }
  },
  "8": {
    "class_type": "VAEDecode",
    "inputs": {
      "samples": ["13", 0],
      "vae": ["31", 0]
    }
  },
  "9": {
    "class_type": "SaveImage",
    "inputs": {
      "filename_prefix": "mcp-image-gen",
      "images": ["8", 0]
    }
  },
  "13": {
    "class_type": "KSampler",
    "inputs": {
      "cfg": 1.0,
      "denoise": 1.0,
      "latent_image": ["27", 0],
      "model": ["32", 0],
      "negative": ["33", 0],
      "positive": ["6", 0],
      "sampler_name": "euler",
      "scheduler": "beta",
      "seed": 42,
      "steps": 4
    }
  },
  "27": {
    "class_type": "EmptySD3LatentImage",
    "inputs": {
      "batch_size": 1,
      "height": 1024,
      "width": 1024
    }
  },
  "30": {
    "class_type": "CLIPLoader",
    "inputs": {
      "clip_name": "qwen_3_4b_heretic.safetensors",
      "type": "flux"
    }
  },
  "31": {
    "class_type": "VAELoader",
    "inputs": {
      "vae_name": "flux2-vae.safetensors"
    }
  },
  "32": {
    "class_type": "UNETLoader",
    "inputs": {
      "unet_name": "flux-2-klein-4b-fp8.safetensors",
      "weight_dtype": "fp8_e4m3fn"
    }
  },
  "33": {
    "class_type": "CLIPTextEncode",
    "inputs": {
      "clip": ["30", 0],
      "text": "NEGATIVE_PLACEHOLDER"
    }
  }
}

Phase 3: server.py — Workflow-Registry

Änderung 1: Workflow-Registry dict (nach _WORKFLOW_PATH)

# Path to the bundled FLUX.1-schnell workflow template
_WORKFLOW_PATH = Path(__file__).parent / "workflows" / "flux_schnell.json"

# Workflow registry: model filename → workflow JSON path
_WORKFLOW_REGISTRY: dict[str, Path] = {
    "flux1-schnell.safetensors": Path(__file__).parent / "workflows" / "flux_schnell.json",
    "flux-2-klein-4b-fp8.safetensors": Path(__file__).parent / "workflows" / "flux2_klein_heretic.json",
}

_DEFAULT_MODEL = "flux1-schnell.safetensors"

Änderung 2: _load_workflow() Hilfsfunktion

def _load_workflow(model: str) -> dict:
    """Load the correct workflow JSON for the requested model.
    
    Falls back to FLUX.1-schnell if model not in registry.
    """
    path = _WORKFLOW_REGISTRY.get(model, _WORKFLOW_PATH)
    if not path.exists():
        raise FileNotFoundError(f"Workflow JSON not found: {path}")
    return json.loads(path.read_text())

Änderung 3: _generate_single() nutzt Registry

Aktueller Code lädt immer _WORKFLOW_PATH. Änderung: _load_workflow(model) aufrufen:

async def _generate_single(
    client: ComfyUIClient,
    prompt: str,
    negative_prompt: str,
    model: str,
    seed: int,
    width: int,
    height: int,
    steps: int,
    output_dir: Path,
    name: str,
) -> tuple[TextContent, ImageContent | None]:
    workflow = _load_workflow(model)  # ← statt json.loads(_WORKFLOW_PATH.read_text())
    # ... rest unchanged

Phase 4: Tests

Neue Tests in mcp/mcp-image-gen/tests/test_server.py:

  1. test_workflow_registry_contains_both_models — Registry hat flux1-schnell + flux2-klein
  2. test_load_workflow_flux1_schnell — lädt flux_schnell.json korrekt
  3. test_load_workflow_flux2_klein — lädt flux2_klein_heretic.json korrekt
  4. test_load_workflow_unknown_model_falls_back — unbekanntes Modell → FLUX.1-schnell
  5. test_generate_image_uses_flux2_workflow — end-to-end Mock mit flux-2-klein-4b-fp8.safetensors

Phase 5: USAGE.md Update

Neuer Abschnitt "FLUX.2 Klein 4B (Heretic)" in mcp/mcp-image-gen/USAGE.md:

  • Download-Befehle für alle 3 neuen Modell-Dateien
  • Erklärung warum Heretic (abliterierter Text-Encoder, KL=0)
  • Beispiel-Aufruf: generate_image("...", model="flux-2-klein-4b-fp8.safetensors")

VRAM-Analyse

Modell VRAM gesamt Passt in 24GB?
FLUX.1-schnell (fp8) ~8GB
FLUX.2 Klein 4B (fp8) + Qwen3-4B ~8.4GB + ~4GB = ~12.4GB
Beide gleichzeitig geladen ~20GB mit Margin

Der RX 7900 XTX mit 24GB VRAM kann beide Modelle komfortabel halten.


Risiken & Mitigationen

Risiko Wahrscheinlichkeit Mitigation
CLIPLoader node nicht verfügbar in ComfyUI niedrig ComfyUI updaten; alternativ custom node
DreamFast-Modell nur als GGUF verfügbar mittel Lockout/qwen3-4b-heretic-zimage GGUF als Fallback
Qwen3-4B braucht anderen node type mittel Live-Test in ComfyUI UI zuerst; workflow anpassen
ROCm + Qwen3-4B Kompatibilität niedrig gleiche ROCm-Umgebung wie FLUX.1-schnell

Entscheidung

Empfehlung: Umsetzen. Minimale Code-Änderungen, kein Breaking Change, klarer Mehrwert.

Der einzige unsichere Punkt ist der genaue ComfyUI-Node-Name für den Qwen3-4B-Loader. Empfohlene Vorgehensweise: Erst in der ComfyUI-Web-UI manuell einen Workflow mit Qwen3-4B aufbauen → JSON exportieren → als flux2_klein_heretic.json speichern. Das garantiert korrekte Node-Namen ohne Guess-Work.