Upload app.py (#1)

- Upload app.py (9cf3d6349d362e183d2d53d432d5375e03d282ac)


Co-authored-by: Michael Lobanov <Michael-Lobanoff@users.noreply.huggingface.co>

app.py

@@ -6,9 +6,6 @@ import torch
 from PIL import Image
 import trimesh
 import random
-import urllib.request
-import io
-import base64
 from transformers import AutoModelForImageSegmentation
 from torchvision import transforms
 from huggingface_hub import hf_hub_download, snapshot_download
@@ -49,24 +46,17 @@ sys.path.append(MV_ADAPTER_CODE_DIR)
 sys.path.append(os.path.join(MV_ADAPTER_CODE_DIR, "scripts"))
 
 HEADER = """
-
 # 🔮 Image to 3D with [TripoSG](https://github.com/VAST-AI-Research/TripoSG)
-
 ## State-of-the-art Open Source 3D Generation Using Large-Scale Rectified Flow Transformers
-
 <p style="font-size: 1.1em;">By <a href="https://www.tripo3d.ai/" style="color: #1E90FF; text-decoration: none; font-weight: bold;">Tripo</a></p>
-
 ## 📋 Quick Start Guide:
 1. **Upload an image** (single object works best)
 2. Click **Generate Shape** to create the 3D mesh
 3. Click **Apply Texture** to add textures
 4. Use **Download GLB** to save your 3D model
 5. Adjust parameters under **Generation Settings** for fine-tuning
-
 Best results come from clean, well-lit images with clear subject isolation. Try it now!
-
 <p style="font-size: 0.9em; margin-top: 10px;">Texture generation powered by <a href="https://github.com/huanngzh/MV-Adapter" style="color: #1E90FF; text-decoration: none;">MV-Adapter</a> - a versatile multi-view adapter for consistent texture generation. Try the <a href="https://huggingface.co/spaces/VAST-AI/MV-Adapter-I2MV-SDXL" style="color: #1E90FF; text-decoration: none;">MV-Adapter demo</a> for multi-view image generation.</p>
-
 """
 
 # # triposg
@@ -122,184 +112,38 @@ def end_session(req: gr.Request):
     save_dir = os.path.join(TMP_DIR, str(req.session_hash))
     shutil.rmtree(save_dir)
 
+def normalize_image(image):
+    if image is None:
+        raise ValueError("Image is None")
+
+    if isinstance(image, Image.Image):
+        return image.convert("RGB")
+
+    if isinstance(image, np.ndarray):
+        return Image.fromarray(image).convert("RGB")
+
+    raise TypeError(f"Unsupported image type: {type(image)}")
+
 def get_random_hex():
     random_bytes = os.urandom(8)
     random_hex = random_bytes.hex()
     return random_hex
 
-def _image_from_value(value):
-    if value is None:
-        return None
-    if isinstance(value, Image.Image):
-        return value
-    if isinstance(value, np.ndarray):
-        return Image.fromarray(value)
-    if isinstance(value, dict):
-        path = value.get("path")
-        url = value.get("url")
-        data = value.get("data")
-        if isinstance(path, str):
-            try:
-                return Image.open(path)
-            except Exception:
-                return None
-        if isinstance(url, str) and url.startswith(("http://", "https://")):
-            try:
-                with urllib.request.urlopen(url) as resp:
-                    data = resp.read()
-                return Image.open(io.BytesIO(data))
-            except Exception:
-                return None
-        if data is not None:
-            try:
-                if isinstance(data, bytes):
-                    return Image.open(io.BytesIO(data))
-                if isinstance(data, str):
-                    if data.startswith("data:image/"):
-                        _, b64 = data.split(",", 1)
-                        return Image.open(io.BytesIO(base64.b64decode(b64)))
-                    return Image.open(io.BytesIO(base64.b64decode(data)))
-            except Exception:
-                return None
-    if hasattr(value, "path"):
-        path = getattr(value, "path", None)
-        if isinstance(path, str):
-            try:
-                return Image.open(path)
-            except Exception:
-                return None
-    if isinstance(value, str) and value.startswith("data:image/"):
-        try:
-            header, b64 = value.split(",", 1)
-            data = base64.b64decode(b64)
-            return Image.open(io.BytesIO(data))
-        except Exception:
-            return None
-    if isinstance(value, str):
-        try:
-            return Image.open(value)
-        except Exception:
-            return None
-    return None
-
-def _build_composite(image_dict):
-    composite_img = _image_from_value(image_dict.get("composite"))
-    if composite_img is not None:
-        return composite_img
-
-    background = _image_from_value(image_dict.get("background"))
-    layers = image_dict.get("layers") or []
-    layer_images = [_image_from_value(layer) for layer in layers]
-    layer_images = [img for img in layer_images if img is not None]
-
-    if background is None and not layer_images:
-        return None
-
-    if background is None and layer_images:
-        background = Image.new("RGBA", layer_images[0].size, (255, 255, 255, 255))
-    elif background is not None:
-        background = background.convert("RGBA")
-
-    base = background
-    for layer_img in layer_images:
-        base = Image.alpha_composite(base, layer_img.convert("RGBA"))
-
-    return base.convert("RGB")
-
-def _to_image_path(image_input, suffix=".png"):
-    if isinstance(image_input, dict):
-        path = image_input.get("path")
-        url = image_input.get("url")
-        data = image_input.get("data")
-        if isinstance(path, str):
-            try:
-                img = Image.open(path)
-                out_path = os.path.join(TMP_DIR, f"input_{get_random_hex()}{suffix}")
-                img.save(out_path)
-                return out_path
-            except Exception:
-                pass
-        if isinstance(url, str) and url.startswith(("http://", "https://")):
-            try:
-                with urllib.request.urlopen(url) as resp:
-                    data = resp.read()
-                out_path = os.path.join(TMP_DIR, f"input_{get_random_hex()}{suffix}")
-                with open(out_path, "wb") as f:
-                    f.write(data)
-                return out_path
-            except Exception:
-                pass
-        if data is not None:
-            try:
-                if isinstance(data, bytes):
-                    raw = data
-                elif isinstance(data, str):
-                    if data.startswith("data:image/"):
-                        _, b64 = data.split(",", 1)
-                        raw = base64.b64decode(b64)
-                    else:
-                        raw = base64.b64decode(data)
-                else:
-                    raw = None
-                if raw:
-                    out_path = os.path.join(TMP_DIR, f"input_{get_random_hex()}{suffix}")
-                    with open(out_path, "wb") as f:
-                        f.write(raw)
-                    return out_path
-            except Exception:
-                pass
-        image = _build_composite(image_input)
-    else:
-        if hasattr(image_input, "path"):
-            path = getattr(image_input, "path", None)
-            if isinstance(path, str):
-                try:
-                    img = Image.open(path)
-                    out_path = os.path.join(TMP_DIR, f"input_{get_random_hex()}{suffix}")
-                    img.save(out_path)
-                    return out_path
-                except Exception:
-                    pass
-        if isinstance(image_input, str) and os.path.isfile(image_input):
-            return image_input
-        image = _image_from_value(image_input)
-    if image is None:
-        return None
-    out_path = os.path.join(TMP_DIR, f"input_{get_random_hex()}{suffix}")
-    image.save(out_path)
-    return out_path
-
-def _describe_image_input(value):
-    try:
-        if isinstance(value, dict):
-            path = value.get("path")
-            url = value.get("url")
-            keys = list(value.keys())
-            return f"dict keys={keys} path={path} url={url}"
-        if hasattr(value, "path"):
-            return f"obj path={getattr(value, 'path', None)}"
-        return f"type={type(value)}"
-    except Exception as e:
-        return f"describe_error={e}"
-
 def get_random_seed(randomize_seed, seed):
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
     return seed
 
 @spaces.GPU(duration=180)
-def run_full(image: object, req: gr.Request):
+def run_full(image: Image, req: gr.Request):
     seed = 0
     num_inference_steps = 50
     guidance_scale = 7.5
     simplify = True
     target_face_num = DEFAULT_FACE_NUMBER
     
-    image_path = _to_image_path(image)
-    if image_path is None:
-        raise ValueError(f"No valid image provided. input={_describe_image_input(image)}")
-
-    image_seg = prepare_image(image_path, bg_color=np.array([1.0, 1.0, 1.0]), rmbg_net=rmbg_net)
+    image = normalize_image(image)
+    image_seg = prepare_image(image, bg_color=np.array([1.0, 1.0, 1.0]), rmbg_net=rmbg_net)
 
     outputs = triposg_pipe(
         image=image_seg,
@@ -359,7 +203,6 @@ def run_full(image: object, req: gr.Request):
         .to(DEVICE)
     )
 
-    image = Image.open(image_path)
     image = remove_bg_fn(image)
     image = preprocess_image(image, height, width)
 
@@ -411,12 +254,13 @@ def run_full(image: object, req: gr.Request):
 
 @spaces.GPU()
 @torch.no_grad()
-def run_segmentation(image: object):
-    image_path = _to_image_path(image)
-    if image_path is None:
-        raise ValueError(f"No valid image provided. input={_describe_image_input(image)}")
-    image = prepare_image(image_path, bg_color=np.array([1.0, 1.0, 1.0]), rmbg_net=rmbg_net)
-    return image
+def run_segmentation(image: Image.Image):
+    image = normalize_image(image)
+    return prepare_image(
+        image,
+        bg_color=np.array([1.0, 1.0, 1.0]),
+        rmbg_net=rmbg_net
+    )
 
 @spaces.GPU(duration=90)
 @torch.no_grad()
@@ -429,6 +273,7 @@ def image_to_3d(
     target_face_num: int,
     req: gr.Request
 ):
+    image = normalize_image(image)
     outputs = triposg_pipe(
         image=image,
         generator=torch.Generator(device=triposg_pipe.device).manual_seed(seed),
@@ -491,10 +336,7 @@ def run_texture(image: Image, mesh_path: str, seed: int, req: gr.Request):
         .to(DEVICE)
     )
 
-    image_path = _to_image_path(image)
-    if image_path is None:
-        raise ValueError("No valid image provided.")
-    image = Image.open(image_path)
+    image = normalize_image(image)
     image = remove_bg_fn(image)
     image = preprocess_image(image, height, width)
 
@@ -545,111 +387,95 @@ def run_texture(image: Image, mesh_path: str, seed: int, req: gr.Request):
     return textured_glb_path
 
 
+
+
+
 with gr.Blocks(title="TripoSG") as demo:
     gr.Markdown(HEADER)
 
-    def debug_echo_image(image_input: object):
-        image = _build_composite(image_input) if isinstance(image_input, dict) else _image_from_value(image_input)
-        if image is None:
-            raise ValueError("No valid image provided.")
-        return image.convert("RGB")
-
-    with gr.Tabs():
-        with gr.Tab("Main"):
+    with gr.Row():
+        with gr.Column():
             with gr.Row():
-                with gr.Column():
-                    with gr.Row():
-                        image_prompts = gr.ImageEditor(label="Input Image", image_mode="RGB")
-                        seg_image = gr.Image(
-                            label="Segmentation Result", type="pil", format="png", interactive=False
-                        )
-
-                    with gr.Accordion("Generation Settings", open=True):
-                        seed = gr.Slider(
-                            label="Seed",
-                            minimum=0,
-                            maximum=MAX_SEED,
-                            step=0,
-                            value=0
-                        )
-                        randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-                        num_inference_steps = gr.Slider(
-                            label="Number of inference steps",
-                            minimum=8,
-                            maximum=50,
-                            step=1,
-                            value=50,
-                        )
-                        guidance_scale = gr.Slider(
-                            label="CFG scale",
-                            minimum=0.0,
-                            maximum=20.0,
-                            step=0.1,
-                            value=7.0,
-                        )
-
-                        with gr.Row():
-                            reduce_face = gr.Checkbox(label="Simplify Mesh", value=True)
-                            target_face_num = gr.Slider(maximum=1000000, minimum=10000, value=DEFAULT_FACE_NUMBER, label="Target Face Number")
-
-                        gen_button = gr.Button("Generate Shape", variant="primary")
-                        gen_texture_button = gr.Button("Apply Texture", interactive=False)
-
-                with gr.Column():
-                    model_output = gr.Model3D(label="Generated GLB", interactive=False)
-                    textured_model_output = gr.Model3D(label="Textured GLB", interactive=False)
+                image_prompts = gr.Image(label="Input Image", type="pil")
+                seg_image = gr.Image(
+                    label="Segmentation Result", type="pil", format="png", interactive=False
+                )
 
-            with gr.Row():
-                examples = gr.Examples(
-                    examples=[
-                        f"{TRIPOSG_CODE_DIR}/assets/example_data/{image}"
-                        for image in os.listdir(f"{TRIPOSG_CODE_DIR}/assets/example_data")
-                    ],
-                    fn=run_full,
-                    inputs=[image_prompts],
-                    outputs=[seg_image, model_output, textured_model_output],
-                    cache_examples=True,
+            with gr.Accordion("Generation Settings", open=True):
+                seed = gr.Slider(
+                    label="Seed",
+                    minimum=0,
+                    maximum=MAX_SEED,
+                    step=0,
+                    value=0
+                )
+                randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+                num_inference_steps = gr.Slider(
+                    label="Number of inference steps",
+                    minimum=8,
+                    maximum=50,
+                    step=1,
+                    value=50,
                 )
+                guidance_scale = gr.Slider(
+                    label="CFG scale",
+                    minimum=0.0,
+                    maximum=20.0,
+                    step=0.1,
+                    value=7.0,
+                )
+
+                with gr.Row():
+                    reduce_face = gr.Checkbox(label="Simplify Mesh", value=True)
+                    target_face_num = gr.Slider(maximum=1000000, minimum=10000, value=DEFAULT_FACE_NUMBER, label="Target Face Number")
+
+                gen_button = gr.Button("Generate Shape", variant="primary")
+                gen_texture_button = gr.Button("Apply Texture", interactive=False)
+
+        with gr.Column():
+            model_output = gr.Model3D(label="Generated GLB", interactive=False)
+            textured_model_output = gr.Model3D(label="Textured GLB", interactive=False)
 
-            gen_button.click(
-                run_segmentation,
-                inputs=[image_prompts],
-                outputs=[seg_image]
-            ).then(
-                get_random_seed,
-                inputs=[randomize_seed, seed],
-                outputs=[seed],
-            ).then(
-                image_to_3d,
-                inputs=[
-                    seg_image,
-                    seed,
-                    num_inference_steps,
-                    guidance_scale,
-                    reduce_face,
-                    target_face_num
-                ],
-                outputs=[model_output]
-            ).then(lambda: gr.Button(interactive=True), outputs=[gen_texture_button])
-
-            gen_texture_button.click(
-                run_texture,
-                inputs=[image_prompts, model_output, seed],
-                outputs=[textured_model_output]
-            )
-
-        with gr.Tab("Test Upload"):
-            test_input = gr.ImageEditor(label="Test Image Input", image_mode="RGB")
-            test_button = gr.Button("Test Input")
-            test_output = gr.Image(label="Composite Output", type="pil")
-
-            test_button.click(
-                debug_echo_image,
-                inputs=[test_input],
-                outputs=[test_output]
-            )
+    with gr.Row():
+        examples = gr.Examples(
+            examples=[
+                f"{TRIPOSG_CODE_DIR}/assets/example_data/{image}"
+                for image in os.listdir(f"{TRIPOSG_CODE_DIR}/assets/example_data")
+            ],
+            fn=run_full,
+            inputs=[image_prompts],
+            outputs=[seg_image, model_output, textured_model_output],
+            cache_examples=True,
+        )
+
+    gen_button.click(
+        run_segmentation,
+        inputs=[image_prompts],
+        outputs=[seg_image]
+    ).then(
+        get_random_seed,
+        inputs=[randomize_seed, seed],
+        outputs=[seed],
+    ).then(
+        image_to_3d,
+        inputs=[
+            seg_image,
+            seed,
+            num_inference_steps,
+            guidance_scale,
+            reduce_face,
+            target_face_num
+        ],
+        outputs=[model_output]
+    ).then(lambda: gr.Button(interactive=True), outputs=[gen_texture_button])
+
+    gen_texture_button.click(
+        run_texture,
+        inputs=[image_prompts, model_output, seed],
+        outputs=[textured_model_output]
+    )
 
     demo.load(start_session)
     demo.unload(end_session)
 
-demo.launch(show_error=True)
+demo.launch()