File size: 3,193 Bytes
e6f80f3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 | """Load FastFill (Qwen3-8B + LoRA adapter + geometry heads) and lay out one room.
pip install torch # for your CUDA / platform first
pip install -r requirements.txt
python predict.py example_condition.json # Qwen/Qwen3-8B from the Hugging Face Hub
python predict.py example_condition.json --qwen /path/to/Qwen3-8B --device cuda
Prints the layout as JSON: for every requested object id, ``target_size_local_m`` [depth, width, height] (m),
``bottom_center_m`` [x, y, z] (m, room frame) and ``yaw_rad`` (the direction its front faces, counter-clockwise from +x).
"""
import argparse
import json
import sys
from dataclasses import replace
from pathlib import Path
HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE / "code"))
import torch # noqa: E402
from fastfill.v2.batch import load_tokenizer # noqa: E402
from fastfill.v2.evaluate import predict_layout # noqa: E402
from fastfill.v2.io import load_checkpoint_config # noqa: E402
from fastfill.v2.model import ModelConfig, StructuredFastFillModel, _qwen_backbone # noqa: E402
CHECKPOINT = HERE / "checkpoint"
def load(qwen: str = "Qwen/Qwen3-8B", device: str = "cuda"):
"""(model, tokenizer, max_length). ``qwen``: the base model, a Hub id or a local directory."""
config = ModelConfig(**json.loads((CHECKPOINT / "model_config.json").read_text()))
config = replace(config, backbone=qwen, local_files_only=Path(qwen).is_dir())
from peft import PeftModel
backbone = _qwen_backbone(replace(config, lora_rank=0)) # the frozen base
backbone = PeftModel.from_pretrained(backbone, CHECKPOINT / "backbone") # FastFill's LoRA adapter
model = StructuredFastFillModel(config, backbone)
state = torch.load(CHECKPOINT / "geometry_model.pt", map_location="cpu", weights_only=True)
missing = model.load_state_dict(state, strict=False)
if missing.unexpected_keys or any(not k.startswith("backbone.") for k in missing.missing_keys):
raise ValueError("geometry heads do not match the model configuration")
tokenizer = load_tokenizer(str(HERE / "tokenizer"), local_files_only=True)
max_length = load_checkpoint_config(CHECKPOINT)["max_length"] or 8192
return model.to(device).eval(), tokenizer, max_length
def main():
ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
ap.add_argument("condition", help="a fastfill.v2 condition JSON (room, objects, constraints)")
ap.add_argument("--qwen", default="Qwen/Qwen3-8B", help="base model: Hub id or local Qwen3-8B directory")
ap.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
ap.add_argument("--decode", default="spread", choices=("spread", "argmax"), help="position grid decoding")
a = ap.parse_args()
model, tokenizer, max_length = load(a.qwen, a.device)
condition = json.loads(Path(a.condition).read_text())
with torch.no_grad():
layout = predict_layout(model, tokenizer, condition, max_length=max_length, device=a.device,
grid_decode=a.decode)
print(json.dumps(layout, indent=1))
if __name__ == "__main__":
main()
|