""" Pixal3D Gradio App Upload an image and generate a 3D mesh. Supports both automatic (MoGe) and fixed camera parameters. """ import os os.environ["no_proxy"] = os.environ.get("no_proxy", "") + ",localhost,127.0.0.1" import torch import tempfile import numpy as np from PIL import Image from torchvision import transforms import gradio as gr from pixal3dpipeline2stage import Pixal3DPipeline2Stage from pixal3dpipeline import Pixal3DPipeline import trimesh from trimesh.visual.material import PBRMaterial from trimesh.transformations import rotation_matrix # Static files directory for model viewer CURRENT_DIR = os.path.dirname(os.path.abspath(__file__)) SAVE_DIR = os.path.join(CURRENT_DIR, "gradio_outputs") # Global pipeline reference pipeline = None rmbg = None def load_pipeline(ckpt_dir="./ckpt", repo_id="Pixal3D/Pixal3D"): """Load all weights at startup.""" global pipeline, rmbg print("Loading Pixal3D 2-Stage pipeline (with MoGe + dense_check)...") pipeline = Pixal3DPipeline2Stage.from_pretrained( ckpt_dir=ckpt_dir, repo_id=repo_id, use_moge=True, use_dense_check=True, ) print("Pipeline loaded!") print("Loading BiRefNet for background removal...") from transformers import AutoModelForImageSegmentation birefnet_model = AutoModelForImageSegmentation.from_pretrained( 'ZhengPeng7/BiRefNet', trust_remote_code=True, ).to("cuda:0") birefnet_model.eval() rmbg = birefnet_model print("BiRefNet loaded!") def remove_background(image_np): """Use BiRefNet to remove background and add alpha channel. Input: numpy array (H, W, 3) RGB Output: numpy array (H, W, 4) RGBA """ pil_img = Image.fromarray(image_np[:, :, :3]).convert('RGB') image_size = (1024, 1024) transform_image = transforms.Compose([ transforms.Resize(image_size), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) input_tensor = transform_image(pil_img).unsqueeze(0).to("cuda:0") with torch.no_grad(): preds = rmbg(input_tensor)[-1].sigmoid().cpu() pred = preds[0].squeeze() pred_pil = transforms.ToPILImage()(pred) mask = pred_pil.resize(pil_img.size) mask = np.array(mask) rgba = np.concatenate([np.array(pil_img), mask[..., None]], axis=-1) return rgba def preprocess_image(image, use_rmbg): """Step 1: process image (background removal or use original), return immediately. use_rmbg=True: run BiRefNet to remove background and generate RGBA use_rmbg=False: directly use the original image (RGB or RGBA), skip background removal """ if image is None: return None if use_rmbg: # Run background removal if rmbg is None: gr.Warning("Background removal model not loaded.") return None processed = remove_background(image) else: # Directly use original image, no background removal processed = image os.makedirs("./gradio_outputs", exist_ok=True) Image.fromarray(processed).save("./gradio_outputs/processed.png") return processed def infer_mesh( processed, use_fixed_camera, camera_angle_x, mesh_scale, dense_steps, dense_guidance_scale, dense_seed, sparse_512_steps, sparse_512_guidance_scale, sparse_1024_steps, sparse_1024_guidance_scale, sparse_seed, dense_threshold, mc_threshold, ): """Step 2: run 3D inference on the already-processed image.""" if processed is None or pipeline is None: return None, None tmp_input = tempfile.NamedTemporaryFile(suffix=".png", delete=False) Image.fromarray(processed).save(tmp_input.name) input_path = tmp_input.name try: if use_fixed_camera: mesh = Pixal3DPipeline.infer( pipeline, image=input_path, camera_angle_x=camera_angle_x, mesh_scale=mesh_scale, dense_steps=int(dense_steps), dense_guidance_scale=dense_guidance_scale, dense_seed=int(dense_seed), sparse_512_steps=int(sparse_512_steps), sparse_512_guidance_scale=sparse_512_guidance_scale, sparse_1024_steps=int(sparse_1024_steps), sparse_1024_guidance_scale=sparse_1024_guidance_scale, sparse_seed=int(sparse_seed), dense_threshold=dense_threshold, mc_threshold=mc_threshold, ) else: mesh = pipeline.infer( image=input_path, mesh_scale=mesh_scale, optimize_mesh_scale=True, target_padding=3, max_optim_iterations=2, dense_steps=int(dense_steps), dense_guidance_scale=dense_guidance_scale, dense_seed=int(dense_seed), sparse_512_steps=int(sparse_512_steps), sparse_512_guidance_scale=sparse_512_guidance_scale, sparse_1024_steps=int(sparse_1024_steps), sparse_1024_guidance_scale=sparse_1024_guidance_scale, sparse_seed=int(sparse_seed), dense_threshold=dense_threshold, mc_threshold=mc_threshold, ) ply_file = tempfile.NamedTemporaryFile(suffix=".ply", delete=False) glb_file = tempfile.NamedTemporaryFile(suffix=".glb", delete=False) ply_path = ply_file.name glb_path = glb_file.name ply_file.close() glb_file.close() mesh.export(ply_path) # Export GLB with PBR material (same as hunyuan_app) material = PBRMaterial(baseColorFactor=[102, 102, 102, 255]) clean_mesh = trimesh.Trimesh(mesh.vertices, mesh.faces) clean_mesh.visual = trimesh.visual.TextureVisuals(material=material) # Rotate mesh to desired view angle (only X rotation needed) rot_x = rotation_matrix(np.radians(-90), [1, 0, 0]) clean_mesh.apply_transform(rot_x) clean_mesh.export(glb_path) return glb_path, ply_path except Exception as e: import traceback traceback.print_exc() return None, None finally: os.unlink(input_path) def build_ui(): # Custom CSS to hide the download button in Model3D custom_css = """ #model3d-viewer button[aria-label="下载"], #model3d-viewer button[aria-label="Download"], #model3d-viewer button[title="下载"], #model3d-viewer button[title="Download"] { display: none !important; } """ with gr.Blocks(title="Pixal3D", theme=gr.themes.Soft(), css=custom_css) as demo: gr.Markdown("# Pixal3D: Pixel-Aligned 3D Generation from Images") with gr.Row(): # Left column: input (scale=1) with gr.Column(scale=1): image_input = gr.Image(label="Input Image", type="numpy", image_mode=None) processed_image = gr.Image( label="Processed Image", image_mode="RGBA", type="numpy", interactive=False, ) use_rmbg = gr.Checkbox( label="Remove Background", value=True, info="Checked: auto remove background via BiRefNet. Unchecked: use original image directly.", ) use_fixed_camera = gr.Checkbox( label="Use Fixed Camera Parameters", value=False, info="If checked, use manually set FOV/distance/mesh_scale instead of MoGe auto-estimation.", ) with gr.Group(visible=False) as fixed_camera_group: gr.Markdown("### Camera Parameters (fixed mode)") camera_angle_x = gr.Number(value=0.2, label="camera_angle_x (rad)", step=0.01) with gr.Group(): gr.Markdown("### Mesh Scale") mesh_scale = gr.Number(value=0.5, label="mesh_scale", step=0.01, info="Initial mesh scale. Fixed mode default: 0.9, Auto mode default: 0.5") with gr.Accordion("Advanced Inference Parameters", open=False): dense_steps = gr.Number(value=50, label="Dense Steps", step=1, precision=0) dense_guidance_scale = gr.Number(value=7.0, label="Dense Guidance Scale", step=0.1) dense_seed = gr.Number(value=0, label="Dense Seed", step=1, precision=0) sparse_512_steps = gr.Number(value=30, label="Sparse 512 Steps", step=1, precision=0) sparse_512_guidance_scale = gr.Number(value=7.0, label="Sparse 512 Guidance Scale", step=0.1) sparse_1024_steps = gr.Number(value=15, label="Sparse 1024 Steps", step=1, precision=0) sparse_1024_guidance_scale = gr.Number(value=7.0, label="Sparse 1024 Guidance Scale", step=0.1) sparse_seed = gr.Number(value=0, label="Sparse Seed", step=1, precision=0) dense_threshold = gr.Number(value=0.1, label="Dense Threshold", step=0.01) mc_threshold = gr.Number(value=0.2, label="MC Threshold", step=0.01) run_btn = gr.Button("Generate 3D Mesh", variant="primary", size="lg") # Right column: output (scale=2) with gr.Column(scale=2): model_viewer = gr.Model3D(label="3D Mesh Preview", interactive=False, clear_color=[1.0, 1.0, 1.0, 1.0], elem_id="model3d-viewer") output_file = gr.File(label="Download .ply") # Toggle fixed camera group visibility and mesh_scale default def on_toggle_fixed(use_fixed): new_scale = 0.9 if use_fixed else 0.5 return gr.update(visible=use_fixed), gr.update(value=new_scale) use_fixed_camera.change( fn=on_toggle_fixed, inputs=[use_fixed_camera], outputs=[fixed_camera_group, mesh_scale], ) # Step 1: preprocess image → show processed image immediately # Step 2: run 3D inference → show mesh and download run_btn.click( fn=preprocess_image, inputs=[image_input, use_rmbg], outputs=[processed_image], ).then( fn=infer_mesh, inputs=[ processed_image, use_fixed_camera, camera_angle_x, mesh_scale, dense_steps, dense_guidance_scale, dense_seed, sparse_512_steps, sparse_512_guidance_scale, sparse_1024_steps, sparse_1024_guidance_scale, sparse_seed, dense_threshold, mc_threshold, ], outputs=[model_viewer, output_file], ) demo.queue(api_open=False) return demo if __name__ == "__main__": import argparse parser = argparse.ArgumentParser() parser.add_argument("--repo_id", type=str, default="TencentARC/Pixal3D-D") args = parser.parse_args() load_pipeline(repo_id=args.repo_id) demo = build_ui() demo.launch( server_name="127.0.0.1", share=True, )