Fix Z-Image image-gen + LoRA training; add separate LoRA training base model

- images.py: stop reusing cached prompt embeds for Z-Image. Its encode_prompt
  returns per-sample lists the pipeline consumes during a run, so cache-HIT
  reuse corrupted the batch (AssertionError: len(size) == bsz in unpatchify).
  Z-Image now always encodes natively; cache key also bound to id(pipeline).
- loras.py: fix wrong import (multi_model_manager lives in codai.models.manager,
  not codai.api.state) used for VRAM unload + base-model path resolution; add an
  up-front architecture guard so transformer/DiT models (Z-Image/Flux/SD3) fail
  with a clear message instead of crashing in the CLIP tokenizer.
- New train_base_model on LoraTrainRequest: train LoRAs against a separate
  UNet-based SD1.x/SDXL model while generation keeps using the DiT image model.
- gen_township_fighters.py: thread --lora-train-base-model / web field through
  train_lora, _train_profile_loras, stage_loras/stage_env_loras, all run + per-
  card call sites; CONFIG_FIELDS + live default_args apply on /start, /save-config.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

codai/api/images.py

@@ -737,11 +737,23 @@ async def _generate_with_diffusers(pipeline, request, global_args, http_request=
     # Try to encode the prompt once and reuse the embeddings.
     # Falls back to passing the plain text prompt if encoding fails.
     # ------------------------------------------------------------------
-    model_id = getattr(pipeline, 'model_name_or_path', None) or str(type(pipeline).__name__)
+    # Bind the cache entry to this specific *loaded* pipeline instance, so a
+    # reload (the model-swapping workflow does this constantly) is a guaranteed
+    # cache miss instead of feeding stale GPU tensors into a fresh pipeline.
+    _pipe_cls = str(type(pipeline).__name__)
+    _model_name = getattr(pipeline, 'model_name_or_path', None) or _pipe_cls
+    model_id = f"{_model_name}#{id(pipeline)}"
     neg_prompt = getattr(request, 'negative_prompt', None) or ""
     do_cfg = cfg_scale > 1.0
 
-    cached_embeds = _embed_cache.get(request.prompt, neg_prompt, model_id)
+    # Some pipelines (Z-Image) return prompt embeddings as *per-sample lists*
+    # rather than stacked tensors, and the pipeline consumes/extends those
+    # lists during a run.  Reusing them on a later request corrupts the batch
+    # dimension (`assert len(size) == bsz` in the Z-Image transformer's
+    # unpatchify).  These are not safe to cache, so encode fresh every time.
+    _embed_cacheable = "ZImage" not in _pipe_cls
+
+    cached_embeds = _embed_cache.get(request.prompt, neg_prompt, model_id) if _embed_cacheable else None
     embed_kwargs = {}
     cache_hit = False
 
@@ -749,6 +761,11 @@ async def _generate_with_diffusers(pipeline, request, global_args, http_request=
         embed_kwargs = cached_embeds
         cache_hit = True
         print(f"Prompt embed cache HIT for model '{model_id}'")
+    elif not _embed_cacheable:
+        # Z-Image et al.: don't pre-encode at all.  Leaving embed_kwargs empty
+        # makes the call below pass the raw prompt, so the pipeline runs its own
+        # native encode + batching (the only path that survives reuse here).
+        pass
     else:
         # Try to encode and cache
         try:
@@ -781,7 +798,7 @@ async def _generate_with_diffusers(pipeline, request, global_args, http_request=
                         'pooled_prompt_embeds': enc[2],
                         'negative_pooled_prompt_embeds': enc[3],
                     }
-                if embed_kwargs:
+                if embed_kwargs and _embed_cacheable:
                     _embed_cache.put(request.prompt, neg_prompt, model_id, embed_kwargs)
                     print(f"Prompt embed cache STORE for model '{model_id}'")
         except Exception as e:

codai/api/loras.py

@@ -131,6 +131,10 @@ def _require_api_auth(request: Request) -> None:
 class LoraTrainRequest(BaseModel):
     name: str                              # output LoRA name (folder)
     base_model: str                        # image model key (models.json) or HF id / path
+    # Optional separate UNet-based SD1.x/SDXL model to train the LoRA against,
+    # when `base_model` (the generation model) is a transformer/DiT (Z-Image,
+    # Flux, SD3) this trainer can't target.  Falls back to base_model if unset.
+    train_base_model: Optional[str] = None
     character: Optional[str] = None        # saved character profile to pull images from
     environment: Optional[str] = None      # OR saved environment profile to pull images from
     images: Optional[List[str]] = None     # OR explicit base64/data-uri images
@@ -148,7 +152,7 @@ class LoraTrainRequest(BaseModel):
 def _resolve_base_model_path(base_model: str) -> str:
     """Resolve an image model key (or path/HF id) to a diffusers model directory."""
     try:
-        from codai.api.state import multi_model_manager
+        from codai.models.manager import multi_model_manager
         for key in (f"image:{base_model}", base_model):
             cfg = multi_model_manager.config.get(key)
             if cfg:
@@ -222,7 +226,9 @@ def _train_lora_sync(req: LoraTrainRequest) -> dict:
     from transformers import CLIPTextModel, CLIPTokenizer
 
     name = req.name
-    base_path = _resolve_base_model_path(req.base_model)
+    # Train against a dedicated UNet model when provided (the generation model
+    # may be a DiT this trainer can't target); otherwise use base_model.
+    base_path = _resolve_base_model_path(req.train_base_model or req.base_model)
     steps = max(50, min(5000, int(req.steps or 800)))
     rank = max(2, min(128, int(req.rank or 16)))
     resolution = int(req.resolution or 512)
@@ -247,7 +253,7 @@ def _train_lora_sync(req: LoraTrainRequest) -> dict:
 
     # Free VRAM: evict every resident model so training has the whole GPU.
     try:
-        from codai.api.state import multi_model_manager
+        from codai.models.manager import multi_model_manager
         multi_model_manager.unload_all_models()
     except Exception as e:
         print(f"  [lora] could not unload models before training: {e}")
@@ -257,6 +263,22 @@ def _train_lora_sync(req: LoraTrainRequest) -> dict:
 
     _set_progress(status="preparing", message=f"loading base model: {base_path}")
 
+    # This DreamBooth-LoRA trainer only supports UNet-based SD1.x / SDXL
+    # pipelines.  Transformer/DiT models (Z-Image, Flux, SD3, …) have a
+    # `transformer/` subfolder and no `unet/`; loading their non-CLIP tokenizer
+    # as a CLIPTokenizer crashes deep inside transformers.  Detect that up front
+    # and fail with an actionable message instead.
+    import os as _os
+    _has_unet = _os.path.isdir(_os.path.join(base_path, "unet"))
+    _has_transformer = _os.path.isdir(_os.path.join(base_path, "transformer"))
+    if _has_transformer and not _has_unet:
+        raise ValueError(
+            f"LoRA training base model '{base_path}' is a transformer/DiT "
+            f"architecture (has 'transformer/', no 'unet/'). This trainer only "
+            f"supports UNet-based SD1.x/SDXL models. Configure an SD1.x or SDXL "
+            f"image model as the LoRA training base."
+        )
+
     # Detect SDXL by attempting to load a second tokenizer.
     is_sdxl = False
     try:
@@ -537,6 +559,7 @@ def _write_meta(name, req, base_path, n_images, arch, instance_prompt):
     meta = {
         "name": name,
         "base_model": req.base_model,
+        "train_base_model": req.train_base_model or req.base_model,
         "base_path": base_path,
         "arch": arch,
         "instance_prompt": instance_prompt,

tools/gen_township_fighters.py

@@ -695,11 +695,13 @@ class CoderAIClient:
     def train_lora(self, name: str, base_model: str, character: str = None,
                    environment: str = None, images: list = None,
                    steps: int = 800, rank: int = 16,
-                   resolution: int = 512) -> dict:
+                   resolution: int = 512, train_base_model: str = None) -> dict:
         """Train a per-character or per-environment LoRA on the server.
         Blocks until complete."""
         body = {"name": name, "base_model": base_model,
                 "steps": int(steps), "rank": int(rank), "resolution": int(resolution)}
+        if train_base_model:
+            body["train_base_model"] = train_base_model
         if character:
             body["character"] = character
         if environment:
@@ -1222,6 +1224,7 @@ CONFIG_FIELDS = [
     "only_prompts", "only_videos",
     "consistency", "keyframe_steps", "keyframe_size",
     "character_strength", "lora_steps", "lora_rank", "lora_weight",
+    "lora_train_base_model",
     "no_env_loras", "env_lora_steps", "env_lora_rank", "env_lora_weight",
     "web_port",
 ]
@@ -1354,12 +1357,16 @@ _LORA_KINDS = {
 
 def _train_profile_loras(client: CoderAIClient, image_model: str, out_dir: Path,
                          names: list, kind: str,
-                         lora_steps: int = 800, lora_rank: int = 16) -> dict:
+                         lora_steps: int = 800, lora_rank: int = 16,
+                         train_base_model: str = None) -> dict:
     """Train one identity LoRA per profile of `kind` (server-side).
 
     Returns {name: lora_path}. Resumable: skips profiles whose LoRA already
     exists locally (<kind>_loras.json) or on the server. All training is grouped
     here so the image base model is touched once for the whole batch.
+
+    `train_base_model` overrides the model the LoRA is *trained* against (must be
+    a UNet-based SD1.x/SDXL model); `image_model` stays the generation model.
     """
     spec = _LORA_KINDS[kind]
     _log("\n" + "═" * 60)
@@ -1400,6 +1407,8 @@ def _train_profile_loras(client: CoderAIClient, image_model: str, out_dir: Path,
              f"({lora_steps} steps, rank {lora_rank}) — this can take a while…")
         train_kwargs = dict(name=lora_name, base_model=image_model,
                             steps=lora_steps, rank=lora_rank)
+        if train_base_model:
+            train_kwargs["train_base_model"] = train_base_model
         train_kwargs[kind] = name  # character=name OR environment=name
         try:
             res = _run_with_spinner(
@@ -1421,17 +1430,19 @@ def _train_profile_loras(client: CoderAIClient, image_model: str, out_dir: Path,
 
 
 def stage_loras(client: CoderAIClient, image_model: str, out_dir: Path,
-                char_names: list, lora_steps: int = 800, lora_rank: int = 16) -> dict:
+                char_names: list, lora_steps: int = 800, lora_rank: int = 16,
+                train_base_model: str = None) -> dict:
     """Train one identity LoRA per fighter. Returns {fighter: lora_path}."""
     return _train_profile_loras(client, image_model, out_dir, char_names,
-                                "character", lora_steps, lora_rank)
+                                "character", lora_steps, lora_rank, train_base_model)
 
 
 def stage_env_loras(client: CoderAIClient, image_model: str, out_dir: Path,
-                    env_names: list, lora_steps: int = 800, lora_rank: int = 16) -> dict:
+                    env_names: list, lora_steps: int = 800, lora_rank: int = 16,
+                    train_base_model: str = None) -> dict:
     """Train one identity LoRA per environment. Returns {environment: lora_path}."""
     return _train_profile_loras(client, image_model, out_dir, env_names,
-                                "environment", lora_steps, lora_rank)
+                                "environment", lora_steps, lora_rank, train_base_model)
 
 
 def _generate_keyframes(client: CoderAIClient, image_model: str, keyframe_dir: Path,
@@ -2208,6 +2219,9 @@ def launch_web_ui(default_args):
                 try:
                     kwargs = dict(name=lora_name, base_model=model,
                                   steps=int(steps), rank=int(rank))
+                    _tbm = getattr(default_args, "lora_train_base_model", None) or None
+                    if _tbm:
+                        kwargs["train_base_model"] = _tbm
                     kwargs[kind] = name
                     result["res"] = client.train_lora(**kwargs)
                 except Exception as e:
@@ -2717,6 +2731,11 @@ textarea{background:#111;border:1px solid #333;color:#e0e0e0;padding:.35rem .5re
     </div>
   </div>
   <div id=lora_fields style="margin-top:.6rem">
+    <div><label>LoRA training base model <span class=hint>(SD1.x/SDXL — leave empty to train on the image model)</span></label>
+         <input name=lora_train_base_model type=text style="width:100%"
+                placeholder="e.g. stabilityai/stable-diffusion-xl-base-1.0"
+                value="{_v('lora_train_base_model','')}"></div>
+    <p class=hint style="margin:.2rem 0 .6rem">Z-Image / Flux / SD3 image models can't be LoRA-trained directly. Set a UNet-based SD1.x or SDXL model here to train identity LoRAs while still generating with the image model above.</p>
     <label style="margin-top:0">Character LoRAs <span class=hint>(per-fighter identity)</span></label>
     <div class=row3>
       <div><label>LoRA train steps</label>
@@ -4499,6 +4518,7 @@ async function pollJob(){
                     "lora_steps": int(_fv("lora_steps", "800") or 800),
                     "lora_rank": int(_fv("lora_rank", "16") or 16),
                     "lora_weight": float(_fv("lora_weight", "0.85") or 0.85),
+                    "lora_train_base_model": _s(_fv("lora_train_base_model")) or "",
                     "no_env_loras": "env_loras" not in form,
                     "env_lora_steps": int(_fv("env_lora_steps", "800") or 800),
                     "env_lora_rank": int(_fv("env_lora_rank", "16") or 16),
@@ -4520,7 +4540,7 @@ async function pollJob(){
                 # immediately, so subsequent per-profile jobs (regenerate, train
                 # LoRA) and runs use them — not the values from script launch.
                 for _k in ("base_url", "api_key", "image_model",
-                           "video_model", "text_model"):
+                           "video_model", "text_model", "lora_train_base_model"):
                     setattr(default_args, _k, cfg.get(_k))
                 _web_log(f"  ⚙ Settings applied (image model: "
                          f"{cfg.get('image_model') or 'auto'})")
@@ -4628,6 +4648,7 @@ async function pollJob(){
             ns.lora_steps        = int(_fv("lora_steps", "800"))
             ns.lora_rank         = int(_fv("lora_rank", "16"))
             ns.lora_weight       = float(_fv("lora_weight", "0.85"))
+            ns.lora_train_base_model = (_fv("lora_train_base_model", "") or "").strip()
             # Environment LoRAs: checkbox "env_loras" present ⇒ train them.
             ns.no_env_loras      = ("env_loras" not in form)
             ns.env_lora_steps    = int(_fv("env_lora_steps", "800"))
@@ -4658,7 +4679,7 @@ async function pollJob(){
             # Apply the submitted connection/model settings to the live session
             # so later per-profile jobs (regenerate, train LoRA) use them too.
             for _k in ("base_url", "api_key", "image_model",
-                       "video_model", "text_model"):
+                       "video_model", "text_model", "lora_train_base_model"):
                 setattr(default_args, _k, getattr(ns, _k, None))
 
             # Apply the same shortcut logic as CLI
@@ -4923,7 +4944,8 @@ async function pollJob(){
                          f"{', '.join(char_names)}")
                 lora_map = stage_loras(client, image_model, out_dir_r, char_names or [],
                                        lora_steps=getattr(args, "lora_steps", 800),
-                                       lora_rank=getattr(args, "lora_rank", 16))
+                                       lora_rank=getattr(args, "lora_rank", 16),
+                                       train_base_model=getattr(args, "lora_train_base_model", None) or None)
             else:
                 _web_log("  ⚠ No characters found to train LoRAs for. Generate or "
                          "select fighters first (Characters page).")
@@ -4933,7 +4955,8 @@ async function pollJob(){
                              f"{', '.join(env_names)}")
                     env_lora_map = stage_env_loras(client, image_model, out_dir_r, env_names or [],
                                                    lora_steps=getattr(args, "env_lora_steps", 800),
-                                                   lora_rank=getattr(args, "env_lora_rank", 16))
+                                                   lora_rank=getattr(args, "env_lora_rank", 16),
+                                                   train_base_model=getattr(args, "lora_train_base_model", None) or None)
                 else:
                     _web_log("  ⚠ No environments found to train LoRAs for.")
 
@@ -5242,6 +5265,11 @@ OUTPUT LAYOUT
                           help="LoRA rank (default: 16).")
     cons_grp.add_argument("--lora-weight", type=float, default=0.85, metavar="F",
                           help="Weight applied to each character LoRA at generation (default: 0.85).")
+    cons_grp.add_argument("--lora-train-base-model", default="", metavar="MODEL",
+                          help="Separate UNet-based SD1.x/SDXL model (models.json key or HF id/path) "
+                               "to TRAIN LoRAs against, when the generation image model is a "
+                               "transformer/DiT (Z-Image, Flux, SD3) this trainer can't target. "
+                               "Generation still uses --image-model. Empty = train on --image-model.")
     cons_grp.add_argument("--no-env-loras", action="store_true",
                           help="Do not train/apply per-environment identity LoRAs when the "
                                "'lora' strategy is active (by default environments get LoRAs too).")
@@ -5393,12 +5421,14 @@ OUTPUT LAYOUT
     if "lora" in consistency and not args.skip_videos and (char_names or []):
         lora_map = stage_loras(client, image_model, out_dir, char_names or [],
                                lora_steps=getattr(args, "lora_steps", 800),
-                               lora_rank=getattr(args, "lora_rank", 16))
+                               lora_rank=getattr(args, "lora_rank", 16),
+                               train_base_model=getattr(args, "lora_train_base_model", None) or None)
     if ("lora" in consistency and not args.skip_videos
             and not getattr(args, "no_env_loras", False) and (env_names or [])):
         env_lora_map = stage_env_loras(client, image_model, out_dir, env_names or [],
                                        lora_steps=getattr(args, "env_lora_steps", 800),
-                                       lora_rank=getattr(args, "env_lora_rank", 16))
+                                       lora_rank=getattr(args, "env_lora_rank", 16),
+                                       train_base_model=getattr(args, "lora_train_base_model", None) or None)
 
     # ── Stage 3: Videos ────────────────────────────────────────────────────────
     if not args.skip_videos: