| """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 |
|
|
| from fastfill.v2.batch import load_tokenizer |
| from fastfill.v2.evaluate import predict_layout |
| from fastfill.v2.io import load_checkpoint_config |
| from fastfill.v2.model import ModelConfig, StructuredFastFillModel, _qwen_backbone |
|
|
| 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)) |
| backbone = PeftModel.from_pretrained(backbone, CHECKPOINT / "backbone") |
| 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() |
|
|
