| """studio_daemon.py - resident model server for LYRA STUDIO. |
| Keeps each specialist loaded in VRAM and executes jobs from BASE/jobs. |
| Writes BASE/results/<id>.json. No time limit.""" |
| import json, os, re, sys, time, traceback |
| from pathlib import Path |
| sys.path.insert(0, "/marimo/ai_studio") |
| import torch |
| from studio_core import (BASE, CHARACTER, VOICE, SEED, DTYPE_NAME, JOBS, JOBS_DONE, RESULTS, LOGS, |
| build_image_prompt, load_refinements, save_refinements) |
|
|
| DTYPE = getattr(torch, DTYPE_NAME) |
| MODELS = {} |
| PAINTER_ID = "Qwen/Qwen-Image" |
| CRITIC_ID = "Qwen/Qwen2.5-VL-7B-Instruct" |
|
|
| def log(m): |
| line = f"[{time.strftime('%H:%M:%S')}] {m}" |
| print(line, flush=True) |
| with open(LOGS/"daemon.log", "a") as f: f.write(line + "\n") |
|
|
| def vram(): |
| free, total = torch.cuda.mem_get_info(); return f"{free/1e9:.0f}/{total/1e9:.0f}GB" |
|
|
| |
| def painter(): |
| if "painter" not in MODELS: |
| from diffusers import QwenImagePipeline |
| log("load Qwen-Image (concept artist) ...") |
| p = QwenImagePipeline.from_pretrained(PAINTER_ID, torch_dtype=DTYPE) |
| p.to("cuda"); p.set_progress_bar_config(disable=True); MODELS["painter"] = p |
| log(f"painter ready {vram()}") |
| return MODELS["painter"] |
|
|
| def brain(): |
| if "brain" not in MODELS: |
| from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor |
| log("load Qwen2.5-VL-7B (critic + director brain) ...") |
| m = Qwen2_5_VLForConditionalGeneration.from_pretrained(CRITIC_ID, torch_dtype=DTYPE, device_map="cuda") |
| pr = AutoProcessor.from_pretrained(CRITIC_ID) |
| MODELS["brain"] = (m, pr); log(f"brain ready {vram()}") |
| return MODELS["brain"] |
|
|
| |
| def _json_block(s): |
| m = re.search(r"\{.*\}", s, re.S) |
| if not m: return None |
| try: return json.loads(m.group(0)) |
| except Exception: |
| try: return json.loads(m.group(0).replace("'", '"')) |
| except Exception: return None |
|
|
| def _gen(model, proc, msgs, images=None, max_new_tokens=900): |
| text = proc.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True) |
| kw = {"text": [text], "return_tensors": "pt"} |
| if images: kw["images"] = images |
| inputs = proc(**kw).to("cuda") |
| with torch.inference_mode(): |
| out = model.generate(**inputs, max_new_tokens=max_new_tokens, do_sample=False) |
| return proc.batch_decode(out[:, inputs.input_ids.shape[1]:], skip_special_tokens=True)[0].strip() |
|
|
| def _render_one(j, refinements, name): |
| pipe = painter() |
| prompt = build_image_prompt(j.get("shot",""), j.get("outfit",""), j.get("mood",""), refinements) |
| g = torch.Generator(device="cuda").manual_seed(j.get("seed", SEED)) |
| w, h = j.get("w",1024), j.get("h",1536) |
| log(f"render {name} {w}x{h} steps={j.get('steps',30)} {vram()}") |
| t0 = time.time() |
| img = pipe(prompt=prompt, negative_prompt=CHARACTER["negative"], width=w, height=h, |
| num_inference_steps=j.get("steps",30), true_cfg_scale=j.get("cfg",4.0), |
| generator=g).images[0] |
| out = BASE/"assets/img"/name; img.save(out) |
| dt = time.time()-t0; log(f"saved {name} in {dt:.1f}s") |
| return {"file": str(out), "prompt": prompt, "seconds": round(dt,1)} |
|
|
| RUBRIC = ("You are the art director of a virtual-idol IP. Be specific and a little ruthless. " |
| "Judge the render against the CHARACTER BIBLE. Return STRICT JSON only, keys: " |
| "score_identity (0-10), score_beauty (0-10), score_brief (0-10), " |
| "problems (list), fixes (list of concrete prompt phrases to ADD to improve the next render), " |
| "verdict (one sentence).") |
|
|
| def _critique_one(name, task, max_new_tokens=900): |
| from PIL import Image |
| model, proc = brain() |
| img = Image.open(BASE/"assets/img"/name).convert("RGB") |
| user = (f"CHARACTER BIBLE:\n{json.dumps(CHARACTER, ensure_ascii=False, indent=1)}\n\n" |
| f"SHOT BRIEF: {task or 'Critique this render of LYRA.'}\n\n{RUBRIC}") |
| reply = _gen(model, proc, [{"role":"user","content":[{"type":"image","image":img}, |
| {"type":"text","text":user}]}], |
| images=[img], max_new_tokens=max_new_tokens) |
| return reply, _json_block(reply) |
|
|
| |
| def job_image(j): |
| rf = j.get("refinements") |
| r = _render_one(j, rf, j["name"]) |
| return {"files": [r["file"]], "prompt": r["prompt"], "seconds": r["seconds"]} |
|
|
| def job_critique(j): |
| reply, parsed = _critique_one(j["image"], j.get("task","")) |
| return {"image": j["image"], "critique": reply, "parsed": parsed} |
|
|
| def job_direct(j): |
| model, proc = brain() |
| log(f"direct: {j['instruction'][:60]}...") |
| reply = _gen(model, proc, [{"role":"user","content":[{"type":"text","text":j["instruction"]}]}], |
| max_new_tokens=j.get("max_new_tokens", 900)) |
| return {"reply": reply} |
|
|
| def job_iterate(j): |
| """render -> critique -> append fixes (learned) -> render again.""" |
| rf = load_refinements() |
| history, learned = [], [] |
| task = j.get("task") or f"Canonical identity shot: {j.get('shot','portrait')}" |
| for r in range(int(j.get("rounds", 2))): |
| name = j["name"] if r == 0 else j["name"].replace(".png", f"_r{r}.png") |
| res = _render_one(j, rf, name) |
| reply, parsed = _critique_one(name, task) |
| fixes = (parsed or {}).get("fixes", []) if isinstance(parsed, dict) else [] |
| scores = {k: (parsed or {}).get(k) for k in ("score_identity","score_beauty","score_brief")} if isinstance(parsed, dict) else {} |
| history.append({"round": r, "file": res["file"], "prompt": res["prompt"], |
| "critique": reply, "parsed": parsed, "scores": scores}) |
| log(f" round {r}: {scores}") |
| if fixes: rf = rf + fixes; learned = learned + fixes |
| if j.get("learn", True) and learned: |
| save_refinements(load_refinements() + learned) |
| log(f"learned {len(learned)} refinements -> {len(load_refinements())} total") |
| return {"history": history, "final": history[-1]["file"], |
| "refinements_file": str(load_refinements())} |
|
|
| def job_evict(j): |
| """Free VRAM so another specialist can use the GPU. The studio shares one card.""" |
| which = j.get("which", "all") |
| names = list(MODELS) if which == "all" else ([which] if isinstance(which, str) else list(which)) |
| for n in names: MODELS.pop(n, None) |
| import gc; gc.collect(); torch.cuda.empty_cache() |
| log(f"evicted {names} -> {vram()}") |
| return {"evicted": names, "vram": vram()} |
|
|
| def job_shutdown(j): |
| Path(LOGS/"daemon.stop").write_text("1"); return {"bye": True} |
|
|
| HANDLERS = {"image": job_image, "critique": job_critique, "direct": job_direct, |
| "iterate": job_iterate, "evict": job_evict, |
| "ping": lambda j: {"pong": True}, "shutdown": job_shutdown} |
|
|
| def main(): |
| (LOGS/"daemon.stop").unlink(missing_ok=True) |
| log(f"daemon up pid={os.getpid()} cuda={torch.cuda.get_device_name(0)} {vram()}") |
| while not (LOGS/"daemon.stop").exists(): |
| found = False |
| for f in sorted(JOBS.glob("*.json")): |
| found = True; jid = f.stem |
| try: |
| j = json.loads(f.read_text()); jid = j.get("id", jid) |
| h = HANDLERS.get(j.get("type")) |
| if h is None: |
| time.sleep(0.5); continue |
| res = h(j) |
| res = {"id": jid, "status": "ok", **(res or {})} |
| except Exception as e: |
| res = {"id": jid, "status": "error", "error": f"{type(e).__name__}: {e}", |
| "trace": traceback.format_exc()[-1500:]} |
| log("FAILED " + res["error"]) |
| tmp = RESULTS/f".{res['id']}.tmp"; tmp.write_text(json.dumps(res, ensure_ascii=False, indent=2)) |
| tmp.rename(RESULTS/f"{res['id']}.json"); f.rename(JOBS_DONE/f.name) |
| log(f"done {res['id']} status={res['status']}") |
| if not found: time.sleep(1) |
|
|
| if __name__ == "__main__": |
| main() |
|
|
