Spaces:
Running
on
Zero
Running
on
Zero
# Set comprehensive OpenMP environment variables to prevent ANY fork() errors | |
# This must be done BEFORE importing any libraries that use OpenMP (torch, numpy, rembg, etc.) | |
import os | |
os.environ["OMP_NUM_THREADS"] = "1" # Limit OpenMP to single thread | |
os.environ["MKL_NUM_THREADS"] = "1" # Intel MKL threading | |
os.environ["NUMEXPR_NUM_THREADS"] = "1" # NumExpr threading | |
os.environ["OPENBLAS_NUM_THREADS"] = "1" # OpenBLAS threading | |
os.environ["VECLIB_MAXIMUM_THREADS"] = "1" # Apple vecLib threading | |
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE" # Avoid Intel MKL errors | |
os.environ["OMP_THREAD_LIMIT"] = "1" # Limit total number of OpenMP threads | |
os.environ["OMP_DISPLAY_ENV"] = "FALSE" # Suppress OpenMP environment display | |
os.environ["KMP_WARNINGS"] = "FALSE" # Suppress KMP warnings | |
os.environ["PYTHONFAULTHANDLER"] = "1" # Better error reporting | |
# Additional settings to prevent subprocess/multiprocessing issues in rembg and other libraries | |
os.environ["TOKENIZERS_PARALLELISM"] = "false" # Disable tokenizer parallelism | |
os.environ["OMP_WAIT_POLICY"] = "PASSIVE" # Use passive waiting | |
os.environ["KMP_INIT_AT_FORK"] = "FALSE" # Don't initialize OpenMP at fork | |
# Hunyuan 3D is licensed under the TENCENT HUNYUAN NON-COMMERCIAL LICENSE AGREEMENT | |
# except for the third-party components listed below. | |
# Hunyuan 3D does not impose any additional limitations beyond what is outlined | |
# in the repsective licenses of these third-party components. | |
# Users must comply with all terms and conditions of original licenses of these third-party | |
# components and must ensure that the usage of the third party components adheres to | |
# all relevant laws and regulations. | |
# For avoidance of doubts, Hunyuan 3D means the large language models and | |
# their software and algorithms, including trained model weights, parameters (including | |
# optimizer states), machine-learning model code, inference-enabling code, training-enabling code, | |
# fine-tuning enabling code and other elements of the foregoing made publicly available | |
# by Tencent in accordance with TENCENT HUNYUAN COMMUNITY LICENSE AGREEMENT. | |
# Apply torchvision compatibility fix before other imports | |
import sys | |
sys.path.insert(0, './hy3dshape') | |
sys.path.insert(0, './hy3dpaint') | |
pythonpath = sys.executable | |
print(pythonpath) | |
try: | |
from torchvision_fix import apply_fix | |
apply_fix() | |
except ImportError: | |
print("Warning: torchvision_fix module not found, proceeding without compatibility fix") | |
except Exception as e: | |
print(f"Warning: Failed to apply torchvision fix: {e}") | |
import os | |
import random | |
import shutil | |
import subprocess | |
import time | |
from glob import glob | |
from pathlib import Path | |
import base64 | |
import json | |
import threading | |
from typing import Dict, Optional, Any | |
from enum import Enum | |
import gradio as gr | |
import torch | |
import trimesh | |
import uvicorn | |
from fastapi import FastAPI, HTTPException, BackgroundTasks, File, Form, UploadFile | |
from fastapi.staticfiles import StaticFiles | |
from fastapi.responses import JSONResponse | |
from pydantic import BaseModel | |
import uuid | |
import numpy as np | |
from PIL import Image | |
import io | |
from hy3dshape.utils import logger | |
from hy3dpaint.convert_utils import create_glb_with_pbr_materials | |
# API Models | |
class JobStatus(Enum): | |
QUEUED = "queued" | |
PROCESSING = "processing" | |
COMPLETED = "completed" | |
FAILED = "failed" | |
class GenerateOptions(BaseModel): | |
enable_pbr: bool = True | |
should_remesh: bool = True | |
should_texture: bool = True # Critical for 3D model quality | |
class JobInfo: | |
def __init__(self, job_id: str): | |
self.job_id = job_id | |
self.status = JobStatus.QUEUED | |
self.progress = 0 | |
self.stage = "queued" | |
self.start_time = time.time() | |
self.end_time = None | |
self.error_message = None | |
self.model_urls = {} | |
self.images = {} | |
self.options = {} | |
# Global job storage | |
jobs: Dict[str, JobInfo] = {} | |
def create_job() -> str: | |
"""Create a new job and return its ID.""" | |
job_id = str(uuid.uuid4()) | |
jobs[job_id] = JobInfo(job_id) | |
return job_id | |
def update_job_status(job_id: str, status: JobStatus, progress: int = None, stage: str = None, error_message: str = None): | |
"""Update job status and progress.""" | |
if job_id in jobs: | |
jobs[job_id].status = status | |
if progress is not None: | |
jobs[job_id].progress = progress | |
if stage is not None: | |
jobs[job_id].stage = stage | |
if error_message is not None: | |
jobs[job_id].error_message = error_message | |
if status in [JobStatus.COMPLETED, JobStatus.FAILED]: | |
jobs[job_id].end_time = time.time() | |
def base64_to_pil_image(base64_string: str) -> Image.Image: | |
"""Convert base64 string to PIL Image.""" | |
try: | |
# Remove data URL prefix if present | |
if base64_string.startswith('data:image'): | |
base64_string = base64_string.split(',')[1] | |
# Ensure we have valid base64 data | |
# Add padding if necessary | |
missing_padding = len(base64_string) % 4 | |
if missing_padding: | |
base64_string += '=' * (4 - missing_padding) | |
# Decode base64 data | |
try: | |
image_data = base64.b64decode(base64_string) | |
except Exception as e: | |
raise ValueError(f"Failed to decode base64 string: {str(e)}") | |
# Ensure we have valid image data | |
if not image_data or len(image_data) == 0: | |
raise ValueError("Empty image data after base64 decoding") | |
# Open as PIL Image | |
image = Image.open(io.BytesIO(image_data)) | |
# Ensure consistent format - convert to RGBA | |
image = image.convert("RGBA") | |
return image | |
except Exception as e: | |
raise HTTPException(status_code=400, detail=f"Invalid image data: {str(e)}") | |
def process_generation_job(job_id: str, images: Dict[str, Image.Image], options: Dict[str, Any]): | |
"""Background task to process generation job.""" | |
global face_reduce_worker, tex_pipeline, HAS_TEXTUREGEN, SAVE_DIR | |
try: | |
update_job_status(job_id, JobStatus.PROCESSING, progress=10, stage="initializing") | |
# Images are already PIL Images | |
pil_images = images | |
# Extract options | |
enable_pbr = options.get("enable_pbr", True) | |
should_remesh = options.get("should_remesh", True) | |
should_texture = options.get("should_texture", True) | |
update_job_status(job_id, JobStatus.PROCESSING, progress=20, stage="preprocessing") | |
# Generate 3D mesh | |
# For non-MV mode, use the front image as the main image, or the first available image | |
main_image = pil_images.get('front') | |
if main_image is None and pil_images: | |
# If no front image, use the first available image | |
main_image = next(iter(pil_images.values())) | |
mesh, main_image, save_folder, stats, seed = _gen_shape( | |
caption=None, | |
image=main_image, | |
mv_image_front=pil_images.get('front'), | |
mv_image_back=pil_images.get('back'), | |
mv_image_left=pil_images.get('left'), | |
mv_image_right=pil_images.get('right'), | |
steps=75, | |
guidance_scale=9.0, | |
seed=1234, | |
octree_resolution=384, | |
check_box_rembg=True, | |
num_chunks=200000, | |
randomize_seed=False, | |
) | |
update_job_status(job_id, JobStatus.PROCESSING, progress=50, stage="shape_generation") | |
# After mesh generation and before exporting, print and store stats | |
try: | |
number_of_faces = mesh.faces.shape[0] if hasattr(mesh, 'faces') else None | |
number_of_vertices = mesh.vertices.shape[0] if hasattr(mesh, 'vertices') else None | |
logger.info(f"Mesh stats: faces={number_of_faces}, vertices={number_of_vertices}") | |
except Exception as e: | |
logger.warning(f"Failed to log mesh stats: {e}") | |
# Print generation parameters for traceability | |
try: | |
logger.info(f"Generation parameters: seed={seed}, steps=75, octree_resolution=384, guidance_scale=9.0, num_chunks=200000, target_face_count=15000") | |
except Exception as e: | |
logger.warning(f"Failed to log generation parameters: {e}") | |
# Export white mesh | |
white_mesh_path = export_mesh(mesh, save_folder, textured=False, type='obj') | |
# Face reduction | |
mesh = face_reduce_worker(mesh) | |
reduced_mesh_path = export_mesh(mesh, save_folder, textured=False, type='obj') | |
update_job_status(job_id, JobStatus.PROCESSING, progress=70, stage="face_reduction") | |
# Texture generation if enabled | |
textured_mesh_path = None | |
if should_texture: | |
if HAS_TEXTUREGEN: | |
try: | |
text_path = os.path.join(save_folder, 'textured_mesh.obj') | |
# Use GPU function for texture generation with lazy initialization | |
try: | |
logger.info(f"Starting texture generation for job {job_id}") | |
# Count available images to adapt texture generation settings | |
num_images = len(pil_images) | |
logger.info(f"Job {job_id} has {num_images} images available") | |
except Exception as e: | |
logger.warning(f"Failed to log texture generation start: {e}") | |
num_images = len(pil_images) if pil_images else 1 | |
# Try texture generation with adaptive settings based on available images | |
textured_mesh_path = generate_texture_lazy_adaptive( | |
mesh_path=reduced_mesh_path, | |
image_path=main_image, | |
output_mesh_path=text_path, | |
num_available_images=num_images | |
) | |
if textured_mesh_path and os.path.exists(textured_mesh_path): | |
try: | |
logger.info(f"Texture generation completed for job {job_id}") | |
except Exception as e: | |
logger.warning(f"Failed to log texture completion: {e}") | |
# Convert to GLB | |
glb_path_textured = os.path.join(save_folder, 'textured_mesh.glb') | |
quick_convert_with_obj2gltf(textured_mesh_path, glb_path_textured) | |
textured_mesh_path = glb_path_textured | |
else: | |
try: | |
logger.warning(f"Texture generation returned None or file doesn't exist for job {job_id}") | |
except Exception as e: | |
logger.warning(f"Failed to log texture warning: {e}") | |
textured_mesh_path = None | |
except Exception as e: | |
logger.error(f"Texture generation failed for job {job_id}: {e}") | |
# Continue without texture - user will get the white mesh | |
textured_mesh_path = None | |
else: | |
try: | |
logger.warning(f"Texture generation requested for job {job_id} but texture pipeline is not available") | |
except Exception as e: | |
logger.warning(f"Failed to log texture unavailable warning: {e}") | |
update_job_status(job_id, JobStatus.PROCESSING, progress=75, stage="texture_generation_unavailable", | |
message="Texture generation is not available - returning mesh without texture") | |
update_job_status(job_id, JobStatus.PROCESSING, progress=90, stage="finalizing") | |
# Prepare model URLs | |
model_urls = {} | |
if textured_mesh_path and os.path.exists(textured_mesh_path): | |
model_urls["glb"] = f"/static/{os.path.relpath(textured_mesh_path, SAVE_DIR)}" | |
else: | |
# Fallback to white mesh | |
white_glb_path = export_mesh(mesh, save_folder, textured=False, type='glb') | |
model_urls["glb"] = f"/static/{os.path.relpath(white_glb_path, SAVE_DIR)}" | |
# Add mesh stats to API output | |
model_urls["number_of_faces"] = number_of_faces | |
model_urls["number_of_vertices"] = number_of_vertices | |
# Update job with results | |
jobs[job_id].model_urls = model_urls | |
update_job_status(job_id, JobStatus.COMPLETED, progress=100, stage="completed") | |
except Exception as e: | |
logger.error(f"Job {job_id} failed: {e}") | |
update_job_status(job_id, JobStatus.FAILED, stage="failed", error_message=str(e)) | |
MAX_SEED = 1e7 | |
ENV = "Huggingface" # "Huggingface" | |
if ENV == 'Huggingface': | |
""" | |
Setup environment for running on Huggingface platform. | |
This block performs the following: | |
- Changes directory to the differentiable renderer folder and runs a shell | |
script to compile the mesh painter. | |
- Installs a custom rasterizer wheel package via pip. | |
Note: | |
This setup assumes the script is running in the Huggingface environment | |
with the specified directory structure. | |
""" | |
import os, spaces, subprocess, sys, shlex | |
from spaces import zero | |
def install_cuda_toolkit(): | |
# CUDA_TOOLKIT_URL = "https://developer.download.nvidia.com/compute/cuda/11.8.0/local_installers/cuda_11.8.0_520.61.05_linux.run" | |
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"], | |
) | |
# Fix: arch_list[-1] += '+PTX'; IndexError: list index out of range | |
os.environ["TORCH_CUDA_ARCH_LIST"] = "8.0;8.6" | |
def prepare_env(): | |
# print('install custom') | |
# os.system(f"cd /home/user/app/hy3dpaint/custom_rasterizer && {pythonpath} -m pip install -e .") | |
# os.system(f"cd /home/user/app/hy3dpaint/packages/custom_rasterizer && pip install -e .") | |
subprocess.run(shlex.split("pip install custom_rasterizer-0.1-cp310-cp310-linux_x86_64.whl"), check=True) | |
print("cd /home/user/app/hy3dpaint/differentiable_renderer/ && bash compile_mesh_painter.sh") | |
os.system("cd /home/user/app/hy3dpaint/DifferentiableRenderer && bash compile_mesh_painter.sh") | |
print("Downloading RealESRGAN model for texture enhancement...") | |
os.makedirs("/home/user/app/hy3dpaint/ckpt", exist_ok=True) | |
os.system("wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth -P /home/user/app/hy3dpaint/ckpt") | |
def check(): | |
import custom_rasterizer | |
print(type(custom_rasterizer)) | |
print(dir(custom_rasterizer)) | |
print(getattr(custom_rasterizer, '__file__', None)) | |
package_dir = None | |
if hasattr(custom_rasterizer, '__file__') and custom_rasterizer.__file__: | |
package_dir = os.path.dirname(custom_rasterizer.__file__) | |
elif hasattr(custom_rasterizer, '__path__'): | |
package_dir = list(custom_rasterizer.__path__)[0] | |
else: | |
raise RuntimeError("Cannot determine package path") | |
print(package_dir) | |
for root, dirs, files in os.walk(package_dir): | |
level = root.replace(package_dir, '').count(os.sep) | |
indent = ' ' * 4 * level | |
print(f"{indent}{os.path.basename(root)}/") | |
subindent = ' ' * 4 * (level + 1) | |
for f in files: | |
print(f"{subindent}{f}") | |
# print(torch.__version__) | |
# install_cuda_toolkit() | |
print(torch.__version__) | |
prepare_env() | |
check() | |
else: | |
""" | |
Define a dummy `spaces` module with a GPU decorator class for local environment. | |
The GPU decorator is a no-op that simply returns the decorated function unchanged. | |
This allows code that uses the `spaces.GPU` decorator to run without modification locally. | |
""" | |
class spaces: | |
class GPU: | |
def __init__(self, duration=60): | |
self.duration = duration | |
def __call__(self, func): | |
return func | |
def get_example_img_list(): | |
""" | |
Load and return a sorted list of example image file paths. | |
Searches recursively for PNG images under the './assets/example_images/' directory. | |
Returns: | |
list[str]: Sorted list of file paths to example PNG images. | |
""" | |
print('Loading example img list ...') | |
return sorted(glob('./assets/example_images/**/*.png', recursive=True)) | |
def get_example_txt_list(): | |
""" | |
Load and return a list of example text prompts. | |
Reads lines from the './assets/example_prompts.txt' file, stripping whitespace. | |
Returns: | |
list[str]: List of example text prompts. | |
""" | |
print('Loading example txt list ...') | |
txt_list = list() | |
for line in open('./assets/example_prompts.txt', encoding='utf-8'): | |
txt_list.append(line.strip()) | |
return txt_list | |
def gen_save_folder(max_size=200): | |
""" | |
Generate a new save folder inside SAVE_DIR, maintaining a maximum number of folders. | |
If the number of existing folders in SAVE_DIR exceeds `max_size`, the oldest folder is removed. | |
Args: | |
max_size (int, optional): Maximum number of folders to keep in SAVE_DIR. Defaults to 200. | |
Returns: | |
str: Path to the newly created save folder. | |
""" | |
os.makedirs(SAVE_DIR, exist_ok=True) | |
dirs = [f for f in Path(SAVE_DIR).iterdir() if f.is_dir()] | |
if len(dirs) >= max_size: | |
oldest_dir = min(dirs, key=lambda x: x.stat().st_ctime) | |
shutil.rmtree(oldest_dir) | |
print(f"Removed the oldest folder: {oldest_dir}") | |
new_folder = os.path.join(SAVE_DIR, str(uuid.uuid4())) | |
os.makedirs(new_folder, exist_ok=True) | |
print(f"Created new folder: {new_folder}") | |
return new_folder | |
# Removed complex PBR conversion functions - using simple trimesh-based conversion | |
def export_mesh(mesh, save_folder, textured=False, type='glb'): | |
""" | |
Export a mesh to a file in the specified folder, optionally including textures. | |
Args: | |
mesh (trimesh.Trimesh): The mesh object to export. | |
save_folder (str): Directory path where the mesh file will be saved. | |
textured (bool, optional): Whether to include textures/normals in the export. Defaults to False. | |
type (str, optional): File format to export ('glb' or 'obj' supported). Defaults to 'glb'. | |
Returns: | |
str: The full path to the exported mesh file. | |
""" | |
if textured: | |
path = os.path.join(save_folder, f'textured_mesh.{type}') | |
else: | |
path = os.path.join(save_folder, f'white_mesh.{type}') | |
if type not in ['glb', 'obj']: | |
mesh.export(path) | |
else: | |
mesh.export(path, include_normals=textured) | |
return path | |
def quick_convert_with_obj2gltf(obj_path: str, glb_path: str) -> bool: | |
# 执行转换 | |
textures = { | |
'albedo': obj_path.replace('.obj', '.jpg'), | |
'metallic': obj_path.replace('.obj', '_metallic.jpg'), | |
'roughness': obj_path.replace('.obj', '_roughness.jpg') | |
} | |
create_glb_with_pbr_materials(obj_path, textures, glb_path) | |
def randomize_seed_fn(seed: int, randomize_seed: bool) -> int: | |
if randomize_seed: | |
seed = random.randint(0, MAX_SEED) | |
return seed | |
def build_model_viewer_html(save_folder, height=660, width=790, textured=False): | |
# Remove first folder from path to make relative path | |
if textured: | |
related_path = f"./textured_mesh.glb" | |
template_name = './assets/modelviewer-textured-template.html' | |
output_html_path = os.path.join(save_folder, f'textured_mesh.html') | |
else: | |
related_path = f"./white_mesh.glb" | |
template_name = './assets/modelviewer-template.html' | |
output_html_path = os.path.join(save_folder, f'white_mesh.html') | |
offset = 50 if textured else 10 | |
with open(os.path.join(CURRENT_DIR, template_name), 'r', encoding='utf-8') as f: | |
template_html = f.read() | |
with open(output_html_path, 'w', encoding='utf-8') as f: | |
template_html = template_html.replace('#height#', f'{height - offset}') | |
template_html = template_html.replace('#width#', f'{width}') | |
template_html = template_html.replace('#src#', f'{related_path}/') | |
f.write(template_html) | |
rel_path = os.path.relpath(output_html_path, SAVE_DIR) | |
iframe_tag = f'<iframe src="/static/{rel_path}" \ | |
height="{height}" width="100%" frameborder="0"></iframe>' | |
print(f'Find html file {output_html_path}, \ | |
{os.path.exists(output_html_path)}, relative HTML path is /static/{rel_path}') | |
return f""" | |
<div style='height: {height}; width: 100%;'> | |
{iframe_tag} | |
</div> | |
""" | |
def _gen_shape( | |
caption=None, | |
image=None, | |
mv_image_front=None, | |
mv_image_back=None, | |
mv_image_left=None, | |
mv_image_right=None, | |
steps=75, | |
guidance_scale=9.0, | |
seed=1234, | |
octree_resolution=384, | |
check_box_rembg=False, | |
num_chunks=200000, | |
randomize_seed: bool = False, | |
): | |
# Check if we're using multi-view mode based on inputs | |
# Only consider non-None AND non-empty images for multi-view detection | |
using_multiview = ((mv_image_front is not None and mv_image_front is not False) or | |
(mv_image_back is not None and mv_image_back is not False) or | |
(mv_image_left is not None and mv_image_left is not False) or | |
(mv_image_right is not None and mv_image_right is not False)) | |
# Single image mode validation | |
if not using_multiview and image is None and caption is None: | |
raise gr.Error("Please provide either a caption or an image.") | |
# Multi-view mode validation and processing | |
if using_multiview: | |
# Check if any valid images were provided | |
has_valid_image = False | |
image = {} | |
if mv_image_front is not None: | |
image['front'] = mv_image_front | |
has_valid_image = True | |
if mv_image_back is not None: | |
image['back'] = mv_image_back | |
has_valid_image = True | |
if mv_image_left is not None: | |
image['left'] = mv_image_left | |
has_valid_image = True | |
if mv_image_right is not None: | |
image['right'] = mv_image_right | |
has_valid_image = True | |
if not has_valid_image: | |
raise gr.Error("Please provide at least one view image.") | |
seed = int(randomize_seed_fn(seed, randomize_seed)) | |
octree_resolution = int(octree_resolution) | |
if caption: print('prompt is', caption) | |
save_folder = gen_save_folder() | |
stats = { | |
'model': { | |
'shapegen': f'{args.model_path}/{args.subfolder}', | |
'texgen': f'{args.texgen_model_path}', | |
}, | |
'params': { | |
'caption': caption, | |
'steps': steps, | |
'guidance_scale': guidance_scale, | |
'seed': seed, | |
'octree_resolution': octree_resolution, | |
'check_box_rembg': check_box_rembg, | |
'num_chunks': num_chunks, | |
} | |
} | |
time_meta = {} | |
if image is None: | |
start_time = time.time() | |
try: | |
image = t2i_worker(caption) | |
except Exception as e: | |
raise gr.Error(f"Text to 3D is disable. \ | |
Please enable it by `python gradio_app.py --enable_t23d`.") | |
time_meta['text2image'] = time.time() - start_time | |
# remove disk io to make responding faster, uncomment at your will. | |
# image.save(os.path.join(save_folder, 'input.png')) | |
# Process images based on whether we're using multi-view mode | |
start_time = time.time() | |
if isinstance(image, dict): # Multi-view mode | |
for k, v in image.items(): | |
if v is not None and (check_box_rembg or v.mode == "RGB"): | |
try: | |
img = rmbg_worker(v.convert('RGB')) | |
image[k] = img | |
except Exception as e: | |
print(f"Error processing {k} view: {e}") | |
# Keep the original image if background removal fails | |
pass | |
else: # Single image mode | |
if image is not None and (check_box_rembg or image.mode == "RGB"): | |
try: | |
image = rmbg_worker(image.convert('RGB')) | |
except Exception as e: | |
print(f"Error removing background: {e}") | |
# Keep the original image if background removal fails | |
pass | |
time_meta['remove background'] = time.time() - start_time | |
# remove disk io to make responding faster, uncomment at your will. | |
# image.save(os.path.join(save_folder, 'rembg.png')) | |
# image to white model | |
start_time = time.time() | |
generator = torch.Generator() | |
generator = generator.manual_seed(int(seed)) | |
outputs = i23d_worker( | |
image=image, | |
num_inference_steps=steps, | |
guidance_scale=guidance_scale, | |
generator=generator, | |
octree_resolution=octree_resolution, | |
num_chunks=num_chunks, | |
output_type='mesh' | |
) | |
time_meta['shape generation'] = time.time() - start_time | |
logger.info("---Shape generation takes %s seconds ---" % (time.time() - start_time)) | |
tmp_start = time.time() | |
mesh = export_to_trimesh(outputs)[0] | |
time_meta['export to trimesh'] = time.time() - tmp_start | |
stats['number_of_faces'] = mesh.faces.shape[0] | |
stats['number_of_vertices'] = mesh.vertices.shape[0] | |
stats['time'] = time_meta | |
# Select the main image for display based on what's available | |
if isinstance(image, dict) and 'front' in image: | |
main_image = image['front'] # Use front view as main image in multi-view mode | |
else: | |
main_image = image # Use the single image in single-image mode | |
return mesh, main_image, save_folder, stats, seed | |
def generation_all( | |
caption=None, | |
image=None, | |
mv_image_front=None, | |
mv_image_back=None, | |
mv_image_left=None, | |
mv_image_right=None, | |
steps=75, | |
guidance_scale=9.0, | |
seed=1234, | |
octree_resolution=384, | |
check_box_rembg=False, | |
num_chunks=200000, | |
randomize_seed: bool = False, | |
): | |
start_time_0 = time.time() | |
mesh, image, save_folder, stats, seed = _gen_shape( | |
caption, | |
image, | |
mv_image_front=mv_image_front, | |
mv_image_back=mv_image_back, | |
mv_image_left=mv_image_left, | |
mv_image_right=mv_image_right, | |
steps=steps, | |
guidance_scale=guidance_scale, | |
seed=seed, | |
octree_resolution=octree_resolution, | |
check_box_rembg=check_box_rembg, | |
num_chunks=num_chunks, | |
randomize_seed=randomize_seed, | |
) | |
path = export_mesh(mesh, save_folder, textured=False) | |
print(path) | |
print('='*40) | |
# tmp_time = time.time() | |
# mesh = floater_remove_worker(mesh) | |
# mesh = degenerate_face_remove_worker(mesh) | |
# logger.info("---Postprocessing takes %s seconds ---" % (time.time() - tmp_time)) | |
# stats['time']['postprocessing'] = time.time() - tmp_time | |
tmp_time = time.time() | |
mesh = face_reduce_worker(mesh) | |
# path = export_mesh(mesh, save_folder, textured=False, type='glb') | |
path = export_mesh(mesh, save_folder, textured=False, type='obj') # 这样操作也会 core dump | |
logger.info("---Face Reduction takes %s seconds ---" % (time.time() - tmp_time)) | |
stats['time']['face reduction'] = time.time() - tmp_time | |
tmp_time = time.time() | |
text_path = os.path.join(save_folder, f'textured_mesh.obj') | |
path_textured = tex_pipeline(mesh_path=path, image_path=image, output_mesh_path=text_path, save_glb=False) | |
logger.info("---Texture Generation takes %s seconds ---" % (time.time() - tmp_time)) | |
stats['time']['texture generation'] = time.time() - tmp_time | |
tmp_time = time.time() | |
# Convert textured OBJ to GLB using obj2gltf with PBR support | |
glb_path_textured = os.path.join(save_folder, 'textured_mesh.glb') | |
conversion_success = quick_convert_with_obj2gltf(path_textured, glb_path_textured) | |
logger.info("---Convert textured OBJ to GLB takes %s seconds ---" % (time.time() - tmp_time)) | |
stats['time']['convert textured OBJ to GLB'] = time.time() - tmp_time | |
stats['time']['total'] = time.time() - start_time_0 | |
model_viewer_html_textured = build_model_viewer_html(save_folder, | |
height=HTML_HEIGHT, | |
width=HTML_WIDTH, textured=True) | |
if args.low_vram_mode: | |
torch.cuda.empty_cache() | |
return ( | |
gr.update(value=path), | |
gr.update(value=glb_path_textured), | |
model_viewer_html_textured, | |
stats, | |
seed, | |
) | |
def generate_texture_lazy_adaptive(mesh_path, image_path, output_mesh_path, num_available_images=1): | |
"""Generate texture for a mesh with adaptive settings based on available images.""" | |
try: | |
# Lazy initialization of texture pipeline inside GPU function | |
from hy3dpaint.textureGenPipeline import Hunyuan3DPaintPipeline, Hunyuan3DPaintConfig | |
# Use the same high-quality settings as the Gradio app | |
max_views = 9 | |
resolution = 768 | |
logger.info(f"Using high quality settings: {max_views} views, {resolution} resolution (same as Gradio app)") | |
conf = Hunyuan3DPaintConfig(max_num_view=max_views, resolution=resolution) | |
conf.realesrgan_ckpt_path = "hy3dpaint/ckpt/RealESRGAN_x4plus.pth" | |
conf.multiview_cfg_path = "hy3dpaint/cfgs/hunyuan-paint-pbr.yaml" | |
conf.custom_pipeline = "hy3dpaint/hunyuanpaintpbr" | |
# Initialize texture pipeline inside GPU function | |
local_tex_pipeline = Hunyuan3DPaintPipeline(conf) | |
# Generate texture with timeout handling | |
try: | |
textured_mesh_path = local_tex_pipeline( | |
mesh_path=mesh_path, | |
image_path=image_path, | |
output_mesh_path=output_mesh_path, | |
save_glb=False | |
) | |
return textured_mesh_path | |
except Exception as texture_error: | |
logger.error(f"Texture generation pipeline failed: {texture_error}") | |
# Try with medium quality settings as fallback | |
try: | |
fallback_views = 4 | |
fallback_resolution = 384 | |
logger.info(f"Trying fallback settings: {fallback_views} views, {fallback_resolution} resolution") | |
conf = Hunyuan3DPaintConfig(max_num_view=fallback_views, resolution=fallback_resolution) | |
conf.realesrgan_ckpt_path = "hy3dpaint/ckpt/RealESRGAN_x4plus.pth" | |
conf.multiview_cfg_path = "hy3dpaint/cfgs/hunyuan-paint-pbr.yaml" | |
conf.custom_pipeline = "hy3dpaint/hunyuanpaintpbr" | |
local_tex_pipeline = Hunyuan3DPaintPipeline(conf) | |
textured_mesh_path = local_tex_pipeline( | |
mesh_path=mesh_path, | |
image_path=image_path, | |
output_mesh_path=output_mesh_path, | |
save_glb=False | |
) | |
return textured_mesh_path | |
except Exception as fallback_error: | |
logger.error(f"Fallback texture generation also failed: {fallback_error}") | |
return None | |
except Exception as e: | |
logger.error(f"Texture generation initialization failed: {e}") | |
return None | |
def generate_texture_lazy(mesh_path, image_path, output_mesh_path): | |
"""Generate texture for a mesh using lazy initialization to avoid CUDA startup issues.""" | |
try: | |
# Lazy initialization of texture pipeline inside GPU function | |
from hy3dpaint.textureGenPipeline import Hunyuan3DPaintPipeline, Hunyuan3DPaintConfig | |
# Use fast settings optimized for 5-minute Hugging Face Spaces limit | |
conf = Hunyuan3DPaintConfig(max_num_view=2, resolution=256) | |
conf.realesrgan_ckpt_path = "hy3dpaint/ckpt/RealESRGAN_x4plus.pth" | |
conf.multiview_cfg_path = "hy3dpaint/cfgs/hunyuan-paint-pbr.yaml" | |
conf.custom_pipeline = "hy3dpaint/hunyuanpaintpbr" | |
# Initialize texture pipeline inside GPU function | |
local_tex_pipeline = Hunyuan3DPaintPipeline(conf) | |
# Generate texture with timeout handling | |
try: | |
textured_mesh_path = local_tex_pipeline( | |
mesh_path=mesh_path, | |
image_path=image_path, | |
output_mesh_path=output_mesh_path, | |
save_glb=False | |
) | |
return textured_mesh_path | |
except Exception as texture_error: | |
logger.error(f"Texture generation pipeline failed: {texture_error}") | |
# Try with even faster settings as fallback | |
try: | |
conf = Hunyuan3DPaintConfig(max_num_view=1, resolution=128) | |
conf.realesrgan_ckpt_path = "hy3dpaint/ckpt/RealESRGAN_x4plus.pth" | |
conf.multiview_cfg_path = "hy3dpaint/cfgs/hunyuan-paint-pbr.yaml" | |
conf.custom_pipeline = "hy3dpaint/hunyuanpaintpbr" | |
local_tex_pipeline = Hunyuan3DPaintPipeline(conf) | |
textured_mesh_path = local_tex_pipeline( | |
mesh_path=mesh_path, | |
image_path=image_path, | |
output_mesh_path=output_mesh_path, | |
save_glb=False | |
) | |
return textured_mesh_path | |
except Exception as fallback_error: | |
logger.error(f"Fallback texture generation also failed: {fallback_error}") | |
return None | |
except Exception as e: | |
logger.error(f"Texture generation initialization failed: {e}") | |
return None | |
def shape_generation( | |
caption=None, | |
image=None, | |
mv_image_front=None, | |
mv_image_back=None, | |
mv_image_left=None, | |
mv_image_right=None, | |
steps=75, | |
guidance_scale=9.0, | |
seed=1234, | |
octree_resolution=384, | |
check_box_rembg=False, | |
num_chunks=200000, | |
randomize_seed: bool = False, | |
): | |
start_time_0 = time.time() | |
mesh, image, save_folder, stats, seed = _gen_shape( | |
caption, | |
image, | |
mv_image_front=mv_image_front, | |
mv_image_back=mv_image_back, | |
mv_image_left=mv_image_left, | |
mv_image_right=mv_image_right, | |
steps=steps, | |
guidance_scale=guidance_scale, | |
seed=seed, | |
octree_resolution=octree_resolution, | |
check_box_rembg=check_box_rembg, | |
num_chunks=num_chunks, | |
randomize_seed=randomize_seed, | |
) | |
stats['time']['total'] = time.time() - start_time_0 | |
mesh.metadata['extras'] = stats | |
path = export_mesh(mesh, save_folder, textured=False) | |
model_viewer_html = build_model_viewer_html(save_folder, height=HTML_HEIGHT, width=HTML_WIDTH) | |
if args.low_vram_mode: | |
torch.cuda.empty_cache() | |
return ( | |
gr.update(value=path), | |
model_viewer_html, | |
stats, | |
seed, | |
) | |
def build_app(): | |
title = 'Hunyuan3D-2: High Resolution Textured 3D Assets Generation' | |
if MV_MODE: | |
title = 'Hunyuan3D-2mv: Image to 3D Generation with 1-4 Views' | |
if 'mini' in args.subfolder: | |
title = 'Hunyuan3D-2mini: Strong 0.6B Image to Shape Generator' | |
title = 'Hunyuan-3D-2.1' | |
if TURBO_MODE: | |
title = title.replace(':', '-Turbo: Fast ') | |
title_html = f""" | |
<div style="font-size: 2em; font-weight: bold; text-align: center; margin-bottom: 5px"> | |
{title} | |
</div> | |
<div align="center"> | |
Tencent Hunyuan3D Team | |
</div> | |
""" | |
custom_css = """ | |
.app.svelte-wpkpf6.svelte-wpkpf6:not(.fill_width) { | |
max-width: 1480px; | |
} | |
.mv-image button .wrap { | |
font-size: 10px; | |
} | |
.mv-image .icon-wrap { | |
width: 20px; | |
} | |
""" | |
with gr.Blocks(theme=gr.themes.Base(), title='Hunyuan-3D-2.1', analytics_enabled=False, css=custom_css) as demo: | |
gr.HTML(title_html) | |
with gr.Row(): | |
with gr.Column(scale=3): | |
with gr.Tabs(selected='tab_img_prompt') as tabs_prompt: | |
with gr.Tab('Image Prompt', id='tab_img_prompt', visible=True) as tab_ip: | |
image = gr.Image(label='Image', type='pil', image_mode='RGBA', height=290) | |
caption = gr.State(None) | |
# with gr.Tab('Text Prompt', id='tab_txt_prompt', visible=HAS_T2I and not MV_MODE) as tab_tp: | |
# caption = gr.Textbox(label='Text Prompt', | |
# placeholder='HunyuanDiT will be used to generate image.', | |
# info='Example: A 3D model of a cute cat, white background.') | |
with gr.Tab('MultiView Prompt', visible=True) as tab_mv: | |
# gr.Label('Please upload at least one front image.') | |
with gr.Row(): | |
mv_image_front = gr.Image(label='Front', type='pil', image_mode='RGBA', height=140, | |
min_width=100, elem_classes='mv-image') | |
mv_image_back = gr.Image(label='Back', type='pil', image_mode='RGBA', height=140, | |
min_width=100, elem_classes='mv-image') | |
with gr.Row(): | |
mv_image_left = gr.Image(label='Left', type='pil', image_mode='RGBA', height=140, | |
min_width=100, elem_classes='mv-image') | |
mv_image_right = gr.Image(label='Right', type='pil', image_mode='RGBA', height=140, | |
min_width=100, elem_classes='mv-image') | |
with gr.Row(): | |
btn = gr.Button(value='Gen Shape', variant='primary', min_width=100) | |
btn_all = gr.Button(value='Gen Textured Shape', | |
variant='primary', | |
visible=True, # Force visible for now, was: HAS_TEXTUREGEN | |
min_width=100) | |
with gr.Group(): | |
file_out = gr.File(label="File", visible=False) | |
file_out2 = gr.File(label="File", visible=False) | |
with gr.Tabs(selected='tab_options' if TURBO_MODE else 'tab_export'): | |
with gr.Tab("Options", id='tab_options', visible=TURBO_MODE): | |
gen_mode = gr.Radio( | |
label='Generation Mode', | |
info='Recommendation: Turbo for most cases, \ | |
Fast for very complex cases, Standard seldom use.', | |
choices=['Turbo', 'Fast', 'Standard'], | |
value='Turbo') | |
decode_mode = gr.Radio( | |
label='Decoding Mode', | |
info='The resolution for exporting mesh from generated vectset', | |
choices=['Low', 'Standard', 'High'], | |
value='Standard') | |
with gr.Tab('Advanced Options', id='tab_advanced_options'): | |
with gr.Row(): | |
check_box_rembg = gr.Checkbox( | |
value=True, | |
label='Remove Background', | |
min_width=100) | |
randomize_seed = gr.Checkbox( | |
label="Randomize seed", | |
value=True, | |
min_width=100) | |
seed = gr.Slider( | |
label="Seed", | |
minimum=0, | |
maximum=MAX_SEED, | |
step=1, | |
value=1234, | |
min_width=100, | |
) | |
with gr.Row(): | |
num_steps = gr.Slider(maximum=100, | |
minimum=1, | |
value=5 if 'turbo' in args.subfolder else 75, | |
step=1, label='Inference Steps') | |
octree_resolution = gr.Slider(maximum=512, | |
minimum=16, | |
value=384, | |
label='Octree Resolution') | |
with gr.Row(): | |
cfg_scale = gr.Number(value=9.0, label='Guidance Scale', min_width=100) | |
num_chunks = gr.Slider(maximum=5000000, minimum=1000, value=8000, | |
label='Number of Chunks', min_width=100) | |
with gr.Tab("Export", id='tab_export'): | |
with gr.Row(): | |
file_type = gr.Dropdown(label='File Type', | |
choices=SUPPORTED_FORMATS, | |
value='glb', min_width=100) | |
reduce_face = gr.Checkbox(label='Simplify Mesh', | |
value=False, min_width=100) | |
export_texture = gr.Checkbox(label='Include Texture', value=False, | |
visible=False, min_width=100) | |
target_face_num = gr.Slider(maximum=1000000, minimum=100, value=15000, | |
label='Target Face Number') | |
with gr.Row(): | |
confirm_export = gr.Button(value="Transform", min_width=100) | |
file_export = gr.DownloadButton(label="Download", variant='primary', | |
interactive=False, min_width=100) | |
with gr.Column(scale=6): | |
with gr.Tabs(selected='gen_mesh_panel') as tabs_output: | |
with gr.Tab('Generated Mesh', id='gen_mesh_panel'): | |
html_gen_mesh = gr.HTML(HTML_OUTPUT_PLACEHOLDER, label='Output') | |
with gr.Tab('Exporting Mesh', id='export_mesh_panel'): | |
html_export_mesh = gr.HTML(HTML_OUTPUT_PLACEHOLDER, label='Output') | |
with gr.Tab('Mesh Statistic', id='stats_panel'): | |
stats = gr.Json({}, label='Mesh Stats') | |
with gr.Column(scale=3 if MV_MODE else 2): | |
with gr.Tabs(selected='tab_img_gallery') as gallery: | |
with gr.Tab('Image to 3D Gallery', | |
id='tab_img_gallery', | |
visible=not MV_MODE) as tab_gi: | |
with gr.Row(): | |
gr.Examples(examples=example_is, inputs=[image], | |
label=None, examples_per_page=18) | |
tab_ip.select(fn=lambda: gr.update(selected='tab_img_gallery'), outputs=gallery) | |
#if HAS_T2I: | |
# tab_tp.select(fn=lambda: gr.update(selected='tab_txt_gallery'), outputs=gallery) | |
btn.click( | |
shape_generation, | |
inputs=[ | |
caption, | |
image, | |
mv_image_front, | |
mv_image_back, | |
mv_image_left, | |
mv_image_right, | |
num_steps, | |
cfg_scale, | |
seed, | |
octree_resolution, | |
check_box_rembg, | |
num_chunks, | |
randomize_seed, | |
], | |
outputs=[file_out, html_gen_mesh, stats, seed] | |
).then( | |
lambda: (gr.update(visible=False, value=False), gr.update(interactive=True), gr.update(interactive=True), | |
gr.update(interactive=False)), | |
outputs=[export_texture, reduce_face, confirm_export, file_export], | |
).then( | |
lambda: gr.update(selected='gen_mesh_panel'), | |
outputs=[tabs_output], | |
) | |
btn_all.click( | |
generation_all, | |
inputs=[ | |
caption, | |
image, | |
mv_image_front, | |
mv_image_back, | |
mv_image_left, | |
mv_image_right, | |
num_steps, | |
cfg_scale, | |
seed, | |
octree_resolution, | |
check_box_rembg, | |
num_chunks, | |
randomize_seed, | |
], | |
outputs=[file_out, file_out2, html_gen_mesh, stats, seed] | |
).then( | |
lambda: (gr.update(visible=True, value=True), gr.update(interactive=False), gr.update(interactive=True), | |
gr.update(interactive=False)), | |
outputs=[export_texture, reduce_face, confirm_export, file_export], | |
).then( | |
lambda: gr.update(selected='gen_mesh_panel'), | |
outputs=[tabs_output], | |
) | |
def on_gen_mode_change(value): | |
if value == 'Turbo': | |
return gr.update(value=5) | |
elif value == 'Fast': | |
return gr.update(value=10) | |
else: | |
return gr.update(value=30) | |
gen_mode.change(on_gen_mode_change, inputs=[gen_mode], outputs=[num_steps]) | |
def on_decode_mode_change(value): | |
if value == 'Low': | |
return gr.update(value=196) | |
elif value == 'Standard': | |
return gr.update(value=256) | |
else: | |
return gr.update(value=384) | |
decode_mode.change(on_decode_mode_change, inputs=[decode_mode], | |
outputs=[octree_resolution]) | |
def on_export_click(file_out, file_out2, file_type, | |
reduce_face, export_texture, target_face_num): | |
if file_out is None: | |
raise gr.Error('Please generate a mesh first.') | |
print(f'exporting {file_out}') | |
print(f'reduce face to {target_face_num}') | |
if export_texture: | |
mesh = trimesh.load(file_out2) | |
save_folder = gen_save_folder() | |
path = export_mesh(mesh, save_folder, textured=True, type=file_type) | |
# for preview | |
save_folder = gen_save_folder() | |
_ = export_mesh(mesh, save_folder, textured=True) | |
model_viewer_html = build_model_viewer_html(save_folder, | |
height=HTML_HEIGHT, | |
width=HTML_WIDTH, | |
textured=True) | |
else: | |
mesh = trimesh.load(file_out) | |
mesh = floater_remove_worker(mesh) | |
mesh = degenerate_face_remove_worker(mesh) | |
if reduce_face: | |
mesh = face_reduce_worker(mesh, target_face_num) | |
save_folder = gen_save_folder() | |
path = export_mesh(mesh, save_folder, textured=False, type=file_type) | |
# for preview | |
save_folder = gen_save_folder() | |
_ = export_mesh(mesh, save_folder, textured=False) | |
model_viewer_html = build_model_viewer_html(save_folder, | |
height=HTML_HEIGHT, | |
width=HTML_WIDTH, | |
textured=False) | |
print(f'export to {path}') | |
return model_viewer_html, gr.update(value=path, interactive=True) | |
confirm_export.click( | |
lambda: gr.update(selected='export_mesh_panel'), | |
outputs=[tabs_output], | |
).then( | |
on_export_click, | |
inputs=[file_out, file_out2, file_type, reduce_face, export_texture, target_face_num], | |
outputs=[html_export_mesh, file_export] | |
) | |
return demo | |
if __name__ == '__main__': | |
import argparse | |
parser = argparse.ArgumentParser() | |
parser.add_argument("--model_path", type=str, default='tencent/Hunyuan3D-2.1') | |
parser.add_argument("--subfolder", type=str, default='hunyuan3d-dit-v2-1') | |
parser.add_argument("--texgen_model_path", type=str, default='tencent/Hunyuan3D-2.1') | |
parser.add_argument('--port', type=int, default=7860) | |
parser.add_argument('--host', type=str, default='0.0.0.0') | |
parser.add_argument('--device', type=str, default='cuda') | |
parser.add_argument('--mc_algo', type=str, default='mc') | |
parser.add_argument('--cache-path', type=str, default='/root/save_dir') | |
parser.add_argument('--enable_t23d', action='store_true') | |
parser.add_argument('--disable_tex', action='store_true') | |
parser.add_argument('--enable_flashvdm', action='store_true') | |
parser.add_argument('--compile', action='store_true') | |
parser.add_argument('--low_vram_mode', action='store_true') | |
args = parser.parse_args() | |
args.enable_flashvdm = False | |
SAVE_DIR = args.cache_path | |
os.makedirs(SAVE_DIR, exist_ok=True) | |
CURRENT_DIR = os.path.dirname(os.path.abspath(__file__)) | |
MV_MODE = True # Force multi-view mode to be enabled | |
TURBO_MODE = 'turbo' in args.subfolder | |
HTML_HEIGHT = 690 if MV_MODE else 650 | |
HTML_WIDTH = 500 | |
HTML_OUTPUT_PLACEHOLDER = f""" | |
<div style='height: {650}px; width: 100%; border-radius: 8px; border-color: #e5e7eb; border-style: solid; border-width: 1px; display: flex; justify-content: center; align-items: center;'> | |
<div style='text-align: center; font-size: 16px; color: #6b7280;'> | |
<p style="color: #8d8d8d;">Welcome to Hunyuan3D!</p> | |
<p style="color: #8d8d8d;">No mesh here.</p> | |
</div> | |
</div> | |
""" | |
INPUT_MESH_HTML = """ | |
<div style='height: 490px; width: 100%; border-radius: 8px; | |
border-color: #e5e7eb; order-style: solid; border-width: 1px;'> | |
</div> | |
""" | |
example_is = get_example_img_list() | |
example_ts = get_example_txt_list() | |
SUPPORTED_FORMATS = ['glb', 'obj', 'ply', 'stl'] | |
HAS_TEXTUREGEN = False | |
if not args.disable_tex: | |
try: | |
print("Initializing texture generation pipeline...") | |
# Apply torchvision fix before importing basicsr/RealESRGAN | |
print("Applying torchvision compatibility fix for texture generation...") | |
try: | |
from torchvision_fix import apply_fix | |
fix_result = apply_fix() | |
if not fix_result: | |
print("Warning: Torchvision fix may not have been applied successfully") | |
else: | |
print("Torchvision fix applied successfully") | |
except ImportError as ie: | |
print(f"Warning: Could not import torchvision_fix: {ie}") | |
except Exception as fix_error: | |
print(f"Warning: Failed to apply torchvision fix: {fix_error}") | |
import traceback | |
traceback.print_exc() | |
# Import texture generation components | |
print("Importing texture generation components...") | |
try: | |
from hy3dpaint.textureGenPipeline import Hunyuan3DPaintPipeline, Hunyuan3DPaintConfig | |
print("Successfully imported texture generation components") | |
except ImportError as ie: | |
print(f"Failed to import texture generation components: {ie}") | |
raise | |
except Exception as e: | |
print(f"Error importing texture generation components: {e}") | |
import traceback | |
traceback.print_exc() | |
raise | |
# Configure texture pipeline | |
print("Configuring texture pipeline...") | |
conf = Hunyuan3DPaintConfig(max_num_view=9, resolution=768) | |
conf.realesrgan_ckpt_path = "hy3dpaint/ckpt/RealESRGAN_x4plus.pth" | |
conf.multiview_cfg_path = "hy3dpaint/cfgs/hunyuan-paint-pbr.yaml" | |
conf.custom_pipeline = "hy3dpaint/hunyuanpaintpbr" | |
# Initialize texture pipeline | |
print("Initializing texture pipeline...") | |
tex_pipeline = Hunyuan3DPaintPipeline(conf) | |
print("Texture pipeline initialized successfully") | |
# Not help much, ignore for now. | |
# if args.compile: | |
# texgen_worker.models['delight_model'].pipeline.unet.compile() | |
# texgen_worker.models['delight_model'].pipeline.vae.compile() | |
# texgen_worker.models['multiview_model'].pipeline.unet.compile() | |
# texgen_worker.models['multiview_model'].pipeline.vae.compile() | |
HAS_TEXTUREGEN = True | |
print("Texture generation is ENABLED - Gen Textured Shape button will be visible") | |
except Exception as e: | |
print(f"Error loading texture generator: {e}") | |
print("Failed to load texture generator.") | |
print('Please try to install requirements by following README.md') | |
import traceback | |
traceback.print_exc() | |
HAS_TEXTUREGEN = False | |
print("Texture generation is DISABLED - Gen Textured Shape button will be hidden") | |
# HAS_T2I = True | |
# if args.enable_t23d: | |
# from hy3dgen.text2image import HunyuanDiTPipeline | |
# t2i_worker = HunyuanDiTPipeline('Tencent-Hunyuan/HunyuanDiT-v1.1-Diffusers-Distilled') | |
# HAS_T2I = True | |
from hy3dshape import FaceReducer, FloaterRemover, DegenerateFaceRemover, MeshSimplifier, \ | |
Hunyuan3DDiTFlowMatchingPipeline | |
from hy3dshape.pipelines import export_to_trimesh | |
from hy3dshape.rembg import BackgroundRemover | |
rmbg_worker = BackgroundRemover() | |
i23d_worker = Hunyuan3DDiTFlowMatchingPipeline.from_pretrained( | |
args.model_path, | |
subfolder=args.subfolder, | |
use_safetensors=False, | |
device=args.device, | |
) | |
if args.enable_flashvdm: | |
mc_algo = 'mc' if args.device in ['cpu', 'mps'] else args.mc_algo | |
i23d_worker.enable_flashvdm(mc_algo=mc_algo) | |
if args.compile: | |
i23d_worker.compile() | |
floater_remove_worker = FloaterRemover() | |
degenerate_face_remove_worker = DegenerateFaceRemover() | |
face_reduce_worker = FaceReducer() | |
# https://discuss.huggingface.co/t/how-to-serve-an-html-file/33921/2 | |
# create a FastAPI app | |
app = FastAPI() | |
# API Endpoints | |
async def generate_3d_model( | |
front: UploadFile = File(None), | |
back: UploadFile = File(None), | |
left: UploadFile = File(None), | |
right: UploadFile = File(None), | |
options: str = Form("{}"), | |
background_tasks: BackgroundTasks = BackgroundTasks() | |
): | |
"""Generate 3D model from images using multipart/form-data.""" | |
try: | |
# Parse options | |
options_dict = json.loads(options) if options else {} | |
generate_options = { | |
"enable_pbr": options_dict.get("enable_pbr", True), | |
"should_remesh": options_dict.get("should_remesh", True), | |
"should_texture": options_dict.get("should_texture", True) | |
} | |
# Process uploaded files | |
images = {} | |
# Validate input - at least one image is required | |
if not front: | |
raise HTTPException(status_code=400, detail="Front image is required") | |
# Process each uploaded file | |
for view, file in { | |
"front": front, | |
"back": back, | |
"left": left, | |
"right": right | |
}.items(): | |
if file and file.filename: | |
# Read file content | |
contents = await file.read() | |
# Convert to PIL Image | |
try: | |
img = Image.open(io.BytesIO(contents)) | |
images[view] = img | |
except Exception as e: | |
raise HTTPException(status_code=400, detail=f"Invalid image for {view}: {str(e)}") | |
# Create job | |
job_id = create_job() | |
# Store job data (store file paths instead of actual images) | |
jobs[job_id].images = {k: "<uploaded_file>" for k in images.keys()} | |
jobs[job_id].options = generate_options | |
# Start background task | |
background_tasks.add_task( | |
process_generation_job, | |
job_id, | |
images, | |
generate_options | |
) | |
return JSONResponse({ | |
"job_id": job_id, | |
"status": "queued" | |
}) | |
except Exception as e: | |
raise HTTPException(status_code=500, detail=str(e)) | |
async def get_job_status(job_id: str): | |
"""Get job status and results.""" | |
if job_id not in jobs: | |
raise HTTPException(status_code=404, detail="Job not found") | |
job = jobs[job_id] | |
response = { | |
"status": job.status.value, | |
"progress": job.progress, | |
"stage": job.stage | |
} | |
if job.status == JobStatus.COMPLETED: | |
# Transform relative URLs to absolute URLs | |
absolute_model_urls = {} | |
# Determine base URL based on environment | |
if ENV == 'Huggingface': | |
base_url = "https://asimfayaz-hunyuan3d-2-1.hf.space" # TODO: Refactor this URL | |
else: | |
# For local development | |
host_for_url = "localhost" if args.host == "0.0.0.0" else args.host | |
base_url = f"http://{host_for_url}:{args.port}" | |
for key, relative_url in job.model_urls.items(): | |
# Handle both string and non-string values | |
if isinstance(relative_url, str): | |
if relative_url.startswith('/'): | |
absolute_model_urls[key] = f"{base_url}{relative_url}" | |
else: | |
absolute_model_urls[key] = relative_url | |
else: | |
# For non-string values (like integers), keep as is | |
absolute_model_urls[key] = relative_url | |
response["model_urls"] = absolute_model_urls | |
elif job.status == JobStatus.FAILED: | |
response["error"] = job.error_message | |
return JSONResponse(response) | |
async def health_check(): | |
"""Health check endpoint.""" | |
return JSONResponse({ | |
"status": "ok", | |
"version": "2.1" | |
}) | |
# create a static directory to store the static files | |
static_dir = Path(SAVE_DIR).absolute() | |
static_dir.mkdir(parents=True, exist_ok=True) | |
app.mount("/static", StaticFiles(directory=static_dir, html=True), name="static") | |
shutil.copytree('./assets/env_maps', os.path.join(static_dir, 'env_maps'), dirs_exist_ok=True) | |
if args.low_vram_mode: | |
torch.cuda.empty_cache() | |
demo = build_app() | |
app = gr.mount_gradio_app(app, demo, path="/") | |
if ENV == 'Huggingface': | |
# for Zerogpu | |
from spaces import zero | |
zero.startup() | |
uvicorn.run(app, host=args.host, port=args.port) | |