| """ |
| mosketch_pipeline.py — the full pipeline as one script with subcommands. |
| |
| Pipeline: |
| 1. Identify objects -> from the semantic file (no Qwen) |
| 2. Classify ARAP vs. not -> `classify` subcommand (one Qwen call, all objects) |
| 3. Narrate + deform -> `narrate` + `deform` subcommands, ONLY for |
| objects marked ARAP in step 2 |
| 4. Render -> `render` subcommand; every object gets real |
| trajectory translation; ARAP objects |
| additionally get deformation on top |
| |
| Run steps individually for debugging, or use `full` to run everything for |
| one sketch in one process — this loads the Qwen model ONCE and reuses it |
| across classify/narrate/deform, instead of loading it 3 separate times. |
| |
| Examples: |
| # step by step |
| python mosketch_pipeline.py classify --model M --caption-file C --sketch-name S --semantic SEM --out deformation.json |
| python mosketch_pipeline.py narrate --model M --caption-file C --sketch-name S --objects dog --out narratives.json |
| python mosketch_pipeline.py deform --model M --svg S.svg --semantic SEM --traj T --narratives narratives.json --deformation deformation.json --out-dir . |
| python mosketch_pipeline.py render --svg S.svg --semantic SEM --traj T --handles-dir . |
| |
| # everything at once, one model load |
| python mosketch_pipeline.py full --model M --caption-file C --sketch-name S --svg S.svg --semantic SEM --traj T --out-dir . |
| """ |
|
|
| import argparse |
| import json |
| import os |
| import re |
| import sys |
|
|
| import numpy as np |
|
|
| from lib import ( |
| load_strokes_from_svg, load_semantic_assignments, filter_strokes, flatten_strokes, |
| load_object, deduplicate_points, build_mesh, nearest_mesh_vertex, |
| auto_select_handles_deduped, object_bbox_size, arap_deform, |
| load_trajectories, bbox_deltas, get_caption, build_stroke_geometry_text, |
| ) |
|
|
| N_KEYFRAMES = 5 |
| MAX_RETRIES = 5 |
| LOW_UTILIZATION_THRESHOLD = 0.30 |
| |
| |
| PLAUSIBILITY_THRESHOLD = 4 |
| FAITHFULNESS_THRESHOLD = 4 |
| |
| QUALITY_THRESHOLD = 4 |
| |
|
|
|
|
| def faithfulness_passed(score): |
| """ |
| faithfulness_score can be a number 1-5, the string "N/A" (no caption |
| was given to compare against, so there's nothing to fail), or missing |
| entirely (treated as NOT passed — can't confirm it's good, so err |
| toward regenerating the narrative rather than assuming it's fine). |
| """ |
| if score is None: |
| return False |
| if isinstance(score, str): |
| return score.strip().upper() == "N/A" |
| if isinstance(score, (int, float)): |
| return score >= FAITHFULNESS_THRESHOLD |
| return False |
|
|
|
|
| def unload_model(model): |
| """Frees GPU memory before loading a different model. Necessary because |
| the narrate/deform steps use a text Qwen model and the judge step uses |
| a separate vision-language model (Qwen3-VL) — on hardware with limited |
| VRAM (this project's RTX A4000, 16GB, already documented as a tight |
| fit for a single model), loading both at once risks the same OOM issue |
| that blocked Wan2.2 integration earlier. Load/unload sequentially |
| instead of assuming both fit simultaneously. |
| |
| CONFIRMED BUG (found on real hardware, invisible to all mocked testing |
| since no real GPU was available to catch it): `del model` here only |
| clears THIS function's own local reference — it does nothing to the |
| caller's variable, which stays alive and keeps the whole model |
| resident in VRAM. torch.cuda.empty_cache() then has nothing to |
| actually free, because the refcount never reaches zero. Fixed by |
| returning None — callers MUST reassign their variable to this return |
| value (e.g. `model = unload_model(model)`), or the bug reappears.""" |
| import gc |
| del model |
| gc.collect() |
| try: |
| import torch |
| if torch.cuda.is_available(): |
| torch.cuda.empty_cache() |
| except ImportError: |
| pass |
| return None |
|
|
| |
| |
| |
| |
| DOG3_CAPTION = ("The person throws a frisbee through the air, and the dog sits poised, " |
| "ready to sprint forward and catch it with its mouth in a swift motion.") |
| DOG3_NARRATIVES = { |
| "dog": [ |
| "the dog is sitting alert, watching the frisbee as it is thrown", |
| "the dog is beginning to rise, weight shifting forward, head reaching toward the frisbee", |
| "the dog is mid-leap, body extended, reaching far forward and up toward the frisbee", |
| "the dog is at the peak of its jump, reaching as far as possible toward the frisbee", |
| "the dog is landing after catching the frisbee, body compacting back down", |
| ], |
| "frisbee": [ |
| "the frisbee has just left the thrower's hand, angled slightly upward", |
| "the frisbee is gliding through the air, tilting slightly as it arcs", |
| "the frisbee is near the peak of its arc, angled toward the dog", |
| "the frisbee is descending toward the dog, tilting down slightly", |
| "the frisbee is at the dog's mouth, being caught", |
| ], |
| } |
|
|
| SVG_PATH_DEFAULT = "/mnt/user-data/uploads/dog3.svg" |
| SEMANTIC_PATH_DEFAULT = "dog3_semantic.txt" |
| DEFAULT_COLOR = "#444444" |
| DEFAULT_LINEWIDTH = 1.1 |
| OBJECT_COLORS = {"dog": "black", "person": "#3F4C57", "frisbee": "#B0463C"} |
| OBJECT_LINEWIDTH = {"dog": 1.1, "person": 1.1, "frisbee": 1.4} |
|
|
|
|
| |
| |
| |
|
|
| def query_qwen(model, tokenizer, prompt, device, max_new_tokens=500, temperature=0.1): |
| import torch |
| messages = [{"role": "user", "content": prompt}] |
| text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) |
| inputs = tokenizer([text], return_tensors="pt").to(device) |
| input_token_count = inputs["input_ids"].shape[1] |
| with torch.no_grad(): |
| output_ids = model.generate(**inputs, max_new_tokens=max_new_tokens, |
| temperature=temperature, do_sample=True) |
| generated = output_ids[0][inputs["input_ids"].shape[1]:] |
| output_token_count = generated.shape[0] |
| response_text = tokenizer.decode(generated, skip_special_tokens=True) |
| return response_text, input_token_count, output_token_count |
|
|
|
|
| def load_qwen_model(model_path): |
| from transformers import AutoModelForCausalLM, AutoTokenizer |
| import torch |
| print(f"loading model from {model_path} ...") |
| tokenizer = AutoTokenizer.from_pretrained(model_path) |
| model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=torch.float16, device_map="auto") |
| device = next(model.parameters()).device |
| print("model loaded.") |
| return model, tokenizer, device |
|
|
|
|
| def parse_json_response(response_text): |
| match = re.search(r"\{.*\}", response_text, re.DOTALL) |
| if not match: |
| raise ValueError("No JSON object found in response:\n" + response_text) |
| return json.loads(match.group(0)) |
|
|
|
|
| |
| |
| |
|
|
| def build_deformation_prompt(caption, object_names): |
| objects_str = ", ".join(f'"{o}"' for o in object_names) |
| lines = "\n".join( |
| f' "{o}": "ARAP" or "TRAJ_ONLY"' + ("," if i < len(object_names) - 1 else "") |
| for i, o in enumerate(object_names) |
| ) |
| return f"""Scene: "{caption}" |
| |
| Objects in this scene: {objects_str} |
| |
| For each object, decide whether representing it correctly needs NON-RIGID DEFORMATION (its body/shape changes — e.g. limbs moving, a neck reaching, a body crouching or leaning) or whether simple RIGID TRANSLATION (the object moves/rotates as a whole, unchanged in shape, or doesn't move at all) is enough. |
| |
| Answer "ARAP" if the object's shape or body configuration changes at any point in the action, even if its overall position doesn't change. Answer "TRAJ_ONLY" if the object is rigid (a vehicle, tool, projectile, furniture, background element) or is simply carried by its own movement without changing shape. |
| |
| Respond with ONLY a JSON object, no other text, in this exact format: |
| {{ |
| {lines} |
| }} |
| """ |
|
|
|
|
| def validate_deformation(parsed, object_names): |
| problems = [] |
| for obj in object_names: |
| if obj not in parsed: |
| problems.append(f"'{obj}' missing from response") |
| continue |
| val = str(parsed[obj]).strip().upper() |
| if val not in ("ARAP", "TRAJ_ONLY"): |
| problems.append(f"'{obj}' has invalid value {parsed[obj]!r}, expected ARAP or TRAJ_ONLY") |
| return problems |
|
|
|
|
| def run_classify(model, tokenizer, device, caption, semantic_path, out_path): |
| assignments = load_semantic_assignments(semantic_path) |
| object_names = list(assignments.keys()) |
| print(f"objects found in {semantic_path}: {object_names}") |
|
|
| prompt = build_deformation_prompt(caption, object_names) |
| print("\n--- CLASSIFY PROMPT ---") |
| print(prompt) |
|
|
| response, in_tok, out_tok = query_qwen(model, tokenizer, prompt, device, |
| max_new_tokens=250, temperature=0.1) |
| print(f"\ntokens: {in_tok} in / {out_tok} out") |
| print("--- RAW RESPONSE ---") |
| print(response) |
|
|
| parsed = parse_json_response(response) |
| problems = validate_deformation(parsed, object_names) |
| print("\n--- PARSED ---") |
| print(json.dumps(parsed, indent=2)) |
| if problems: |
| print("--- VALIDATION PROBLEMS ---") |
| for p in problems: |
| print(f" - {p}") |
|
|
| arap_objs = [o for o in object_names if str(parsed.get(o, "")).strip().upper() == "ARAP"] |
| traj_only_objs = [o for o in object_names if o not in arap_objs] |
| print(f"\nARAP: {arap_objs}") |
| print(f"TRAJ_ONLY: {traj_only_objs}") |
|
|
| with open(out_path, "w") as f: |
| json.dump(parsed, f, indent=2) |
| print(f"wrote {out_path}") |
| return parsed, arap_objs |
|
|
|
|
| def build_objects_info(svg_path, semantic_path, arap_objects): |
| """ |
| Standalone version of the mesh/joint setup previously embedded inside |
| run_deform — factored out so the unified narrate+deform+judge retry |
| loop can build this ONCE before the loop (mesh/joints never change |
| between attempts) instead of recomputing it every attempt. |
| """ |
| objects_info = {} |
| for obj_name in arap_objects: |
| points, slices = load_object(obj_name, svg_path, semantic_path) |
| bbox_size = object_bbox_size(points) |
| unique_points, p2u = deduplicate_points(points, tol=0.35) |
| tri, edges = build_mesh(unique_points) |
| strokes = filter_strokes(load_strokes_from_svg(svg_path), |
| load_semantic_assignments(semantic_path)[obj_name]) |
| joints, anchor_idx, handle_idxs, joint_mesh_indices = auto_select_handles_deduped( |
| strokes, unique_points, k=4) |
|
|
| if len(handle_idxs) == 0: |
| print(f"WARNING: '{obj_name}' has no independent handles after dedup, skipping") |
| continue |
|
|
| print(f"'{obj_name}': {len(joints)} joints, anchor={anchor_idx}, handles={handle_idxs}, " |
| f"mesh_indices={joint_mesh_indices}, bbox_size={bbox_size:.1f}") |
|
|
| objects_info[obj_name] = { |
| "joints": joints, "anchor_idx": anchor_idx, "handle_idxs": handle_idxs, |
| "joint_mesh_indices": joint_mesh_indices, "bbox_size": bbox_size, "strokes": strokes, |
| "points": points, "slices": slices, "unique_points": unique_points, |
| "p2u": p2u, "edges": edges, |
| } |
| return objects_info |
|
|
|
|
| def build_deformed_stroke_geometry_text(object_info, kf_targets, n_points=2): |
| """ |
| Reconstructs what a previous attempt's ACTUAL DEFORMED shape looked |
| like, as compact text — used to replace the "attached image of the |
| previous attempt" memory in svg.py, which has no images at all. |
| |
| kf_targets: {joint_i_str: [x, y]} for ONE keyframe, in the same |
| "joint_i" key format apply_deform_clip_and_write saves (i.e. one |
| entry of deform_outputs[obj_name][kf_key]). |
| |
| Runs the SAME ARAP solve used for real rendering, then formats the |
| resulting deformed stroke points the same way build_stroke_geometry_text |
| formats the rest pose — so a retry sees the previous attempt's actual |
| resulting SHAPE in text, not just the isolated handle-joint numbers. |
| """ |
| import re as _re |
| unique_points = object_info["unique_points"] |
| edges = object_info["edges"] |
| p2u = object_info["p2u"] |
| points = object_info["points"] |
| slices = object_info["slices"] |
| anchor_idx = object_info["anchor_idx"] |
| joints = object_info["joints"] |
| joint_mesh_indices = object_info["joint_mesh_indices"] |
|
|
| handle_mesh_indices, handle_targets = [], [] |
| for key, target in kf_targets.items(): |
| m = _re.match(r"joint_(\d+)", key) |
| if not m: |
| continue |
| joint_i = int(m.group(1)) |
| if joint_i >= len(joint_mesh_indices): |
| continue |
| handle_mesh_indices.append(joint_mesh_indices[joint_i]) |
| handle_targets.append(target) |
| |
| handle_mesh_indices.append(joint_mesh_indices[anchor_idx]) |
| handle_targets.append(joints[anchor_idx].tolist()) |
|
|
| if not handle_mesh_indices: |
| return None |
|
|
| deformed_unique = arap_deform(unique_points, edges, handle_mesh_indices, |
| np.array(handle_targets), iterations=10) |
| deformed_points = deformed_unique[p2u] |
|
|
| lines = [] |
| for i, (start, end) in enumerate(slices): |
| stroke_pts = deformed_points[start:end] |
| idxs = np.linspace(0, len(stroke_pts) - 1, n_points).astype(int) |
| pts = stroke_pts[idxs] |
| pts_str = " -> ".join(f"({x:.0f},{y:.0f})" for x, y in pts) |
| lines.append(f" stroke_{i}: {pts_str}") |
| return "\n".join(lines) |
|
|
|
|
| def render_rest_pose_multi(object_names, svg_path, semantic_path, out_path): |
| """ |
| Renders multiple objects' ORIGINAL strokes together, at their real |
| positions in the source SVG (no deformation, no trajectory translation) |
| — this is the "what does the sketch actually look like" image every |
| attempt is grounded in, so Qwen can see what's actually drawable |
| (e.g. whether the dog's back legs even exist as strokes) instead of |
| only reasoning from the caption's text description. |
| """ |
| import matplotlib.pyplot as plt |
| fig, ax = plt.subplots(figsize=(6, 6)) |
| for name in object_names: |
| points, slices = load_object(name, svg_path, semantic_path) |
| for start, end in slices: |
| seg = points[start:end] |
| ax.plot(seg[:, 0], seg[:, 1], color="black", linewidth=1.2) |
| ax.invert_yaxis() |
| ax.set_aspect("equal") |
| ax.set_title("rest pose") |
| fig.savefig(out_path, dpi=150, bbox_inches="tight") |
| plt.close(fig) |
| return out_path |
|
|
|
|
| def build_combined_object_section(object_name, joints, anchor_idx, handle_idxs, bbox_size, |
| previous_narrative=None, feedback=None, n_keyframes=N_KEYFRAMES, |
| freeze_narrative=False, strokes=None, previous_shape_text=None): |
| """ |
| freeze_narrative: if True, previous_narrative is used as a FIXED target |
| pose description (the object's narrative already passed faithfulness — |
| only the numeric deformation needs to improve, not the story). If |
| False (default), the narrative is regenerated fresh, informed by |
| feedback (and, in the image-based pipeline variants, the attached |
| image(s) — this variant, svg.py, has NO images at all, relying |
| entirely on text: stroke geometry + previous numeric targets). |
| |
| strokes: if given, the object's actual stroke points are included as |
| TEXT (not just the rendered image) — added specifically because |
| multimodal LLMs can under-attend to image content relative to text; |
| this gives the same geometric information in a text-native form the |
| model is more likely to actually use. Kept deliberately sparse (2 |
| points per stroke, start+end only) since a complex object can have |
| 60+ strokes — measured on real dog3 data: 2 points/stroke costs |
| ~570 tokens for a 64-stroke object vs ~2650 for the full 12 |
| points/stroke used internally for the ARAP mesh. |
| |
| previous_shape_text: pre-formatted text (from build_deformed_stroke_geometry_text, |
| computed one level up where the full mesh/ARAP data is available) showing |
| the previous attempt's ACTUAL DEFORMED STROKE positions at 2 representative |
| keyframes — not raw joint numbers, the reconstructed resulting SHAPE, in the |
| same stroke-point text format as the rest pose. This is svg.py's replacement |
| for showing Qwen the previous attempt's rendered images — since there are no |
| images here at all, without this the model has zero way to know what its own |
| previous output actually looked like on a retry. |
| """ |
| cap = round(bbox_size * 0.25, 1) |
| joint_lines = "\n".join( |
| f' - joint_{i}: rest position (x={joints[i][0]:.1f}, y={joints[i][1]:.1f})' |
| + (" <-- ANCHOR, must stay at or near this position in EVERY keyframe" if i == anchor_idx else "") |
| for i in range(len(joints)) |
| ) |
| handle_list_str = ', joint_'.join(str(i) for i in handle_idxs) |
|
|
| geometry_block = "" |
| if strokes: |
| geometry_text = build_stroke_geometry_text(strokes, n_points=2) |
| geometry_block = ( |
| f"\n This object's ACTUAL drawn strokes (start -> end point of each stroke, same coordinate " |
| f"space as the joints above) — use this to know exactly what is and isn't actually drawn, don't " |
| f"invent motion for parts that have no strokes here:\n{geometry_text}\n" |
| ) |
|
|
| previous_targets_block = "" |
| if previous_shape_text: |
| previous_targets_block = ( |
| f"\n Your previous attempt's ACTUAL DEFORMED SHAPE at two representative keyframes " |
| f"(reconstructed from your own previous target positions — same stroke-point format as " |
| f"the rest pose above, so you can directly compare what changed):\n{previous_shape_text}\n" |
| ) |
|
|
| if freeze_narrative and previous_narrative: |
| pose_lines = "\n".join(f" kf{i}: {desc}" for i, desc in enumerate(previous_narrative)) |
| feedback_line = f'\n This pose story already matches the intended action — it is FIXED, do not change it. ' \ |
| f'Only the numeric target positions need to improve.' \ |
| + (f' Previous attempt was judged: "{feedback}"' if feedback else "") + \ |
| "\n Adjust specifically from your previous target positions shown above where the " \ |
| "critique applies — this is a revision of known values, not a fresh guess." |
| return f"""Object: "{object_name}" |
| Joints: |
| {joint_lines} |
| {geometry_block} |
| Target pose across all {n_keyframes} keyframes (FIXED, already correct — do not rewrite): |
| {pose_lines} |
| {feedback_line} |
| {previous_targets_block} |
| For non-anchor joints (joint_{handle_list_str}), do not move more than {cap} pixels from REST in any keyframe. IMPORTANT: these {n_keyframes} keyframes are SPARSE anchor points spanning the ENTIRE action, NOT consecutive video frames — a large, dramatic difference between consecutive keyframes is NORMAL and EXPECTED, not an error; the actual in-between motion will be generated separately later by a different model. Positions should progress in a DIRECTIONALLY COHERENT way (don't make real progress toward the action and then have a LATER keyframe randomly revert backward without the narrative describing a reason to — e.g. only "landing"/"settling" should move back toward rest). Small, timid, barely-different positions between keyframes are themselves a mistake, not a safe choice.""" |
|
|
| previous_block = "" |
| if previous_narrative or feedback: |
| parts = [] |
| if previous_narrative: |
| parts.append(f"Your previous narrative attempt was:\n{json.dumps(previous_narrative, indent=2)}") |
| if feedback: |
| parts.append(f'That attempt was judged and received this critique: "{feedback}"') |
| parts.append("Revise BOTH the narrative and the target positions to address the critique — don't just " |
| "reword the narrative superficially while leaving the same underlying numeric problem.") |
| previous_block = "\n " + "\n ".join(parts) + "\n" |
| previous_block += previous_targets_block |
|
|
| return f"""Object: "{object_name}" |
| Joints: |
| {joint_lines} |
| {geometry_block} |
| {previous_block} |
| For non-anchor joints (joint_{handle_list_str}), do not move more than {cap} pixels from REST in any keyframe. IMPORTANT: these {n_keyframes} keyframes are SPARSE anchor points spanning the ENTIRE action, NOT consecutive video frames — a large, dramatic difference between consecutive keyframes is NORMAL and EXPECTED, not an error; the actual in-between motion will be generated separately later by a different model. Positions should progress in a DIRECTIONALLY COHERENT way (don't make real progress toward the action and then have a LATER keyframe randomly revert backward without the narrative describing a reason to — e.g. only "landing"/"settling" should move back toward rest). Small, timid, barely-different positions between keyframes are themselves a mistake, not a safe choice.""" |
|
|
|
|
| DOG3_COMBINED_FEWSHOT_EXAMPLE = { |
| "dog": { |
| "narrative": [ |
| "the dog is sitting alert, watching the frisbee as it is thrown", |
| "the dog is beginning to rise, weight shifting forward, head reaching toward the frisbee", |
| "the dog is mid-leap, body extended, reaching far forward and up toward the frisbee", |
| "the dog is at the peak of its jump, reaching as far as possible toward the frisbee", |
| "the dog is landing after catching the frisbee, body compacting back down", |
| ], |
| |
| |
| |
| |
| |
| |
| |
| "targets": { |
| "kf0": {"joint_1": [178.3, 136.8], "joint_2": [228.8, 194.9], "joint_3": [189.3, 151.3]}, |
| "kf1": {"joint_1": [168.0, 127.0], "joint_2": [232.0, 191.0], "joint_3": [184.0, 144.0]}, |
| "kf2": {"joint_1": [160.0, 120.0], "joint_2": [237.0, 186.0], "joint_3": [177.0, 137.0]}, |
| "kf3": {"joint_1": [159.0, 119.0], "joint_2": [240.0, 183.0], "joint_3": [174.0, 134.0]}, |
| "kf4": {"joint_1": [168.0, 128.0], "joint_2": [231.0, 192.0], "joint_3": [185.0, 146.0]}, |
| }, |
| } |
| } |
|
|
|
|
| def build_combined_narrate_deform_prompt(objects_info, caption, previous_narratives=None, feedback=None, |
| n_keyframes=N_KEYFRAMES, is_retry=False, few_shot=True, |
| freeze_narrative=False, previous_deform_outputs=None): |
| |
| |
| |
| |
| |
| SHAPE_MEMORY_KEYFRAMES = [2, 4] |
|
|
| sections, example_parts = [], [] |
| for name, info in objects_info.items(): |
| prev_narrative_for_obj = (previous_narratives or {}).get(name) |
|
|
| previous_shape_text = None |
| prev_deform_for_obj = (previous_deform_outputs or {}).get(name) |
| if prev_deform_for_obj: |
| shape_parts = [] |
| for kf in SHAPE_MEMORY_KEYFRAMES: |
| kf_key = f"kf{kf}" |
| if kf_key not in prev_deform_for_obj: |
| continue |
| shape_text = build_deformed_stroke_geometry_text(info, prev_deform_for_obj[kf_key], n_points=2) |
| if shape_text: |
| shape_parts.append(f" -- {kf_key} --\n{shape_text}") |
| if shape_parts: |
| previous_shape_text = "\n".join(shape_parts) |
|
|
| sections.append(build_combined_object_section( |
| name, info["joints"], info["anchor_idx"], info["handle_idxs"], info["bbox_size"], |
| previous_narrative=prev_narrative_for_obj, feedback=feedback, n_keyframes=n_keyframes, |
| freeze_narrative=freeze_narrative, strokes=info.get("strokes"), |
| previous_shape_text=previous_shape_text, |
| )) |
| kf_examples = ",\n".join( |
| " \"kf%d\": {%s}" % (kf, ", ".join(f'"joint_{i}": [x, y]' for i in info["handle_idxs"])) |
| for kf in range(n_keyframes) |
| ) |
| if freeze_narrative: |
| example_parts.append( |
| f' "{name}": {{\n' |
| f' "targets": {{\n{kf_examples}\n }}\n' |
| f' }}' |
| ) |
| else: |
| example_parts.append( |
| f' "{name}": {{\n' |
| f' "narrative": [<{n_keyframes} short pose description strings, one per keyframe>],\n' |
| f' "targets": {{\n{kf_examples}\n }}\n' |
| f' }}' |
| ) |
|
|
| all_sections = "\n\n".join(sections) |
| example_json = "{\n" + ",\n".join(example_parts) + "\n}" |
|
|
| image_context = "" |
| if is_retry: |
| image_context = ("This is a RETRY. There are no images — everything you need is given as TEXT below: " |
| "each object's rest-pose stroke geometry, and (where available) your previous " |
| "attempt's actual deformed shape at two representative keyframes.") |
| else: |
| image_context = ("There are no images in this prompt — the object's rest-pose stroke geometry is " |
| "given below as TEXT. Use it to understand exactly what strokes exist and are " |
| "available to move; do not invent motion for body parts that aren't actually drawn.") |
|
|
| fewshot_block = "" |
| if few_shot: |
| fewshot_json = json.dumps(DOG3_COMBINED_FEWSHOT_EXAMPLE, indent=2) |
| fewshot_block = f"""Example — for the scene "{DOG3_CAPTION}", a good answer looks like: |
| {fewshot_json} |
| |
| Notice: each narrative keyframe reads as a distinct, substantially different stage of the action — not a near-duplicate of its neighbor, and not a small incremental change from it. The joint targets BUILD UP smoothly (kf0 -> kf1 -> kf2 -> kf3 each moving further than the last) and only settle back toward rest at the FINAL keyframe, matching the narrative's "landing" moment — no keyframe overshoots and then has a later keyframe revert back toward rest without a narrative reason to. Match this style and this kind of numeric consistency for the new scene below. |
| |
| """ |
|
|
| if freeze_narrative: |
| output_instruction = ( |
| 'For EACH object above, the narrative/pose story is already fixed (shown above) — ' |
| 'produce ONLY:\n' |
| ' "targets": target (x, y) positions for its non-anchor joints, at every keyframe, ' |
| 'consistent with the fixed pose story above.' |
| ) |
| else: |
| output_instruction = ( |
| "For EACH object above, produce BOTH:\n" |
| ' 1. "narrative": a plain-English pose description for each keyframe. REMEMBER: these are ' |
| "SPARSE keyframes spanning the WHOLE action, not consecutive video frames — each description " |
| "should be a meaningfully, substantially different stage of the action from its neighbors, not " |
| "a small incremental change. Write these like 5 distinct captions for 5 different moments spread " |
| "across an entire action, not like 5 near-duplicate snapshots a split-second apart. Under 20 " |
| "words each.\n" |
| ' 2. "targets": target (x, y) positions for its non-anchor joints, at every keyframe, ' |
| "consistent with your own narrative." |
| ) |
|
|
| return f"""{fewshot_block}You are directing a {n_keyframes}-keyframe animated sequence for a hand-drawn sketch, viewed from the side. Coordinate system: x increases rightward, y increases DOWNWARD. |
| |
| IMPORTANT: these {n_keyframes} keyframes are SPARSE anchor points sampled across the ENTIRE action from start to finish — NOT consecutive video frames. Think of them like 5 widely-spaced snapshots of a whole motion, not neighboring frames a fraction of a second apart. Large, dramatic pose changes between consecutive keyframes are normal and expected; a separate model will generate the actual in-between motion frames later. Do not treat these like near-continuous animation frames. |
| |
| Scene: "{caption}" |
| |
| {image_context} |
| |
| {all_sections} |
| |
| {output_instruction} |
| |
| Consider objects together (e.g. a dog reaching toward a frisbee should be spatially consistent with the frisbee's own position) and consider each object's OWN sequence together — these {n_keyframes} keyframes are SPARSE anchor points spanning the WHOLE action, not consecutive video frames, so large differences between consecutive keyframes are expected and correct, not something to avoid. The many actual in-between motion frames will be generated separately later. Only avoid a keyframe making real progress and then a LATER keyframe randomly reverting backward without the narrative describing why. |
| |
| Respond with ONLY one JSON object, no other text, in this exact format: |
| {example_json} |
| """ |
|
|
|
|
| def run_combined_narrate_deform(model, processor, images, prompt): |
| """ |
| Same multi-image calling pattern as vlm_judge.run_judge, but svg.py |
| calls this with images=[] (or None) on every call — no images at |
| all, text-only. UNCONFIRMED against a real model: whether |
| processor(images=[]) behaves identically to omitting the images kwarg |
| entirely. Handled defensively here rather than assumed, since this |
| can't be tested without real hardware. |
| """ |
| import torch |
| has_images = bool(images) |
| content = [{"type": "image", "image": img} for img in images] if has_images else [] |
| content.append({"type": "text", "text": prompt}) |
| messages = [{"role": "user", "content": content}] |
|
|
| text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) |
| if has_images: |
| inputs = processor(text=[text], images=images, return_tensors="pt").to(model.device) |
| else: |
| inputs = processor(text=[text], return_tensors="pt").to(model.device) |
|
|
| with torch.no_grad(): |
| output_ids = model.generate(**inputs, max_new_tokens=500 * N_KEYFRAMES, temperature=0.1, do_sample=True) |
|
|
| generated = output_ids[:, inputs["input_ids"].shape[1]:] |
| response = processor.batch_decode(generated, skip_special_tokens=True)[0] |
| return response |
|
|
|
|
| def build_svg_judge_prompt(sketch_name, all_keyframes_geometry_text, caption=None, |
| dino_stagnation=None, dino_temporal=None, clip_scores=None, |
| cap_utilization=None, n_keyframes=N_KEYFRAMES): |
| """ |
| Text-only judge prompt — reconstructs each keyframe's ACTUAL deformed |
| stroke geometry (same function used for retry memory) and gives ALL |
| 5 keyframes to the judge as text, instead of showing it the rendered |
| images. Kept separate from vlm_judge.build_judge_prompt (which stays |
| image-based) so pipe.py/pipeline1.py are never affected by this. |
| |
| HONEST CAVEAT, not glossed over: the QUALITY criterion below ("does |
| this look like a clean line drawing vs garbled/noisy") is inherently |
| a pixel-level, visual question. Assessing it from coordinate text |
| alone is a fundamentally harder, more indirect task than looking at |
| the actual rendering — this is the part of the experiment most |
| likely to perform worse than the image-based judge, not a solved |
| problem. |
| """ |
| caption_block = "" |
| if caption: |
| caption_block = f'\nThe sketch is supposed to depict: "{caption}"\n' |
|
|
| metrics_lines = [] |
| if dino_stagnation is not None: |
| sims = ", ".join(f"kf{i}={s:.3f}" for i, s in enumerate(dino_stagnation["per_keyframe_similarity"])) |
| metrics_lines.append(f"- DINOv2 similarity to the PREVIOUS attempt's RENDERED images, per keyframe " |
| f"(1.0 = identical): {sims} (mean {dino_stagnation['mean_similarity']:.3f})") |
| if dino_temporal is not None: |
| temp_str = ", ".join(f"kf{i}->kf{i+1}={s:.3f}" for i, s in enumerate(dino_temporal)) |
| metrics_lines.append(f"- DINOv2 similarity between CONSECUTIVE keyframes' RENDERED images: {temp_str}") |
| if cap_utilization is not None: |
| util_str = ", ".join(f"{obj}={frac*100:.0f}%" for obj, frac in cap_utilization.items()) |
| metrics_lines.append(f"- Movement allowance used, per object: {util_str}") |
| if clip_scores is not None: |
| clip_str = ", ".join(f"kf{i}={s:.3f}" for i, s in enumerate(clip_scores["per_keyframe_clip_score"])) |
| metrics_lines.append(f"- CLIP image-caption similarity, per keyframe (from the RENDERED images): " |
| f"{clip_str} (mean {clip_scores['mean_clip_score']:.3f})") |
| metrics_block = "" |
| if metrics_lines: |
| metrics_block = ("\nObjective measurements computed from the actual rendered images (these ARE " |
| "image-derived even though you are not shown the images directly — use them to " |
| "inform your reasoning):\n" + "\n".join(metrics_lines) + "\n") |
|
|
| return f"""You are judging a sequence of {n_keyframes} keyframes generated for an animated sketch named "{sketch_name}". You are NOT shown images — instead, each keyframe's actual deformed stroke geometry is given below as TEXT (start -> end point of each stroke, in pixel coordinates, y increases DOWNWARD). |
| |
| IMPORTANT: these are SPARSE keyframes sampled across the ENTIRE action from start to finish — NOT consecutive video frames. Large, dramatic differences between consecutive keyframes are normal and correct. |
| {caption_block}{metrics_block} |
| {all_keyframes_geometry_text} |
| |
| Evaluate the sequence on these three criteria, reasoning from the coordinate data above: |
| |
| 1. PLAUSIBILITY (1-5): Does the sequence of stroke positions, taken as sparse waypoints across the whole action, describe a coherent and physically believable progression? 5 = each keyframe is a sensible, meaningfully different stage of the action, in a believable order. 1 = static (coordinates barely change between keyframes) or the positions describe something physically impossible/nonsensical. |
| |
| 2. FAITHFULNESS (1-5, or "N/A" if no caption was given above): Does the described motion match what the caption says should be happening? 5 = clearly matches, 1 = unrelated to the description. |
| |
| 3. QUALITY (1-5): Judging ONLY from the coordinate data (you cannot see the actual rendered image) — do the stroke positions look geometrically coherent, e.g. no wildly self-intersecting or degenerate configurations that would likely render as visual noise? 5 = coordinates look clean and coherent, 1 = coordinates suggest severe geometric distortion. NOTE: this criterion is fundamentally harder to assess without seeing the actual rendering — be appropriately uncertain rather than overconfident. |
| |
| Respond with ONLY a JSON object, no other text, in this exact format: |
| {{ |
| "plausibility_score": <1-5>, |
| "plausibility_notes": "<one sentence>", |
| "faithfulness_score": <1-5 or "N/A">, |
| "faithfulness_notes": "<one sentence>", |
| "quality_score": <1-5>, |
| "quality_notes": "<one sentence>", |
| "overall_verdict": "<one or two sentence summary>" |
| }} |
| """ |
|
|
|
|
| def run_svg_judge(model, processor, prompt): |
| """Text-only judge call — same no-images pattern as run_combined_narrate_deform.""" |
| import torch |
| messages = [{"role": "user", "content": [{"type": "text", "text": prompt}]}] |
| text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) |
| inputs = processor(text=[text], return_tensors="pt").to(model.device) |
|
|
| with torch.no_grad(): |
| output_ids = model.generate(**inputs, max_new_tokens=400, temperature=0.2, do_sample=True) |
|
|
| generated = output_ids[:, inputs["input_ids"].shape[1]:] |
| response = processor.batch_decode(generated, skip_special_tokens=True)[0] |
| return response |
|
|
|
|
| def build_all_keyframes_geometry_text(objects_info, deform_outputs, n_keyframes=N_KEYFRAMES): |
| """ |
| For EVERY object and EVERY keyframe of the CURRENT attempt, reconstruct |
| the actual deformed stroke geometry as text — unlike the 2-keyframe |
| retry memory, judging needs the FULL sequence to assess the whole |
| story, not just a sample. Real, measured cost: ~2,866 tokens for ONE |
| 64-stroke object across all 5 keyframes — this multiplies per object, |
| so multi-object scenes get expensive fast. No trimming applied here; |
| if this needs to be cut down for a specific scene, that's the next |
| thing to adjust. |
| """ |
| sections = [] |
| for obj_name, info in objects_info.items(): |
| obj_targets = deform_outputs.get(obj_name) |
| if not obj_targets: |
| continue |
| kf_parts = [] |
| for kf in range(n_keyframes): |
| kf_key = f"kf{kf}" |
| if kf_key not in obj_targets: |
| continue |
| shape_text = build_deformed_stroke_geometry_text(info, obj_targets[kf_key], n_points=2) |
| if shape_text: |
| kf_parts.append(f" -- {kf_key} --\n{shape_text}") |
| if kf_parts: |
| sections.append(f'Object "{obj_name}":\n' + "\n".join(kf_parts)) |
| return "\n\n".join(sections) |
|
|
|
|
| def apply_deform_clip_and_write(parsed, objects_info, out_dir, sketch_name, n_keyframes=N_KEYFRAMES, |
| frozen_narratives=None): |
| """ |
| Shared post-processing for the combined call's "targets" section: |
| same hard-clip logic as the old run_deform, applied here instead. |
| Returns (narratives_dict, deform_outputs_dict, cap_utilization_dict). |
| |
| frozen_narratives: {obj_name: [...]} — used as a fallback when the |
| response doesn't include a "narrative" key for an object, which |
| happens when freeze_narrative=True was used in the prompt (the model |
| was never asked to produce one, so its absence is expected, not an |
| error — carry the frozen one forward instead of losing it). |
| |
| cap_utilization: {obj_name: mean_fraction_of_cap_used} — CONFIRMED on |
| real hardware (horsecar5's person) that Qwen can propose displacement |
| well within the movement cap without ever being told it did so — the |
| handle selection and clipping were both working correctly, but the |
| actual output was too timid to be visible (e.g. a leg moving only 18% |
| of its allowed range). The hard clip only ever catches OVER the cap; |
| nothing previously caught UNDER-using it. This surfaces that as an |
| explicit number so it can be fed back to Qwen directly. |
| """ |
| narratives_out = {} |
| deform_outputs = {} |
| utilization_by_obj = {} |
|
|
| for obj_name, info in objects_info.items(): |
| if obj_name not in parsed: |
| print(f" WARNING: '{obj_name}' missing from response entirely, skipping") |
| continue |
| obj_result = parsed[obj_name] |
| utilization_by_obj[obj_name] = [] |
|
|
| if "narrative" in obj_result: |
| narratives_out[obj_name] = obj_result["narrative"] |
| elif frozen_narratives and obj_name in frozen_narratives: |
| narratives_out[obj_name] = frozen_narratives[obj_name] |
| else: |
| print(f" WARNING: '{obj_name}' has no narrative in response and no frozen narrative " |
| f"to fall back to — narratives.json will be missing this object") |
|
|
| deform_outputs[obj_name] = {} |
| targets = obj_result.get("targets", {}) |
| for kf in range(n_keyframes): |
| kf_key = f"kf{kf}" |
| if kf_key not in targets: |
| print(f" WARNING: '{obj_name}' missing {kf_key} targets, skipping this frame") |
| continue |
| kf_result = targets[kf_key] |
| out = {} |
| joint_targets_this_kf = {} |
| for name, target in kf_result.items(): |
| m = re.match(r"joint_(\d+)", name) |
| if not m: |
| print(f" WARNING: unexpected key '{name}' for '{obj_name}' {kf_key}, skipping") |
| continue |
| joint_i = int(m.group(1)) |
| if joint_i >= len(info["joint_mesh_indices"]): |
| print(f" WARNING: '{obj_name}' joint_{joint_i} out of range, skipping") |
| continue |
| mesh_idx = info["joint_mesh_indices"][joint_i] |
|
|
| rest = np.array(info["joints"][joint_i]) |
| cap = round(info["bbox_size"] * 0.25, 1) |
| target_arr = np.array(target, dtype=float) |
| disp = target_arr - rest |
| dist = np.linalg.norm(disp) |
| utilization_by_obj[obj_name].append(min(dist / cap, 1.0) if cap > 0 else 0.0) |
| if dist > cap: |
| clipped = rest + disp / dist * cap |
| print(f" CLIPPED '{obj_name}' {kf_key} joint_{joint_i}: requested {dist:.1f}px " |
| f"(cap {cap}px) -> clipped to {cap}px, direction preserved") |
| target = clipped.tolist() |
|
|
| out[str(mesh_idx)] = target |
| joint_targets_this_kf[f"joint_{joint_i}"] = target |
| anchor_mesh_idx = info["joint_mesh_indices"][info["anchor_idx"]] |
| out[str(anchor_mesh_idx)] = info["joints"][info["anchor_idx"]].tolist() |
| deform_outputs[obj_name][kf_key] = joint_targets_this_kf |
|
|
| out_path = os.path.join(out_dir, f"qwen_{sketch_name}_{obj_name}_kf{kf}.json") |
| with open(out_path, "w") as f: |
| json.dump(out, f, indent=2) |
| print(f" wrote {out_path}") |
|
|
| cap_utilization = {} |
| for obj_name, fractions in utilization_by_obj.items(): |
| if fractions: |
| mean_frac = sum(fractions) / len(fractions) |
| cap_utilization[obj_name] = mean_frac |
| print(f" '{obj_name}': mean cap utilization = {mean_frac*100:.0f}% " |
| f"(across {len(fractions)} joint-keyframe pairs)") |
|
|
| return narratives_out, deform_outputs, cap_utilization |
| |
| |
| |
|
|
| def run_render(handles_dir, svg_path, semantic_path, traj_path, |
| out_path=None, frames_dir=None): |
| """ |
| out_path: if given, ALSO saves the combined strip image (all 5 keyframes |
| side by side) here, outside frames_dir. Optional — pass None |
| to keep output confined to frames_dir only. |
| frames_dir: if given, saves each keyframe as its own individual PNG |
| (kf0.png ... kf4.png) plus the combined strip, named after |
| the sketch itself ({sketch_name}.png), all inside this one |
| folder. |
| """ |
| import matplotlib.pyplot as plt |
|
|
| sketch_name = os.path.splitext(os.path.basename(svg_path))[0] |
| real_trajectories = load_trajectories(traj_path) |
| object_names = list(real_trajectories.keys()) |
|
|
| object_data = {} |
| for name in object_names: |
| points, slices = load_object(name, svg_path, semantic_path) |
| dx_vals, dy_vals = bbox_deltas(real_trajectories[name]) |
| unique_points, p2u = deduplicate_points(points, tol=0.35) |
| tri, edges = build_mesh(unique_points) |
| object_data[name] = { |
| "points": points, "slices": slices, "dx": dx_vals, "dy": dy_vals, |
| "unique_points": unique_points, "p2u": p2u, "edges": edges, |
| } |
|
|
| if frames_dir: |
| os.makedirs(frames_dir, exist_ok=True) |
|
|
| xmin, xmax, ymin, ymax = 0, 260, 60, 230 |
|
|
| |
| |
| keyframe_lines = [] |
| for kf in range(N_KEYFRAMES): |
| lines_this_kf = {} |
| for name in object_names: |
| od = object_data[name] |
| handles_path = os.path.join(handles_dir, f"qwen_{sketch_name}_{name}_kf{kf}.json") |
|
|
| if os.path.exists(handles_path): |
| with open(handles_path) as f: |
| spec = json.load(f) |
| handle_indices = [int(k) for k in spec.keys()] |
| handle_targets = np.array([spec[k] for k in spec.keys()]) |
| n_verts = len(od["unique_points"]) |
| bad = [i for i in handle_indices if i >= n_verts] |
| if bad: |
| print(f" ERROR: {handles_path} has out-of-bounds indices {bad} for '{name}' " |
| f"({n_verts} mesh vertices) — likely from a DIFFERENT sketch's mesh. " |
| f"Falling back to translation-only.") |
| deformed_points = od["points"] |
| mode = "translation-only (handles file failed validation)" |
| else: |
| deformed_unique = arap_deform(od["unique_points"], od["edges"], |
| handle_indices, handle_targets, iterations=10) |
| deformed_points = deformed_unique[od["p2u"]] |
| mode = "ARAP" |
| else: |
| deformed_points = od["points"] |
| mode = "translation-only (no handles file found)" |
|
|
| moved = deformed_points + np.array([od["dx"][kf], od["dy"][kf]]) |
| lines_this_kf[name] = [moved[start:end] for start, end in od["slices"]] |
|
|
| print(f"kf{kf} '{name}': {mode}, points_after_move_range=" |
| f"x[{moved[:,0].min():.1f},{moved[:,0].max():.1f}] " |
| f"y[{moved[:,1].min():.1f},{moved[:,1].max():.1f}]") |
|
|
| keyframe_lines.append(lines_this_kf) |
|
|
| if frames_dir: |
| fig_i, ax_i = plt.subplots(figsize=(6, 5.5)) |
| for name, segs in lines_this_kf.items(): |
| for seg in segs: |
| ax_i.plot(seg[:, 0], seg[:, 1], |
| color=OBJECT_COLORS.get(name, DEFAULT_COLOR), |
| linewidth=OBJECT_LINEWIDTH.get(name, DEFAULT_LINEWIDTH)) |
| ax_i.set_xlim(xmin, xmax) |
| ax_i.set_ylim(ymax, ymin) |
| ax_i.set_aspect("equal") |
| ax_i.set_title(f"{sketch_name} — kf{kf}", fontsize=12, fontweight="bold") |
| frame_path = os.path.join(frames_dir, f"kf{kf}.png") |
| fig_i.savefig(frame_path, dpi=140, bbox_inches="tight") |
| plt.close(fig_i) |
| print(f" wrote {frame_path}") |
|
|
| |
| fig, axes = plt.subplots(1, N_KEYFRAMES, figsize=(24, 5)) |
| for kf in range(N_KEYFRAMES): |
| ax = axes[kf] |
| for name, segs in keyframe_lines[kf].items(): |
| for seg in segs: |
| ax.plot(seg[:, 0], seg[:, 1], |
| color=OBJECT_COLORS.get(name, DEFAULT_COLOR), |
| linewidth=OBJECT_LINEWIDTH.get(name, DEFAULT_LINEWIDTH)) |
| ax.set_xlim(xmin, xmax) |
| ax.set_ylim(ymax, ymin) |
| ax.set_aspect("equal") |
| ax.set_title(f"kf{kf}", fontsize=13, fontweight="bold") |
|
|
| plt.tight_layout() |
|
|
| if out_path: |
| plt.savefig(out_path, dpi=140, bbox_inches="tight") |
| print(f"wrote {out_path}") |
|
|
| if frames_dir: |
| combined_frame_path = os.path.join(frames_dir, f"{sketch_name}.png") |
| plt.savefig(combined_frame_path, dpi=140, bbox_inches="tight") |
| print(f"wrote {combined_frame_path}") |
|
|
| if not out_path and not frames_dir: |
| print("WARNING: neither out_path nor frames_dir given, combined strip image not saved anywhere") |
|
|
| plt.close(fig) |
|
|
|
|
| |
| |
| |
| |
| |
| |
|
|
| |
| |
| |
| SVG_DIR_DEFAULT = "/user/HS400/rk01499/my_scratch/sketch/data/raw/60sketches/svg" |
| PROCESSED_DIR_DEFAULT = "/user/HS400/rk01499/my_scratch/sketch/data/processed" |
| CAPTION_FILE_DEFAULT = "/user/HS400/rk01499/my_scratch/sketch/data/raw/60sketches/caption.txt" |
| MODEL_PATH_DEFAULT = "/user/HS400/rk01499/my_scratch/models/qwen2.5-7b/" |
|
|
|
|
| def resolve_sketch_paths(sketch, svg_dir, processed_dir, caption_file): |
| """ |
| sketch: either a bare sketch name ("dog9") or a path to its SVG |
| ("/path/to/dog9.svg") — either way, everything else (semantic, |
| traj, caption) is derived from the same naming convention used |
| across this dataset: {name}.svg, {name}/{name}_semantic.txt, |
| {name}/{name}_traj.txt, and a lookup in one shared caption.txt. |
| """ |
| name = os.path.splitext(os.path.basename(sketch))[0] |
| svg_path = sketch if sketch.endswith(".svg") else os.path.join(svg_dir, f"{name}.svg") |
| semantic_path = os.path.join(processed_dir, name, f"{name}_semantic.txt") |
| traj_path = os.path.join(processed_dir, name, f"{name}_traj.txt") |
|
|
| missing = [p for p in [svg_path, semantic_path, traj_path, caption_file] if not os.path.exists(p)] |
| if missing: |
| raise SystemExit( |
| f"Could not find these expected files for sketch '{name}':\n " + |
| "\n ".join(missing) + |
| "\n\nIf your directory layout differs from the default, pass --svg-dir / " |
| "--processed-dir / --caption-file explicitly." |
| ) |
|
|
| caption = get_caption(caption_file, name) |
| return name, svg_path, semantic_path, traj_path, caption |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser( |
| description="Run the full sketch deformation pipeline for one image. " |
| "The only required input is the sketch — everything else " |
| "(semantic assignments, trajectory, caption) is looked up " |
| "automatically from the standard dataset layout.") |
| ap.add_argument("sketch", type=str, |
| help="sketch name (e.g. 'dog9') or path to its .svg file") |
| ap.add_argument("--model", type=str, default=MODEL_PATH_DEFAULT) |
| ap.add_argument("--svg-dir", type=str, default=SVG_DIR_DEFAULT) |
| ap.add_argument("--processed-dir", type=str, default=PROCESSED_DIR_DEFAULT) |
| ap.add_argument("--caption-file", type=str, default=CAPTION_FILE_DEFAULT) |
| ap.add_argument("--out-dir", type=str, default=".") |
| ap.add_argument("--no-fewshot", action="store_true", |
| help="disable the dog3 few-shot example in narrate (for A/B comparison)") |
| ap.add_argument("--deform-only", action="store_true", |
| help="run classify + ONE narrate+deform attempt + render, then STOP — no judge, " |
| "no retries. Prints cap utilization directly and saves the render, so you can " |
| "inspect raw generation quality without the judge's assessment as a confound.") |
| args = ap.parse_args() |
|
|
| name, svg_path, semantic_path, traj_path, caption = resolve_sketch_paths( |
| args.sketch, args.svg_dir, args.processed_dir, args.caption_file) |
| print(f"sketch: {name}") |
| print(f" svg: {svg_path}") |
| print(f" semantic: {semantic_path}") |
| print(f" traj: {traj_path}") |
| print(f" caption: {caption!r}") |
|
|
| os.makedirs(args.out_dir, exist_ok=True) |
| json_dir = os.path.join(args.out_dir, "json", name) |
| os.makedirs(json_dir, exist_ok=True) |
| print(f" json output dir: {json_dir}") |
|
|
| print("\n########## STEP 1: CLASSIFY ##########") |
| classify_model, classify_tokenizer, classify_device = load_qwen_model(args.model) |
| deformation, arap_objects = run_classify( |
| classify_model, classify_tokenizer, classify_device, caption, semantic_path, |
| os.path.join(json_dir, f"{name}_deformation.json")) |
| classify_model = unload_model(classify_model) |
|
|
| if arap_objects: |
| import vlm_judge |
|
|
| temp_dir = os.path.join(args.out_dir, "temp", name) |
| objects_info = build_objects_info(svg_path, semantic_path, arap_objects) |
| if not objects_info: |
| print("No objects with valid handles found after mesh setup. Nothing to do.") |
| objects_info = None |
|
|
| rest_pose_image_path = os.path.join(json_dir, f"{name}_rest_pose.png") |
| render_rest_pose_multi(arap_objects, svg_path, semantic_path, rest_pose_image_path) |
| |
| |
| |
|
|
| import dino_similarity |
| dino_model, dino_processor = dino_similarity.load_dino_model() |
|
|
| import clip_score |
| clip_model, clip_processor = clip_score.load_clip_model() |
|
|
| feedback = None |
| previous_narratives = None |
| previous_temp_dir = None |
| freeze_narrative = False |
| consecutive_stagnant = 0 |
| previous_deform_outputs = None |
| |
| final_verdict = None |
| winning_attempt = None |
| all_attempts_summary = [] |
|
|
| for attempt in range(1, MAX_RETRIES + 1): |
| if not objects_info: |
| break |
| print(f"\n########## ATTEMPT {attempt}/{MAX_RETRIES} ##########") |
|
|
| attempt_json_dir = os.path.join(json_dir, "attempts", f"attempt_{attempt}") |
| attempt_temp_dir = os.path.join(temp_dir, "attempts", f"attempt_{attempt}") |
| os.makedirs(attempt_json_dir, exist_ok=True) |
|
|
| |
| |
| |
| |
| |
| is_retry = attempt > 1 |
|
|
| prompt = build_combined_narrate_deform_prompt( |
| objects_info, caption, previous_narratives=previous_narratives, |
| feedback=feedback, is_retry=is_retry, few_shot=not args.no_fewshot, |
| freeze_narrative=freeze_narrative, previous_deform_outputs=previous_deform_outputs) |
|
|
| print(f"\n---------- STEP 2+3: NARRATE+DEFORM (attempt {attempt}, " |
| f"NO IMAGES, text-only" |
| f"{', narrative FROZEN' if freeze_narrative else ''}) ----------") |
| vlm_model, vlm_processor = vlm_judge.load_vlm("Qwen/Qwen3-VL-4B-Instruct") |
| response = run_combined_narrate_deform(vlm_model, vlm_processor, [], prompt) |
| try: |
| parsed = parse_json_response(response) |
| except (ValueError, json.JSONDecodeError) as e: |
| print(f"FAILED TO PARSE: {e}\nraw: {response}") |
| vlm_model = unload_model(vlm_model) |
| feedback = "the previous attempt's output could not be parsed; produce valid JSON in the exact requested format" |
| all_attempts_summary.append({"attempt": attempt, "plausibility_score": None, "note": "narrate+deform parse failed"}) |
| continue |
|
|
| narratives_this_attempt, deform_outputs_this_attempt, cap_utilization_this_attempt = apply_deform_clip_and_write( |
| parsed, objects_info, attempt_json_dir, name, |
| frozen_narratives=previous_narratives if freeze_narrative else None) |
| previous_deform_outputs = deform_outputs_this_attempt |
| with open(os.path.join(attempt_json_dir, f"{name}_narratives.json"), "w") as f: |
| json.dump(narratives_this_attempt, f, indent=2) |
|
|
| print(f"\n---------- STEP 4: RENDER (attempt {attempt}, no Qwen) ----------") |
| run_render(attempt_json_dir, svg_path, semantic_path, traj_path, frames_dir=attempt_temp_dir) |
|
|
| if args.deform_only: |
| print(f"\n########## --deform-only: STOPPING after attempt 1, no judge ##########") |
| print(f"cap_utilization (raw, unfiltered by any threshold):") |
| for obj, frac in (cap_utilization_this_attempt or {}).items(): |
| print(f" {obj}: {frac*100:.1f}% of allowed movement used") |
| print(f"\nInspect the actual render directly at: {attempt_temp_dir}") |
| print(f"(kf0.png ... kf4.png, plus the combined strip)") |
| sys.exit(0) |
|
|
| stagnation_result = None |
| if previous_temp_dir: |
| print(f"\n---------- STAGNATION CHECK (attempt {attempt} vs attempt {attempt - 1}) ----------") |
| prev_dino_images = vlm_judge.load_keyframe_images(previous_temp_dir) |
| curr_dino_images = vlm_judge.load_keyframe_images(attempt_temp_dir) |
| stagnation_result = dino_similarity.stagnation_score( |
| dino_model, dino_processor, prev_dino_images, curr_dino_images) |
| stagnant = dino_similarity.is_stagnant(stagnation_result) |
| print(f"mean attempt-to-attempt similarity: {stagnation_result['mean_similarity']:.4f} " |
| f"({'STAGNANT' if stagnant else 'changed'})") |
| consecutive_stagnant = consecutive_stagnant + 1 if stagnant else 0 |
|
|
| print(f"\n---------- TEMPORAL CONSISTENCY (attempt {attempt}, diagnostic only) ----------") |
| temporal_images = vlm_judge.load_keyframe_images(attempt_temp_dir) |
| temporal_result = dino_similarity.temporal_consistency(dino_model, dino_processor, temporal_images) |
|
|
| clip_result = None |
| if caption: |
| print(f"\n---------- CLIP SCORE (attempt {attempt}) ----------") |
| clip_images = vlm_judge.load_keyframe_images(attempt_temp_dir) |
| clip_result = clip_score.compute_sequence_clip_scores(clip_model, clip_processor, clip_images, caption) |
|
|
| print(f"\n---------- STEP 5: JUDGE (attempt {attempt}, SVG-based, no images) ----------") |
| |
| |
| all_kf_geometry_text = build_all_keyframes_geometry_text(objects_info, deform_outputs_this_attempt) |
| judge_prompt = build_svg_judge_prompt( |
| name, all_kf_geometry_text, caption, dino_stagnation=stagnation_result, |
| dino_temporal=temporal_result, clip_scores=clip_result, |
| cap_utilization=cap_utilization_this_attempt) |
| judge_response = run_svg_judge(vlm_model, vlm_processor, judge_prompt) |
| vlm_model = unload_model(vlm_model) |
|
|
| try: |
| verdict = vlm_judge.parse_judge_response(judge_response) |
| except (ValueError, json.JSONDecodeError) as e: |
| print(f"JUDGE FAILED TO PARSE: {e}\nraw: {judge_response}") |
| print("Treating as a failed attempt, retrying without specific feedback.") |
| feedback = "the previous attempt's evaluation could not be parsed; try a clearer, more varied pose progression" |
| previous_narratives = narratives_this_attempt |
| previous_temp_dir = attempt_temp_dir |
| freeze_narrative = False |
| all_attempts_summary.append({"attempt": attempt, "plausibility_score": None, "note": "judge parse failed"}) |
| continue |
|
|
| problems = vlm_judge.validate_judge_response(verdict) |
| print("\n--- JUDGE VERDICT ---") |
| print(json.dumps(verdict, indent=2)) |
| if problems: |
| for p in problems: |
| print(f" VALIDATION PROBLEM: {p}") |
|
|
| with open(os.path.join(attempt_json_dir, f"{name}_judge_verdict.json"), "w") as f: |
| json.dump(verdict, f, indent=2) |
|
|
| final_verdict = verdict |
| winning_attempt = attempt |
| score = verdict.get("plausibility_score") |
| faith_score = verdict.get("faithfulness_score") |
| quality_score = verdict.get("quality_score") |
| all_attempts_summary.append({"attempt": attempt, "plausibility_score": score, |
| "plausibility_notes": verdict.get("plausibility_notes"), |
| "faithfulness_score": faith_score, |
| "faithfulness_notes": verdict.get("faithfulness_notes"), |
| "quality_score": quality_score, |
| "quality_notes": verdict.get("quality_notes"), |
| "dino_stagnant": dino_similarity.is_stagnant(stagnation_result) if stagnation_result else None}) |
|
|
| plausibility_ok = isinstance(score, (int, float)) and score >= PLAUSIBILITY_THRESHOLD |
| faithfulness_ok = faithfulness_passed(faith_score) |
| quality_ok = isinstance(quality_score, (int, float)) and quality_score >= QUALITY_THRESHOLD |
| print(f"\nplausibility_score = {score} (threshold = {PLAUSIBILITY_THRESHOLD}, " |
| f"{'PASS' if plausibility_ok else 'FAIL'})") |
| print(f"faithfulness_score = {faith_score} (threshold = {FAITHFULNESS_THRESHOLD}, " |
| f"{'PASS' if faithfulness_ok else 'FAIL'})") |
| print(f"quality_score = {quality_score} (threshold = {QUALITY_THRESHOLD}, " |
| f"{'PASS' if quality_ok else 'FAIL'})") |
|
|
| if plausibility_ok and faithfulness_ok and quality_ok: |
| print(f"All three thresholds met on attempt {attempt} — stopping.") |
| break |
|
|
| if consecutive_stagnant >= 2: |
| print(f"\nDINOv2 confirmed NO real change across {consecutive_stagnant} consecutive attempts " |
| f"(attempt {attempt} vs {attempt-1}, and {attempt-1} vs {attempt-2}) — further retries " |
| f"are very unlikely to help. Stopping early and keeping this attempt's result rather " |
| f"than burning through the remaining {MAX_RETRIES - attempt} attempts.") |
| break |
|
|
| previous_narratives = narratives_this_attempt |
| previous_temp_dir = attempt_temp_dir |
| |
| |
| |
| freeze_narrative = faithfulness_ok |
|
|
| if attempt == MAX_RETRIES: |
| print("Below threshold on the final attempt — no retry left, skipping feedback construction.") |
| else: |
| plausibility_note = verdict.get("plausibility_notes", "the pose progression needs to look more plausible") |
| faithfulness_note = verdict.get("faithfulness_notes") |
| quality_note = verdict.get("quality_notes") |
| notes = [f"plausibility: {plausibility_note}"] |
| if not faithfulness_ok and faithfulness_note and faithfulness_note != "N/A": |
| notes.append(f"faithfulness to the intended action: {faithfulness_note}") |
| if not quality_ok and quality_note: |
| notes.append(f"rendering quality: {quality_note}") |
| feedback = " | ".join(notes) |
|
|
| |
| |
| |
| |
| |
| |
| if stagnation_result and dino_similarity.is_stagnant(stagnation_result): |
| feedback = (f"CRITICAL: your last attempt was measured as nearly IDENTICAL to the one " |
| f"before it (DINOv2 similarity {stagnation_result['mean_similarity']:.3f}) — " |
| f"you are NOT making real changes. You MUST produce substantially different " |
| f"target positions this time, not a superficial rewording. Original feedback: {feedback}") |
| elif cap_utilization_this_attempt: |
| |
| |
| |
| |
| |
| low_objs = {obj: frac for obj, frac in cap_utilization_this_attempt.items() |
| if frac < LOW_UTILIZATION_THRESHOLD} |
| if low_objs: |
| low_str = ", ".join(f"{obj} used only {frac*100:.0f}%" for obj, frac in low_objs.items()) |
| feedback = (f"CRITICAL: {low_str} of their allowed movement range — this is too timid " |
| f"to be visible. You MUST commit to larger, more decisive displacement " |
| f"(up to the stated cap) for these objects. Original feedback: {feedback}") |
|
|
| print(f"Below threshold — retrying with feedback: {feedback!r} " |
| f"(narrative will be {'FROZEN' if freeze_narrative else 'regenerated'})") |
| else: |
| print(f"\nReached MAX_RETRIES ({MAX_RETRIES}) without meeting both thresholds. " |
| f"Using the last attempt's result.") |
|
|
| |
| |
| print("\n########## ATTEMPT SUMMARY ##########") |
| for a in all_attempts_summary: |
| print(f" attempt {a['attempt']}: plausibility={a.get('plausibility_score')} " |
| f"faithfulness={a.get('faithfulness_score')} quality={a.get('quality_score')} " |
| f"dino_stagnant={a.get('dino_stagnant')} " |
| f"— {a.get('plausibility_notes') or a.get('note', '')}") |
|
|
| if winning_attempt: |
| import shutil |
| winning_json = os.path.join(json_dir, "attempts", f"attempt_{winning_attempt}") |
| winning_temp = os.path.join(temp_dir, "attempts", f"attempt_{winning_attempt}") |
| for f in os.listdir(winning_json): |
| shutil.copy2(os.path.join(winning_json, f), os.path.join(json_dir, f)) |
| for f in os.listdir(winning_temp): |
| src = os.path.join(winning_temp, f) |
| if os.path.isfile(src): |
| shutil.copy2(src, os.path.join(temp_dir, f)) |
| print(f"\ncopied winning attempt ({winning_attempt}) to the top-level " |
| f"json/{name}/ and temp/{name}/ locations") |
| print(f"all {len(all_attempts_summary)} attempts preserved under " |
| f"json/{name}/attempts/ and temp/{name}/attempts/ for comparison") |
| else: |
| print("\nNo ARAP objects — skipping narrate/deform/judge entirely.") |
| temp_dir = os.path.join(args.out_dir, "temp", name) |
| print("\n########## RENDER (no Qwen) ##########") |
| run_render(json_dir, svg_path, semantic_path, traj_path, frames_dir=temp_dir) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
