fix(texture): reuse pipeline for IP adapter

common.py

@@ -138,11 +138,7 @@ elif os.getenv("GRADIO_APP").startswith("textto3d"):
 elif os.getenv("GRADIO_APP") == "texture_edit":
     DELIGHT = DelightingModel()
     IMAGESR_MODEL = ImageRealESRGAN(outscale=4)
-    PIPELINE_IP = build_texture_gen_pipe(
-        base_ckpt_dir="./weights",
-        ip_adapt_scale=0.7,
-        device="cuda",
-    )
+    PIPELINE_HAS_IP_ADAPTER = False
     PIPELINE = build_texture_gen_pipe(
         base_ckpt_dir="./weights",
         ip_adapt_scale=0,
@@ -671,27 +667,62 @@ def generate_texture_mvimages(
     sub_idxs: tuple[tuple[int]] = ((0, 1, 2), (3, 4, 5)),
     req: gr.Request = None,
 ) -> list[str]:
+    global PIPELINE, PIPELINE_HAS_IP_ADAPTER
+
     output_root = os.path.join(TMP_DIR, str(req.session_hash))
     use_ip_adapter = True if ip_img_path and ip_adapt_scale > 0 else False
-    PIPELINE_IP.set_ip_adapter_scale([ip_adapt_scale])
-    img_save_paths = infer_pipe(
-        index_file=f"{output_root}/condition/index.json",
-        controlnet_cond_scale=controlnet_cond_scale,
-        guidance_scale=guidance_scale,
-        strength=strength,
-        num_inference_steps=num_inference_steps,
-        ip_adapt_scale=ip_adapt_scale,
-        ip_img_path=ip_img_path,
-        uid=uid,
-        prompt=prompt,
-        save_dir=f"{output_root}/multi_view",
-        sub_idxs=sub_idxs,
-        pipeline=PIPELINE_IP if use_ip_adapter else PIPELINE,
-        seed=seed,
-    )
 
-    gc.collect()
-    torch.cuda.empty_cache()
+    if PIPELINE is None:
+        PIPELINE = build_texture_gen_pipe(
+            base_ckpt_dir="./weights",
+            ip_adapt_scale=0,
+            device="cuda",
+        )
+
+    if use_ip_adapter and not PIPELINE_HAS_IP_ADAPTER:
+        logger.info("Load IP adapter into default texture pipeline")
+        if hasattr(PIPELINE.unet, "encoder_hid_proj"):
+            PIPELINE.unet.text_encoder_hid_proj = (
+                PIPELINE.unet.encoder_hid_proj
+            )
+        PIPELINE.load_ip_adapter(
+            "./weights/Kolors-IP-Adapter-Plus",
+            subfolder="",
+            weight_name=["ip_adapter_plus_general.bin"],
+        )
+        PIPELINE_HAS_IP_ADAPTER = True
+
+    if PIPELINE_HAS_IP_ADAPTER:
+        PIPELINE.set_ip_adapter_scale(
+            [ip_adapt_scale if use_ip_adapter else 0.0]
+        )
+
+    try:
+        img_save_paths = infer_pipe(
+            index_file=f"{output_root}/condition/index.json",
+            controlnet_cond_scale=controlnet_cond_scale,
+            guidance_scale=guidance_scale,
+            strength=strength,
+            num_inference_steps=num_inference_steps,
+            ip_adapt_scale=ip_adapt_scale if use_ip_adapter else 0.0,
+            ip_img_path=ip_img_path if use_ip_adapter else None,
+            uid=uid,
+            prompt=prompt,
+            save_dir=f"{output_root}/multi_view",
+            sub_idxs=sub_idxs,
+            pipeline=PIPELINE,
+            seed=seed,
+        )
+    finally:
+        if use_ip_adapter and PIPELINE_HAS_IP_ADAPTER:
+            logger.info("Unload IP adapter from default texture pipeline")
+            if hasattr(PIPELINE, "unload_ip_adapter"):
+                PIPELINE.unload_ip_adapter()
+            else:
+                PIPELINE = None
+            PIPELINE_HAS_IP_ADAPTER = False
+        gc.collect()
+        torch.cuda.empty_cache()
 
     return img_save_paths + img_save_paths
 

embodied_gen/scripts/render_mv.py

@@ -130,11 +130,22 @@ def infer_pipe(
             device=device,
         )
 
-    if ip_adapt_scale > 0 and ip_img_path is not None and len(ip_img_path) > 0:
+    has_ip_adapter = any(
+        hasattr(attn_processor, "to_k_ip")
+        for attn_processor in pipeline.unet.attn_processors.values()
+    )
+    use_ip_adapter = (
+        ip_adapt_scale > 0 and ip_img_path is not None and len(ip_img_path) > 0
+    )
+    effective_ip_adapt_scale = ip_adapt_scale if use_ip_adapter else 0.0
+    if use_ip_adapter:
         ip_image = Image.open(ip_img_path).convert("RGB")
         ip_image = ip_image.resize(target_hw[::-1])
         ip_image = [ip_image]
-        pipeline.set_ip_adapter_scale([ip_adapt_scale])
+        pipeline.set_ip_adapter_scale([effective_ip_adapt_scale])
+    elif has_ip_adapter:
+        ip_image = [Image.new("RGB", target_hw[::-1], color=(0, 0, 0))]
+        pipeline.set_ip_adapter_scale([0.0])
     else:
         ip_image = None
 
@@ -182,7 +193,7 @@ def infer_pipe(
         [str(item) for sublist in sub_idxs for item in sublist]
     )
     save_path = os.path.join(
-        save_dir, f"sample_idx{str(sub_idxs)}_ip{ip_adapt_scale}.jpg"
+        save_dir, f"sample_idx{str(sub_idxs)}_ip{effective_ip_adapt_scale}.jpg"
     )
     make_image_grid(grid_image, row_num, col_num).save(save_path)
     logger.info(f"Visualize in {save_path}")