import torch import numpy as np from PIL import Image import trimesh import tempfile from typing import Union, Optional, Dict, Any from pathlib import Path import os import logging import random import time import threading from huggingface_hub import snapshot_download import shutil # Set up detailed logging for 3D generation logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) class TimeoutError(Exception): """Custom timeout exception""" pass class Hunyuan3DGenerator: """3D model generation using Hunyuan3D-2.1 directly""" def __init__(self, device: str = "cuda"): logger.info(f"🔧 Initializing Hunyuan3DGenerator with device: {device}") self.device = device if torch.cuda.is_available() else "cpu" logger.info(f"🔧 Final device selection: {self.device}") self.model = None self.preprocessor = None # Model configuration self.model_id = "tencent/Hunyuan3D-2.1" self.model_path = None # Generation parameters self.num_inference_steps = 30 # Reduced for faster generation self.guidance_scale = 7.5 self.resolution = 256 # 3D resolution # Timeout configuration self.generation_timeout = 180 # 3 minutes timeout for local generation # Use full model since we have enough RAM logger.info(f"🔧 Using full Hunyuan3D-2.1 model") logger.info(f"⏱️ Generation timeout set to: {self.generation_timeout} seconds") def _check_vram(self) -> bool: """Check if we have enough VRAM for full model""" logger.info("🔍 Checking VRAM availability...") if not torch.cuda.is_available(): logger.info("❌ CUDA not available") return False try: vram = torch.cuda.get_device_properties(0).total_memory vram_gb = vram / (1024 * 1024 * 1024) logger.info(f"🔍 Available VRAM: {vram_gb:.2f} GB") # Need at least 12GB for full model has_enough = vram > 12 * 1024 * 1024 * 1024 logger.info(f"🔍 Has enough VRAM (>12GB): {has_enough}") return has_enough except Exception as e: logger.error(f"❌ Error checking VRAM: {e}") return False def load_model(self): """Load Hunyuan3D model and run necessary setup""" if self.model is None: logger.info("🚀 Starting Hunyuan3D model loading and setup...") try: import subprocess import sys import os def run_setup_command(command, cwd): logger.info(f"Running command: {' '.join(command)} in {cwd}") try: process = subprocess.run( command, check=True, capture_output=True, text=True, cwd=cwd ) logger.info(f"✅ Command successful.") if process.stdout: logger.info(f"STDOUT:\n{process.stdout}") if process.stderr: logger.warning(f"STDERR:\n{process.stderr}") except subprocess.CalledProcessError as e: logger.error(f"❌ Command failed with exit code {e.returncode}") logger.error(f"STDOUT:\n{e.stdout}") logger.error(f"STDERR:\n{e.stderr}") raise # Re-raise the exception to halt execution and see the error # Download model repository if not already present logger.info(f"📥 Downloading Hunyuan3D repository from {self.model_id}...") self.model_path = snapshot_download( repo_id=self.model_id, repo_type="space", cache_dir="./models/hunyuan3d_cache" ) logger.info(f"✅ Model repository downloaded to: {self.model_path}") # # List the contents of the downloaded directory for debugging # logger.info(f"🔍 Listing contents of {self.model_path}...") # run_setup_command(['ls', '-R'], cwd=self.model_path) # --- Installation and Compilation --- logger.info("🔧 Running Hunyuan3D setup scripts with detailed logging...") # 1. Install requirements from the model's specific requirements file # requirements_path = os.path.join(self.model_path, 'requirements_hunyuan3d.txt') # if os.path.exists(requirements_path): # pip_command = [ # sys.executable, '-m', 'pip', 'install', '-r', requirements_path, # '--extra-index-url', 'https://mirrors.cloud.tencent.com/pypi/simple/', # '--extra-index-url', 'https://mirrors.aliyun.com/pypi/simple' # ] # run_setup_command(pip_command, cwd=self.model_path) # 2. Install custom rasterizer dependencies (torch) # logger.info("Installing torch, torchvision, torchaudio...") # pip_command_torch = [sys.executable, '-m', 'pip', 'install', 'torch==2.5.1', 'torchvision==0.20.1', 'torchaudio==2.5.1', '--index-url', 'https://download.pytorch.org/whl/cu124'] # run_setup_command(pip_command_torch, cwd=self.model_path) # 3. Install custom rasterizer rasterizer_path = os.path.join(self.model_path, 'hy3dpaint', 'packages', 'custom_rasterizer') if os.path.exists(rasterizer_path): pip_command_rasterizer = [sys.executable, '-m', 'pip', 'install', '--no-build-isolation', '-e', '.'] run_setup_command(pip_command_rasterizer, cwd=rasterizer_path) # 4. Compile mesh painter renderer_path = os.path.join(self.model_path, 'hy3dpaint', 'DifferentiableRenderer') compile_script_path = os.path.join(renderer_path, 'compile_mesh_painter.sh') if os.path.exists(compile_script_path): bash_command = ['bash', compile_script_path] run_setup_command(bash_command, cwd=renderer_path) logger.info("✅ Hunyuan3D setup completed successfully.") # --- Pipeline Initialization --- logger.info("✈️ Initializing Hunyuan3D pipelines...") # Add subdirectories to Python path sys.path.insert(0, os.path.join(self.model_path, 'hy3dshape')) sys.path.insert(0, os.path.join(self.model_path, 'hy3dpaint')) # Import the correct pipelines from hy3dshape.pipelines import Hunyuan3DDiTFlowMatchingPipeline from textureGenPipeline import Hunyuan3DPaintPipeline, Hunyuan3DPaintConfig # Instantiate pipelines logger.info("Instantiating shape pipeline...") self.shape_pipeline = Hunyuan3DDiTFlowMatchingPipeline.from_pretrained( self.model_path, torch_dtype=torch.bfloat16 ).to(self.device) logger.info("Instantiating paint pipeline...") paint_config = Hunyuan3DPaintConfig(max_num_view=8, resolution=1024, pbr_optimization=True) self.paint_pipeline = Hunyuan3DPaintPipeline(paint_config) self.model = "direct_model" logger.info("✅ Hunyuan3D pipelines loaded successfully.") except Exception as e: logger.error(f"❌ Failed to set up Hunyuan3D pipeline: {e}", exc_info=True) logger.warning("🔄 Falling back to simplified 3D generation...") self.model = "simplified" def image_to_3d(self, image: Union[str, Image.Image, np.ndarray], remove_background: bool = True, texture_resolution: int = 1024) -> Union[str, trimesh.Trimesh]: """Convert 2D image to 3D model using local Hunyuan3D""" logger.info("🎯 Starting image-to-3D conversion process...") logger.info(f"🎯 Input type: {type(image)}") logger.info(f"🎯 Remove background: {remove_background}") logger.info(f"🎯 Texture resolution: {texture_resolution}") try: # Load model if needed logger.info("🔍 Checking if model needs loading...") if self.model is None: logger.info("📦 Model not loaded, initiating loading...") self.load_model() else: logger.info("✅ Model already loaded") # Prepare image logger.info("🖼️ Preparing input image...") if isinstance(image, str): logger.info(f"🖼️ Loading image from path: {image}") image = Image.open(image) elif isinstance(image, np.ndarray): logger.info("🖼️ Converting numpy array to PIL Image") image = Image.fromarray(image) # Ensure image is PIL Image if not isinstance(image, Image.Image): logger.error("❌ Invalid image type") raise ValueError("Image must be PIL Image, numpy array, or path string") logger.info(f"🖼️ Image mode: {image.mode}, size: {image.size}") # Process based on model type if self.model == "direct_model": logger.info("🌐 Using direct Hunyuan3D model for 3D generation...") return self._generate_with_direct_model(image, remove_background, texture_resolution) elif self.model == "simplified": logger.info("🔄 Using simplified Hunyuan3D generation...") return self._generate_simplified_3d(image) else: # Fallback to simple 3D generation logger.info("🔄 Using fallback 3D generation...") return self._generate_fallback_3d(image) except Exception as e: logger.error(f"❌ 3D generation error: {e}") logger.error(f"❌ Error type: {type(e).__name__}") logger.info("🔄 Falling back to simple 3D generation...") return self._generate_fallback_3d(image) def _generate_with_direct_model(self, image: Image.Image, remove_background: bool, texture_resolution: int) -> str: """Generate 3D model using the official Hunyuan3D pipelines""" try: # Remove background if requested if remove_background: logger.info("🎭 Removing background...") image = self._remove_background(image) # Save image to a temporary file, as pipelines expect a path temp_image_path = self._save_temp_image(image) # 1. Generate the untextured mesh logger.info("🔲 Generating 3D shape with Hunyuan3DDiTFlowMatchingPipeline...") # The pipeline returns a list of meshes, we take the first one mesh_untextured_path = self.shape_pipeline( image=temp_image_path, num_inference_steps=self.num_inference_steps, guidance_scale=self.guidance_scale, seed=random.randint(1, 10000) )[0] logger.info(f"✅ Untextured mesh saved to: {mesh_untextured_path}") # 2. Generate the texture for the mesh logger.info("🎨 Generating texture with Hunyuan3DPaintPipeline...") mesh_textured_path = self.paint_pipeline( mesh_path=mesh_untextured_path, image_path=temp_image_path, guidance_scale=self.guidance_scale, seed=random.randint(1, 10000) ) logger.info(f"✅ Textured mesh saved to: {mesh_textured_path}") # 3. Save the final output to a consistent location output_path = self._save_output_mesh(mesh_textured_path) logger.info(f"✅ 3D model generation successful. Final model at: {output_path}") return output_path except Exception as e: logger.error(f"❌ Direct model generation failed: {e}", exc_info=True) raise def _generate_simplified_3d(self, image: Image.Image) -> str: """Generate 3D using simplified approach with PyTorch operations""" logger.info("🔧 Using simplified 3D generation with PyTorch...") try: # Convert image to tensor import torchvision.transforms as transforms transform = transforms.Compose([ transforms.Resize((256, 256)), transforms.ToTensor(), transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) ]) image_tensor = transform(image).unsqueeze(0).to(self.device) # Create a depth map from the image logger.info("📏 Generating depth map...") # Simple depth estimation based on image brightness gray_image = image.convert('L') depth_array = np.array(gray_image.resize((64, 64))) / 255.0 # Apply some smoothing and scaling from scipy.ndimage import gaussian_filter depth_array = gaussian_filter(depth_array, sigma=2) depth_array = depth_array * 0.3 + 0.1 # Scale depth # Generate mesh from depth map logger.info("🔲 Creating mesh from depth map...") mesh = self._depthmap_to_mesh(depth_array, image) # Save mesh output_path = self._save_mesh(mesh) logger.info(f"✅ Simplified 3D model generated: {output_path}") return output_path except Exception as e: logger.error(f"❌ Simplified generation failed: {e}") return self._generate_fallback_3d(image) def _depthmap_to_mesh(self, depth_map: np.ndarray, texture_image: Image.Image) -> trimesh.Trimesh: """Convert depth map to textured 3D mesh""" h, w = depth_map.shape # Create vertices with texture coordinates vertices = [] faces = [] vertex_colors = [] # Resize texture to match depth map texture_resized = texture_image.resize((w, h)) texture_array = np.array(texture_resized) # Create vertex grid with colors for i in range(h): for j in range(w): x = (j - w/2) / w * 2 y = (i - h/2) / h * 2 z = depth_map[i, j] vertices.append([x, y, z]) # Add vertex color from texture if len(texture_array.shape) == 3: color = texture_array[i, j, :3] else: color = [texture_array[i, j]] * 3 vertex_colors.append(color) # Create faces (two triangles per grid square) for i in range(h-1): for j in range(w-1): v1 = i * w + j v2 = v1 + 1 v3 = v1 + w v4 = v3 + 1 faces.append([v1, v2, v3]) faces.append([v2, v4, v3]) vertices = np.array(vertices) faces = np.array(faces) vertex_colors = np.array(vertex_colors, dtype=np.uint8) # Create mesh with vertex colors mesh = trimesh.Trimesh( vertices=vertices, faces=faces, vertex_colors=vertex_colors ) # Apply smoothing mesh = mesh.smoothed() # Add a base to make it more stable base_vertices, base_faces = self._create_base(vertices, w, h) base_mesh = trimesh.Trimesh(vertices=base_vertices, faces=base_faces) # Combine mesh with base mesh = trimesh.util.concatenate([mesh, base_mesh]) return mesh def _create_base(self, vertices: np.ndarray, w: int, h: int) -> tuple: """Create a base for the mesh""" base_z = vertices[:, 2].min() - 0.1 base_vertices = [] base_faces = [] # Get boundary vertices - fix the indexing boundary_indices = [] # Top edge (excluding corners) for j in range(1, w-1): boundary_indices.append(j) # Right edge (including top-right corner) for i in range(h): boundary_indices.append(i * w + w - 1) # Bottom edge (excluding bottom-right corner, going right to left) for j in range(w-2, 0, -1): boundary_indices.append((h-1) * w + j) # Left edge (including bottom-left corner, going bottom to top) for i in range(h-1, -1, -1): boundary_indices.append(i * w) # Remove duplicate indices (first and last should not be the same) if boundary_indices and boundary_indices[0] == boundary_indices[-1]: boundary_indices = boundary_indices[:-1] # Create base vertices start_idx = len(vertices) for idx in boundary_indices: if idx < len(vertices): # Safety check v = vertices[idx].copy() v[2] = base_z base_vertices.append(v) if not base_vertices: # If no base vertices were created, return empty arrays return np.array([]), np.array([]) # Create center vertex center = np.mean(base_vertices, axis=0) base_vertices.append(center) center_idx = len(base_vertices) - 1 # Create base faces for i in range(len(boundary_indices)): next_i = (i + 1) % len(boundary_indices) base_faces.append([ i, next_i, center_idx ]) return np.array(base_vertices), np.array(base_faces) def _remove_background(self, image: Image.Image) -> Image.Image: """Remove background from image""" try: # Try using rembg if available from rembg import remove return remove(image) except: # Fallback: simple background removal # Convert to RGBA image = image.convert("RGBA") # Simple white background removal datas = image.getdata() new_data = [] for item in datas: # Remove white-ish backgrounds if item[0] > 230 and item[1] > 230 and item[2] > 230: new_data.append((255, 255, 255, 0)) else: new_data.append(item) image.putdata(new_data) return image def _generate_fallback_3d(self, image: Union[Image.Image, np.ndarray]) -> str: """Generate fallback 3D model when main model fails""" # Create a simple 3D representation based on image if isinstance(image, np.ndarray): image = Image.fromarray(image) elif isinstance(image, str): image = Image.open(image) # Analyze image for basic shape image_array = np.array(image.resize((64, 64))) # Create height map from image brightness gray = np.mean(image_array, axis=2) if len(image_array.shape) == 3 else image_array height_map = gray / 255.0 # Create mesh from height map mesh = self._heightmap_to_mesh(height_map) # Save and return path return self._save_mesh(mesh) def _heightmap_to_mesh(self, heightmap: np.ndarray) -> trimesh.Trimesh: """Convert heightmap to 3D mesh""" h, w = heightmap.shape # Create vertices vertices = [] faces = [] # Create vertex grid for i in range(h): for j in range(w): x = (j - w/2) / w * 2 y = (i - h/2) / h * 2 z = heightmap[i, j] * 0.5 vertices.append([x, y, z]) # Create faces for i in range(h-1): for j in range(w-1): # Two triangles per grid square v1 = i * w + j v2 = v1 + 1 v3 = v1 + w v4 = v3 + 1 faces.append([v1, v2, v3]) faces.append([v2, v4, v3]) vertices = np.array(vertices) faces = np.array(faces) # Create mesh mesh = trimesh.Trimesh(vertices=vertices, faces=faces) # Apply smoothing mesh = mesh.smoothed() return mesh def _save_mesh(self, mesh: trimesh.Trimesh) -> str: """Save mesh to file""" # Create temporary file with tempfile.NamedTemporaryFile(suffix='.glb', delete=False) as tmp: mesh_path = tmp.name # Export mesh mesh.export(mesh_path) return mesh_path def _save_temp_image(self, image: Image.Image) -> str: """Save PIL image to temporary file""" with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as tmp: image_path = tmp.name # Save image image.save(image_path, 'PNG') logger.info(f"💾 Saved temp image to: {image_path}") return image_path def _save_output_mesh(self, source_mesh_path: str) -> str: """Copy generated mesh to our output location""" # Create output directory if it doesn't exist output_dir = "/tmp/hunyuan3d_output" os.makedirs(output_dir, exist_ok=True) # Generate unique filename timestamp = tempfile.mktemp().split('/')[-1] output_filename = f"hunyuan3d_mesh_{timestamp}.glb" output_path = os.path.join(output_dir, output_filename) # Copy the file shutil.copy2(source_mesh_path, output_path) logger.info(f"📁 Copied mesh from {source_mesh_path} to {output_path}") return output_path def text_to_3d(self, text_prompt: str) -> str: """Generate 3D model from text description""" # First generate image, then convert to 3D # This would require image generator integration raise NotImplementedError("Text to 3D requires image generation first") def to(self, device: str): """Update device preference""" self.device = device logger.info(f"🔧 Device preference updated to: {device}") def __del__(self): """Cleanup when object is destroyed""" if hasattr(self, 'model') and self.model not in [None, "fallback_mode", "simplified"]: del self.model if torch.cuda.is_available(): torch.cuda.empty_cache()