Upload demo/real_world_pipeline.py with huggingface_hub

demo/real_world_pipeline.py

@@ -24,7 +24,7 @@ from demo.infer.run_caption_bbox_infer import (  # noqa: E402
     draw_boxes,
     infer_caption_bbox,
 )
-from demo.hf_repo_assets import build_repo_asset_overrides, get_stage2_model_repo_id  # noqa: E402
+from demo.hf_repo_assets import build_repo_asset_overrides  # noqa: E402
 from demo.infer.vlm_bbox_inference import get_model_and_processor  # noqa: E402
 from infer.common_infer import initialize_pipeline  # noqa: E402
 from infer.infer import build_run_save_dir, get_real_boxes, load_adapter_image  # noqa: E402
@@ -35,9 +35,11 @@ DEFAULT_REAL_CONFIG_PATH = PROJECT_ROOT / "infer" / "infer.yaml"
 DEFAULT_WORK_DIR = PROJECT_ROOT / "demo" / "outputs" / "real_world_demo"
 DEFAULT_RUN_NAME = "step_120000"
 DEFAULT_TARGET_SIZE = 1024
+DEFAULT_MODEL_REPO_ID = "SynLayers/Bbox-caption-8b"
 
 _BBOX_CACHE: dict[str, object] = {"model_path": None, "model": None, "processor": None}
 _REAL_CACHE: dict[str, object] = {"key": None, "pipeline": None, "transp_vae": None}
+RELEASE_BBOX_AFTER_CAPTION = os.environ.get("SYNLAYERS_RELEASE_BBOX_AFTER_CAPTION", "0") == "1"
 
 
 def slugify(text: str) -> str:
@@ -59,9 +61,8 @@ DEFAULT_DECOMP_CKPT_ROOT = Path(
     resolve_existing_path(
         os.environ.get("SYNLAYERS_DECOMP_CKPT_ROOT"),
         PROJECT_ROOT / "SynLayers_ckpt" / "step_120000",
-        "/project/llmsvgen/share/data/kmw_layered_checkpoint/SynLayers_ckpt/step_120000",
     )
-    or "/project/llmsvgen/share/data/kmw_layered_checkpoint/SynLayers_ckpt/step_120000"
+    or PROJECT_ROOT / "SynLayers_ckpt" / "step_120000"
 )
 
 
@@ -148,7 +149,8 @@ def build_runtime_config(
     seed: int | None = None,
 ) -> dict:
     config = load_config(str(config_path))
-    repo_overrides = build_repo_asset_overrides(get_stage2_model_repo_id())
+    model_repo = os.environ.get("SYNLAYERS_MODEL_REPO") or DEFAULT_MODEL_REPO_ID
+    repo_overrides = build_repo_asset_overrides(model_repo)
     decomp_ckpt_root = Path(
         os.environ.get("SYNLAYERS_DECOMP_CKPT_ROOT")
         or repo_overrides.get("decomp_ckpt_root")
@@ -368,7 +370,7 @@ def run_real_world_pipeline(
     bbox_model = (
         bbox_model
         or os.environ.get("SYNLAYERS_BBOX_MODEL")
-        or os.environ.get("SYNLAYERS_BBOX_MODEL_REPO")
+        or os.environ.get("SYNLAYERS_MODEL_REPO")
         or DEFAULT_BBOX_MODEL
     )
     config_path = Path(config_path or os.environ.get("SYNLAYERS_REAL_CONFIG", str(DEFAULT_REAL_CONFIG_PATH)))
@@ -405,7 +407,8 @@ def run_real_world_pipeline(
     bbox_jsonl = write_bbox_jsonl(record, run_dir / "caption_bbox_infer.jsonl")
     bbox_vis_path = run_dir / "bbox_vis" / f"{normalized_sample_name}_vis.png"
     draw_boxes(prepared_image_path, bboxes, bbox_vis_path)
-    release_bbox_bundle()
+    if RELEASE_BBOX_AFTER_CAPTION:
+        release_bbox_bundle()
 
     config = build_runtime_config(
         config_path=config_path,
@@ -466,4 +469,4 @@ def main():
 
 
 if __name__ == "__main__":
-    main()
+    main()