Patrick Plate
ea0c5d39c4
- 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
301 lines
8.9 KiB
Markdown
301 lines
8.9 KiB
Markdown
# 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.
|