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,