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

mcp/mcp-image-gen/src/server.py

@@ -39,8 +39,14 @@ COMFYUI_DIR = Path(
 # Maximum number of images allowed in a single batch call
 MAX_COUNT = 10
 
-# 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
+# This allows us to support multiple models (FLUX.1-schnell + FLUX.2 Klein with Heretic encoder)
+_WORKFLOW_REGISTRY: dict[str, Path] = {
+    "flux1-schnell.safetensors": Path(__file__).parent / "workflows" / "flux_schnell.json",
+    "flux-2-klein-4b.safetensors": Path(__file__).parent / "workflows" / "flux2_klein_heretic.json",
+}
+
+_DEFAULT_MODEL = "flux1-schnell.safetensors"
 
 
 # ---------------------------------------------------------------------------
@@ -181,21 +187,37 @@ class ComfyUIClient:
             return resp.content
 
     async def get_models(self) -> list[str]:
-        """Return the list of available checkpoint model filenames."""
-        async with httpx.AsyncClient(timeout=10.0) as client:
-            resp = await client.get(
-                f"{self.base_url}/object_info/CheckpointLoaderSimple"
-            )
-            resp.raise_for_status()
-            data = resp.json()
-            # ComfyUI returns: {"CheckpointLoaderSimple": {"input": {"required": {"ckpt_name": [["model1.safetensors", ...], ...]}}}}
-            node_info = data.get("CheckpointLoaderSimple", {})
-            ckpt_list = (
-                node_info.get("input", {})
-                .get("required", {})
-                .get("ckpt_name", [[]])[0]
-            )
-            return ckpt_list if isinstance(ckpt_list, list) else []
+        """Return the list of available checkpoint model filenames.
+
+        Combines models known to ComfyUI with our internal registry
+        (including FLUX.2 Klein with Heretic encoder).
+        """
+        models = set()
+
+        # Get models from ComfyUI
+        try:
+            async with httpx.AsyncClient(timeout=10.0) as client:
+                resp = await client.get(
+                    f"{self.base_url}/object_info/CheckpointLoaderSimple"
+                )
+                resp.raise_for_status()
+                data = resp.json()
+                node_info = data.get("CheckpointLoaderSimple", {})
+                ckpt_list = (
+                    node_info.get("input", {})
+                    .get("required", {})
+                    .get("ckpt_name", [[]])[0]
+                )
+                if isinstance(ckpt_list, list):
+                    models.update(ckpt_list)
+        except Exception:
+            # ComfyUI not reachable — fall back to registry only
+            pass
+
+        # Add our registered models
+        models.update(_WORKFLOW_REGISTRY.keys())
+
+        return sorted(list(models))
 
 
 # ---------------------------------------------------------------------------
@@ -209,13 +231,20 @@ def build_flux_workflow(
     height: int,
     steps: int,
     seed: int,
-    model: str,
+    model: str = _DEFAULT_MODEL,
 ) -> dict:
-    """Build a ComfyUI API-format workflow dict for FLUX.1-schnell text-to-image.
+    """Build a ComfyUI API-format workflow dict for the requested model.
 
-    This is a pure function — no I/O, fully testable.
+    Supports:
+    - "flux1-schnell.safetensors" (original)
+    - "flux-2-klein-4b-fp8.safetensors" (with Heretic-abliterated Qwen3-4B text encoder)
+
+    Falls back to FLUX.1-schnell if model is unknown.
+    This is a pure function — no I/O outside the registry, fully testable.
     """
-    with open(_WORKFLOW_PATH) as f:
+    workflow_path = _WORKFLOW_REGISTRY.get(model, _WORKFLOW_REGISTRY[_DEFAULT_MODEL])
+
+    with open(workflow_path) as f:
         wf = json.load(f)
     wf = copy.deepcopy(wf)
 
@@ -277,18 +306,13 @@ async def _generate_single(
 ) -> list:
     """Generate a single image and return [TextContent, ImageContent] or [TextContent] on error.
 
-    Args:
-        client:              ComfyUIClient instance.
-        prompt:              Positive text prompt.
-        negative_prompt:     Negative text prompt.
-        width / height:      Image dimensions.
-        steps:               Inference steps.
-        seed:                Seed value (-1 = random).
-        model:               ComfyUI model filename.
-        resolved_output_dir: Resolved output directory Path.
-        name:                User-supplied name prefix (unsanitized).
-        label:               Human-readable label for TextContent prefix (e.g. "[lumen 1/3]").
+    Supports two models:
+    - flux1-schnell.safetensors (default, fast 4-step)
+    - flux-2-klein-4b.safetensors (with Heretic-abliterated Qwen3-4B text encoder — no refusals)
     """
+    if model not in _WORKFLOW_REGISTRY:
+        model = _DEFAULT_MODEL
+        logger.warning("Unknown model %s, falling back to %s", model, _DEFAULT_MODEL)
     # Build and submit workflow
     try:
         workflow = build_flux_workflow(