Upload demo/real_world_pipeline.py with huggingface_hub

demo/real_world_pipeline.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import argparse
+import gc
 import json
 import os
 import re
@@ -95,6 +96,19 @@ def load_bbox_bundle(model_path: str):
     return model, processor
 
 
+def release_bbox_bundle():
+    model = _BBOX_CACHE.get("model")
+    processor = _BBOX_CACHE.get("processor")
+    if model is not None:
+        del model
+    if processor is not None:
+        del processor
+    _BBOX_CACHE.update({"model_path": None, "model": None, "processor": None})
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
+
 def load_real_bundle(config: dict):
     key = (
         config.get("pretrained_model_name_or_path"),
@@ -390,6 +404,7 @@ 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()
 
     config = build_runtime_config(
         config_path=config_path,