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 diffusers import DiffusionPipeline
# Import and setup SPAR3D
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
# Constants
COND_WIDTH = 512
COND_HEIGHT = 512
COND_DISTANCE = 2.2
COND_FOVY = 0.591627
BACKGROUND_COLOR = [0.5, 0.5, 0.5]
# Initialize models
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)
# Initialize FLUX model
dtype = torch.bfloat16
flux_pipe = DiffusionPipeline.from_pretrained(
"black-forest-labs/FLUX.1-schnell",
torch_dtype=dtype
).to(device)
# Initialize camera parameters
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:
"""Create an RGBA image from RGB image and optional mask."""
rgba_image = rgb_image.convert('RGBA')
if mask is not None:
# Ensure mask is 2D before converting to alpha
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]:
"""Prepare image batch for model input."""
# Resize and convert input image to numpy array
resized_image = input_image.resize((COND_WIDTH, COND_HEIGHT))
img_array = np.array(resized_image).astype(np.float32) / 255.0
# Extract RGB and alpha channels
if img_array.shape[-1] == 4: # RGBA
rgb = img_array[..., :3]
mask = img_array[..., 3:4]
else: # RGB
rgb = img_array
mask = np.ones((*img_array.shape[:2], 1), dtype=np.float32)
# Convert to tensors while keeping channel-last format
rgb = torch.from_numpy(rgb).float() # [H, W, 3]
mask = torch.from_numpy(mask).float() # [H, W, 1]
# Create background blend (match channel-last format)
bg_tensor = torch.tensor(BACKGROUND_COLOR).view(1, 1, 3) # [1, 1, 3]
# Blend RGB with background using mask (all in channel-last format)
rgb_cond = torch.lerp(bg_tensor, rgb, mask) # [H, W, 3]
# Move channels to correct dimension and add batch dimension
# Important: For SPAR3D image tokenizer, we need [B, H, W, C] format
rgb_cond = rgb_cond.unsqueeze(0) # [1, H, W, 3]
mask = mask.unsqueeze(0) # [1, H, W, 1]
# Create the batch dictionary
batch = {
"rgb_cond": rgb_cond, # [1, H, W, 3]
"mask_cond": mask, # [1, H, W, 1]
"c2w_cond": c2w_cond.unsqueeze(0), # [1, 4, 4]
"intrinsic_cond": intrinsic.unsqueeze(0), # [1, 3, 3]
"intrinsic_normed_cond": intrinsic_normed_cond.unsqueeze(0), # [1, 3, 3]
}
for k, v in batch.items():
print(f"[debug] {k} final shape:", v.shape)
return batch
def forward_model(batch, system, guidance_scale=3.0, seed=0, device="cuda"):
"""Process batch through model and generate point cloud."""
batch_size = batch["rgb_cond"].shape[0]
assert batch_size == 1, f"Expected batch size 1, got {batch_size}"
# Generate point cloud tokens
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 points
sample_iter = system.sampler.sample_batch_progressive(
batch_size,
cond_tokens,
guidance_scale=guidance_scale,
device=device
)
# Get final samples
for x in sample_iter:
samples = x["xstart"]
pc_cond = samples.permute(0, 2, 1).float()
# Normalize point cloud
pc_cond = spar3d_utils.normalize_pc_bbox(pc_cond)
# Subsample to 512 points
pc_cond = pc_cond[:, torch.randperm(pc_cond.shape[1])[:512]]
return pc_cond
def generate_and_process_3d(prompt: str) -> tuple[str | None, Image.Image | None]:
"""Generate image from prompt and convert to 3D model."""
width: int = 1024
height: int = 1024
# Generate random seed
seed = np.random.randint(0, np.iinfo(np.int32).max)
try:
# Set random seeds
torch.manual_seed(seed)
np.random.seed(seed)
# Generate image using FLUX
generator = torch.Generator(device=device).manual_seed(seed)
generated_image = flux_pipe(
prompt=prompt,
width=width,
height=height,
num_inference_steps=4,
generator=generator,
guidance_scale=0.0
).images[0]
rgb_image = generated_image.convert('RGB')
# bg_remover returns a PIL Image already, no need to convert
no_bg_image = bg_remover.process(rgb_image)
print(f"[debug] no_bg_image type: {type(no_bg_image)}, mode: {no_bg_image.mode}")
# Convert to RGBA if not already
rgba_image = no_bg_image.convert('RGBA')
print(f"[debug] rgba_image mode: {rgba_image.mode}")
processed_image = spar3d_utils.foreground_crop(
rgba_image,
crop_ratio=1.3,
newsize=(COND_WIDTH, COND_HEIGHT),
no_crop=False
)
# Show the processed image alpha channel for debugging
alpha = np.array(processed_image)[:, :, 3]
print(f"[debug] Alpha channel stats - min: {alpha.min()}, max: {alpha.max()}, unique: {np.unique(alpha)}")
# Prepare batch for processing
batch = create_batch(processed_image)
batch = {k: v.to(device) for k, v in batch.items()}
# Generate point cloud
pc_cond = forward_model(
batch,
spar3d_model,
guidance_scale=3.0,
seed=seed,
device=device
)
batch["pc_cond"] = pc_cond
# Generate mesh
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, # texture_resolution
remesh="none",
vertex_count=-1,
estimate_illumination=True
)
trimesh_mesh = trimesh_mesh[0]
# Export to GLB
temp_dir = tempfile.mkdtemp()
output_path = os.path.join(temp_dir, 'output.glb')
trimesh_mesh.export(output_path, file_type="glb", include_normals=True)
return output_path
except Exception as e:
print(f"Error during generation: {str(e)}")
import traceback
traceback.print_exc()
return None
# Create Gradio app using Blocks
with gr.Blocks() as demo:
gr.Markdown("# Text to 3D")
gr.Markdown("This space is based on [Stable Point-Aware 3D](https://huggingface.co/spaces/stabilityai/stable-point-aware-3d) by Stability AI.")
with gr.Row():
prompt_input = gr.Text(
label="Enter your prompt",
placeholder="eg. isometric 3D castle"
)
with gr.Row():
generate_btn = gr.Button("Generate", variant="primary")
with gr.Row():
model_output = gr.Model3D(
label="Generated .GLB model",
clear_color=[0.0, 0.0, 0.0, 0.0],
)
# Event handler
generate_btn.click(
fn=generate_and_process_3d,
inputs=[prompt_input],
outputs=[model_output],
api_name="generate"
)
if __name__ == "__main__":
demo.queue().launch()