update

common.py

@@ -170,6 +170,25 @@ def _diag_pipeline(name: str, pipeline: object) -> None:
     )
 
 
+def _move_pipeline(pipeline: object, device: str) -> None:
+    started = time.monotonic()
+    _diag(
+        "PIPELINE_MOVE_START",
+        pipeline_type=type(pipeline).__name__,
+        target_device=device,
+    )
+    pipeline.to(device)
+    image_encoder = getattr(pipeline, "image_encoder", None)
+    if image_encoder is not None:
+        image_encoder.to(device)
+    _diag(
+        "PIPELINE_MOVE_READY",
+        pipeline_type=type(pipeline).__name__,
+        target_device=device,
+        elapsed_seconds=round(time.monotonic() - started, 3),
+    )
+
+
 def _install_zerogpu_diagnostics() -> None:
     try:
         from spaces.zero import client as zero_client
@@ -296,30 +315,32 @@ elif os.getenv("GRADIO_APP").startswith("textto3d"):
         elapsed_seconds=round(time.monotonic() - _pipeline_init_started, 3),
     )
     text_model_dir = "weights/Kolors"
-    _pipeline_img_ip_started = time.monotonic()
-    _diag("IMPORT_MODEL_START", name="PIPELINE_IMG_IP", model="Kolors-IP")
-    PIPELINE_IMG_IP = build_text2img_ip_pipeline(text_model_dir, ref_scale=0.3)
+    _pipeline_img_started = time.monotonic()
     _diag(
-        "IMPORT_MODEL_READY",
-        name="PIPELINE_IMG_IP",
-        model_type=type(PIPELINE_IMG_IP).__name__,
-        elapsed_seconds=round(time.monotonic() - _pipeline_img_ip_started, 3),
+        "IMPORT_MODEL_START",
+        name="PIPELINE_IMG",
+        model="Kolors",
+        device="cpu",
     )
-    _pipeline_img_started = time.monotonic()
-    _diag("IMPORT_MODEL_START", name="PIPELINE_IMG", model="Kolors")
-    PIPELINE_IMG = build_text2img_pipeline(text_model_dir)
+    PIPELINE_IMG = build_text2img_pipeline(
+        text_model_dir,
+        device="cpu",
+    )
+    _PIPELINE_IMG_IP = None
     _diag(
         "IMPORT_MODEL_READY",
         name="PIPELINE_IMG",
         model_type=type(PIPELINE_IMG).__name__,
+        device=str(getattr(PIPELINE_IMG, "device", None)),
         elapsed_seconds=round(time.monotonic() - _pipeline_img_started, 3),
     )
     _diag(
         "IMPORT_MODELS_COMPLETE",
-        global_cuda_pipeline_count=3,
+        global_cuda_pipeline_count=1,
+        global_cpu_pipeline_count=1,
         note=(
-            "Sam3D plus two full Kolors pipelines are resident before "
-            "ZeroGPU dispatch"
+            "Only Sam3D is CUDA-resident; the standard Kolors pipeline "
+            "stays on CPU and IP-Adapter is loaded only when requested"
         ),
     )
     SEG_CHECKER = ImageSegChecker(GPT_CLIENT)
@@ -818,6 +839,8 @@ def text2image_fn(
     n_sample: int = 3,
     req: gr.Request = None,
 ) -> list[str]:
+    global _PIPELINE_IMG_IP
+
     started = time.monotonic()
     _diag(
         "WORKER_ENTER",
@@ -848,32 +871,77 @@ def text2image_fn(
             os.makedirs(output_root, exist_ok=True)
         _diag("OUTPUT_DIRECTORY_READY", output_root=output_root)
 
-        pipeline_name = "PIPELINE_IMG" if ip_image is None else "PIPELINE_IMG_IP"
-        pipeline = PIPELINE_IMG if ip_image is None else PIPELINE_IMG_IP
-        _diag_pipeline(pipeline_name, pipeline)
-        if ip_image is not None:
+        if ip_image is None:
+            pipeline_name = "PIPELINE_IMG"
+            pipeline = PIPELINE_IMG
+        else:
+            if _PIPELINE_IMG_IP is None:
+                ip_pipeline_started = time.monotonic()
+                _diag(
+                    "IP_PIPELINE_INIT_START",
+                    model="Kolors-IP",
+                    device="cpu",
+                )
+                _PIPELINE_IMG_IP = build_text2img_ip_pipeline(
+                    "weights/Kolors",
+                    ref_scale=ip_adapt_scale,
+                    device="cpu",
+                )
+                _diag(
+                    "IP_PIPELINE_INIT_READY",
+                    model_type=type(_PIPELINE_IMG_IP).__name__,
+                    elapsed_seconds=round(
+                        time.monotonic() - ip_pipeline_started, 3
+                    ),
+                )
+            pipeline_name = "PIPELINE_IMG_IP"
+            pipeline = _PIPELINE_IMG_IP
             _diag("IP_ADAPTER_SCALE_START", scale=ip_adapt_scale)
             pipeline.set_ip_adapter_scale([ip_adapt_scale])
             _diag("IP_ADAPTER_SCALE_READY", scale=ip_adapt_scale)
+        _diag_pipeline(pipeline_name, pipeline)
 
-        generation_started = time.monotonic()
-        _diag("TEXT2IMAGE_GENERATION_START", pipeline=pipeline_name)
-        images = text2img_gen(
-            prompt=prompt,
-            n_sample=n_sample,
-            guidance_scale=guidance_scale,
-            pipeline=pipeline,
-            ip_image=ip_image,
-            image_wh=image_wh,
-            infer_step=infer_step,
-            seed=seed,
-        )
-        _diag(
-            "TEXT2IMAGE_GENERATION_READY",
-            image_count=len(images),
-            elapsed_seconds=round(time.monotonic() - generation_started, 3),
-        )
-        _diag_cuda("after_text2image_generation")
+        pipeline_move_attempted = False
+        try:
+            pipeline_move_attempted = True
+            _move_pipeline(pipeline, "cuda")
+            _diag_pipeline(pipeline_name, pipeline)
+            _diag_cuda("after_kolors_move_to_cuda")
+
+            generation_started = time.monotonic()
+            _diag("TEXT2IMAGE_GENERATION_START", pipeline=pipeline_name)
+            images = text2img_gen(
+                prompt=prompt,
+                n_sample=n_sample,
+                guidance_scale=guidance_scale,
+                pipeline=pipeline,
+                ip_image=ip_image,
+                image_wh=image_wh,
+                infer_step=infer_step,
+                seed=seed,
+            )
+            _diag(
+                "TEXT2IMAGE_GENERATION_READY",
+                image_count=len(images),
+                elapsed_seconds=round(
+                    time.monotonic() - generation_started, 3
+                ),
+            )
+            _diag_cuda("after_text2image_generation")
+        finally:
+            if pipeline_move_attempted:
+                try:
+                    _move_pipeline(pipeline, "cpu")
+                except Exception as exc:
+                    _diag(
+                        "PIPELINE_MOVE_CPU_EXCEPTION",
+                        exception_type=type(exc).__name__,
+                        exception=str(exc)[:2000],
+                    )
+                gc.collect()
+                if torch.cuda.is_available():
+                    torch.cuda.empty_cache()
+                _diag_cuda("after_kolors_move_to_cpu")
 
         for idx, image in enumerate(images):
             preprocess_started = time.monotonic()

embodied_gen/models/text_model.py

@@ -59,11 +59,15 @@ PROMPT_APPEND = (
 PROMPT_KAPPEND = "Single {object}, in the center of the image, white background, 3D style, best quality"
 
 
-def download_kolors_weights(local_dir: str = "weights/Kolors") -> None:
+def download_kolors_weights(
+    local_dir: str = "weights/Kolors",
+    include_ip_adapter: bool = False,
+) -> None:
     """Downloads Kolors model weights from HuggingFace.
 
     Args:
-        local_dir (str, optional): Local directory to store weights.
+        local_dir: Local directory to store weights.
+        include_ip_adapter: Whether to also download IP-Adapter weights.
     """
     logger.info(f"Download kolors weights from huggingface...")
     os.makedirs(local_dir, exist_ok=True)
@@ -79,6 +83,9 @@ def download_kolors_weights(local_dir: str = "weights/Kolors") -> None:
         check=True,
     )
 
+    if not include_ip_adapter:
+        return
+
     ip_adapter_path = f"{local_dir}/../Kolors-IP-Adapter-Plus"
     subprocess.run(
         [
@@ -114,7 +121,7 @@ def build_text2img_ip_pipeline(
         pipe = build_text2img_ip_pipeline("weights/Kolors", ref_scale=0.3)
         ```
     """
-    download_kolors_weights(ckpt_dir)
+    download_kolors_weights(ckpt_dir, include_ip_adapter=True)
 
     text_encoder = ChatGLMModel.from_pretrained(
         f"{ckpt_dir}/text_encoder", torch_dtype=torch.float16