# 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)
```bash
# 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
```bash
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)
```bash
# 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`](mcp/mcp-image-gen/src/workflows/flux2_klein_heretic.json)

FLUX.2 Klein verwendet andere ComfyUI-Nodes als FLUX.1-schnell:
- `DualCLIPLoader` → `CLIPLoader` (einzelner Qwen-Encoder)
- `UNETLoader` mit `diffusion_models/` Pfad statt `checkpoints/`
- `EmptySD3LatentImage` → gleich (kompatibel)
- `KSampler` → gleich aber `sampler_name: "euler"`, `scheduler: "beta"`, `steps: 4`

```json
{
  "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`)

```python
# 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

```python
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:

```python
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`](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`](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.