Delete generator.py

generator.py

@@ -1,409 +0,0 @@
-from __future__ import annotations
-
-import math
-import tempfile
-from dataclasses import dataclass
-from pathlib import Path
-from typing import Generator
-
-import numpy as np
-import trimesh
-from scipy import ndimage
-from skimage import measure
-
-from llm_parser import DEFAULT_LOCAL_MODEL, parse_prompt_with_local_llm
-from model_runtime import TARGET_OMNI_MODEL, ensure_target_model_cached
-from parser import PromptSpec, parse_prompt
-
-
-@dataclass
-class BuildArtifacts:
-    ply_path: str
-    glb_path: str
-    summary: dict
-
-
-SCALE_FACTORS = {
-    "small": 1.0,
-    "medium": 1.35,
-    "large": 1.85,
-}
-
-
-def _sample_box_surface(center, size, density: int, label: int) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
-    cx, cy, cz = center
-    sx, sy, sz = size
-    n = max(4, density)
-    u = np.linspace(-0.5, 0.5, n)
-    vv = np.linspace(-0.5, 0.5, n)
-    pts = []
-    normals = []
-    labels = []
-    for ax in (-1, 1):
-        x = np.full((n, n), cx + ax * sx / 2)
-        y, z = np.meshgrid(u * sy + cy, vv * sz + cz)
-        pts.append(np.column_stack([x.ravel(), y.ravel(), z.ravel()]))
-        normals.append(np.tile([ax, 0, 0], (n * n, 1)))
-        labels.append(np.full(n * n, label))
-    for ay in (-1, 1):
-        y = np.full((n, n), cy + ay * sy / 2)
-        x, z = np.meshgrid(u * sx + cx, vv * sz + cz)
-        pts.append(np.column_stack([x.ravel(), y.ravel(), z.ravel()]))
-        normals.append(np.tile([0, ay, 0], (n * n, 1)))
-        labels.append(np.full(n * n, label))
-    for az in (-1, 1):
-        z = np.full((n, n), cz + az * sz / 2)
-        x, y = np.meshgrid(u * sx + cx, vv * sy + cy)
-        pts.append(np.column_stack([x.ravel(), y.ravel(), z.ravel()]))
-        normals.append(np.tile([0, 0, az], (n * n, 1)))
-        labels.append(np.full(n * n, label))
-    return np.vstack(pts), np.vstack(normals), np.concatenate(labels)
-
-
-def _sample_ellipsoid_surface(center, radii, density: int, label: int) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
-    cx, cy, cz = center
-    rx, ry, rz = radii
-    nu = max(16, density * 3)
-    nv = max(10, density * 2)
-    u = np.linspace(0, 2 * math.pi, nu, endpoint=False)
-    v = np.linspace(-math.pi / 2, math.pi / 2, nv)
-    uu, vv = np.meshgrid(u, v)
-    x = cx + rx * np.cos(vv) * np.cos(uu)
-    y = cy + ry * np.cos(vv) * np.sin(uu)
-    z = cz + rz * np.sin(vv)
-    pts = np.column_stack([x.ravel(), y.ravel(), z.ravel()])
-    normals = np.column_stack([
-        (x - cx).ravel() / max(rx, 1e-6),
-        (y - cy).ravel() / max(ry, 1e-6),
-        (z - cz).ravel() / max(rz, 1e-6),
-    ])
-    normals /= np.linalg.norm(normals, axis=1, keepdims=True) + 1e-8
-    labels = np.full(len(pts), label)
-    return pts, normals, labels
-
-
-def _sample_cylinder_surface(center, radius, length, axis: str, density: int, label: int) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
-    cx, cy, cz = center
-    nt = max(18, density * 4)
-    nl = max(6, density)
-    theta = np.linspace(0, 2 * math.pi, nt, endpoint=False)
-    line = np.linspace(-length / 2, length / 2, nl)
-    tt, ll = np.meshgrid(theta, line)
-    if axis == "x":
-        x = cx + ll
-        y = cy + radius * np.cos(tt)
-        z = cz + radius * np.sin(tt)
-        normals = np.column_stack([np.zeros(x.size), np.cos(tt).ravel(), np.sin(tt).ravel()])
-    elif axis == "y":
-        x = cx + radius * np.cos(tt)
-        y = cy + ll
-        z = cz + radius * np.sin(tt)
-        normals = np.column_stack([np.cos(tt).ravel(), np.zeros(x.size), np.sin(tt).ravel()])
-    else:
-        x = cx + radius * np.cos(tt)
-        y = cy + radius * np.sin(tt)
-        z = cz + ll
-        normals = np.column_stack([np.cos(tt).ravel(), np.sin(tt).ravel(), np.zeros(x.size)])
-    pts = np.column_stack([x.ravel(), y.ravel(), z.ravel()])
-    labels = np.full(len(pts), label)
-    return pts, normals, labels
-
-
-def export_point_cloud_as_ply(points: np.ndarray, labels: np.ndarray, path: str) -> str:
-    colors = np.array([
-        [170, 170, 180],
-        [120, 180, 255],
-        [255, 190, 120],
-        [180, 180, 255],
-        [255, 120, 120],
-        [200, 255, 180],
-        [255, 255, 180],
-    ], dtype=np.uint8)
-    c = colors[labels % len(colors)]
-    pc = trimesh.points.PointCloud(vertices=points, colors=c)
-    pc.export(path)
-    return path
-
-
-def export_mesh_as_glb(mesh: trimesh.Trimesh, path: str) -> str:
-    mesh.visual.vertex_colors = np.tile(np.array([[185, 190, 200, 255]], dtype=np.uint8), (len(mesh.vertices), 1))
-    mesh.export(path)
-    return path
-
-
-def _resolve_spec(prompt: str, parser_mode: str, model_id: str | None = None) -> tuple[PromptSpec, str]:
-    parser_mode = (parser_mode or "heuristic").strip().lower()
-    if parser_mode.startswith("local"):
-        spec = parse_prompt_with_local_llm(prompt, model_id=model_id or DEFAULT_LOCAL_MODEL)
-        return spec, f"local_llm:{model_id or DEFAULT_LOCAL_MODEL}"
-    return parse_prompt(prompt), "heuristic"
-
-
-def _iter_part_specs(spec: PromptSpec, detail: int):
-    scale = SCALE_FACTORS[spec.scale]
-    density = max(6, detail)
-
-    hull_len = 2.8 * scale
-    hull_w = 1.2 * scale
-    hull_h = 0.8 * scale
-
-    if spec.hull_style == "rounded":
-        yield "Hull", *_sample_ellipsoid_surface((0.0, 0.0, 0.0), (hull_len / 2, hull_w / 2, hull_h / 2), density, 0)
-    elif spec.hull_style == "sleek":
-        p1, n1, l1 = _sample_ellipsoid_surface((0.12 * scale, 0.0, 0.0), (hull_len / 2.3, hull_w / 2.8, hull_h / 2.6), density, 0)
-        p2, n2, l2 = _sample_box_surface((-0.15 * scale, 0.0, -0.02 * scale), (hull_len * 0.52, hull_w * 0.5, hull_h * 0.55), max(4, density // 2), 0)
-        yield "Hull", np.vstack([p1, p2]), np.vstack([n1, n2]), np.concatenate([l1, l2])
-    else:
-        yield "Hull", *_sample_box_surface((0.0, 0.0, 0.0), (hull_len, hull_w, hull_h), density, 0)
-
-    cockpit_center = (hull_len / 2 - hull_len * spec.cockpit_ratio * 0.8, 0.0, hull_h * 0.14)
-    yield "Cockpit", *_sample_ellipsoid_surface(cockpit_center, (hull_len * spec.cockpit_ratio, hull_w * 0.22, hull_h * 0.24), max(4, density // 2), 1)
-
-    if spec.cargo_ratio > 0.16:
-        cargo_center = (-hull_len * 0.18, 0.0, -hull_h * 0.06)
-        cargo_size = (hull_len * spec.cargo_ratio, hull_w * 0.76, hull_h * 0.6)
-        yield "Cargo bay", *_sample_box_surface(cargo_center, cargo_size, max(4, density // 2), 2)
-
-    if spec.wing_span > 0:
-        wing_length = hull_len * 0.34
-        wing_width = hull_w * 0.18
-        wing_height = hull_h * 0.08
-        yoff = hull_w * 0.45 + wing_width * 0.6
-        wing_parts = []
-        wing_normals = []
-        wing_labels = []
-        for side in (-1, 1):
-            wc = (-0.1 * scale, side * yoff, -0.04 * scale)
-            pp, pn, pl = _sample_box_surface(wc, (wing_length, wing_width, wing_height), max(6, density // 3), 3)
-            wing_parts.append(pp)
-            wing_normals.append(pn)
-            wing_labels.append(pl)
-        yield "Wings", np.vstack(wing_parts), np.vstack(wing_normals), np.concatenate(wing_labels)
-
-    engine_radius = 0.14 * scale if spec.object_type != "fighter" else 0.1 * scale
-    engine_length = 0.48 * scale
-    engine_y_positions = np.linspace(-hull_w * 0.32, hull_w * 0.32, spec.engine_count)
-    engine_parts = []
-    engine_normals = []
-    engine_labels = []
-    for ypos in engine_y_positions:
-        ec = (-hull_len / 2 + engine_length * 0.3, ypos, 0.0)
-        pp, pn, pl = _sample_cylinder_surface(ec, engine_radius, engine_length, "x", max(6, density // 3), 4)
-        engine_parts.append(pp)
-        engine_normals.append(pn)
-        engine_labels.append(pl)
-    yield "Engines", np.vstack(engine_parts), np.vstack(engine_normals), np.concatenate(engine_labels)
-
-    if spec.fin_height > 0:
-        fin_center = (-hull_len * 0.25, 0.0, hull_h * 0.42)
-        fin_size = (hull_len * 0.18, hull_w * 0.1, hull_h * max(spec.fin_height, 0.12))
-        yield "Fin", *_sample_box_surface(fin_center, fin_size, max(6, density // 3), 5)
-
-    if spec.landing_gear:
-        gear_x = np.array([-hull_len * 0.18, hull_len * 0.12])
-        gear_y = np.array([-hull_w * 0.28, hull_w * 0.28])
-        gear_parts = []
-        gear_normals = []
-        gear_labels = []
-        for gx in gear_x:
-            for gy in gear_y:
-                gc = (gx, gy, -hull_h * 0.45)
-                pp, pn, pl = _sample_cylinder_surface(gc, 0.04 * scale, 0.22 * scale, "z", max(5, density // 5), 6)
-                gear_parts.append(pp)
-                gear_normals.append(pn)
-                gear_labels.append(pl)
-        yield "Landing gear", np.vstack(gear_parts), np.vstack(gear_normals), np.concatenate(gear_labels)
-
-
-def iter_blueprint_session(
-    prompt: str,
-    detail: int = 24,
-    parser_mode: str = "heuristic",
-    model_id: str | None = None,
-) -> Generator[dict, None, dict]:
-    prompt = (prompt or "").strip()
-    if not prompt:
-        raise ValueError("Enter a prompt first.")
-
-    out_dir = Path(tempfile.mkdtemp(prefix="particle_blueprint_session_"))
-    yield {"status": "Parsing prompt and planning shape…", "stage_index": 0, "stage_count": 1, "session_dir": str(out_dir)}
-
-    spec, parser_backend = _resolve_spec(prompt, parser_mode=parser_mode, model_id=model_id)
-    stages = list(_iter_part_specs(spec, detail=detail))
-
-    all_points = []
-    all_normals = []
-    all_labels = []
-
-    for idx, (stage_name, points, normals, labels) in enumerate(stages, start=1):
-        if spec.asymmetry > 0 and stage_name in {"Hull", "Cockpit", "Cargo bay"}:
-            mask = points[:, 1] > 0
-            points = points.copy()
-            points[mask, 2] += spec.asymmetry * np.sin(points[mask, 0] * 2.0)
-
-        all_points.append(points)
-        all_normals.append(normals)
-        all_labels.append(labels)
-
-        merged_points = np.vstack(all_points).astype(np.float32)
-        merged_normals = np.vstack(all_normals).astype(np.float32)
-        merged_labels = np.concatenate(all_labels).astype(np.int32)
-
-        preview_path = str(out_dir / f"blueprint_stage_{idx:02d}.ply")
-        export_point_cloud_as_ply(merged_points, merged_labels, preview_path)
-
-        summary = {
-            "prompt": prompt,
-            "parser_backend": parser_backend,
-            "spec": spec.to_dict(),
-            "stage": stage_name,
-            "stage_index": idx,
-            "stage_count": len(stages),
-            "point_count": int(len(merged_points)),
-        }
-        yield {
-            "status": f"{stage_name} added ({idx}/{len(stages)})",
-            "blueprint_path": preview_path,
-            "summary": summary,
-            "stage_index": idx,
-            "stage_count": len(stages),
-            "session_dir": str(out_dir),
-        }
-
-    final_points = np.vstack(all_points).astype(np.float32)
-    final_normals = np.vstack(all_normals).astype(np.float32)
-    final_labels = np.concatenate(all_labels).astype(np.int32)
-
-    npz_path = str(out_dir / "blueprint_data.npz")
-    np.savez_compressed(npz_path, points=final_points, normals=final_normals, labels=final_labels)
-
-    final_ply = str(out_dir / "blueprint_final.ply")
-    export_point_cloud_as_ply(final_points, final_labels, final_ply)
-
-    state = {
-        "prompt": prompt,
-        "parser_backend": parser_backend,
-        "spec": spec.to_dict(),
-        "point_count": int(len(final_points)),
-        "session_dir": str(out_dir),
-        "npz_path": npz_path,
-        "blueprint_path": final_ply,
-        "target_model": TARGET_OMNI_MODEL,
-    }
-    yield {
-        "status": "Blueprint ready. Inspect it, then run mesh generation when happy.",
-        "blueprint_path": final_ply,
-        "summary": {
-            **state,
-            "stage": "complete",
-        },
-        "stage_index": len(stages),
-        "stage_count": len(stages),
-        "state": state,
-        "session_dir": str(out_dir),
-    }
-    return state
-
-
-def points_to_mesh(points: np.ndarray, pitch: float = 0.08, padding: int = 5, sigma: float = 1.2, level: float = 0.11) -> trimesh.Trimesh:
-    mins = points.min(axis=0) - padding * pitch
-    maxs = points.max(axis=0) + padding * pitch
-    dims = np.ceil((maxs - mins) / pitch).astype(int) + 1
-    dims = np.clip(dims, 24, 192)
-
-    grid = np.zeros(tuple(dims.tolist()), dtype=np.float32)
-    coords = ((points - mins) / pitch).astype(int)
-    coords = np.clip(coords, 0, dims - 1)
-    np.add.at(grid, (coords[:, 0], coords[:, 1], coords[:, 2]), 1.0)
-
-    grid = ndimage.gaussian_filter(grid, sigma=sigma)
-    verts, faces, normals, _ = measure.marching_cubes(grid, level=level)
-    verts = verts * pitch + mins
-
-    mesh = trimesh.Trimesh(vertices=verts, faces=faces, vertex_normals=normals, process=True)
-    mesh.update_faces(mesh.nondegenerate_faces())
-    mesh.update_faces(mesh.unique_faces())
-    mesh.remove_unreferenced_vertices()
-    try:
-        mesh.fill_holes()
-    except Exception:
-        pass
-    try:
-        trimesh.smoothing.filter_humphrey(mesh, iterations=2)
-    except Exception:
-        pass
-    return mesh
-
-
-def iter_meshify_session(
-    state: dict,
-    voxel_pitch: float = 0.08,
-    use_target_model_cache: bool = True,
-) -> Generator[dict, None, dict]:
-    if not state or not state.get("npz_path"):
-        raise ValueError("Generate a blueprint first.")
-
-    data = np.load(state["npz_path"])
-    points = data["points"].astype(np.float32)
-    labels = data["labels"].astype(np.int32)
-    session_dir = Path(state["session_dir"])
-
-    model_note = None
-    if use_target_model_cache:
-        yield {"status": f"Preparing target model cache for {TARGET_OMNI_MODEL}…"}
-        model_cache = ensure_target_model_cached(TARGET_OMNI_MODEL)
-        model_note = model_cache["message"]
-        yield {"status": model_note}
-
-    yield {"status": "Converting blueprint into a watertight mesh…"}
-    mesh = points_to_mesh(points, pitch=voxel_pitch)
-
-    yield {"status": "Exporting GLB…"}
-    glb_path = str(session_dir / "mesh_final.glb")
-    export_mesh_as_glb(mesh, glb_path)
-
-    summary = {
-        **state,
-        "mesh_backend": "local_voxel_mesher",
-        "target_model_cached": bool(use_target_model_cache),
-        "target_model": TARGET_OMNI_MODEL,
-        "target_model_note": model_note,
-        "vertex_count": int(len(mesh.vertices)),
-        "face_count": int(len(mesh.faces)),
-        "bounds": mesh.bounds.round(3).tolist(),
-        "voxel_pitch": voxel_pitch,
-        "mesh_path": glb_path,
-    }
-    yield {
-        "status": "Mesh ready.",
-        "mesh_path": glb_path,
-        "summary": summary,
-        "mesh_file": glb_path,
-    }
-    return summary
-
-
-# Backward-compatible helper for older single-click flow.
-def run_pipeline(
-    prompt: str,
-    detail: int = 24,
-    voxel_pitch: float = 0.08,
-    parser_mode: str = "heuristic",
-    model_id: str | None = None,
-) -> BuildArtifacts:
-    final_state = None
-    final_summary = None
-    blueprint_path = None
-    for update in iter_blueprint_session(prompt, detail=detail, parser_mode=parser_mode, model_id=model_id):
-        blueprint_path = update.get("blueprint_path", blueprint_path)
-        final_state = update.get("state", final_state)
-        final_summary = update.get("summary", final_summary)
-    mesh_summary = None
-    mesh_path = None
-    if final_state is None:
-        raise RuntimeError("Blueprint generation failed.")
-    for update in iter_meshify_session(final_state, voxel_pitch=voxel_pitch, use_target_model_cache=False):
-        mesh_path = update.get("mesh_path", mesh_path)
-        mesh_summary = update.get("summary", mesh_summary)
-    summary = mesh_summary or final_summary or {}
-    return BuildArtifacts(ply_path=blueprint_path or "", glb_path=mesh_path or "", summary=summary)