gradio_app.py

import spaces
import os
import tempfile
from typing import Any
import torch
import numpy as np
from PIL import Image
import gradio as gr
import trimesh
from transparent_background import Remover
from pathlib import Path
import subprocess
import uuid

def install_cuda_toolkit():
    CUDA_TOOLKIT_URL = "https://developer.download.nvidia.com/compute/cuda/12.2.0/local_installers/cuda_12.2.0_535.54.03_linux.run"
    CUDA_TOOLKIT_FILE = "/tmp/%s" % os.path.basename(CUDA_TOOLKIT_URL)
    subprocess.call(["wget", "-q", CUDA_TOOLKIT_URL, "-O", CUDA_TOOLKIT_FILE])
    subprocess.call(["chmod", "+x", CUDA_TOOLKIT_FILE])
    subprocess.call([CUDA_TOOLKIT_FILE, "--silent", "--toolkit"])

    os.environ["CUDA_HOME"] = "/usr/local/cuda"
    os.environ["PATH"] = "%s/bin:%s" % (os.environ["CUDA_HOME"], os.environ["PATH"])
    os.environ["LD_LIBRARY_PATH"] = "%s/lib:%s" % (
        os.environ["CUDA_HOME"],
        "" if "LD_LIBRARY_PATH" not in os.environ else os.environ["LD_LIBRARY_PATH"],
    )
    os.environ["TORCH_CUDA_ARCH_LIST"] = "8.0;8.6"

install_cuda_toolkit()

os.system("USE_CUDA=1 pip install -vv --no-build-isolation ./texture_baker ./uv_unwrapper")
import spar3d.utils as spar3d_utils
from spar3d.system import SPAR3D

COND_WIDTH = 512
COND_HEIGHT = 512
COND_DISTANCE = 2.2
COND_FOVY = 0.591627
BACKGROUND_COLOR = [0.5, 0.5, 0.5]
OUTPUT_DIR = "output"
os.makedirs(OUTPUT_DIR, exist_ok=True)

device = spar3d_utils.get_device()
bg_remover = Remover()
spar3d_model = SPAR3D.from_pretrained(
    "stabilityai/stable-point-aware-3d",
    config_name="config.yaml",
    weight_name="model.safetensors"
).eval().to(device)

c2w_cond = spar3d_utils.default_cond_c2w(COND_DISTANCE)
intrinsic, intrinsic_normed_cond = spar3d_utils.create_intrinsic_from_fov_rad(
    COND_FOVY, COND_HEIGHT, COND_WIDTH
)

def create_rgba_image(rgb_image: Image.Image, mask: np.ndarray = None) -> Image.Image:
    rgba_image = rgb_image.convert('RGBA')
    if mask is not None:
        if len(mask.shape) > 2:
            mask = mask.squeeze()
        alpha = Image.fromarray((mask * 255).astype(np.uint8))
        rgba_image.putalpha(alpha)
    return rgba_image

def create_batch(input_image: Image.Image) -> dict[str, Any]:
    resized_image = input_image.resize((COND_WIDTH, COND_HEIGHT))
    img_array = np.array(resized_image).astype(np.float32) / 255.0

    if img_array.shape[-1] == 4:
        rgb = img_array[..., :3]
        mask = img_array[..., 3:4]
    else:
        rgb = img_array
        mask = np.ones((*img_array.shape[:2], 1), dtype=np.float32)
    
    rgb = torch.from_numpy(rgb).float()
    mask = torch.from_numpy(mask).float()
    bg_tensor = torch.tensor(BACKGROUND_COLOR).view(1, 1, 3)
    rgb_cond = torch.lerp(bg_tensor, rgb, mask)
    rgb_cond = rgb_cond.unsqueeze(0)
    mask = mask.unsqueeze(0)
    
    batch = {
        "rgb_cond": rgb_cond,
        "mask_cond": mask,
        "c2w_cond": c2w_cond.unsqueeze(0),
        "intrinsic_cond": intrinsic.unsqueeze(0),
        "intrinsic_normed_cond": intrinsic_normed_cond.unsqueeze(0),
    }
    
    return batch

def forward_model(batch, system, guidance_scale=3.0, seed=0, device="cuda"):
    batch_size = batch["rgb_cond"].shape[0]
    assert batch_size == 1, f"Expected batch size 1, got {batch_size}"
    
    try:
        cond_tokens = system.forward_pdiff_cond(batch)
    except Exception as e:
        print("\n[ERROR] Failed in forward_pdiff_cond:")
        print(e)
        print("\nInput tensor properties:")
        print("rgb_cond dtype:", batch["rgb_cond"].dtype)
        print("rgb_cond device:", batch["rgb_cond"].device)
        print("rgb_cond requires_grad:", batch["rgb_cond"].requires_grad)
        raise
    
    sample_iter = system.sampler.sample_batch_progressive(
        batch_size,
        cond_tokens,
        guidance_scale=guidance_scale,
        device=device
    )
    
    for x in sample_iter:
        samples = x["xstart"]
    
    pc_cond = samples.permute(0, 2, 1).float()
    pc_cond = spar3d_utils.normalize_pc_bbox(pc_cond)
    pc_cond = pc_cond[:, torch.randperm(pc_cond.shape[1])[:512]]
    return pc_cond

@spaces.GPU
@torch.inference_mode()
def generate_and_process_3d(image: Image.Image) -> str:
    seed = np.random.randint(0, np.iinfo(np.int32).max)
    
    try:
        rgb_image = image.convert('RGB')
        no_bg_image = bg_remover.process(rgb_image)
        rgba_image = no_bg_image.convert('RGBA')
        
        processed_image = spar3d_utils.foreground_crop(
            rgba_image,
            crop_ratio=1.3,
            newsize=(COND_WIDTH, COND_HEIGHT),
            no_crop=False
        )
        
        batch = create_batch(processed_image)
        batch = {k: v.to(device) for k, v in batch.items()}
        pc_cond = forward_model(
            batch,
            spar3d_model,
            guidance_scale=3.0,
            seed=seed,
            device=device
        )
        batch["pc_cond"] = pc_cond

        with torch.no_grad():
            with torch.autocast(device_type='cuda' if torch.cuda.is_available() else 'cpu', dtype=torch.bfloat16):
                trimesh_mesh, _ = spar3d_model.generate_mesh(
                    batch,
                    1024,
                    remesh="none",
                    vertex_count=-1,
                    estimate_illumination=True
                )
                trimesh_mesh = trimesh_mesh[0]

        unique_id = str(uuid.uuid4())
        filename = f'model_{unique_id}.glb'
        output_path = os.path.join(OUTPUT_DIR, filename)
        trimesh_mesh.export(output_path, file_type="glb", include_normals=True)
        public_url = f"https://john6666-image-to-3d-test.hf.space/file={output_path}"
        
        return public_url
        
    except Exception as e:
        print(f"Error during generation: {str(e)}")
        import traceback
        traceback.print_exc()
        return None

# Create Gradio interface
with gr.Blocks() as demo:
    with gr.Row():
        input_img = gr.Image(
            type="pil",
            label=None,  # Remove the label
            show_label=False, # Further remove label
            sources="upload",
            image_mode="RGBA",
            elem_id="hidden-upload" # Add an ID for CSS targeting
        )
    
    # Make all output components invisible
    with gr.Row(visible=False):
        model_url = gr.Textbox(label="Model URL")
    
    input_img.upload(
        fn=generate_and_process_3d,
        inputs=[input_img],
        outputs=[model_url],
        api_name="generate"
    )
    
if __name__ == "__main__":
    demo.queue().launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=True,
        ssr_mode=False,
        allowed_paths=[Path(OUTPUT_DIR).resolve()]
    )