Skip to content

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

Township UI: profile management, regenerate/upload refs, env parity 

Web UI (tools/gen_township_fighters.py):
- Characters and Environments management pages: edit every meta field
  (data-driven form, not a fixed set), remove a profile, delete individual
  reference images — synced to CoderAI (DELETE/PATCH).
- Regenerate references: background job generates N new images guided by the
  kept references (IP-Adapter) and APPENDS them, preserving non-deleted ones.
- Upload your own reference image files (magic-byte validated, appended).
- Clear run banner + status-pill label indicating which run/step is executing.
- Config: -s/--save and -c/--config; web Save button writes server-side and,
  when launched with -c, defaults to that same config path.

Environment parity / consistency:
- Per-environment identity LoRAs (loras.py: environment field, env image
  resolver, "place" instance prompt; script stage_env_loras, env_loras.json,
  applied to keyframes + video alongside fighter LoRAs).
- Environment IP-Adapter: environment_profiles on the image request
  (imagerequest.py) resolved into the IP-Adapter reference set (images.py),
  so environment regeneration can match kept references.
Co-Authored-By: 's avatarClaude Opus 4.8 <noreply@anthropic.com>
parent f5cb44e1
Hide whitespace changes
Inline Side-by-side
Showing with 1813 additions and 144 deletions
codai/api/images.py
View file @ 278fab3e
......@@ -51,6 +51,14 @@ from codai.api.state import get_load_mode
import hashlib as _hashlib
import threading as _threading
# Serializes all diffusers from_pretrained() calls.
# huggingface_hub acquires per-repo .lock files during from_pretrained; running
# two from_pretrained calls concurrently (or one alongside snapshot_download on
# the same repo) causes a filelock deadlock that hangs the process indefinitely.
# A single threading.Lock here ensures only one pipeline loads at a time.
_DIFFUSERS_LOAD_LOCK = _threading.Lock()
class _PromptEmbedCache:
"""Single-entry LRU cache for diffusers prompt embeddings."""
......@@ -323,21 +331,29 @@ def _disable_safety_checker(pipe):
return pipe
def _load_diffusers_pipeline(model_name: str, global_args):
def _load_diffusers_pipeline(model_name: str, global_args, model_config: dict = None):
"""
Try to load a model using the diffusers library.
Returns the loaded pipeline or None if diffusers can't handle this model.
Raises Exception if loading fails for other reasons.
Per-model configuration (model_config) is the source of truth and takes
precedence over CLI/global args for precision, offload, quantization, etc.
"""
from diffusers import StableDiffusionPipeline, StableDiffusionXLPipeline, DiffusionPipeline
import torch
# Check for --no-ram mode
no_ram = getattr(global_args, 'no_ram', False) if global_args else False
# Determine precision from CLI argument (--image-precision)
precision = getattr(global_args, 'image_precision', 'f32') or 'f32'
_mc = model_config or {}
def _cfg(key, default=None):
"""Read a value from the per-model configuration only (source of truth)."""
v = _mc.get(key)
return v if v is not None else default
# All loading parameters come from the per-model configuration.
no_ram = bool(_cfg('no_ram', False))
precision = _cfg('precision', 'f32') or 'f32'
precision_map = {
'bf16': torch.bfloat16,
'f32': torch.float32,
......@@ -354,14 +370,45 @@ def _load_diffusers_pipeline(model_name: str, global_args):
else:
print(f"Using precision: {precision} ({dtype})")
# Check if CPU offload is requested via CLI
use_sequential_offload = getattr(global_args, 'image_cpu_offload', False)
# CPU offload comes from the per-model configuration: an explicit
# cpu_offload flag, or an offload_strategy that implies CPU offloading.
_offload_strategy = _mc.get('offload_strategy')
use_sequential_offload = bool(
_mc.get('cpu_offload')
or (_offload_strategy in ('cpu', 'sequential', 'model', 'disk'))
)
# Quantization (per-model config). Builds a diffusers quantization config
# applied per-component so 4-bit/8-bit image models use less VRAM. Per-model
# 'component_quantization' overrides win; otherwise the global flag applies
# to all heavy components (backbone + text encoders).
from codai.models.hf_loading import (
build_pipeline_quant_config, build_gguf_pipeline_components)
_img_quant_config, _img_quant_desc = build_pipeline_quant_config(model_name, _mc, dtype)
if _img_quant_config is not None:
print(f"Image quantization: {_img_quant_desc}")
_img_gguf_components, _img_gguf_desc = build_gguf_pipeline_components(model_name, _mc, dtype)
if _img_gguf_components:
print(f"Image GGUF components: {_img_gguf_desc}")
# --no-ram mode: never use CPU offload
if no_ram and use_sequential_offload:
print("--no-ram mode: ignoring --image-cpu-offload, forcing full GPU loading")
use_sequential_offload = False
# Refuse to load a model that is currently being downloaded — the HF hub
# file lock on the same repo would deadlock the process.
try:
from codai.admin.routes import get_active_download_model_ids
active_downloads = get_active_download_model_ids()
if model_name in active_downloads:
raise RuntimeError(
f"Model '{model_name}' is currently being downloaded. "
"Wait for the download to finish before loading it."
)
except ImportError:
pass
# =====================================================================
# --no-ram mode: load directly on GPU, no CPU RAM fallback
# =====================================================================
......@@ -370,20 +417,31 @@ def _load_diffusers_pipeline(model_name: str, global_args):
print(f"--no-ram mode: loading diffusers model directly on {cuda_device}")
try:
_xtra = {}
if _img_quant_config is not None:
_xtra['quantization_config'] = _img_quant_config
if _img_gguf_components:
_xtra.update(_img_gguf_components)
with _DIFFUSERS_LOAD_LOCK:
try:
pipeline = StableDiffusionXLPipeline.from_pretrained(
model_name,
torch_dtype=dtype,
use_safetensors=True,
**_xtra,
)
except Exception:
pipeline = DiffusionPipeline.from_pretrained(
model_name,
torch_dtype=dtype,
use_safetensors=True,
**_xtra,
)
try:
pipeline = StableDiffusionXLPipeline.from_pretrained(
model_name,
torch_dtype=dtype,
use_safetensors=True,
)
pipeline = pipeline.to(cuda_device)
except Exception:
pipeline = DiffusionPipeline.from_pretrained(
model_name,
torch_dtype=dtype,
use_safetensors=True,
)
pipeline = pipeline.to(cuda_device)
if _img_quant_config is None:
raise # only quantized pipelines may reject .to()
print(f"--no-ram: Diffusers model loaded on {cuda_device}")
return pipeline
except Exception as e:
......@@ -395,20 +453,39 @@ def _load_diffusers_pipeline(model_name: str, global_args):
# =====================================================================
# Proactive VRAM eviction before first load attempt
# =====================================================================
# HF repo IDs have no local file, so their size is unknown; we can't
# pre-compute needed_gb. Instead, evict if other models are loaded and
# free VRAM is below 15% of total (almost certainly an OOM on attempt 1).
# Evict only the minimum needed so models that fit together can coexist.
# Prefer the model's configured/estimated VRAM need; only fall back to the
# blunt "free < 15%" heuristic when the size is genuinely unknown.
if torch.cuda.is_available():
try:
from codai.models.manager import multi_model_manager as _mmm
if _mmm.models:
_free, _total = torch.cuda.mem_get_info()
if _total > 0 and (_free / _total) < 0.15:
print(f"Low VRAM ({_free/1e9:.1f} GB free of {_total/1e9:.1f} GB) with "
f"{len(_mmm.models)} model(s) loaded — evicting before load attempt")
_mmm.unload_all_models()
except Exception:
pass
_free_gb = _free / 1e9
# Needed VRAM for this model (config used_vram_gb, with quant/offload
# factors applied) — 0 when it can't be determined.
_key = None
for _k in (model_key, model_name, f"image:{model_name}"):
if _k in _mmm.config:
_key = _k
break
_need_gb = _mmm._get_model_used_vram_gb(_key or model_name, model_name)
if _need_gb > 0:
if _free_gb < _need_gb:
print(f"Image model needs {_need_gb:.1f} GB, {_free_gb:.1f} GB free "
f"— evicting the minimum to fit (others may coexist)")
_mmm._evict_models_for_vram(_need_gb)
else:
print(f"Image model needs {_need_gb:.1f} GB, {_free_gb:.1f} GB free "
f"— no eviction needed (coexisting with loaded models)")
elif _total > 0 and (_free / _total) < 0.15:
# Size unknown and VRAM nearly full — evict LRU one at a time
# until we clear ~25% headroom, instead of nuking everything.
print(f"Low VRAM ({_free_gb:.1f} GB free of {_total/1e9:.1f} GB), "
f"unknown model size — evicting LRU to free headroom")
_mmm._evict_models_for_vram(_total * 0.25 / 1e9)
except Exception as _ee:
print(f" Proactive eviction skipped: {_ee}")
# =====================================================================
# Standard loading path (with OOM fallback)
......@@ -422,22 +499,46 @@ def _load_diffusers_pipeline(model_name: str, global_args):
try:
load_attempt += 1
print(f"Loading attempt {load_attempt}/{max_attempts}...")
# Try to load as Stable Diffusion XL first, then generic DiffusionPipeline
try:
pipeline = StableDiffusionXLPipeline.from_pretrained(
model_name,
torch_dtype=dtype,
use_safetensors=True,
)
except Exception:
# Try generic diffusion pipeline (supports custom pipelines like ZImagePipeline)
pipeline = DiffusionPipeline.from_pretrained(
model_name,
torch_dtype=dtype,
use_safetensors=True,
)
# Acquire the global load lock before any from_pretrained call.
# This prevents concurrent HF hub file-lock conflicts (e.g. when
# another pipeline or snapshot_download holds the same .lock file).
with _DIFFUSERS_LOAD_LOCK:
# Re-check download conflict inside the lock — a download may
# have started between our first check and acquiring the lock.
try:
from codai.admin.routes import get_active_download_model_ids
if model_name in get_active_download_model_ids():
raise RuntimeError(
f"Model '{model_name}' started downloading while waiting "
"for the load lock. Wait for the download to finish."
)
except ImportError:
pass
# Inject per-model quantization config when configured.
_xtra = {}
if _img_quant_config is not None:
_xtra['quantization_config'] = _img_quant_config
if _img_gguf_components:
_xtra.update(_img_gguf_components)
# Try to load as Stable Diffusion XL first, then generic DiffusionPipeline
try:
pipeline = StableDiffusionXLPipeline.from_pretrained(
model_name,
torch_dtype=dtype,
use_safetensors=True,
**_xtra,
)
except Exception:
# Try generic diffusion pipeline (supports custom pipelines like ZImagePipeline)
pipeline = DiffusionPipeline.from_pretrained(
model_name,
torch_dtype=dtype,
use_safetensors=True,
**_xtra,
)
# Apply memory optimizations based on attempt
if torch.cuda.is_available():
if load_attempt >= 2:
......@@ -445,8 +546,17 @@ def _load_diffusers_pipeline(model_name: str, global_args):
print("Enabling attention slicing for lower VRAM usage...")
if hasattr(pipeline, 'enable_attention_slicing'):
pipeline.enable_attention_slicing()
if load_attempt >= 3 or use_sequential_offload:
if _img_quant_config is not None:
# Quantized (bitsandbytes) pipelines are already placed on GPU
# by from_pretrained and cannot be moved with .to(); only the
# non-quantized components need an explicit device move.
print("Quantized pipeline — placing non-quantized components on GPU")
try:
pipeline = pipeline.to("cuda")
except Exception:
pass # bitsandbytes components stay where loaded
elif load_attempt >= 3 or use_sequential_offload:
# Third attempt or offload requested: enable sequential CPU offload
print("Enabling sequential CPU offload for lower VRAM usage...")
if hasattr(pipeline, 'enable_sequential_cpu_offload'):
......@@ -484,6 +594,85 @@ def _load_diffusers_pipeline(model_name: str, global_args):
return pipeline
async def _apply_vae_override(pipeline, vae_model_id: str):
"""Swap the pipeline's VAE with an alternate model (diffusers only)."""
try:
import torch
from diffusers import AutoencoderKL
dtype = next(pipeline.parameters()).dtype if hasattr(pipeline, 'parameters') else torch.float16
vae = AutoencoderKL.from_pretrained(vae_model_id, torch_dtype=dtype)
vae = vae.to(pipeline.device)
pipeline.vae = vae
_log.info("VAE override applied: %s", vae_model_id)
except Exception as e:
_log.warning("Could not load VAE override %s: %s", vae_model_id, e)
def _ensure_ip_adapter_loaded(pipeline) -> bool:
"""Lazily load IP-Adapter weights matching the pipeline architecture.
Returns True if IP-Adapter is loaded and ready (so the caller can pass
ip_adapter_image), False if this pipeline type isn't supported. The result
is cached on the pipeline so repeated requests don't reload.
"""
# Already loaded (or already known-unsupported) for this pipeline instance.
flag = getattr(pipeline, '_coderai_ip_state', None)
if flag == 'loaded':
return True
if flag == 'unsupported':
return False
if not hasattr(pipeline, 'load_ip_adapter'):
try:
pipeline._coderai_ip_state = 'unsupported'
except Exception:
pass
return False
# Detect SDXL vs SD1.5 by the presence of a second text encoder.
is_sdxl = hasattr(pipeline, 'text_encoder_2') and getattr(pipeline, 'text_encoder_2', None) is not None
cls_name = type(pipeline).__name__.lower()
if 'xl' in cls_name:
is_sdxl = True
attempts = []
if is_sdxl:
attempts.append(("h94/IP-Adapter", "sdxl_models", "ip-adapter_sdxl.bin"))
else:
attempts.append(("h94/IP-Adapter", "models", "ip-adapter_sd15.bin"))
for repo, subfolder, weight_name in attempts:
try:
pipeline.load_ip_adapter(repo, subfolder=subfolder, weight_name=weight_name)
pipeline._coderai_ip_state = 'loaded'
_log.info("IP-Adapter loaded: %s/%s/%s", repo, subfolder, weight_name)
return True
except Exception as e:
_log.warning("IP-Adapter load failed (%s/%s): %s", repo, weight_name, e)
try:
pipeline._coderai_ip_state = 'unsupported'
except Exception:
pass
return False
def _apply_loras(pipeline, loras):
"""Load and activate LoRA weights on a diffusers pipeline."""
try:
names = []
weights = []
for i, lora in enumerate(loras):
name = lora.name or f"lora_{i}"
pipeline.load_lora_weights(lora.model, adapter_name=name)
names.append(name)
weights.append(float(lora.weight if lora.weight is not None else 1.0))
if names:
pipeline.set_adapters(names, weights)
_log.info("LoRA weights applied: %s", names)
except Exception as e:
_log.warning("Could not apply LoRA weights: %s", e)
async def _generate_with_diffusers(pipeline, request, global_args, http_request=None):
"""Generate images using a diffusers pipeline (with prompt-embedding cache)."""
import torch
......@@ -493,6 +682,14 @@ async def _generate_with_diffusers(pipeline, request, global_args, http_request=
if getattr(request, 'disable_safety_checker', False):
_disable_safety_checker(pipeline)
# Apply optional per-request VAE override
if getattr(request, 'vae_model', None):
await _apply_vae_override(pipeline, request.vae_model)
# Apply optional per-request LoRA weights
if getattr(request, 'loras', None):
_apply_loras(pipeline, request.loras)
# Determine size
width, height = 1024, 1024
if request.size:
......@@ -513,7 +710,9 @@ async def _generate_with_diffusers(pipeline, request, global_args, http_request=
try:
if hasattr(pipeline, 'enable_attention_slicing'):
pipeline.enable_attention_slicing(slice_size="auto")
if hasattr(pipeline, 'enable_vae_slicing'):
if hasattr(pipeline, 'vae') and hasattr(pipeline.vae, 'enable_slicing'):
pipeline.vae.enable_slicing()
elif hasattr(pipeline, 'enable_vae_slicing'):
pipeline.enable_vae_slicing()
except Exception as e:
print(f"Warning: Could not enable memory optimizations: {e}")
......@@ -554,13 +753,18 @@ async def _generate_with_diffusers(pipeline, request, global_args, http_request=
# Try to encode and cache
try:
if hasattr(pipeline, 'encode_prompt'):
enc = pipeline.encode_prompt(
prompt=request.prompt,
device=pipeline.device,
num_images_per_prompt=1,
do_classifier_free_guidance=do_cfg,
negative_prompt=neg_prompt or None,
)
import inspect as _inspect
_ep_params = set(_inspect.signature(pipeline.encode_prompt).parameters)
_ep_kwargs = {"prompt": request.prompt}
if "device" in _ep_params:
_ep_kwargs["device"] = pipeline.device
if "num_images_per_prompt" in _ep_params:
_ep_kwargs["num_images_per_prompt"] = 1
if "do_classifier_free_guidance" in _ep_params:
_ep_kwargs["do_classifier_free_guidance"] = do_cfg
if "negative_prompt" in _ep_params:
_ep_kwargs["negative_prompt"] = neg_prompt or None
enc = pipeline.encode_prompt(**_ep_kwargs)
# enc is a tuple; length varies by pipeline type
if len(enc) == 2:
# SD 1.x: (prompt_embeds, negative_prompt_embeds)
......@@ -586,6 +790,12 @@ async def _generate_with_diffusers(pipeline, request, global_args, http_request=
def _step_cb(pipe, step_index, timestep, callback_kwargs):
_progress_step(step_index + 1)
# Mid-generation thermal checkpoint: pause between denoise steps if too hot.
try:
from codai.models.thermal import checkpoint as _thermal_checkpoint
_thermal_checkpoint(context="image-gen")
except Exception:
pass
return callback_kwargs
# Resolve character references (saved profiles + inline images)
......@@ -600,6 +810,16 @@ async def _generate_with_diffusers(pipeline, request, global_args, http_request=
except Exception:
pass
# Environment profiles feed the SAME IP-Adapter reference set, so a
# regenerated location keyframe can match the references kept on disk.
try:
env_profiles = getattr(request, 'environment_profiles', None) or []
if env_profiles:
from codai.api.environments import resolve_environment_profiles
char_images += resolve_environment_profiles(env_profiles)
except Exception:
pass
# Build call kwargs
if embed_kwargs:
call_kwargs = dict(
......@@ -625,21 +845,27 @@ async def _generate_with_diffusers(pipeline, request, global_args, http_request=
callback_on_step_end=_step_cb,
)
# Inject IP-Adapter images if character references provided
# Inject IP-Adapter images if character references provided. The pipeline
# must have IP-Adapter *weights* loaded first — _ensure_ip_adapter_loaded
# lazily downloads + loads the right checkpoint for the architecture.
if char_images and hasattr(pipeline, 'set_ip_adapter_scale'):
try:
strength = getattr(request, 'character_strength', 0.6) or 0.6
ref_imgs = []
for ref in char_images:
from PIL import Image as PILImage
if ref.startswith('data:'):
_, b64 = ref.split(',', 1)
raw = base64.b64decode(b64)
else:
raw = base64.b64decode(ref)
ref_imgs.append(PILImage.open(io.BytesIO(raw)).convert('RGB'))
pipeline.set_ip_adapter_scale(strength)
call_kwargs['ip_adapter_image'] = ref_imgs[0] if len(ref_imgs) == 1 else ref_imgs
if _ensure_ip_adapter_loaded(pipeline):
strength = getattr(request, 'character_strength', 0.6) or 0.6
ref_imgs = []
for ref in char_images:
from PIL import Image as PILImage
if ref.startswith('data:'):
_, b64 = ref.split(',', 1)
raw = base64.b64decode(b64)
else:
raw = base64.b64decode(ref)
ref_imgs.append(PILImage.open(io.BytesIO(raw)).convert('RGB'))
pipeline.set_ip_adapter_scale(strength)
call_kwargs['ip_adapter_image'] = ref_imgs[0] if len(ref_imgs) == 1 else ref_imgs
else:
print("Note: IP-Adapter weights unavailable for this pipeline — "
"relying on prompt/LoRA for character consistency")
except Exception as _ip_err:
print(f"Warning: IP-Adapter injection failed ({_ip_err}), continuing without character refs")
......@@ -959,9 +1185,19 @@ async def create_image_generation(request: ImageGenerationRequest, http_request:
# Step 1: Ask the manager to resolve the model and manage VRAM
# =====================================================================
_progress_loading(request.model or "image")
model_info = multi_model_manager.request_model(
# Reserve VRAM for any per-request LoRA adapters so eviction frees enough
# headroom for base weights + adapters before the pipeline loads.
_lora_extra_gb = 0.0
if getattr(request, 'loras', None):
try:
_lora_extra_gb = multi_model_manager._lora_vram_gb(request.loras)
except Exception:
_lora_extra_gb = 0.0
model_info = await asyncio.to_thread(
multi_model_manager.request_model,
requested_model=request.model,
model_type="image"
model_type="image",
extra_vram_gb=_lora_extra_gb,
)
# Check if the model was rejected as not allowed
......@@ -1030,14 +1266,22 @@ async def create_image_generation(request: ImageGenerationRequest, http_request:
if not is_gguf:
try:
print(f"Loading diffusers model: {model_name}")
pipeline = _load_diffusers_pipeline(model_name, global_args)
_diff_cfg = (multi_model_manager.config.get(model_key)
or multi_model_manager.config.get(model_name) or {})
_vram_before = multi_model_manager.vram_before_load()
pipeline = await asyncio.to_thread(
_load_diffusers_pipeline, model_name, global_args, _diff_cfg)
if pipeline is not None:
# Cache the loaded pipeline in the manager
multi_model_manager.add_model(model_key, pipeline)
multi_model_manager.current_model_key = model_key
try:
multi_model_manager.record_vram_delta(model_key, _vram_before)
except Exception:
pass
print(f"Loaded diffusers model: {model_name}")
return await _generate_with_diffusers(pipeline, request, global_args, http_request)
except ImportError as e:
......@@ -1054,19 +1298,24 @@ async def create_image_generation(request: ImageGenerationRequest, http_request:
# For GGUF models or URLs, resolve the model path through the cache
resolved_path = model_name
if is_gguf or model_name.startswith('http://') or model_name.startswith('https://'):
resolved_path = multi_model_manager.load_model(model_name)
resolved_path = await asyncio.to_thread(multi_model_manager.load_model, model_name)
if not resolved_path:
raise Exception(f"Failed to resolve model path: {model_name}")
# Only use sd.cpp if we have a local file path
if resolved_path and os.path.isfile(resolved_path):
cfg = multi_model_manager.config.get(model_key) or multi_model_manager.config.get(model_name) or {}
sd_model = _load_sdcpp_model(resolved_path, global_args, model_config=cfg)
_vram_before = multi_model_manager.vram_before_load()
sd_model = await asyncio.to_thread(_load_sdcpp_model, resolved_path, global_args, model_config=cfg)
if sd_model is not None:
# Cache the loaded model in the manager
multi_model_manager.add_model(model_key, sd_model)
multi_model_manager.current_model_key = model_key
try:
multi_model_manager.record_vram_delta(model_key, _vram_before)
except Exception:
pass
print(f"Loaded sd.cpp model: {model_name}")
return await _generate_with_sdcpp(sd_model, request, global_args, http_request)
......@@ -1155,7 +1404,8 @@ def _load_img2img_pipeline(model_name: str, global_args):
for attempt in range(3):
try:
pipe = PipeClass.from_pretrained(model_name, torch_dtype=torch_dtype)
with _DIFFUSERS_LOAD_LOCK:
pipe = PipeClass.from_pretrained(model_name, torch_dtype=torch_dtype)
pipe = pipe.to(device)
if attempt >= 1:
pipe.enable_attention_slicing()
......@@ -1189,7 +1439,8 @@ async def create_image_edit(request: ImageEditRequest, http_request: Request = N
raise HTTPException(status_code=400, detail="image is required")
_progress_loading(request.model or "image")
model_info = multi_model_manager.request_model(request.model, model_type="image")
model_info = await asyncio.to_thread(
multi_model_manager.request_model, request.model, model_type="image")
model_name = model_info.get('model_name')
if not model_name:
err = model_info.get('error', f"Model '{request.model}' not found or not registered")
......@@ -1294,7 +1545,8 @@ def _load_inpaint_pipeline(model_name: str, global_args):
PClass = DiffusionPipeline
for attempt in range(3):
try:
pipe = PClass.from_pretrained(model_name, torch_dtype=torch_dtype)
with _DIFFUSERS_LOAD_LOCK:
pipe = PClass.from_pretrained(model_name, torch_dtype=torch_dtype)
pipe = pipe.to(device)
if attempt >= 1:
pipe.enable_attention_slicing()
......@@ -1323,7 +1575,8 @@ async def create_image_inpaint(request: ImageInpaintRequest, http_request: Reque
if not request.image or not request.mask:
raise HTTPException(status_code=400, detail="image and mask are required")
_progress_loading(request.model or "image")
model_info = multi_model_manager.request_model(request.model, model_type="image")
model_info = await asyncio.to_thread(
multi_model_manager.request_model, request.model, model_type="image")
model_name = model_info.get('model_name')
if not model_name:
raise HTTPException(status_code=404, detail=model_info.get('error', 'Model not found'))
......@@ -1432,7 +1685,8 @@ async def create_image_upscale(request: ImageUpscaleRequest, http_request: Reque
"""Upscale an image using Real-ESRGAN or PIL LANCZOS fallback."""
global global_args
_progress_loading(request.model or "image")
model_info = multi_model_manager.request_model(request.model, model_type="image")
model_info = await asyncio.to_thread(
multi_model_manager.request_model, request.model, model_type="image")
model_name = model_info.get('model_name') or request.model
model_key = f"upscale:{model_name}"
upscaler = multi_model_manager.models.get(model_key)
......@@ -1465,11 +1719,17 @@ class ImageDepthRequest(BaseModel):
extra = "allow"
def _load_depth_model(model_name: str, global_args):
def _load_depth_model(model_name: str, global_args, model_config: dict = None):
device = _derive_diffusers_device(global_args)
try:
from transformers import pipeline as hf_pipeline
pipe = hf_pipeline("depth-estimation", model=model_name, device=device)
from codai.models.hf_loading import pipeline_device_kwargs
pk = pipeline_device_kwargs(model_config)
# device and device_map are mutually exclusive in HF pipeline.
if 'device_map' in pk:
pipe = hf_pipeline("depth-estimation", model=model_name, **pk)
else:
pipe = hf_pipeline("depth-estimation", model=model_name, device=device, **pk)
return ('transformers', pipe)
except Exception:
pass
......@@ -1542,9 +1802,11 @@ async def create_image_depth(request: ImageDepthRequest, http_request: Request =
model_key = f"depth:{model_name}"
depth_model = multi_model_manager.models.get(model_key)
if depth_model is None:
_sp_cfg = (multi_model_manager.config.get(f"spatial:{model_name}")
or multi_model_manager.config.get(model_name) or {})
try:
depth_model = await asyncio.get_event_loop().run_in_executor(
None, _load_depth_model, model_name, global_args)
None, _load_depth_model, model_name, global_args, _sp_cfg)
except Exception as e:
raise HTTPException(status_code=500, detail=f"Failed to load depth model: {e}")
multi_model_manager.models[model_key] = depth_model
......@@ -1572,19 +1834,28 @@ class ImageSegmentRequest(BaseModel):
extra = "allow"
def _load_segmentation_model(model_name: str, global_args):
def _load_segmentation_model(model_name: str, global_args, model_config: dict = None):
device = _derive_diffusers_device(global_args)
from codai.models.hf_loading import build_from_pretrained_kwargs, pipeline_device_kwargs
try:
from transformers import SamModel, SamProcessor
import torch
model = SamModel.from_pretrained(model_name).to(device)
fp = build_from_pretrained_kwargs(model_config)
model = SamModel.from_pretrained(model_name, **fp)
# Quantized/offloaded models are already placed; only plain models move.
if 'quantization_config' not in fp and 'device_map' not in fp:
model = model.to(device)
processor = SamProcessor.from_pretrained(model_name)
return ('sam', (model, processor, device))
except Exception:
pass
try:
from transformers import pipeline as hf_pipeline
pipe = hf_pipeline("image-segmentation", model=model_name, device=device)
pk = pipeline_device_kwargs(model_config)
if 'device_map' in pk:
pipe = hf_pipeline("image-segmentation", model=model_name, **pk)
else:
pipe = hf_pipeline("image-segmentation", model=model_name, device=device, **pk)
return ('transformers', pipe)
except Exception as e:
raise RuntimeError(f"Cannot load segmentation model: {e}")
......@@ -1637,9 +1908,11 @@ async def create_image_segment(request: ImageSegmentRequest, http_request: Reque
model_key = f"segment:{model_name}"
seg_model = multi_model_manager.models.get(model_key)
if seg_model is None:
_sp_cfg = (multi_model_manager.config.get(f"spatial:{model_name}")
or multi_model_manager.config.get(model_name) or {})
try:
seg_model = await asyncio.get_event_loop().run_in_executor(
None, _load_segmentation_model, model_name, global_args)
None, _load_segmentation_model, model_name, global_args, _sp_cfg)
except Exception as e:
raise HTTPException(status_code=500, detail=f"Failed to load segmentation model: {e}")
multi_model_manager.models[model_key] = seg_model
......
codai/api/loras.py 0 → 100644
View file @ 278fab3e
# CoderAI - OpenAI-compatible API server
# Copyright (C) 2026 Stefy Lanza <stefy@nexlab.net>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program. If not, see <https://www.gnu.org/licenses/>.
"""
LoRA training endpoints.
Train a small per-character (or per-style) LoRA adapter from a handful of
reference images, then apply it to BOTH image and video diffusion pipelines for
consistent identity across models.
POST /v1/loras/train – train a LoRA from a saved character profile or images
GET /v1/loras – list trained LoRAs
GET /v1/loras/progress – training progress (for the active job)
GET /v1/loras/{name} – info about one trained LoRA
DELETE /v1/loras/{name} – delete a trained LoRA
All model execution stays server-side. Training runs in-process so it can share
the model manager's VRAM (it evicts resident models first) and honour the global
thermal-protection checkpoints.
"""
import base64
import io
import json
import os
import threading
import time
from typing import List, Optional
from fastapi import APIRouter, Depends, HTTPException, Request
from pydantic import BaseModel, ConfigDict
from codai.platform_paths import default_loras_dir
router = APIRouter()
_LORAS_DIR: Optional[str] = None
# Single-job training progress (training is VRAM-heavy; we run one at a time).
_progress_lock = threading.Lock()
_progress = {
"active": False,
"name": None,
"step": 0,
"total": 0,
"status": "idle", # idle | preparing | training | saving | done | error
"message": "",
"started_at": 0.0,
"path": None,
}
_train_lock = threading.Lock()
def set_global_args(args):
global _LORAS_DIR
base = getattr(args, 'file_path', None)
if base and os.path.isdir(base):
root = base
elif base:
root = os.path.dirname(base)
else:
root = None
_LORAS_DIR = os.path.join(root, 'loras') if root else str(default_loras_dir())
os.makedirs(_LORAS_DIR, exist_ok=True)
def _loras_dir() -> str:
if _LORAS_DIR:
os.makedirs(_LORAS_DIR, exist_ok=True)
return _LORAS_DIR
d = str(default_loras_dir())
os.makedirs(d, exist_ok=True)
return d
def _lora_dir(name: str) -> str:
return os.path.join(_loras_dir(), name)
def _lora_weight_file(name: str) -> Optional[str]:
"""Return the path to the trained weights file for a LoRA, if present."""
d = _lora_dir(name)
for fn in ("pytorch_lora_weights.safetensors", "pytorch_lora_weights.bin"):
p = os.path.join(d, fn)
if os.path.isfile(p):
return p
return None
def _require_api_auth(request: Request) -> None:
"""Raise 401 if auth is enabled and the request carries no valid credential."""
try:
from codai.admin import routes as _admin_routes
sm = _admin_routes.session_manager
except Exception:
return
if sm is None:
return
auth = request.headers.get("authorization", "")
if auth.lower().startswith("bearer "):
token = auth[7:].strip()
if sm.verify_token(token):
return
cookie = request.cookies.get("session", "")
if cookie.endswith(".MUST_CHANGE"):
cookie = cookie[:-12]
if cookie and sm.validate_session(cookie):
return
raise HTTPException(
status_code=401,
detail={"message": "Invalid API key. Provide a valid Bearer token.",
"type": "invalid_request_error", "code": "invalid_api_key"},
)
# ── Pydantic models ───────────────────────────────────────────────────────────
class LoraTrainRequest(BaseModel):
name: str # output LoRA name (folder)
base_model: str # image model key (models.json) or HF id / path
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
instance_prompt: Optional[str] = None # e.g. "a photo of sks man"; auto from name if None
steps: Optional[int] = 800 # training steps (balanced default)
rank: Optional[int] = 16 # LoRA rank
learning_rate: Optional[float] = 1e-4
resolution: Optional[int] = 512
seed: Optional[int] = 42
model_config = ConfigDict(extra="allow")
# ── Base-model resolution ─────────────────────────────────────────────────────
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
for key in (f"image:{base_model}", base_model):
cfg = multi_model_manager.config.get(key)
if cfg:
for k in ('path', 'model_path', 'model', 'diffusers_path'):
v = cfg.get(k)
if v and isinstance(v, str):
return v
except Exception:
pass
# Treat as a direct path or HF repo id.
return base_model
def _decode_image(ref: str):
from PIL import Image as PILImage
if ref.startswith('data:'):
ref = ref.split(',', 1)[1]
raw = base64.b64decode(ref)
return PILImage.open(io.BytesIO(raw)).convert('RGB')
def _gather_images(req: LoraTrainRequest):
"""Return a list of PIL images from the character profile and/or inline images."""
imgs = []
if req.character:
try:
from codai.api.characters import resolve_character_profiles
for b64 in resolve_character_profiles([req.character]):
try:
imgs.append(_decode_image(b64))
except Exception:
pass
except Exception:
pass
if req.environment:
try:
from codai.api.environments import resolve_environment_profiles
for b64 in resolve_environment_profiles([req.environment]):
try:
imgs.append(_decode_image(b64))
except Exception:
pass
except Exception:
pass
for ref in (req.images or []):
try:
imgs.append(_decode_image(ref))
except Exception:
pass
return imgs
def _set_progress(**kw):
with _progress_lock:
_progress.update(kw)
# ── Training core ─────────────────────────────────────────────────────────────
def _train_lora_sync(req: LoraTrainRequest) -> dict:
"""Run a DreamBooth-style LoRA training in-process. Returns {name, path}."""
import torch
import torch.nn.functional as F
from diffusers import (
AutoencoderKL, DDPMScheduler, UNet2DConditionModel,
)
from diffusers.optimization import get_scheduler
from diffusers.utils import convert_state_dict_to_diffusers
from peft import LoraConfig as PeftLoraConfig
from peft.utils import get_peft_model_state_dict
from transformers import CLIPTextModel, CLIPTokenizer
name = req.name
base_path = _resolve_base_model_path(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)
lr = float(req.learning_rate or 1e-4)
seed = int(req.seed if req.seed is not None else 42)
_set_progress(active=True, name=name, step=0, total=steps,
status="preparing", message="loading reference images",
started_at=time.time(), path=None)
images = _gather_images(req)
if not images:
raise HTTPException(status_code=400,
detail="No training images (provide `character` or `images`)")
if req.instance_prompt:
instance_prompt = req.instance_prompt
elif req.environment and not req.character:
instance_prompt = f"a photo of {name} place"
else:
instance_prompt = f"a photo of {name} person"
# Free VRAM: evict every resident model so training has the whole GPU.
try:
from codai.api.state import multi_model_manager
multi_model_manager.unload_all_models()
except Exception as e:
print(f" [lora] could not unload models before training: {e}")
device = "cuda" if torch.cuda.is_available() else "cpu"
weight_dtype = torch.float32 # train in fp32 for stability
_set_progress(status="preparing", message=f"loading base model: {base_path}")
# Detect SDXL by attempting to load a second tokenizer.
is_sdxl = False
try:
from transformers import CLIPTokenizer as _CT
_CT.from_pretrained(base_path, subfolder="tokenizer_2")
is_sdxl = True
except Exception:
is_sdxl = False
if is_sdxl:
result = _train_sdxl(req, base_path, images, instance_prompt,
steps, rank, resolution, lr, seed, device)
else:
result = _train_sd15(req, base_path, images, instance_prompt,
steps, rank, resolution, lr, seed, device)
return result
def _make_dataset(images, tokenizers, text_encoders, instance_prompt,
resolution, vae, device, weight_dtype, is_sdxl):
"""Pre-encode latents + text embeddings for every reference image once."""
import torch
from torchvision import transforms
tfm = transforms.Compose([
transforms.Resize(resolution, interpolation=transforms.InterpolationMode.BILINEAR),
transforms.CenterCrop(resolution),
transforms.ToTensor(),
transforms.Normalize([0.5], [0.5]),
])
latents_list = []
with torch.no_grad():
for img in images:
px = tfm(img).unsqueeze(0).to(device, dtype=vae.dtype)
lat = vae.encode(px).latent_dist.sample() * vae.config.scaling_factor
latents_list.append(lat.to(weight_dtype).cpu())
return latents_list
def _train_sd15(req, base_path, images, instance_prompt,
steps, rank, resolution, lr, seed, device):
import torch
import torch.nn.functional as F
from diffusers import AutoencoderKL, DDPMScheduler, UNet2DConditionModel
from diffusers.utils import convert_state_dict_to_diffusers
from diffusers import StableDiffusionPipeline
from peft import LoraConfig as PeftLoraConfig
from peft.utils import get_peft_model_state_dict
from transformers import CLIPTextModel, CLIPTokenizer
name = req.name
g = torch.Generator(device=device).manual_seed(seed)
tokenizer = CLIPTokenizer.from_pretrained(base_path, subfolder="tokenizer")
text_encoder = CLIPTextModel.from_pretrained(base_path, subfolder="text_encoder").to(device)
vae = AutoencoderKL.from_pretrained(base_path, subfolder="vae").to(device)
unet = UNet2DConditionModel.from_pretrained(base_path, subfolder="unet").to(device)
noise_scheduler = DDPMScheduler.from_pretrained(base_path, subfolder="scheduler")
vae.requires_grad_(False)
text_encoder.requires_grad_(False)
unet.requires_grad_(False)
lora_cfg = PeftLoraConfig(
r=rank, lora_alpha=rank, init_lora_weights="gaussian",
target_modules=["to_k", "to_q", "to_v", "to_out.0"],
)
unet.add_adapter(lora_cfg)
lora_params = [p for p in unet.parameters() if p.requires_grad]
optimizer = torch.optim.AdamW(lora_params, lr=lr)
# Pre-encode latents and the (single) instance-prompt embedding.
latents_list = _make_dataset(images, [tokenizer], [text_encoder], instance_prompt,
resolution, vae, device, torch.float32, is_sdxl=False)
with torch.no_grad():
tok = tokenizer(instance_prompt, padding="max_length",
max_length=tokenizer.model_max_length, truncation=True,
return_tensors="pt").input_ids.to(device)
encoder_hidden_states = text_encoder(tok)[0]
_set_progress(status="training", message="training (SD1.5)")
unet.train()
n = len(latents_list)
for step in range(steps):
latents = latents_list[step % n].to(device)
noise = torch.randn_like(latents)
bsz = latents.shape[0]
timesteps = torch.randint(0, noise_scheduler.config.num_train_timesteps,
(bsz,), device=device).long()
noisy = noise_scheduler.add_noise(latents, noise, timesteps)
model_pred = unet(noisy, timesteps, encoder_hidden_states).sample
if noise_scheduler.config.prediction_type == "v_prediction":
target = noise_scheduler.get_velocity(latents, noise, timesteps)
else:
target = noise
loss = F.mse_loss(model_pred.float(), target.float(), reduction="mean")
loss.backward()
torch.nn.utils.clip_grad_norm_(lora_params, 1.0)
optimizer.step()
optimizer.zero_grad()
if step % 10 == 0 or step == steps - 1:
_set_progress(step=step + 1, message=f"step {step+1}/{steps} loss={loss.item():.4f}")
# Mid-training thermal checkpoint (pauses if CPU/GPU too hot).
try:
from codai.models.thermal import checkpoint as _thermal_checkpoint
_thermal_checkpoint(context="lora-train", throttle_seconds=2.0)
except Exception:
pass
_set_progress(status="saving", message="saving LoRA weights")
save_dir = _lora_dir(name)
os.makedirs(save_dir, exist_ok=True)
unet_lora = convert_state_dict_to_diffusers(get_peft_model_state_dict(unet))
StableDiffusionPipeline.save_lora_weights(save_directory=save_dir,
unet_lora_layers=unet_lora,
safe_serialization=True)
_write_meta(name, req, base_path, len(images), "sd15", instance_prompt)
# Release training tensors.
del unet, vae, text_encoder, optimizer, latents_list
try:
torch.cuda.empty_cache()
except Exception:
pass
path = _lora_weight_file(name) or save_dir
_set_progress(active=False, status="done", message="done", path=path)
return {"name": name, "path": path}
def _train_sdxl(req, base_path, images, instance_prompt,
steps, rank, resolution, lr, seed, device):
import torch
import torch.nn.functional as F
from diffusers import AutoencoderKL, DDPMScheduler, UNet2DConditionModel
from diffusers import StableDiffusionXLPipeline
from diffusers.utils import convert_state_dict_to_diffusers
from peft import LoraConfig as PeftLoraConfig
from peft.utils import get_peft_model_state_dict
from transformers import CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer
from torchvision import transforms
name = req.name
g = torch.Generator(device=device).manual_seed(seed)
tokenizer_1 = CLIPTokenizer.from_pretrained(base_path, subfolder="tokenizer")
tokenizer_2 = CLIPTokenizer.from_pretrained(base_path, subfolder="tokenizer_2")
text_encoder_1 = CLIPTextModel.from_pretrained(base_path, subfolder="text_encoder").to(device)
text_encoder_2 = CLIPTextModelWithProjection.from_pretrained(base_path, subfolder="text_encoder_2").to(device)
vae = AutoencoderKL.from_pretrained(base_path, subfolder="vae").to(device)
unet = UNet2DConditionModel.from_pretrained(base_path, subfolder="unet").to(device)
noise_scheduler = DDPMScheduler.from_pretrained(base_path, subfolder="scheduler")
for m in (vae, text_encoder_1, text_encoder_2, unet):
m.requires_grad_(False)
lora_cfg = PeftLoraConfig(
r=rank, lora_alpha=rank, init_lora_weights="gaussian",
target_modules=["to_k", "to_q", "to_v", "to_out.0"],
)
unet.add_adapter(lora_cfg)
lora_params = [p for p in unet.parameters() if p.requires_grad]
optimizer = torch.optim.AdamW(lora_params, lr=lr)
tfm = transforms.Compose([
transforms.Resize(resolution, interpolation=transforms.InterpolationMode.BILINEAR),
transforms.CenterCrop(resolution),
transforms.ToTensor(),
transforms.Normalize([0.5], [0.5]),
])
# Pre-encode latents.
latents_list = []
with torch.no_grad():
for img in images:
px = tfm(img).unsqueeze(0).to(device, dtype=vae.dtype)
lat = vae.encode(px).latent_dist.sample() * vae.config.scaling_factor
latents_list.append(lat.float().cpu())
# SDXL text conditioning: concat hidden states + pooled embeds from encoder 2.
with torch.no_grad():
ids_1 = tokenizer_1(instance_prompt, padding="max_length",
max_length=tokenizer_1.model_max_length, truncation=True,
return_tensors="pt").input_ids.to(device)
ids_2 = tokenizer_2(instance_prompt, padding="max_length",
max_length=tokenizer_2.model_max_length, truncation=True,
return_tensors="pt").input_ids.to(device)
enc1 = text_encoder_1(ids_1, output_hidden_states=True)
enc2 = text_encoder_2(ids_2, output_hidden_states=True)
# penultimate hidden states
prompt_embeds = torch.cat([enc1.hidden_states[-2], enc2.hidden_states[-2]], dim=-1)
pooled = enc2[0] # text_embeds (pooled) from projection encoder
add_time_ids = torch.tensor(
[[resolution, resolution, 0, 0, resolution, resolution]],
device=device, dtype=prompt_embeds.dtype,
)
_set_progress(status="training", message="training (SDXL)")
unet.train()
n = len(latents_list)
for step in range(steps):
latents = latents_list[step % n].to(device)
noise = torch.randn_like(latents)
bsz = latents.shape[0]
timesteps = torch.randint(0, noise_scheduler.config.num_train_timesteps,
(bsz,), device=device).long()
noisy = noise_scheduler.add_noise(latents, noise, timesteps)
added = {"text_embeds": pooled, "time_ids": add_time_ids}
model_pred = unet(noisy, timesteps, prompt_embeds, added_cond_kwargs=added).sample
if noise_scheduler.config.prediction_type == "v_prediction":
target = noise_scheduler.get_velocity(latents, noise, timesteps)
else:
target = noise
loss = F.mse_loss(model_pred.float(), target.float(), reduction="mean")
loss.backward()
torch.nn.utils.clip_grad_norm_(lora_params, 1.0)
optimizer.step()
optimizer.zero_grad()
if step % 10 == 0 or step == steps - 1:
_set_progress(step=step + 1, message=f"step {step+1}/{steps} loss={loss.item():.4f}")
try:
from codai.models.thermal import checkpoint as _thermal_checkpoint
_thermal_checkpoint(context="lora-train", throttle_seconds=2.0)
except Exception:
pass
_set_progress(status="saving", message="saving LoRA weights")
save_dir = _lora_dir(name)
os.makedirs(save_dir, exist_ok=True)
unet_lora = convert_state_dict_to_diffusers(get_peft_model_state_dict(unet))
StableDiffusionXLPipeline.save_lora_weights(save_directory=save_dir,
unet_lora_layers=unet_lora,
safe_serialization=True)
_write_meta(name, req, base_path, len(images), "sdxl", instance_prompt)
del unet, vae, text_encoder_1, text_encoder_2, optimizer, latents_list
try:
torch.cuda.empty_cache()
except Exception:
pass
path = _lora_weight_file(name) or save_dir
_set_progress(active=False, status="done", message="done", path=path)
return {"name": name, "path": path}
def _write_meta(name, req, base_path, n_images, arch, instance_prompt):
meta = {
"name": name,
"base_model": req.base_model,
"base_path": base_path,
"arch": arch,
"instance_prompt": instance_prompt,
"steps": req.steps,
"rank": req.rank,
"resolution": req.resolution,
"num_images": n_images,
"created_at": int(time.time()),
}
try:
with open(os.path.join(_lora_dir(name), "meta.json"), "w") as f:
json.dump(meta, f, indent=2)
except Exception:
pass
# ── Endpoints ─────────────────────────────────────────────────────────────────
@router.post("/v1/loras/train")
async def train_lora(req: LoraTrainRequest, _auth=Depends(_require_api_auth)):
"""Train a LoRA from a saved character profile or supplied images (blocking)."""
import asyncio
if not req.name or '/' in req.name or '..' in req.name:
raise HTTPException(status_code=400, detail="Invalid LoRA name")
if not req.base_model:
raise HTTPException(status_code=400, detail="base_model is required")
if not _train_lock.acquire(blocking=False):
raise HTTPException(status_code=409, detail="A LoRA training job is already running")
try:
try:
result = await asyncio.to_thread(_train_lora_sync, req)
except HTTPException:
raise
except Exception as e:
import traceback
traceback.print_exc()
_set_progress(active=False, status="error", message=str(e))
raise HTTPException(status_code=500, detail=f"LoRA training failed: {e}")
finally:
_train_lock.release()
return {"ok": True, **result}
@router.get("/v1/loras/progress")
async def lora_progress():
with _progress_lock:
return dict(_progress)
@router.get("/v1/loras")
async def list_loras(_auth=Depends(_require_api_auth)):
out = []
d = _loras_dir()
if os.path.isdir(d):
for name in sorted(os.listdir(d)):
wf = _lora_weight_file(name)
if not wf:
continue
meta = {}
mp = os.path.join(_lora_dir(name), "meta.json")
if os.path.isfile(mp):
try:
with open(mp) as f:
meta = json.load(f)
except Exception:
pass
out.append({"name": name, "path": wf, **meta})
return {"loras": out}
@router.get("/v1/loras/{name}")
async def get_lora(name: str, _auth=Depends(_require_api_auth)):
wf = _lora_weight_file(name)
if not wf:
raise HTTPException(status_code=404, detail=f"LoRA '{name}' not found")
meta = {}
mp = os.path.join(_lora_dir(name), "meta.json")
if os.path.isfile(mp):
try:
with open(mp) as f:
meta = json.load(f)
except Exception:
pass
return {"name": name, "path": wf, **meta}
@router.delete("/v1/loras/{name}")
async def delete_lora(name: str, _auth=Depends(_require_api_auth)):
d = _lora_dir(name)
if not os.path.isdir(d):
raise HTTPException(status_code=404, detail=f"LoRA '{name}' not found")
import shutil
shutil.rmtree(d)
return {"ok": True, "name": name}
codai/pydantic/imagerequest.py
View file @ 278fab3e
......@@ -21,6 +21,13 @@ from typing import Dict, List, Optional
from pydantic import BaseModel, ConfigDict
class LoraConfig(BaseModel):
model: str
weight: float = 1.0
name: Optional[str] = None
model_config = ConfigDict(extra="allow")
class ImageGenerationRequest(BaseModel):
model: str
prompt: str
......@@ -36,10 +43,15 @@ class ImageGenerationRequest(BaseModel):
disable_safety_checker: Optional[bool] = False
negative_prompt: Optional[str] = None
# Per-request component overrides
vae_model: Optional[str] = None # Override the VAE for this request
loras: Optional[List[LoraConfig]] = None # Additional LoRA weights for this request
# Character consistency
character_profiles: Optional[List[str]] = None # saved profile names
character_references: Optional[List[str]] = None # inline base64 images
character_strength: Optional[float] = 0.6 # IP-Adapter scale
environment_profiles: Optional[List[str]] = None # saved environment profile names (IP-Adapter)
model_config = ConfigDict(extra="allow")
......@@ -47,4 +59,4 @@ class ImageGenerationRequest(BaseModel):
class ImageGenerationResponse(BaseModel):
created: int
data: List[Dict]
model_config = ConfigDict(extra="allow")
\ No newline at end of file
model_config = ConfigDict(extra="allow")
tools/gen_township_fighters.py
View file @ 278fab3e
......@@ -503,6 +503,40 @@ class CoderAIClient:
r.raise_for_status()
return r.json()
def _delete(self, path: str) -> dict:
r = self.session.delete(f"{self.base}{path}", timeout=30)
r.raise_for_status()
return r.json()
def _patch(self, path: str, body: dict) -> dict:
r = self.session.patch(f"{self.base}{path}", json=body, timeout=60)
if not r.ok:
raise RuntimeError(f"PATCH {path} → {r.status_code}: {r.text[:400]}")
return r.json()
def delete_profile(self, kind: str, name: str) -> dict:
plural = "characters" if kind == "character" else "environments"
return self._delete(f"/v1/{plural}/{name}")
def patch_profile(self, kind: str, name: str, description: str = None,
remove_indices: list = None, add_images: list = None) -> dict:
plural = "characters" if kind == "character" else "environments"
body = {}
if description is not None:
body["description"] = description
if remove_indices:
body["remove_indices"] = remove_indices
if add_images:
# Each entry may be a data-uri/base64 str or a {data,label} dict.
imgs = []
for j, im in enumerate(add_images):
if isinstance(im, dict):
imgs.append(im)
else:
imgs.append({"data": im, "label": f"regen_{j:02d}"})
body["add_images"] = imgs
return self._patch(f"/v1/{plural}/{name}", body)
def list_models(self) -> list:
return self._get("/v1/models").get("data", [])
......@@ -561,7 +595,7 @@ class CoderAIClient:
character_profiles: list = None,
loras: list = None, character_strength: float = 0.7,
size: str = "512x512", steps: int = 28,
seed: int = None) -> bytes:
seed: int = None, environment_profiles: list = None) -> bytes:
"""Generate a single still image (used for keyframes). Returns PNG bytes."""
w, h = size.split("x")
body = {
......@@ -572,6 +606,9 @@ class CoderAIClient:
if character_profiles:
body["character_profiles"] = list(character_profiles)
body["character_strength"] = character_strength
if environment_profiles:
body["environment_profiles"] = list(environment_profiles)
body["character_strength"] = character_strength
if loras:
body["loras"] = loras
if seed is not None:
......@@ -588,13 +625,17 @@ class CoderAIClient:
return base64.b64decode(raw)
def train_lora(self, name: str, base_model: str, character: str = None,
images: list = None, steps: int = 800, rank: int = 16,
environment: str = None, images: list = None,
steps: int = 800, rank: int = 16,
resolution: int = 512) -> dict:
"""Train a per-character LoRA on the server. Blocks until complete."""
"""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 character:
body["character"] = character
if environment:
body["environment"] = environment
if images:
body["images"] = images
return self._post("/v1/loras/train", body)
......@@ -1108,6 +1149,7 @@ CONFIG_FIELDS = [
"only_prompts", "only_videos",
"consistency", "keyframe_steps", "keyframe_size",
"character_strength", "lora_steps", "lora_rank", "lora_weight",
"no_env_loras", "env_lora_steps", "env_lora_rank", "env_lora_weight",
"web_port",
]
......@@ -1190,18 +1232,38 @@ def _lora_specs_for(fighters: list, lora_map: dict, weight: float) -> list:
return specs
def stage_loras(client: CoderAIClient, image_model: str, out_dir: Path,
char_names: list, lora_steps: int = 800, lora_rank: int = 16) -> dict:
"""Train one identity LoRA per fighter (server-side). Returns {fighter: lora_path}.
def _env_lora_specs_for(env: str, env_lora_map: dict, weight: float) -> list:
"""Build the `loras` request entry for the environment used in a clip."""
if not env:
return []
path = (env_lora_map or {}).get(env)
if path:
return [{"model": path, "weight": float(weight), "name": f"env_{env}"}]
return []
Resumable: skips fighters whose LoRA already exists locally (loras.json) or on
the server. All training is grouped here so the image base model is touched
once for the whole batch.
# Per-kind LoRA training parameters: server name prefix, local cache file,
# the train_lora keyword used to pull reference images, and a friendly label.
_LORA_KINDS = {
"character": {"prefix": "fighter_", "file": "loras.json", "label": "Character"},
"environment": {"prefix": "env_", "file": "env_loras.json", "label": "Environment"},
}
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:
"""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.
"""
spec = _LORA_KINDS[kind]
_log("\n" + "═" * 60)
_log(" STAGE — Character LoRA training")
_log(f" STAGE — {spec['label']} LoRA training")
_log("═" * 60)
lora_file = out_dir / "loras.json"
lora_file = out_dir / spec["file"]
lora_map = {}
if lora_file.exists():
try:
......@@ -1218,28 +1280,29 @@ def stage_loras(client: CoderAIClient, image_model: str, out_dir: Path,
try:
lora_file.write_text(json.dumps(lora_map, indent=2))
except Exception as e:
_log(f" ⚠ could not save loras.json: {e}")
_log(f" ⚠ could not save {spec['file']}: {e}")
for i, name in enumerate(char_names, 1):
lora_name = f"fighter_{name}"
for i, name in enumerate(names, 1):
lora_name = f"{spec['prefix']}{name}"
# Already trained and recorded?
cur = lora_map.get(name)
if cur and Path(cur).exists():
_log(f" [{i}/{len(char_names)}] {name}: reusing trained LoRA")
_log(f" [{i}/{len(names)}] {name}: reusing trained LoRA")
continue
if lora_name in existing and existing[lora_name]:
lora_map[name] = existing[lora_name]
_save()
_log(f" [{i}/{len(char_names)}] {name}: found existing LoRA on server")
_log(f" [{i}/{len(names)}] {name}: found existing LoRA on server")
continue
_log(f" [{i}/{len(char_names)}] {name}: training LoRA "
_log(f" [{i}/{len(names)}] {name}: training LoRA "
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)
train_kwargs[kind] = name # character=name OR environment=name
try:
res = _run_with_spinner(
f"training LoRA '{name}'",
client.train_lora,
name=lora_name, base_model=image_model, character=name,
steps=lora_steps, rank=lora_rank,
f"training {kind} LoRA '{name}'",
client.train_lora, **train_kwargs,
)
path = res.get("path")
if path:
......@@ -1251,14 +1314,29 @@ def stage_loras(client: CoderAIClient, image_model: str, out_dir: Path,
except Exception as e:
_log(f" ✗ LoRA training failed for {name}: {e}")
_log(f"\n LoRAs ready: {len(lora_map)}/{len(char_names)}")
_log(f"\n {spec['label']} LoRAs ready: {len(lora_map)}/{len(names)}")
return lora_map
def stage_loras(client: CoderAIClient, image_model: str, out_dir: Path,
char_names: list, lora_steps: int = 800, lora_rank: int = 16) -> 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)
def stage_env_loras(client: CoderAIClient, image_model: str, out_dir: Path,
env_names: list, lora_steps: int = 800, lora_rank: int = 16) -> 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)
def _generate_keyframes(client: CoderAIClient, image_model: str, keyframe_dir: Path,
fight_plan: list, outcome_plan: list, consistency: set,
lora_map: dict, char_strength: float, keyframe_steps: int,
keyframe_size: str, lora_weight: float):
keyframe_size: str, lora_weight: float,
env_lora_map: dict = None, env_lora_weight: float = 0.8):
"""Generate one keyframe still per clip (image model). Saved as PNG keyed by
the clip's output stem so the render phase can pick them up as init images.
Resumable: existing PNGs are kept."""
......@@ -1284,7 +1362,10 @@ def _generate_keyframes(client: CoderAIClient, image_model: str, keyframe_dir: P
skipped += 1
continue
profiles = list(fighters) if use_ip else None
loras = _lora_specs_for(fighters, lora_map, lora_weight) if use_lora else None
loras = None
if use_lora:
loras = (_lora_specs_for(fighters, lora_map, lora_weight)
+ _env_lora_specs_for(env, env_lora_map, env_lora_weight)) or None
kf_prompt = prompt
if env:
kf_prompt = f"[{env} location] " + kf_prompt
......@@ -1315,7 +1396,8 @@ def stage_videos(client: CoderAIClient, video_model: str, out_dir: Path,
consistency: set = None, image_model: str = None,
lora_map: dict = None, char_strength: float = 0.7,
keyframe_steps: int = 28, keyframe_size: str = "512x512",
lora_weight: float = 0.85, keyframes_only: bool = False):
lora_weight: float = 0.85, keyframes_only: bool = False,
env_lora_map: dict = None, env_lora_weight: float = 0.8):
_log("\n" + "═" * 60)
_log(" STAGE 3 — Videos")
_log("═" * 60)
......@@ -1338,7 +1420,8 @@ def stage_videos(client: CoderAIClient, video_model: str, out_dir: Path,
_generate_keyframes(client, image_model, keyframe_dir,
saved.get("fight_plan", []), saved.get("outcome_plan", []),
consistency or {"prompt", "keyframe"}, lora_map or {},
char_strength, keyframe_steps, keyframe_size, lora_weight)
char_strength, keyframe_steps, keyframe_size, lora_weight,
env_lora_map=env_lora_map or {}, env_lora_weight=env_lora_weight)
return
consistency = consistency or {"prompt"}
......@@ -1377,14 +1460,16 @@ def stage_videos(client: CoderAIClient, video_model: str, out_dir: Path,
if use_keyframe and image_model:
_generate_keyframes(client, image_model, keyframe_dir,
fight_plan, outcome_plan, consistency, lora_map,
char_strength, keyframe_steps, keyframe_size, lora_weight)
char_strength, keyframe_steps, keyframe_size, lora_weight,
env_lora_map=env_lora_map or {}, env_lora_weight=env_lora_weight)
# Jump straight to Phase 3 (rendering) below.
return _stage_videos_render(
client, video_model, video_dir, fight_plan, outcome_plan,
total_matches, total_outcomes, fps, clip_delay,
consistency=consistency, lora_map=lora_map,
keyframe_dir=keyframe_dir if use_keyframe else None,
lora_weight=lora_weight)
lora_weight=lora_weight,
env_lora_map=env_lora_map or {}, env_lora_weight=env_lora_weight)
# =========================================================================
# PHASE 1 — PLAN every clip up front (no API calls).
......@@ -1520,25 +1605,28 @@ def stage_videos(client: CoderAIClient, video_model: str, out_dir: Path,
if use_keyframe and image_model:
_generate_keyframes(client, image_model, keyframe_dir,
fight_plan, outcome_plan, consistency, lora_map,
char_strength, keyframe_steps, keyframe_size, lora_weight)
char_strength, keyframe_steps, keyframe_size, lora_weight,
env_lora_map=env_lora_map or {}, env_lora_weight=env_lora_weight)
return _stage_videos_render(
client, video_model, video_dir, fight_plan, outcome_plan,
total_matches, total_outcomes, fps, clip_delay,
consistency=consistency, lora_map=lora_map,
keyframe_dir=keyframe_dir if use_keyframe else None,
lora_weight=lora_weight)
lora_weight=lora_weight,
env_lora_map=env_lora_map or {}, env_lora_weight=env_lora_weight)
def _stage_videos_render(client, video_model, video_dir, fight_plan, outcome_plan,
total_matches, total_outcomes, fps, clip_delay,
consistency=None, lora_map=None, keyframe_dir=None,
lora_weight=0.85):
lora_weight=0.85, env_lora_map=None, env_lora_weight=0.8):
"""PHASE 3 — render ALL videos from pre-written prompts (video model stays loaded)."""
_log("\n ── Phase B — rendering all videos (video model) ──")
render_start = time.monotonic()
consistency = consistency or {"prompt"}
lora_map = lora_map or {}
env_lora_map = env_lora_map or {}
use_lora = "lora" in consistency
def _keyframe_bytes(stem: str):
......@@ -1555,7 +1643,10 @@ def _stage_videos_render(client, video_model, video_dir, fight_plan, outcome_pla
def _render(label, prompt, profiles, env, nf, out_path, stem=None, fighters=None):
"""Render one clip; returns (ok, duration_or_None, fatal)."""
init_image = _keyframe_bytes(stem) if stem else None
loras = _lora_specs_for(fighters or profiles or [], lora_map, lora_weight) if use_lora else None
loras = None
if use_lora:
loras = (_lora_specs_for(fighters or profiles or [], lora_map, lora_weight)
+ _env_lora_specs_for(env, env_lora_map, env_lora_weight)) or None
try:
mp4 = _run_with_spinner(
label, client.generate_video_clip,
......@@ -1710,6 +1801,7 @@ def launch_web_ui(default_args):
_state = {
"running": False,
"done": False,
"current": "", # label of the run currently/last executing
"log_lines": [], # all lines so far (for late-joining SSE clients)
"abort": threading.Event(),
"jobs": {}, # job_id -> {status, progress, output, error}
......@@ -1832,6 +1924,108 @@ def launch_web_ui(default_args):
except Exception as exc:
_fail(str(exc))
def _next_ref_path(base: Path, ext: str = ".png") -> Path:
"""Return the next free ref_NN path in a profile folder. The index is
unique across extensions so ref_00.png and ref_00.jpg can't coexist."""
i = 0
while True:
if not any((base / f"ref_{i:02d}{e}").exists()
for e in (".png", ".jpg", ".jpeg", ".webp", ".gif")):
return base / f"ref_{i:02d}{ext}"
i += 1
def _run_regen_job(job_id: str, kind: str, name: str, count: int, guide: bool):
"""Generate `count` NEW reference images for a profile and APPEND them,
preserving every existing (non-deleted) image. Runs server-side image
generation; updates job progress for the profile page to poll."""
with _jobs_lock:
_state["jobs"][job_id] = {"status": "running", "progress": 3,
"output": None, "error": None,
"_msg": "starting…", "added": 0}
def _prog(pct, msg=""):
with _jobs_lock:
_state["jobs"][job_id]["progress"] = pct
if msg:
_state["jobs"][job_id]["_msg"] = msg
def _fail(msg):
with _jobs_lock:
_state["jobs"][job_id].update({"status": "error", "error": msg})
try:
base = out_dir / (kind + "s") / name
meta = {}
try:
meta = json.loads((base / "meta.json").read_text())
except Exception:
pass
# Build a generation prompt from the saved profile.
prompt = (meta.get("prompt") or meta.get("description") or name).strip()
if kind == "environment":
size = "768x512"
else:
size = "512x512"
client = CoderAIClient(default_args.base_url,
getattr(default_args, "api_key", None))
_prog(8, "selecting image model…")
model = getattr(default_args, "image_model", None)
if not model:
try:
model = pick_model(client, "image", None)
except Exception as e:
_fail(f"no image model available: {e}")
return
# Guide new images with the surviving references via IP-Adapter so
# regenerated refs match the ones the user kept. Characters use
# character_profiles; environments use environment_profiles.
char_p = [name] if (guide and kind == "character") else None
env_p = [name] if (guide and kind == "environment") else None
base.mkdir(parents=True, exist_ok=True)
added_uris = []
for k in range(count):
_prog(int(10 + 80 * k / max(1, count)),
f"generating image {k+1}/{count}…")
try:
img = client.generate_image(
prompt=prompt, model=model,
character_profiles=char_p, environment_profiles=env_p,
character_strength=0.7,
size=size, steps=28, seed=random.randint(0, 2**31),
)
except Exception as e:
_web_log(f" ✗ regen image {k+1}/{count} for {name} failed: {e}")
continue
out_png = _next_ref_path(base)
out_png.write_bytes(img)
added_uris.append("data:image/png;base64," +
base64.b64encode(img).decode())
if not added_uris:
_fail("no images were generated")
return
# Append to the CoderAI server profile too (best-effort), so video
# and keyframe generation that resolves this profile sees them.
_prog(94, "syncing new images to CoderAI…")
synced = True
try:
client.patch_profile(kind, name, add_images=added_uris)
except Exception:
synced = False
with _jobs_lock:
_state["jobs"][job_id].update({
"status": "done", "progress": 100,
"added": len(added_uris), "synced": synced,
"_msg": f"added {len(added_uris)} image(s)",
})
except Exception as exc:
_fail(str(exc))
def _web_log(msg: str):
"""Override _log so output goes to both stdout and the web log queue."""
print(msg, flush=True)
......@@ -1886,6 +2080,7 @@ input[type=checkbox]{width:auto;accent-color:#f5a623;margin-right:.3rem}
height:340px;overflow-y:auto;font-family:monospace;font-size:.78rem;
line-height:1.55;white-space:pre-wrap;word-break:break-all}
#log-box .info{color:#9ad89a}#log-box .warn{color:#f5c842}#log-box .err{color:#e07070}
#log-box .head{color:#f5a623;font-weight:700}
.status-pill{display:inline-block;padding:.2rem .55rem;border-radius:10px;
font-size:.72rem;font-weight:700}
.status-idle{background:#333;color:#888}
......@@ -1920,11 +2115,28 @@ input[type=checkbox]{width:auto;accent-color:#f5a623;margin-right:.3rem}
.progress-fill{height:100%;background:#f5a623;border-radius:4px;transition:width .4s}
.job-status{font-size:.78rem;margin-top:.4rem;min-height:1.2rem}
.job-status.done{color:#7ed87e}.job-status.error{color:#e07070}
/* profile editor */
textarea{background:#111;border:1px solid #333;color:#e0e0e0;padding:.35rem .5rem;
border-radius:4px;width:100%;font-size:.85rem;font-family:inherit;
resize:vertical;min-height:3rem}
.pf-head{display:flex;justify-content:space-between;align-items:center;gap:.6rem}
.pf-name{font-weight:700;color:#f5a623;font-size:1.05rem}
.pf-thumbs{display:flex;gap:.4rem;flex-wrap:wrap;margin:.5rem 0}
.pf-thumb{position:relative;width:92px;height:92px}
.pf-thumb img{width:92px;height:92px;object-fit:cover;border-radius:4px;background:#111}
.pf-thumb-del{position:absolute;top:2px;right:2px;background:rgba(192,57,43,.92);color:#fff;
border:none;border-radius:3px;cursor:pointer;font-size:.7rem;
width:18px;height:18px;line-height:1;padding:0}
.pf-thumb-del:hover{background:#c0392b}
.pf-status{font-size:.76rem;color:#7ed87e;min-height:1.1rem;margin-left:.5rem}
.pf-actions{display:flex;gap:.5rem;align-items:center;margin-top:.7rem}
"""
def _page(title, body, active="run"):
nav_items = [
("run", "/", "▶ Run"),
("characters", "/characters", "👤 Characters"),
("environments", "/environments", "🏞 Environments"),
("gallery", "/gallery", "🎬 Gallery"),
]
nav = "".join(
......@@ -1940,9 +2152,16 @@ input[type=checkbox]{width:auto;accent-color:#f5a623;margin-right:.3rem}
</body></html>"""
def _run_page_html(args_ns):
import json as _json
def _v(attr, default=""): return getattr(args_ns, attr, default)
def _c(attr): return " checked" if getattr(args_ns, attr, False) else ""
# If the script was launched with -c/--config, the Save button defaults
# to that same path so saving overwrites the loaded config file.
_cfg_arg = getattr(args_ns, "config", None)
_save_default = os.path.abspath(_cfg_arg) if _cfg_arg else "township_config.json"
_save_default_js = _json.dumps(_save_default)
char_mode = ("reuse" if _v("reuse_fighters") else
"fighters" if _v("fighters") else
"skip" if _v("skip_characters") else "generate")
......@@ -2076,7 +2295,8 @@ input[type=checkbox]{width:auto;accent-color:#f5a623;margin-right:.3rem}
</div>
</div>
<div id=lora_fields style="margin-top:.6rem">
<div class=row>
<label style="margin-top:0">Character LoRAs <span class=hint>(per-fighter identity)</span></label>
<div class=row3>
<div><label>LoRA train steps</label>
<input name=lora_steps type=number min=100 max=3000 step=50 value="{_v('lora_steps', 800)}"></div>
<div><label>LoRA rank</label>
......@@ -2084,6 +2304,17 @@ input[type=checkbox]{width:auto;accent-color:#f5a623;margin-right:.3rem}
<div><label>LoRA weight <span class=hint>(at generation)</span></label>
<input name=lora_weight type=number min=0 max=2 step=0.05 value="{_v('lora_weight', 0.85)}"></div>
</div>
<div style="margin-top:.6rem">
<label><input type=checkbox name=env_loras{"" if _v('no_env_loras') else " checked"}> Also train per-environment LoRAs <span class=hint>(lock each location’s look)</span></label>
</div>
<div class=row3 style="margin-top:.4rem">
<div><label>Env LoRA train steps</label>
<input name=env_lora_steps type=number min=100 max=3000 step=50 value="{_v('env_lora_steps', 800)}"></div>
<div><label>Env LoRA rank</label>
<input name=env_lora_rank type=number min=2 max=128 value="{_v('env_lora_rank', 16)}"></div>
<div><label>Env LoRA weight <span class=hint>(at generation)</span></label>
<input name=env_lora_weight type=number min=0 max=2 step=0.05 value="{_v('env_lora_weight', 0.8)}"></div>
</div>
</div>
</div>
......@@ -2139,11 +2370,13 @@ async function runStep(step){{
if(j.error){{ appendLog('✗ '+j.error); return; }}
setStatus(true,false);
startSSE();
setTimeout(refreshStatus, 500);
}}
async function saveConfig(){{
// Save the current options to a file ON THE SERVER (where the script runs).
// Relative paths are written inside the output directory.
const def = 'township_config.json';
// Relative paths are written inside the output directory. When launched with
// -c/--config, the default below is that same config path (overwrite-in-place).
const def = {_save_default_js};
const path = prompt('Save configuration to file (relative paths go inside the output dir):', def);
if(path === null) return; // cancelled
const fd = new FormData(document.getElementById('run-form'));
......@@ -2162,6 +2395,8 @@ function clearLog(){{ document.getElementById('log-box').innerHTML=''; }}
let _es = null;
function colorLine(t){{
const low = t.toLowerCase();
if(t.indexOf('▶')!==-1 || t.trim().startsWith('━'))
return '<span class=head>'+escHtml(t)+'</span>';
if(low.includes('✗')||low.includes('error')||low.includes('oom')||low.includes('fatal'))
return '<span class=err>'+escHtml(t)+'</span>';
if(low.includes('✓')||low.includes('loaded')||low.includes('saved')||low.includes('done'))
......@@ -2178,28 +2413,32 @@ function appendLog(t){{
box.innerHTML += colorLine(t)+'\\n';
box.scrollTop=box.scrollHeight;
}}
function setStatus(running, done){{
function setStatus(running, done, label){{
const pill=document.getElementById('status-pill');
const startBtn=document.getElementById('start-btn');
const stopBtn=document.getElementById('stop-btn');
if(done){{ pill.className='status-pill status-done'; pill.textContent='Done'; }}
else if(running){{ pill.className='status-pill status-run'; pill.textContent='Running…'; }}
const lbl = label ? (' — '+label) : '';
if(running){{ pill.className='status-pill status-run'; pill.textContent='Running…'+lbl; }}
else if(done){{ pill.className='status-pill status-done'; pill.textContent='Done'+lbl; }}
else{{ pill.className='status-pill status-idle'; pill.textContent='Idle'; }}
startBtn.style.display = running ? 'none' : '';
stopBtn.style.display = running ? '' : 'none';
if(!running && _es){{ _es.close(); _es=null; }}
}}
function refreshStatus(){{
fetch('/status').then(r=>r.json()).then(d=>setStatus(d.running,d.done,d.label)).catch(()=>{{}});
}}
function startSSE(){{
if(_es){{ _es.close(); }}
_es = new EventSource('/stream');
_es.onmessage = e => appendLog(e.data);
_es.onerror = () => {{
setTimeout(()=>fetch('/status').then(r=>r.json()).then(d=>setStatus(d.running,d.done)),1000);
setTimeout(()=>fetch('/status').then(r=>r.json()).then(d=>setStatus(d.running,d.done,d.label)),1000);
}};
}}
function pollStatus(){{
fetch('/status').then(r=>r.json()).then(d=>{{
setStatus(d.running,d.done);
setStatus(d.running,d.done,d.label);
if(d.running) setTimeout(pollStatus,3000);
}}).catch(()=>setTimeout(pollStatus,5000));
}}
......@@ -2211,6 +2450,7 @@ document.getElementById('run-form').onsubmit = async function(e){{
if(j.error){{ appendLog('✗ '+j.error); return; }}
setStatus(true,false);
startSSE();
setTimeout(refreshStatus, 500);
}};
async function stopRun(){{
await fetch('/stop',{{method:'POST'}});
......@@ -2218,12 +2458,231 @@ async function stopRun(){{
// Restore state on page load
toggleConsFields();
fetch('/status').then(r=>r.json()).then(d=>{{
setStatus(d.running,d.done);
setStatus(d.running,d.done,d.label);
if(d.running) startSSE();
d.log.forEach(l=>appendLog(l));
}});
</script>"""
def _list_profiles(kind: str) -> list:
"""Return locally-saved profiles of a kind with their meta + image files."""
base = out_dir / (kind + "s")
out = []
if not base.exists():
return out
for d in sorted(base.iterdir()):
if not d.is_dir():
continue
mp = d / "meta.json"
if not mp.exists():
continue
try:
meta = json.loads(mp.read_text())
except Exception:
meta = {}
imgs = sorted(
p.name for p in d.iterdir()
if p.suffix.lower() in (".png", ".jpg", ".jpeg", ".webp")
)
out.append({"name": d.name, "meta": meta, "images": imgs})
return out
def _profiles_html(kind: str):
import html as _html
label = "Characters" if kind == "character" else "Environments"
profiles = _list_profiles(kind)
def esc(v):
return _html.escape(str(v if v is not None else ""), quote=True)
cards = []
for p in profiles:
name = p["name"]
meta = p["meta"]
thumbs = "".join(
f'<div class=pf-thumb>'
f'<img src="/media/{kind}s/{esc(name)}/{esc(img)}" loading=lazy alt="{esc(img)}">'
f'<button class=pf-thumb-del title="Delete this image" '
f'onclick="delImg(\'{kind}\',\'{esc(name)}\',\'{esc(img)}\')">✕</button>'
f'</div>'
for img in p["images"]
) or '<span class=hint>No reference images.</span>'
gender_html = ""
if kind == "character":
gender_html = (
f'<div><label>Gender</label>'
f'<input type=text data-field=gender value="{esc(meta.get("gender",""))}"></div>'
)
# Both kinds can guide regeneration with their kept references via
# IP-Adapter (characters → character_profiles, envs → environment_profiles).
guide_html = (
'<label style="margin:0;font-size:.78rem;display:inline-flex;align-items:center;gap:.25rem">'
'<input type=checkbox data-regen=guide checked style="width:auto"> match kept refs</label>'
)
# Any OTHER scalar fields present in meta.json become editable inputs
# too, so every field of a profile can be edited (not just the fixed
# set). Bookkeeping / already-rendered keys are excluded.
_shown = {"name", "region", "gender", "description", "prompt",
"images", "image_count", "created_at", "created"}
extra_rows = []
for fk, fval in meta.items():
if fk in _shown or isinstance(fval, (dict, list)):
continue
extra_rows.append(
f'<div><label>{esc(fk)}</label>'
f'<input type=text data-field="{esc(fk)}" value="{esc(fval)}"></div>'
)
extra_html = (f'<div class=row3 style="margin-top:.4rem">{"".join(extra_rows)}</div>'
if extra_rows else "")
cards.append(
f'<div class=card id="pf-{kind}-{esc(name)}">'
f' <div class=pf-head>'
f' <span class=pf-name>{esc(name)}</span>'
f' <span class=hint>{len(p["images"])} image(s)</span>'
f' </div>'
f' <div class=pf-thumbs>{thumbs}</div>'
f' <div class=row>'
f' <div><label>Region</label>'
f'<input type=text data-field=region value="{esc(meta.get("region",""))}"></div>'
f' {gender_html}'
f' </div>'
f' <label>Description <span class=hint>(synced to CoderAI)</span></label>'
f' <textarea data-field=description rows=2>{esc(meta.get("description",""))}</textarea>'
f' <label>Prompt</label>'
f' <textarea data-field=prompt rows=3>{esc(meta.get("prompt",""))}</textarea>'
f' {extra_html}'
f' <div class=pf-actions>'
f' <button class="btn btn-primary" style="font-size:.82rem;padding:.35rem .9rem" '
f'onclick="saveProfile(\'{kind}\',\'{esc(name)}\')">💾 Save</button>'
f' <button class="btn btn-danger" style="font-size:.82rem;padding:.35rem .9rem" '
f'onclick="delProfile(\'{kind}\',\'{esc(name)}\')">🗑 Remove</button>'
f' <span class=pf-status></span>'
f' </div>'
f' <div class=pf-actions style="border-top:1px solid #222;padding-top:.6rem;margin-top:.6rem">'
f' <label style="margin:0;font-size:.78rem">Add <input type=number data-regen=count '
f'value=4 min=1 max=8 style="width:54px;display:inline-block"> new ref(s)</label>'
f' {guide_html}'
f' <button class="btn btn-secondary" style="font-size:.82rem;padding:.35rem .9rem" '
f'onclick="regenProfile(\'{kind}\',\'{esc(name)}\')">♻ Regenerate references</button>'
f' <span class=pf-regen-status style="font-size:.76rem;color:#7ea8f7"></span>'
f' </div>'
f' <div class=pf-actions style="padding-top:.5rem">'
f' <label style="margin:0;font-size:.78rem">Or upload your own:</label>'
f' <input type=file data-upload=files accept="image/*" multiple '
f'style="font-size:.76rem;width:auto;flex:1;min-width:160px">'
f' <button class="btn btn-secondary" style="font-size:.82rem;padding:.35rem .9rem" '
f'onclick="uploadRefs(\'{kind}\',\'{esc(name)}\')">⬆ Upload references</button>'
f' <span class=pf-upload-status style="font-size:.76rem;color:#7ea8f7"></span>'
f' </div>'
f'</div>'
)
if cards:
inner = "".join(cards)
else:
inner = (f'<div class=card style="color:#666">No {label.lower()} found in '
f'<code>{esc(str(out_dir))}</code> yet. Generate some from the Run page first.</div>')
script = """
<script>
async function saveProfile(kind,name){
const root=document.getElementById('pf-'+kind+'-'+name);
const st=root.querySelector('.pf-status');
st.style.color='#aaa'; st.textContent='Saving…';
const fd=new FormData();
fd.append('kind',kind); fd.append('name',name);
root.querySelectorAll('[data-field]').forEach(el=>fd.append(el.getAttribute('data-field'),el.value));
try{
const r=await fetch('/profile/save',{method:'POST',body:fd});
const j=await r.json();
if(j.error){st.style.color='#e07070'; st.textContent='✗ '+j.error; return;}
st.style.color='#7ed87e'; st.textContent='✓ Saved'+(j.synced?' (synced to CoderAI)':'');
}catch(e){st.style.color='#e07070'; st.textContent='✗ '+e;}
}
async function delProfile(kind,name){
if(!confirm('Remove "'+name+'" and all its images? This deletes the local profile'
+' and removes it from CoderAI. This cannot be undone.'))return;
const fd=new FormData(); fd.append('kind',kind); fd.append('name',name);
const r=await fetch('/profile/delete',{method:'POST',body:fd});
const j=await r.json();
if(j.error){alert('Delete failed: '+j.error); return;}
const el=document.getElementById('pf-'+kind+'-'+name);
if(el) el.remove();
}
async function delImg(kind,name,file){
if(!confirm('Delete image "'+file+'"?'))return;
const fd=new FormData(); fd.append('kind',kind); fd.append('name',name); fd.append('file',file);
const r=await fetch('/profile/delete-image',{method:'POST',body:fd});
const j=await r.json();
if(j.error){alert('Delete failed: '+j.error); return;}
location.reload();
}
async function regenProfile(kind,name){
const root=document.getElementById('pf-'+kind+'-'+name);
const st=root.querySelector('.pf-regen-status');
const cnt=root.querySelector('[data-regen=count]');
const guideEl=root.querySelector('[data-regen=guide]');
const count=Math.max(1,Math.min(8,parseInt(cnt&&cnt.value||'4',10)||4));
const fd=new FormData();
fd.append('kind',kind); fd.append('name',name); fd.append('count',count);
fd.append('guide', (guideEl && guideEl.checked) ? '1' : '0');
st.style.color='#aaa'; st.textContent='Starting…';
let j;
try{
const r=await fetch('/profile/regenerate',{method:'POST',body:fd});
j=await r.json();
}catch(e){ st.style.color='#e07070'; st.textContent='✗ '+e; return; }
if(j.error){ st.style.color='#e07070'; st.textContent='✗ '+j.error; return; }
const jobId=j.job_id;
st.style.color='#7ea8f7';
const poll=async()=>{
let d;
try{ d=await (await fetch('/job/'+jobId)).json(); }
catch(e){ setTimeout(poll,1500); return; }
const pct=d.progress||0;
if(d.status==='running'){ st.textContent='⏳ '+(d._msg||('working… '+pct+'%')); setTimeout(poll,1200); }
else if(d.status==='done'){
st.style.color='#7ed87e';
st.textContent='✓ added '+(d.added||0)+' image(s)'+(d.synced===false?' (local only)':'')+' — reloading…';
setTimeout(()=>location.reload(),900);
} else {
st.style.color='#e07070'; st.textContent='✗ '+(d.error||'failed');
}
};
setTimeout(poll,800);
}
async function uploadRefs(kind,name){
const root=document.getElementById('pf-'+kind+'-'+name);
const inp=root.querySelector('[data-upload=files]');
const st=root.querySelector('.pf-upload-status');
if(!inp||!inp.files||!inp.files.length){
st.style.color='#e07070'; st.textContent='Choose image file(s) first'; return;
}
const fd=new FormData();
fd.append('kind',kind); fd.append('name',name);
for(const f of inp.files) fd.append('files',f);
st.style.color='#aaa'; st.textContent='Uploading '+inp.files.length+' file(s)…';
try{
const r=await fetch('/profile/upload-image',{method:'POST',body:fd});
const j=await r.json();
if(j.error){ st.style.color='#e07070'; st.textContent='✗ '+j.error; return; }
st.style.color='#7ed87e';
st.textContent='✓ added '+j.added+(j.rejected?(' ('+j.rejected+' skipped)'):'')
+(j.synced===false?' (local only)':'')+' — reloading…';
setTimeout(()=>location.reload(),800);
}catch(e){ st.style.color='#e07070'; st.textContent='✗ '+e; }
}
</script>"""
return (f'<div style="display:flex;justify-content:space-between;align-items:center">'
f'<h1>{label}</h1>'
f'<a href="/{kind}s" class="btn btn-secondary" style="font-size:.8rem">↻ Refresh</a></div>'
f'<p class=hint style="margin-bottom:.8rem">Edit a profile’s fields and Save, or '
f'Remove it entirely. Changes apply to the local output folder and are synced to CoderAI.</p>'
f'{inner}{script}')
def _gallery_html(out_path: Path):
sections = []
......@@ -2426,6 +2885,14 @@ async function pollJob(){
html = _page("Run", _run_page_html(default_args), "run")
self._send(200, "text/html; charset=utf-8", html)
elif path == "/characters":
html = _page("Characters", _profiles_html("character"), "characters")
self._send(200, "text/html; charset=utf-8", html)
elif path == "/environments":
html = _page("Environments", _profiles_html("environment"), "environments")
self._send(200, "text/html; charset=utf-8", html)
elif path == "/gallery":
html = _page("Gallery", _gallery_html(out_dir), "gallery")
self._send(200, "text/html; charset=utf-8", html)
......@@ -2435,6 +2902,7 @@ async function pollJob(){
payload = _j.dumps({
"running": _state["running"],
"done": _state["done"],
"label": _state.get("current", ""),
"log": _state["log_lines"][-200:],
})
self._send(200, "application/json", payload)
......@@ -2541,6 +3009,208 @@ async function pollJob(){
self._send(200, "application/json", _j.dumps({"ok": True}))
return
if path == "/profile/regenerate":
import json as _j, uuid as _u
clen = int(self.headers.get("Content-Length", 0))
raw = self.rfile.read(clen)
ctype = self.headers.get("Content-Type", "")
if "multipart/form-data" in ctype:
boundary = ctype.split("boundary=")[-1].strip().encode()
form = _parse_multipart(raw, boundary)
else:
form = dict(urllib.parse.parse_qsl(raw.decode(errors="replace")))
def _fv(k, default=""):
v = form.get(k)
if v is None: return default
return v if isinstance(v, str) else v.decode(errors="replace")
kind = _fv("kind"); name = _fv("name")
if (kind not in ("character", "environment") or not name
or "/" in name or "\\" in name or ".." in name):
self._send(400, "application/json",
_j.dumps({"error": "invalid kind/name"}))
return
try:
count = max(1, min(8, int(_fv("count", "4") or 4)))
except ValueError:
count = 4
guide = _fv("guide", "1") not in ("0", "false", "no", "")
job_id = _u.uuid4().hex[:12]
threading.Thread(target=_run_regen_job,
args=(job_id, kind, name, count, guide),
daemon=True).start()
self._send(200, "application/json", _j.dumps({"job_id": job_id}))
return
if path == "/profile/upload-image":
# Append user-uploaded image files as new references, preserving
# all existing ones.
import json as _j
clen = int(self.headers.get("Content-Length", 0))
raw = self.rfile.read(clen)
ctype = self.headers.get("Content-Type", "")
if "multipart/form-data" not in ctype:
self._send(400, "application/json",
_j.dumps({"error": "expected multipart/form-data"}))
return
boundary = ctype.split("boundary=")[-1].strip().encode()
fields, files = _parse_multipart_full(raw, boundary)
kind = fields.get("kind", "")
name = fields.get("name", "")
if (kind not in ("character", "environment") or not name
or "/" in name or "\\" in name or ".." in name):
self._send(400, "application/json",
_j.dumps({"error": "invalid kind/name"}))
return
def _img_ext(data: bytes):
if data[:4] == b"\x89PNG": return ".png", "image/png"
if data[:2] == b"\xff\xd8": return ".jpg", "image/jpeg"
if data[:4] == b"RIFF" and data[8:12] == b"WEBP": return ".webp", "image/webp"
if data[:6] in (b"GIF87a", b"GIF89a"): return ".gif", "image/gif"
return None, None
base = out_dir / (kind + "s") / name
base.mkdir(parents=True, exist_ok=True)
added_uris, rejected = [], 0
for f in files:
data = f.get("data") or b""
ext, mime = _img_ext(data)
if not ext:
rejected += 1
continue
out_p = _next_ref_path(base, ext)
try:
out_p.write_bytes(data)
except Exception:
rejected += 1
continue
added_uris.append("data:%s;base64,%s" % (
mime, __import__("base64").b64encode(data).decode()))
if not added_uris:
self._send(400, "application/json",
_j.dumps({"error": "no valid image files uploaded"}))
return
synced = True
try:
client = CoderAIClient(default_args.base_url,
getattr(default_args, "api_key", None))
client.patch_profile(kind, name, add_images=added_uris)
except Exception:
synced = False
self._send(200, "application/json",
_j.dumps({"ok": True, "added": len(added_uris),
"rejected": rejected, "synced": synced}))
return
if path in ("/profile/save", "/profile/delete", "/profile/delete-image"):
import json as _j
import shutil as _shutil
clen = int(self.headers.get("Content-Length", 0))
raw = self.rfile.read(clen)
ctype = self.headers.get("Content-Type", "")
if "multipart/form-data" in ctype:
boundary = ctype.split("boundary=")[-1].strip().encode()
form = _parse_multipart(raw, boundary)
else:
form = dict(urllib.parse.parse_qsl(raw.decode(errors="replace")))
def _fv(k, default=""):
v = form.get(k)
if v is None: return default
return v if isinstance(v, str) else v.decode(errors="replace")
kind = _fv("kind")
name = _fv("name")
# Reject anything that could escape the profile directory.
if (kind not in ("character", "environment") or not name
or "/" in name or "\\" in name or ".." in name):
self._send(400, "application/json",
_j.dumps({"error": "invalid kind/name"}))
return
base = out_dir / (kind + "s") / name
client = CoderAIClient(default_args.base_url,
getattr(default_args, "api_key", None))
if path == "/profile/delete":
try:
if base.exists():
_shutil.rmtree(base)
except Exception as e:
self._send(500, "application/json",
_j.dumps({"error": f"cannot delete local: {e}"}))
return
synced = True
try:
client.delete_profile(kind, name)
except Exception:
synced = False
self._send(200, "application/json",
_j.dumps({"ok": True, "synced": synced}))
return
if path == "/profile/delete-image":
file = _fv("file")
if not file or "/" in file or "\\" in file or ".." in file:
self._send(400, "application/json",
_j.dumps({"error": "invalid file"}))
return
imgs = sorted(
p.name for p in base.iterdir()
if p.is_file() and p.suffix.lower() in (".png", ".jpg", ".jpeg", ".webp")
) if base.exists() else []
idx = imgs.index(file) if file in imgs else None
fp = base / file
try:
if fp.exists():
fp.unlink()
except Exception as e:
self._send(500, "application/json",
_j.dumps({"error": f"cannot delete image: {e}"}))
return
if idx is not None:
try:
client.patch_profile(kind, name, remove_indices=[idx])
except Exception:
pass
self._send(200, "application/json", _j.dumps({"ok": True}))
return
# /profile/save — update local meta.json (+ sync description to server)
meta_path = base / "meta.json"
meta = {}
try:
meta = json.loads(meta_path.read_text())
except Exception:
pass
# Write back every submitted field (the editor exposes all of
# them), skipping control + bookkeeping keys.
_reserved = {"kind", "name", "images", "image_count",
"created_at", "created"}
for fk in form.keys():
if fk in _reserved:
continue
meta[fk] = _fv(fk)
meta["name"] = name
try:
base.mkdir(parents=True, exist_ok=True)
meta_path.write_text(json.dumps(meta, indent=2))
except Exception as e:
self._send(500, "application/json",
_j.dumps({"error": f"cannot save: {e}"}))
return
synced = False
try:
client.patch_profile(kind, name, description=meta.get("description", ""))
synced = True
except Exception:
pass
self._send(200, "application/json",
_j.dumps({"ok": True, "synced": synced}))
return
if path == "/save-config":
# Write the submitted options to a config file ON THE SERVER
# (the machine running this script), reusable later via --config.
......@@ -2589,6 +3259,10 @@ 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),
"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),
"env_lora_weight": float(_fv("env_lora_weight", "0.8") or 0.8),
"skip_characters": cm == "skip",
"reuse_fighters": cm == "reuse",
"fighters": _s(_fv("fighters")) if cm == "fighters" else None,
......@@ -2706,6 +3380,11 @@ 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"))
# 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"))
ns.env_lora_rank = int(_fv("env_lora_rank", "16"))
ns.env_lora_weight = float(_fv("env_lora_weight", "0.8"))
# char mode
cm = _fv("char_mode", "generate")
ns.skip_characters = (cm == "skip")
......@@ -2763,10 +3442,43 @@ async function pollJob(){
elif step == "videos":
ns.skip_characters = True; ns.skip_environments = True; ns.only_videos = True
# Human-readable label for what's being executed, shown as a banner
# at the top of the (freshly cleared) log so it's always obvious
# which run/step is currently in progress.
_STEP_LABELS = {
"characters": "Step 1 · Generate Characters",
"environments": "Step 2 · Generate Environments",
"prompts": "Step 3 · Write Video Prompts",
"loras": "Step 4 · Train Character LoRAs",
"keyframes": "Step 5 · Generate Keyframes",
"videos": "Step 6 · Render Videos",
}
if step:
run_label = _STEP_LABELS.get(step, f"Step · {step}")
elif ns.only_characters:
run_label = "Full Run · Characters only"
elif ns.only_environments:
run_label = "Full Run · Environments only"
elif ns.only_assets:
run_label = "Full Run · Assets only (characters + environments)"
elif ns.only_prompts:
run_label = "Full Run · Prompts only"
elif ns.only_videos:
run_label = "Full Run · Videos only (render from saved prompts)"
else:
run_label = "Full Run · All stages"
_state["abort"].clear()
_state["log_lines"].clear()
_state["done"] = False
_state["running"] = True
_state["current"] = run_label
_bar = "━" * 58
_web_log(_bar)
_web_log(f"▶ {run_label} [{time.strftime('%H:%M:%S')}]")
_web_log(f" consistency: {ns.consistency} output: {ns.out_dir}")
_web_log(_bar)
def _run_thread():
import sys as _sys
......@@ -2779,7 +3491,7 @@ async function pollJob(){
finally:
_state["running"] = False
_state["done"] = True
_web_log("✓ Run complete.")
_web_log(f"✓ {run_label} — complete.")
with _sse_lock:
for q in list(_sse_clients):
try: q.put(None)
......@@ -2809,6 +3521,40 @@ async function pollJob(){
result[name] = value.decode(errors="replace")
return result
def _parse_multipart_full(body: bytes, boundary: bytes):
"""Parse multipart/form-data preserving raw bytes for file parts.
Returns (fields, files) where fields maps name→str and files is a list
of {'name', 'filename', 'data': bytes} for parts with a filename.
"""
fields, files = {}, []
delimiter = b"--" + boundary
for part in body.split(delimiter)[1:]:
if part[:2] == b"\r\n":
part = part[2:]
if part.strip() in (b"", b"--"):
continue
if b"\r\n\r\n" not in part:
continue
header_raw, _, value = part.partition(b"\r\n\r\n")
if value.endswith(b"\r\n"):
value = value[:-2]
headers_text = header_raw.decode(errors="replace")
name = filename = None
for hdr_line in headers_text.splitlines():
if "Content-Disposition" in hdr_line:
if 'name="' in hdr_line:
name = hdr_line.split('name="')[1].split('"')[0]
if 'filename="' in hdr_line:
filename = hdr_line.split('filename="')[1].split('"')[0]
if name is None:
continue
if filename is not None:
files.append({"name": name, "filename": filename, "data": value})
else:
fields[name] = value.decode(errors="replace")
return fields, files
def _run_main_with_args(args):
"""Run the full generation pipeline with a pre-built args Namespace."""
# Identical logic to main() after parse_args(), driven by args.
......@@ -2897,21 +3643,31 @@ async function pollJob(){
only_loras = getattr(args, "only_loras", False)
only_keyframes = getattr(args, "only_keyframes", False)
# Load any previously-trained LoRA map from disk so keyframe/video steps
# can reuse it without retraining.
lora_map = {}
lora_file = out_dir_r / "loras.json"
if lora_file.exists():
try:
lora_map = json.loads(lora_file.read_text()) or {}
except Exception:
lora_map = {}
# Load any previously-trained LoRA maps from disk so keyframe/video steps
# can reuse them without retraining (characters + environments).
def _load_map(fname):
fp = out_dir_r / fname
if fp.exists():
try:
return json.loads(fp.read_text()) or {}
except Exception:
return {}
return {}
lora_map = _load_map("loras.json")
env_lora_map = _load_map("env_loras.json")
_no_env_loras = getattr(args, "no_env_loras", False)
_env_lora_weight = getattr(args, "env_lora_weight", 0.8)
# Train LoRAs when requested (full run with lora strategy, or the LoRA step).
if "lora" in consistency and (char_names or []) and (only_loras or not args.skip_videos):
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))
if ("lora" in consistency and not _no_env_loras and (env_names or [])
and (only_loras or not args.skip_videos)):
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))
if only_loras:
_web_log("\n✓ LoRA step complete.")
......@@ -2927,6 +3683,7 @@ async function pollJob(){
keyframe_steps=getattr(args, "keyframe_steps", 28),
keyframe_size=getattr(args, "keyframe_size", "512x512"),
lora_weight=getattr(args, "lora_weight", 0.85),
env_lora_map=env_lora_map, env_lora_weight=_env_lora_weight,
keyframes_only=True,
)
_web_log("\n✓ Keyframe step complete.")
......@@ -2949,6 +3706,7 @@ async function pollJob(){
keyframe_steps=getattr(args, "keyframe_steps", 28),
keyframe_size=getattr(args, "keyframe_size", "512x512"),
lora_weight=getattr(args, "lora_weight", 0.85),
env_lora_map=env_lora_map, env_lora_weight=_env_lora_weight,
)
_web_log("\n✓ Done.")
......@@ -3212,6 +3970,15 @@ 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("--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).")
cons_grp.add_argument("--env-lora-steps", type=int, default=800, metavar="N",
help="Training steps per environment LoRA (default: 800).")
cons_grp.add_argument("--env-lora-rank", type=int, default=16, metavar="N",
help="Environment LoRA rank (default: 16).")
cons_grp.add_argument("--env-lora-weight", type=float, default=0.8, metavar="F",
help="Weight applied to each environment LoRA at generation (default: 0.8).")
parser.add_argument("--cli-mode", action="store_true",
help="Run in CLI mode (default when --cli-mode is present). "
......@@ -3344,12 +4111,18 @@ OUTPUT LAYOUT
else:
env_names = stage_environments(client, image_model, out_dir, region_filter=args.region)
# ── Stage 2.5: Character LoRA training (image base model) ───────────────────
# ── Stage 2.5: LoRA training (image base model) — characters + environments ─
lora_map = {}
env_lora_map = {}
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))
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))
# ── Stage 3: Videos ────────────────────────────────────────────────────────
if not args.skip_videos:
......@@ -3371,6 +4144,8 @@ OUTPUT LAYOUT
keyframe_steps=getattr(args, "keyframe_steps", 28),
keyframe_size=getattr(args, "keyframe_size", "512x512"),
lora_weight=getattr(args, "lora_weight", 0.85),
env_lora_map=env_lora_map,
env_lora_weight=getattr(args, "env_lora_weight", 0.8),
)
_log("\n✓ Done.")
......
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