| |
| """Queue MOSS-Image prompts against the local ComfyUI (http://127.0.0.1:8188). |
| |
| python3 moss.py "a gorilla in a nightclub..." --seeds 1 2 3 |
| python3 moss.py -f prompts.txt --enhance --lora r256 |
| """ |
| import argparse, copy, json, os, sys, time, urllib.request |
|
|
| HERE = os.path.dirname(os.path.abspath(__file__)) |
| URL = "http://127.0.0.1:8188" |
| OUT = os.path.join(HERE, "..", "..", "output") |
| MODELS = { |
| "base": "moss_image_base_bf16.safetensors", |
| "sft": "moss_image_sft_bf16.safetensors", |
| } |
| LORAS = { |
| "b21": "moss_image_sft_lora_blocks_21_to_49_rank_512_bf16.safetensors", |
| "r256": "moss_image_sft_lora_rank_256_bf16.safetensors", |
| "r512": "moss_image_sft_lora_rank_512_bf16.safetensors", |
| } |
|
|
|
|
| def call(path, data=None): |
| req = urllib.request.Request(URL + path, data=json.dumps(data).encode() if data is not None else None, |
| headers={"Content-Type": "application/json"}) |
| with urllib.request.urlopen(req) as r: |
| return json.load(r) |
|
|
|
|
| def build(base, prompt, seed, a): |
| g = copy.deepcopy(base) |
| g["1"]["inputs"]["unet_name"] = MODELS.get(a.model, a.model) |
| if a.swap_layers is not None or a.swap_other: |
| |
| g["40"] = {"class_type": "UNETLoader", "inputs": {"unet_name": MODELS.get(a.donor, a.donor), "weight_dtype": "default"}} |
| g["41"] = {"class_type": "MOSSBlockSwap", "inputs": {"model": ["1", 0], "donor": ["40", 0], "layers": a.swap_layers or "", |
| "parts": a.swap_parts, "include_other": a.swap_other, "ratio": a.swap_ratio}} |
| g["16"]["inputs"]["model"] = ["41", 0] |
| del g["30"] |
| elif a.lora == "none": |
| g["16"]["inputs"]["model"] = ["1", 0] |
| del g["30"] |
| else: |
| g["30"]["inputs"]["lora_name"] = LORAS.get(a.lora, a.lora) |
| g["30"]["inputs"]["strength_model"] = a.strength |
| if a.lora2: |
| g["31"] = {"class_type": "LoraLoaderModelOnly", "inputs": {"model": g["16"]["inputs"]["model"], |
| "lora_name": LORAS.get(a.lora2, a.lora2), "strength_model": a.strength2}} |
| g["16"]["inputs"]["model"] = ["31", 0] |
| if a.enhance: |
| g["17"]["inputs"]["prompt"] = prompt |
| g["17"]["inputs"]["sampling_mode.seed"] = seed |
| else: |
| g["5"]["inputs"]["text"] = prompt |
| for k in ("17", "18", "19"): |
| del g[k] |
| if a.probe: |
| g["42"] = {"class_type": "MOSSStripeProbe", "inputs": {"model": g["16"]["inputs"]["model"], "head_from": MODELS.get(a.model, a.model), |
| "tag": f"{a.probe}_s{seed}", "width": a.width, "height": a.height}} |
| g["16"]["inputs"]["model"] = ["42", 0] |
| if a.no_tsr: |
| g["2"]["inputs"]["model"] = g["16"]["inputs"]["model"] |
| del g["16"] |
| if a.no_slg: |
| for k in ("20", "22"): |
| g[k]["inputs"]["model"] = g["2"]["inputs"]["model"] |
| del g["2"] |
| if a.sampler != "er_sde": |
| g["21"] = {"class_type": "KSamplerSelect", "inputs": {"sampler_name": a.sampler}} |
| if a.negative is not None: |
| g["6"]["inputs"]["text"] = a.negative |
| g["20"]["inputs"].update(noise_seed=seed, cfg=a.cfg) |
| g["22"]["inputs"]["steps"] = a.steps |
| g["7"]["inputs"].update(width=a.width, height=a.height) |
| g["10"]["inputs"]["filename_prefix"] = "moss/" + a.prefix |
| if a.window_steps: |
| window(g, a) |
| return g |
|
|
|
|
| def window(g, a): |
| """Apply --window-lora only for sampling steps a..b (1-indexed) via ComfyUI hook keyframes on both conds. |
| --window-steps "1-N" applies it for the whole run; strength 0 everywhere is a no-op control.""" |
| lo, hi = (int(v) for v in a.window_steps.split("-")) |
| n = a.steps |
| g["70"] = {"class_type": "CreateHookLoraModelOnly", "inputs": {"lora_name": a.window_lora, "strength_model": a.window_strength}} |
| kf = ([(0.0, 0.0)] if lo > 1 else []) + [((lo - 1) / n, 1.0)] + ([(hi / n, 0.0)] if hi < n else []) |
| prev = None |
| for i, (pct, mult) in enumerate(kf): |
| node = {"class_type": "CreateHookKeyframe", "inputs": {"strength_mult": mult, "start_percent": pct}} |
| if prev: |
| node["inputs"]["prev_hook_kf"] = [prev, 0] |
| prev = str(71 + i) |
| g[prev] = node |
| g["75"] = {"class_type": "SetHookKeyframes", "inputs": {"hooks": ["70", 0], "hook_kf": [prev, 0]}} |
| for cid, src in (("76", "5"), ("77", "6")): |
| g[cid] = {"class_type": "ConditioningSetProperties", "inputs": {"cond_NEW": [src, 0], "strength": 1.0, "set_cond_area": "default", "hooks": ["75", 0]}} |
| g["20"]["inputs"]["positive"], g["20"]["inputs"]["negative"] = ["76", 0], ["77", 0] |
|
|
|
|
| def main(): |
| p = argparse.ArgumentParser() |
| p.add_argument("prompt", nargs="*") |
| p.add_argument("-f", "--file", help="one prompt per paragraph (blank-line separated)") |
| p.add_argument("--seeds", type=int, nargs="+", default=[16]) |
| p.add_argument("--enhance", action="store_true", help="rewrite the prompt with Qwen3.5 first") |
| p.add_argument("--model", default="sft", help="base | sft | <filename>") |
| p.add_argument("--lora", default="b21", help="b21 | r256 | none | <filename>") |
| p.add_argument("--strength", type=float, default=-0.5) |
| p.add_argument("--donor", default="sft", help="model whose exact weights are swapped in (with --swap-layers)") |
| p.add_argument("--swap-layers", help='layers to take from --donor, e.g. "21-41" or "3,4"; enables block swap') |
| p.add_argument("--swap-parts", default="all", help="all | comma list of qkv,o,ffn_up_gate,ffn_down,norm") |
| p.add_argument("--swap-other", action="store_true", help="also swap non-layer weights (embedders, modulation, head)") |
| p.add_argument("--swap-ratio", type=float, default=1.0) |
| p.add_argument("--probe", help="record per-layer stripe spectra to output/moss/probe/<probe>_s<seed>.npz") |
| p.add_argument("--window-steps", help='e.g. "5-12": apply --window-lora only for these sampling steps (1-indexed)') |
| p.add_argument("--window-lora", default="moss_image_sft_banding_fix_L37-41_bf16.safetensors") |
| p.add_argument("--window-strength", type=float, default=1.0) |
| p.add_argument("--lora2", help="second LoRA chained after the first, e.g. --lora r512 --strength 1 --lora2 b21 --strength2 -0.5") |
| p.add_argument("--strength2", type=float, default=1.0) |
| p.add_argument("--no-tsr", action="store_true") |
| p.add_argument("--no-slg", action="store_true") |
| p.add_argument("--sampler", default="er_sde", help="er_sde (workflow's SamplerER_SDE) or any KSampler sampler name") |
| p.add_argument("--steps", type=int, default=25) |
| p.add_argument("--cfg", type=float, default=3.0) |
| p.add_argument("--width", type=int, default=1024) |
| p.add_argument("--height", type=int, default=1024) |
| p.add_argument("--negative") |
| p.add_argument("--prefix", default="moss") |
| a = p.parse_args() |
|
|
| prompts = [" ".join(a.prompt)] if a.prompt else [] |
| if a.file: |
| prompts += [x.strip() for x in open(a.file).read().split("\n\n") if x.strip()] |
| if not prompts: |
| p.error("give a prompt or -f file") |
|
|
| base = json.load(open(os.path.join(HERE, "moss_image_t2i_api.json"))) |
| jobs = [] |
| for prompt in prompts: |
| for seed in a.seeds: |
| r = call("/prompt", {"prompt": build(base, prompt, seed, a)}) |
| if r.get("node_errors"): |
| sys.exit(json.dumps(r["node_errors"], indent=1)) |
| jobs.append((r["prompt_id"], prompt, seed)) |
| print(f"queued {len(jobs)}", flush=True) |
|
|
| for pid, prompt, seed in jobs: |
| while True: |
| h = call(f"/history/{pid}").get(pid) |
| if h and h.get("status", {}).get("completed") is not None: |
| break |
| time.sleep(1) |
| if h["status"].get("status_str") != "success": |
| print(f"FAILED seed {seed}: {json.dumps(h['status'].get('messages', [])[-1:])[:800]}") |
| continue |
| outs = h["outputs"] |
| for img in outs.get("10", {}).get("images", []): |
| print(os.path.normpath(os.path.join(OUT, img["subfolder"], img["filename"])), f"seed={seed}") |
| if "19" in outs: |
| print(" enhanced:", outs["19"]["text"][0][:1500]) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
