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()] )