Use local RMBG model, fix device placement, improve tqdm progress tracking

app.py

@@ -10,13 +10,10 @@ import numpy as np
 import base64
 import io
 import json
-import tempfile
 from datetime import datetime
 from typing import *
 from PIL import Image
 
-from gradio_client import Client as GradioClient, handle_file as gradio_handle_file
-
 import threading
 try:
     import nest_asyncio
@@ -140,7 +137,6 @@ def init_models():
         pipeline.image_cond_model_shape_1024 = build_image_cond_model(IMAGE_COND_CONFIGS["shape_1024"])
         pipeline.image_cond_model_tex_1024 = build_image_cond_model(IMAGE_COND_CONFIGS["tex_1024"])
         
-        pipeline.rembg_model = None  # Use remote BRIA-RMBG-2.0 instead
         pipeline.low_vram = False
         pipeline.cuda()
         
@@ -167,52 +163,6 @@ def init_models():
             'courtyard': EnvMap(torch.tensor(cv2.cvtColor(cv2.imread(os.path.join(_base, 'assets/hdri/courtyard.exr'), cv2.IMREAD_UNCHANGED), cv2.COLOR_BGR2RGB), dtype=torch.float32, device='cuda')),
         }
 
-# ============================================================================
-# Remote Background Removal (same as Microsoft TRELLIS.2 official)
-# ============================================================================
-
-rmbg_client = GradioClient("briaai/BRIA-RMBG-2.0")
-
-def remove_background_remote(input: Image.Image) -> Image.Image:
-    """Remove background using remote BRIA-RMBG-2.0 Space (no local GPU needed)."""
-    with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as f:
-        input = input.convert('RGB')
-        input.save(f.name)
-        output = rmbg_client.predict(gradio_handle_file(f.name), api_name="/image")[0][0]
-        result = Image.open(output)
-        os.unlink(f.name)
-        return result
-
-def preprocess_image_remote(input: Image.Image, bg_color: tuple = (0, 0, 0)) -> Image.Image:
-    """Preprocess image using remote rembg (no GPU required)."""
-    # If has alpha channel, use it directly
-    has_alpha = False
-    if input.mode == 'RGBA':
-        alpha = np.array(input)[:, :, 3]
-        if not np.all(alpha == 255):
-            has_alpha = True
-    max_size = max(input.size)
-    scale = min(1, 1024 / max_size)
-    if scale < 1:
-        input = input.resize((int(input.width * scale), int(input.height * scale)), Image.Resampling.LANCZOS)
-    if has_alpha:
-        output = input
-    else:
-        output = remove_background_remote(input)
-    output_np = np.array(output)
-    alpha = output_np[:, :, 3]
-    bbox = np.argwhere(alpha > 0.8 * 255)
-    bbox = np.min(bbox[:, 1]), np.min(bbox[:, 0]), np.max(bbox[:, 1]), np.max(bbox[:, 0])
-    center = (bbox[0] + bbox[2]) / 2, (bbox[1] + bbox[3]) / 2
-    size = max(bbox[2] - bbox[0], bbox[3] - bbox[1])
-    size = int(size * 1)
-    bbox = center[0] - size // 2, center[1] - size // 2, center[0] + size // 2, center[1] + size // 2
-    output = output.crop(bbox)
-    output = np.array(output).astype(np.float32) / 255
-    output = output[:, :, :3] * output[:, :, 3:4] + np.array(bg_color) / 255 * (1 - output[:, :, 3:4])
-    output = Image.fromarray((output * 255).astype(np.uint8))
-    return output
-
 # ============================================================================
 # Utilities
 # ============================================================================
@@ -395,9 +345,11 @@ async def progress_sse(request: Request):
     return StreamingResponse(event_stream(), media_type="text/event-stream")
 
 @app.api()
+@spaces.GPU(duration=30)
 def preprocess(image: FileData) -> FileData:
+    init_models()
     img = Image.open(image["path"])
-    processed = preprocess_image_remote(img)
+    processed = pipeline.preprocess_image(img)
     out_path = os.path.join(TMP_DIR, f"preprocessed_{int(time.time()*1000)}.png")
     processed.save(out_path)
     return FileData(path=out_path)
@@ -540,4 +492,4 @@ if __name__ == "__main__":
     # Pre-initialize models before launching the server
     init_models()
     
-    app.launch(show_error=True, share=True)
+    app.launch(show_error=True, share=True)

trellis2/pipelines/pixal3d_image_to_3d.py

@@ -120,7 +120,7 @@ class Pixal3DImageTo3DPipeline(Pipeline):
         pipeline.image_cond_model_shape_1024 = None
         pipeline.image_cond_model_tex_1024 = None
 
-        pipeline.rembg_model = None  # Skip local RMBG loading; use remote client instead
+        pipeline.rembg_model = getattr(rembg, args['rembg_model']['name'])(**args['rembg_model']['args'])
         
         pipeline.low_vram = args.get('low_vram', True)
         pipeline.default_pipeline_type = args.get('default_pipeline_type', '1024_cascade')