Skip to content

C
coderai
Commit f21de22e authored by Stefy Lanza (nextime / spora )'s avatar Stefy Lanza (nextime / spora )
Browse files
Options

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: 's avatarClaude Opus 4.8 <noreply@anthropic.com>
parent 2c736da4
Hide whitespace changes
Inline Side-by-side
Showing with 88 additions and 18 deletions
codai/api/images.py
View file @ f21de22e
......@@ -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
View file @ f21de22e
......@@ -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
View file @ f21de22e
......@@ -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:
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment