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
This commit is contained in:
Patrick Plate
@@ -0,0 +1,300 @@
|
||||
# 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.
|
||||
Reference in New Issue
Block a user