Spaces:
Running
on
Zero
Running
on
Zero
# --- Environment Variables Used --- | |
# DISABLE_ZEROGPU: Set to 'true' or '1' to disable @spaces.GPU decorator (for Hugging Face Spaces). | |
# TRIPOSG_CODE_PATH: Absolute path to a local directory containing the checked-out TripoSG repository (scribble branch). | |
# GITHUB_TOKEN: A GitHub token used for cloning the TripoSG repo if TRIPOSG_CODE_PATH is not provided. | |
# WEIGHTS_PATH: Absolute path to a local directory containing the TripoSG-scribble model weights. | |
# HF_TOKEN: A Hugging Face Hub token used for downloading weights/models if local paths (WEIGHTS_PATH, WD14_CONVNEXT_PATH) are not provided. | |
# WD14_CONVNEXT_PATH: Absolute path to a local directory containing the WD14 ConvNeXT tagger model.onnx and selected_tags.csv. | |
# ---------------------------------- | |
import gradio as gr | |
import os | |
import sys | |
import subprocess | |
from huggingface_hub import snapshot_download, HfFolder, hf_hub_download | |
import random # Import random for seed generation | |
import re # For WD14 tag processing | |
import cv2 # For WD14 preprocessing | |
import pandas as pd # For WD14 tags | |
from onnxruntime import InferenceSession # For WD14 model | |
from typing import Mapping, Tuple, Dict # Type hints | |
# --- Repo Setup --- | |
DEFAULT_REPO_DIR = "./TripoSG-repo" # Directory to clone into if not using local path | |
REPO_GIT_URL = "github.com/VAST-AI-Research/TripoSG.git" # Base URL without schema/token | |
BRANCH = "scribble" | |
code_source_path = None | |
# Option 1: Use local path if TRIPOSG_CODE_PATH env var is set | |
local_code_path = os.environ.get("TRIPOSG_CODE_PATH") | |
if local_code_path: | |
print(f"Attempting to use local code path specified by TRIPOSG_CODE_PATH: {local_code_path}") | |
# Basic check: does it exist and seem like a git repo (has .git)? | |
if os.path.isdir(local_code_path) and os.path.isdir(os.path.join(local_code_path, ".git")): | |
code_source_path = os.path.abspath(local_code_path) | |
print(f"Using local TripoSG code directory: {code_source_path}") | |
# You might want to add a check here to verify the branch is correct, e.g.: | |
# try: | |
# current_branch = subprocess.run(["git", "rev-parse", "--abbrev-ref", "HEAD"], cwd=code_source_path, check=True, capture_output=True, text=True).stdout.strip() | |
# if current_branch != BRANCH: | |
# print(f"Warning: Local repo is on branch '{current_branch}', expected '{BRANCH}'. Attempting checkout...") | |
# subprocess.run(["git", "checkout", BRANCH], cwd=code_source_path, check=True) | |
# except Exception as e: | |
# print(f"Warning: Could not verify or checkout branch '{BRANCH}' in {code_source_path}: {e}") | |
else: | |
print(f"Warning: TRIPOSG_CODE_PATH '{local_code_path}' not found or not a valid git repository directory. Falling back to cloning.") | |
# Option 2: Clone from GitHub (if local path not used or invalid) | |
if not code_source_path: | |
repo_url_to_clone = f"https://{REPO_GIT_URL}" | |
github_token = os.environ.get("GITHUB_TOKEN") | |
if github_token: | |
print("Using GITHUB_TOKEN for repository cloning.") | |
repo_url_to_clone = f"https://{github_token}@{REPO_GIT_URL}" | |
else: | |
print("No GITHUB_TOKEN found. Using public HTTPS for cloning.") | |
repo_target_dir = os.path.abspath(DEFAULT_REPO_DIR) | |
if not os.path.exists(repo_target_dir): | |
print(f"Cloning TripoSG repository ({BRANCH} branch) into {repo_target_dir}...") | |
try: | |
subprocess.run(["git", "clone", "--branch", BRANCH, "--depth", "1", repo_url_to_clone, repo_target_dir], check=True) | |
code_source_path = repo_target_dir | |
print("Repository cloned successfully.") | |
except subprocess.CalledProcessError as e: | |
print(f"Error cloning repository: {e}") | |
print("Please ensure the URL is correct, the branch '{BRANCH}' exists, and you have access rights (or provide a GITHUB_TOKEN).") | |
sys.exit(1) | |
except Exception as e: | |
print(f"An unexpected error occurred during cloning: {e}") | |
sys.exit(1) | |
else: | |
print(f"Directory {repo_target_dir} already exists. Assuming it contains the correct code/branch.") | |
# Optional: Add checks here like git pull or verifying the branch | |
code_source_path = repo_target_dir | |
if not code_source_path: | |
print("Error: Could not determine TripoSG code source path.") | |
sys.exit(1) | |
# Add repo to Python path | |
sys.path.insert(0, code_source_path) # Use the determined absolute path | |
print(f"Added {code_source_path} to sys.path") | |
# --- End Repo Setup --- | |
# --- ZeroGPU Setup --- | |
DISABLE_ZEROGPU = os.environ.get("DISABLE_ZEROGPU", "false").lower() in ("true", "1", "t") | |
ENABLE_ZEROGPU = not DISABLE_ZEROGPU | |
print(f"ZeroGPU Enabled: {ENABLE_ZEROGPU}") | |
# --- End ZeroGPU Setup --- | |
if ENABLE_ZEROGPU: | |
import spaces # Import spaces for ZeroGPU | |
from PIL import Image | |
import numpy as np | |
import torch | |
from triposg.pipelines.pipeline_triposg_scribble import TripoSGScribblePipeline | |
import tempfile | |
# --- Weight Loading Logic --- | |
HF_TOKEN = os.environ.get("HF_TOKEN") | |
if HF_TOKEN: | |
HfFolder.save_token(HF_TOKEN) | |
HUGGING_FACE_REPO_ID = "VAST-AI/TripoSG-scribble" | |
DEFAULT_CACHE_PATH = "./pretrained_weights/TripoSG-scribble" | |
# Option 1: Use local path if WEIGHTS_PATH env var is set | |
local_weights_path = os.environ.get("WEIGHTS_PATH") | |
model_load_path = None | |
if local_weights_path: | |
print(f"Attempting to load weights from local path specified by WEIGHTS_PATH: {local_weights_path}") | |
if os.path.isdir(local_weights_path): | |
model_load_path = local_weights_path | |
print(f"Using local weights directory: {model_load_path}") | |
else: | |
print(f"Warning: WEIGHTS_PATH '{local_weights_path}' not found or not a directory. Falling back to Hugging Face download.") | |
# Option 2: Download from Hugging Face (if local path not used or invalid) | |
if not model_load_path: | |
hf_token = os.environ.get("HF_TOKEN") | |
print(f"Attempting to download weights from Hugging Face repo: {HUGGING_FACE_REPO_ID}") | |
if hf_token: | |
print("Using Hugging Face token for download.") | |
auth_token = hf_token | |
else: | |
print("No Hugging Face token found. Attempting public download.") | |
auth_token = None | |
try: | |
model_load_path = snapshot_download( | |
repo_id=HUGGING_FACE_REPO_ID, | |
local_dir=DEFAULT_CACHE_PATH, | |
local_dir_use_symlinks=False, # Recommended for Spaces | |
token=auth_token, | |
# revision="main" # Specify branch/commit if needed | |
) | |
print(f"Weights downloaded/cached to: {model_load_path}") | |
except Exception as e: | |
print(f"Error downloading weights from Hugging Face: {e}") | |
print("Please ensure the repository exists and is accessible, or provide a valid WEIGHTS_PATH.") | |
sys.exit(1) # Exit if weights cannot be loaded | |
# Load the pipeline using the determined path | |
print(f"Loading pipeline from: {model_load_path}") | |
pipe = TripoSGScribblePipeline.from_pretrained(model_load_path) | |
pipe.to(dtype=torch.float16, device="cuda") | |
print("Pipeline loaded.") | |
# --- End Weight Loading Logic --- | |
# Create a white background image and a transparent layer for drawing | |
canvas_width, canvas_height = 512, 512 | |
initial_background = Image.new("RGB", (canvas_width, canvas_height), color="white") | |
initial_layer = Image.new("RGBA", (canvas_width, canvas_height), color=(0, 0, 0, 0)) # Transparent layer | |
# Prepare the initial value dictionary for ImageEditor | |
initial_value = { | |
"background": initial_background, | |
"layers": [initial_layer], # Add the transparent layer | |
"composite": None | |
} | |
# --- ZeroGPU Setup --- | |
# ... existing ZeroGPU setup ... | |
MAX_SEED = np.iinfo(np.int32).max | |
def get_random_seed(): | |
return random.randint(0, MAX_SEED) | |
# --- WD14 Helper Functions --- | |
def make_square(img, target_size): | |
old_size = img.shape[:2] | |
desired_size = max(old_size) | |
desired_size = max(desired_size, target_size) | |
delta_w = desired_size - old_size[1] | |
delta_h = desired_size - old_size[0] | |
top, bottom = delta_h // 2, delta_h - (delta_h // 2) | |
left, right = delta_w // 2, delta_w - (delta_w // 2) | |
color = [255, 255, 255] # White padding | |
return cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color) | |
def smart_resize(img, size): | |
if img.shape[0] > size: | |
img = cv2.resize(img, (size, size), interpolation=cv2.INTER_AREA) | |
elif img.shape[0] < size: | |
img = cv2.resize(img, (size, size), interpolation=cv2.INTER_CUBIC) | |
return img | |
RE_SPECIAL = re.compile(r'([\()])') | |
# --- WD14 Tagger Class --- | |
class WaifuDiffusionInterrogator: | |
def __init__( | |
self, | |
repo: str, | |
model_filename='model.onnx', | |
tags_filename='selected_tags.csv', | |
local_model_dir: str | None = None # Added local path option | |
) -> None: | |
self.__repo = repo | |
self.__model_filename = model_filename | |
self.__tags_filename = tags_filename | |
self.__local_model_dir = local_model_dir | |
self.__initialized = False | |
self._model = None | |
self._tags = None | |
def _init(self) -> None: | |
if self.__initialized: | |
return | |
model_path = None | |
tags_path = None | |
if self.__local_model_dir: | |
print(f"WD14: Attempting to load from local directory: {self.__local_model_dir}") | |
potential_model_path = os.path.join(self.__local_model_dir, self.__model_filename) | |
potential_tags_path = os.path.join(self.__local_model_dir, self.__tags_filename) | |
if os.path.exists(potential_model_path) and os.path.exists(potential_tags_path): | |
model_path = potential_model_path | |
tags_path = potential_tags_path | |
print("WD14: Found local model and tags file.") | |
else: | |
print("WD14: Local files not found. Falling back to Hugging Face download.") | |
if model_path is None or tags_path is None: | |
print(f"WD14: Downloading from repo: {self.__repo}") | |
hf_token = os.environ.get("HF_TOKEN") # Reuse HF token if available | |
try: | |
model_path = hf_hub_download(self.__repo, filename=self.__model_filename, token=hf_token) | |
tags_path = hf_hub_download(self.__repo, filename=self.__tags_filename, token=hf_token) | |
print("WD14: Download complete.") | |
except Exception as e: | |
print(f"WD14: Error downloading from Hugging Face: {e}") | |
# Decide how to handle this - maybe raise error or disable tagging? | |
# For now, we'll let it fail later if model is None | |
return # Cannot initialize | |
try: | |
self._model = InferenceSession(str(model_path)) | |
self._tags = pd.read_csv(tags_path) | |
self.__initialized = True | |
print("WD14: Tagger initialized successfully.") | |
except Exception as e: | |
print(f"WD14: Error initializing ONNX session or reading tags: {e}") | |
def _calculation(self, image: Image.Image) -> pd.DataFrame | None: | |
self._init() | |
if not self._model or self._tags is None: | |
print("WD14: Tagger not initialized.") | |
return None | |
_, height, _, _ = self._model.get_inputs()[0].shape | |
image = image.convert('RGBA') | |
new_image = Image.new('RGBA', image.size, 'WHITE') | |
new_image.paste(image, mask=image) | |
image = new_image.convert('RGB') | |
image = np.asarray(image) | |
image = image[:, :, ::-1] | |
image = make_square(image, height) | |
image = smart_resize(image, height) | |
image = image.astype(np.float32) | |
image = np.expand_dims(image, 0) | |
input_name = self._model.get_inputs()[0].name | |
label_name = self._model.get_outputs()[0].name | |
confidence = self._model.run([label_name], {input_name: image})[0] | |
full_tags = self._tags[['name', 'category']].copy() | |
full_tags['confidence'] = confidence[0] | |
return full_tags | |
def interrogate(self, image: Image.Image) -> Tuple[Dict[str, float], Dict[str, float]] | None: | |
full_tags = self._calculation(image) | |
if full_tags is None: | |
return None | |
ratings = dict(full_tags[full_tags['category'] == 9][['name', 'confidence']].values) | |
tags = dict(full_tags[full_tags['category'] != 9][['name', 'confidence']].values) | |
return ratings, tags | |
# --- Instantiate WD14 Tagger --- | |
WD14_CONVNEXT_REPO = 'SmilingWolf/wd-v1-4-convnext-tagger' | |
wd14_local_path = os.environ.get("WD14_CONVNEXT_PATH") | |
wd14_tagger = WaifuDiffusionInterrogator(repo=WD14_CONVNEXT_REPO, local_model_dir=wd14_local_path) | |
# --- Helper to format tags --- | |
def format_wd14_tags(tags: Dict[str, float], threshold: float = 0.35) -> str: | |
filtered_tags = { | |
tag: score for tag, score in tags.items() | |
if score >= threshold and "background" not in tag and tag not in {"monochrome", "greyscale", "no_humans", "comic", "solo"} | |
} | |
print(filtered_tags) | |
# Sort by score descending, then alphabetically | |
tags_pairs = sorted(filtered_tags.items(), key=lambda x: (-x[1], x[0])) | |
text_items = [tag.replace('_', ' ') for tag, score in tags_pairs] | |
return ', '.join(text_items) | |
# Apply decorator conditionally | |
def generate_3d(scribble_image_dict, prompt, scribble_confidence, text_confidence, seed): | |
print("Generating 3D model...") | |
input_prompt = prompt # Keep track of original prompt for return on early exit | |
if scribble_image_dict is None or scribble_image_dict.get("composite") is None: | |
print("No scribble image provided.") | |
return None, input_prompt # Return None for model, original prompt | |
# --- Prompt Handling --- | |
input_prompt = prompt.strip() | |
if not input_prompt: | |
print("Prompt is empty, attempting WD14 tagging...") | |
try: | |
# Get the user drawing (black on white) for tagging | |
user_drawing_img = Image.fromarray(scribble_image_dict["composite"]).convert("RGB") | |
tag_results = wd14_tagger.interrogate(user_drawing_img) | |
if tag_results: | |
ratings, tags = tag_results | |
generated_prompt = format_wd14_tags(tags) # Use default threshold | |
if generated_prompt: | |
print(f"WD14 generated prompt: {generated_prompt}") | |
input_prompt = generated_prompt | |
else: | |
print("WD14 tagging did not produce tags above threshold.") | |
input_prompt = "3d object" # Fallback prompt | |
else: | |
print("WD14 tagging failed or tagger not initialized.") | |
input_prompt = "3d object" # Fallback prompt | |
except Exception as e: | |
print(f"Error during WD14 tagging: {e}") | |
input_prompt = "3d object" # Fallback prompt | |
else: | |
print(f"Using user provided prompt: {input_prompt}") | |
# --- End Prompt Handling --- | |
# --- Seed Handling --- | |
current_seed = int(seed) | |
print(f"Using seed: {current_seed}") | |
# --- End Seed Handling --- | |
# --- Image Preprocessing for TripoSG --- | |
# Get the composite image again (safer in case dict is modified) | |
# The composite might be RGBA if a layer was involved, ensure RGB for processing | |
image_for_triposg = Image.fromarray(scribble_image_dict["composite"]).convert("RGB") | |
# Preprocess the image: invert colors (black on white -> white on black) | |
image_np = np.array(image_for_triposg) | |
processed_image_np = 255 - image_np | |
processed_image = Image.fromarray(processed_image_np) | |
print("Image preprocessed for TripoSG.") | |
# --- End Image Preprocessing --- | |
# --- Generator Setup --- | |
generator = torch.Generator(device='cuda').manual_seed(current_seed) | |
# --- End Generator Setup --- | |
# --- Run Pipeline --- | |
print("Running pipeline...") | |
try: | |
out = pipe( | |
processed_image, | |
prompt=input_prompt, # Use the potentially generated prompt | |
num_tokens=512, # Default value from example | |
guidance_scale=0, # Default value from example | |
num_inference_steps=16, # Default value from example | |
attention_kwargs={ | |
"cross_attention_scale": text_confidence, | |
"cross_attention_2_scale": scribble_confidence | |
}, | |
generator=generator, | |
use_flash_decoder=False, # Default value from example | |
dense_octree_depth=8, # Default value from example | |
hierarchical_octree_depth=8 # Default value from example | |
) | |
print("Pipeline finished.") | |
except Exception as e: | |
print(f"Error during pipeline execution: {e}") | |
return None, input_prompt # Return None for model, the prompt used | |
# --- End Run Pipeline --- | |
# --- Save Output --- | |
if out.meshes and len(out.meshes) > 0: | |
# Create a temporary file with .glb extension | |
with tempfile.NamedTemporaryFile(suffix=".glb", delete=False) as tmpfile: | |
output_path = tmpfile.name | |
out.meshes[0].export(output_path) | |
print(f"Mesh saved to temporary file: {output_path}") | |
return output_path, input_prompt # Return model path and the prompt used | |
else: | |
print("Pipeline did not generate any meshes.") | |
return None, input_prompt # Return None for model, the prompt used | |
# --- End Save Output --- | |
# Create the Gradio interface | |
with gr.Blocks() as demo: | |
gr.Markdown("# TripoSG Scribble!!") | |
gr.Markdown(""" | |
### [GitHub](https://github.com/VAST-AI-Research/TripoSG) | [Paper](https://arxiv.org/abs/2502.06608) | [Project Page](https://yg256li.github.io/TripoSG-Page/) | |
### Fast 3D shape prototyping with simple scribble and text prompt. Presented by [Tripo](https://www.tripo3d.ai/). | |
- For local deployment, simply clone this space, set up the environment and run with DISABLE_ZEROGPU=1. | |
- Feel free to tune the scribble confidence to balance fidelity and alignment :) | |
""") | |
with gr.Row(): | |
with gr.Column(scale=1): | |
image_input = gr.ImageEditor( | |
label="Scribble Input (Draw Black on White)", | |
value=initial_value, | |
image_mode="RGB", | |
brush=gr.Brush(default_color="#000000", color_mode="fixed", default_size=4), | |
interactive=True, | |
eraser=gr.Brush(default_color="#FFFFFF", color_mode="fixed", default_size=20), | |
canvas_size=(canvas_width, canvas_height), | |
fixed_canvas=True, | |
height=canvas_height + 128, | |
) | |
with gr.Column(scale=1): | |
with gr.Row(): | |
prompt_input = gr.Textbox(label="Prompt", placeholder="e.g., a cat", scale=3) | |
seed_input = gr.Number(label="Seed", value=0, precision=0, scale=1) | |
with gr.Row(): # Add row for sliders | |
confidence_input = gr.Slider(minimum=0.0, maximum=1.0, value=0.4, step=0.05, label="Scribble Confidence") | |
prompt_confidence_input = gr.Slider(minimum=0.0, maximum=1.0, value=1.0, step=0.05, label="Prompt Confidence") | |
with gr.Row(): | |
submit_button = gr.Button("Generate 3D Model", variant="primary", scale=1) | |
lucky_button = gr.Button("I'm Feeling Lucky", scale=1) | |
model_output = gr.Model3D(label="Generated 3D Model", interactive=False, height=384) | |
# Define the inputs for the main generation function | |
gen_inputs = [image_input, prompt_input, confidence_input, prompt_confidence_input, seed_input] # Added text_confidence_input | |
submit_button.click( | |
fn=generate_3d, | |
inputs=gen_inputs, | |
outputs=[model_output, prompt_input] # Add prompt_input to outputs | |
) | |
# Define inputs for the lucky button (same as main button for the final call) | |
lucky_gen_inputs = [image_input, prompt_input, confidence_input, prompt_confidence_input, seed_input] # Added text_confidence_input | |
lucky_button.click( | |
fn=get_random_seed, | |
inputs=[], | |
outputs=[seed_input] | |
).then( | |
fn=generate_3d, | |
inputs=lucky_gen_inputs, | |
outputs=[model_output, prompt_input] # Add prompt_input to outputs | |
) | |
# Launch with queue enabled if using ZeroGPU | |
print("Launching Gradio interface...") | |
demo.launch(share=False, server_name="0.0.0.0") | |
print("Gradio interface launched.") | |