bennyguo
minor adjustment
bd1714d
# --- 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
@spaces.GPU() if ENABLE_ZEROGPU else lambda func: func
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.")