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
# --- HF_TOKEN INTEGRATION ---
HF_TOKEN = os.environ.get("HF_TOKEN")
if not HF_TOKEN:
raise ValueError(
"HF_TOKEN environment variable must be set to access gated models."
)
# ----------------------------
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()
# --- HF_TOKEN is not neeeded ---- just check that HF_TOKEN exists---
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://rgndgn-i3d.hf.space/gradio_api/file={Path(output_path).resolve()}"
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:
input_img = gr.Image(
type="pil",
label=None, # Remove the label
show_label=False, # Further remove label
sources="upload",
image_mode="RGBA",
width=40,
elem_id="hidden-upload" # Add an ID for CSS targeting
)
# Make textbox visible but hide it with CSS
model_url = gr.Textbox(
label="Model URL",
elem_id="model-url-output", # Add this for CSS targeting
show_copy_button=True,
)
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()]
)