Reverb commited on
Commit
8138769
·
1 Parent(s): 85a0e9b

Fix: preprocess image on remote TRELLIS.2 Space before calling image_to_3d

Browse files

image_to_3d uses preprocess_image=False — it expects the image to already
have background removed and be cropped to the object bounding box. In the
UI flow this is done by the image_prompt.upload() event handler before
generation runs. When calling via gradio_client we bypass that handler,
sending a raw RGB image that causes the model to raise RuntimeError.

Fix: add a Step 0 that calls /preprocess_image on the remote Space first.
This runs background removal (BRIA-RMBG via the Space's own rmbg_client)
and crops to bbox, producing the RGBA image image_to_3d expects.
Falls back to raw image if preprocessing endpoint is unavailable.

Files changed (1) hide show
  1. src/stages/stage1_generate.py +25 -3
src/stages/stage1_generate.py CHANGED
@@ -84,10 +84,28 @@ def generate_trellis(
84
 
85
  client = _get_trellis_client()
86
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
87
  try:
88
- from gradio_client import handle_file
89
  state, _html = client.predict(
90
- image=handle_file(image_path),
91
  seed=int(seed),
92
  resolution=preset["resolution"],
93
  ss_guidance_strength=7.5,
@@ -105,7 +123,11 @@ def generate_trellis(
105
  api_name="/image_to_3d",
106
  )
107
  except Exception as e:
108
- return f"❌ TRELLIS.2 generation failed:\n```\n{type(e).__name__}: {e}\n```"
 
 
 
 
109
 
110
  try:
111
  hp_src, _ = client.predict(
 
84
 
85
  client = _get_trellis_client()
86
 
87
+ from gradio_client import handle_file
88
+
89
+ # Step 0: preprocess on the remote Space (background removal + bbox crop).
90
+ # image_to_3d expects a preprocessed RGBA image — in the UI flow this is
91
+ # done by the image_prompt.upload() event before generation is triggered.
92
+ # When calling via API we must do it explicitly or the model crashes.
93
+ try:
94
+ preprocessed_path = client.predict(
95
+ handle_file(image_path),
96
+ api_name="/preprocess_image",
97
+ )
98
+ gen_image = handle_file(
99
+ preprocessed_path if isinstance(preprocessed_path, str) else image_path
100
+ )
101
+ except Exception:
102
+ # If remote preprocessing is unavailable, send the raw image and hope
103
+ # the image already has a transparent background.
104
+ gen_image = handle_file(image_path)
105
+
106
  try:
 
107
  state, _html = client.predict(
108
+ image=gen_image,
109
  seed=int(seed),
110
  resolution=preset["resolution"],
111
  ss_guidance_strength=7.5,
 
123
  api_name="/image_to_3d",
124
  )
125
  except Exception as e:
126
+ return (
127
+ f"❌ TRELLIS.2 generation failed:\n```\n{type(e).__name__}: {e}\n```\n\n"
128
+ f"Tip: make sure your image has a clean subject on a white/transparent "
129
+ f"background for best results."
130
+ )
131
 
132
  try:
133
  hp_src, _ = client.predict(