models/model_3d_generator.py

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()