generators/hunyuan.py

"""Hunyuan3D-2.1 3D model generation."""

# CRITICAL: Import spaces BEFORE torch/CUDA packages
import spaces

import torch
from pathlib import Path
from gradio_client import Client, handle_file
import httpx
from tenacity import retry, stop_after_attempt, wait_exponential, retry_if_exception_type

from core.config import HUNYUAN_SETTINGS, QualityPreset
from utils.memory import MemoryManager


class HunyuanGenerator:
    """Generates 3D models using Hunyuan3D-2.1."""
    
    def __init__(self):
        self.memory_manager = MemoryManager()
    
    @retry(
        stop=stop_after_attempt(3),
        wait=wait_exponential(multiplier=1, min=4, max=10),
        retry=retry_if_exception_type((httpx.TimeoutException, httpx.NetworkError))
    )
    def _call_api(self, client: Client, **kwargs):
        """Call Hunyuan3D API with automatic retry."""
        return client.predict(**kwargs)
    
    @spaces.GPU(duration=90)
    def generate(
        self,
        image_path: Path,
        preset: QualityPreset,
        output_dir: Path
    ) -> Path:
        """Generate 3D model from 2D image."""
        try:
            print(f"[Hunyuan3D] Generating 3D model: {preset.name} quality")
            print(f"[Hunyuan3D] Input image: {image_path}")
            print(f"[Hunyuan3D] Settings: steps={preset.hunyuan_steps}, guidance={preset.hunyuan_guidance}, octree={preset.octree_resolution}")
            
            # Validate input image exists
            if not image_path.exists():
                raise FileNotFoundError(f"Input image not found: {image_path}")
            
            # Connect to API
            print(f"[Hunyuan3D] Connecting to {HUNYUAN_SETTINGS['space_id']}...")
            client = Client(
                HUNYUAN_SETTINGS["space_id"],
                httpx_kwargs={
                    "timeout": httpx.Timeout(
                        HUNYUAN_SETTINGS["timeout"],
                        connect=HUNYUAN_SETTINGS["connect_timeout"]
                    )
                }
            )
            print(f"[Hunyuan3D] Connected successfully")
            
            # Call API (returns tuple: file, output, mesh_stats, seed)
            print(f"[Hunyuan3D] Calling API with parameters...")
            result = self._call_api(
                client,
                image=handle_file(str(image_path)),
                mv_image_front=None,
                mv_image_back=None,
                mv_image_left=None,
                mv_image_right=None,
                steps=preset.hunyuan_steps,
                guidance_scale=preset.hunyuan_guidance,
                seed=1234,
                octree_resolution=preset.octree_resolution,
                check_box_rembg=True,
                num_chunks=preset.num_chunks,
                randomize_seed=True,
                api_name="/shape_generation"
            )
            print(f"[Hunyuan3D] API call completed")
            
            # Extract GLB path from tuple response
            # API returns: (file, output, mesh_stats, seed)
            print(f"[Hunyuan3D] Raw result type: {type(result)}")
            print(f"[Hunyuan3D] Raw result length: {len(result) if isinstance(result, (tuple, list)) else 'N/A'}")
            
            if not isinstance(result, tuple):
                raise ValueError(
                    f"Unexpected result type from Hunyuan3D API: {type(result)}. "
                    f"Expected tuple of (file, output, mesh_stats, seed)."
                )
            
            if len(result) != 4:
                raise ValueError(
                    f"Unexpected result length from Hunyuan3D API: {len(result)}. "
                    f"Expected 4 elements (file, output, mesh_stats, seed), got {len(result)}."
                )
            
            # Extract GLB file path (first element)
            file_data, html_output, mesh_stats, used_seed = result
            
            print(f"[Hunyuan3D] file_data type: {type(file_data)}")
            print(f"[Hunyuan3D] mesh_stats: {mesh_stats}")
            print(f"[Hunyuan3D] used_seed: {used_seed}")
            
            # Extract path from file_data
            if file_data is None:
                raise ValueError(
                    "Hunyuan3D API returned None for file. "
                    "This usually means the generation failed on the server side. "
                    "Possible causes:\n"
                    "  - Invalid image input\n"
                    "  - API timeout\n"
                    "  - Server overload\n"
                    "Try again with a different image or quality setting."
                )
            
            # Handle different file_data formats
            if isinstance(file_data, dict):
                print(f"[Hunyuan3D] file_data is dict with keys: {file_data.keys()}")
                if 'path' in file_data:
                    glb_path = file_data['path']
                elif 'value' in file_data:
                    glb_path = file_data['value']
                elif 'name' in file_data:
                    glb_path = file_data['name']
                else:
                    raise ValueError(
                        f"Unexpected dict format from Hunyuan3D API. "
                        f"Keys: {list(file_data.keys())}"
                    )
            elif isinstance(file_data, str):
                glb_path = file_data
            else:
                raise ValueError(
                    f"Unexpected file_data type: {type(file_data)}. "
                    f"Expected dict or str."
                )
            
            print(f"[Hunyuan3D] Extracted GLB path: {glb_path}")
            
            # Validate path exists
            if not glb_path or glb_path == "None":
                raise ValueError(
                    "Hunyuan3D API returned invalid path. "
                    "The generation may have failed on the server side."
                )
            
            if not Path(glb_path).exists():
                raise ValueError(
                    f"GLB file not found at path: {glb_path}. "
                    f"The file may not have been generated or saved correctly."
                )
            
            print(f"[Hunyuan3D] Model generated: {glb_path}")
            
            # Cleanup
            del client
            import gc
            gc.collect()
            torch.cuda.empty_cache()
            
            return Path(glb_path)
            
        except Exception as e:
            import traceback
            error_details = traceback.format_exc()
            print(f"[Hunyuan3D] ERROR: {e}")
            print(f"[Hunyuan3D] Full traceback:\n{error_details}")
            
            # Provide helpful error message based on error type
            error_str = str(e).lower()
            
            if "quota" in error_str or "zerogpu" in error_str:
                raise RuntimeError(
                    f"⚠️ Hunyuan3D Space is out of GPU quota.\n"
                    f"This is a limitation of the free Hunyuan3D-2.1 Space.\n\n"
                    f"Solutions:\n"
                    f"1. Wait for quota reset (resets daily)\n"
                    f"2. Try again in a few hours\n"
                    f"3. Use a different time of day (less traffic)\n\n"
                    f"Note: Your L4 GPU is only used for FLUX generation.\n"
                    f"Hunyuan3D runs on an external space with quota limits."
                ) from e
            elif "list index out of range" in str(e) or "unexpected result" in error_str:
                raise ValueError(
                    f"❌ Hunyuan3D API returned empty result.\n"
                    f"This usually means:\n"
                    f"1. The Hunyuan3D Space is overloaded\n"
                    f"2. GPU quota exhausted\n"
                    f"3. Invalid image input\n\n"
                    f"Try again in a few minutes."
                ) from e
            elif "timeout" in error_str:
                raise TimeoutError(
                    f"⏱️ Hunyuan3D generation timed out.\n"
                    f"Try using a lower quality preset (Fast or Balanced)."
                ) from e
            elif "not found" in error_str or "404" in error_str:
                raise RuntimeError(
                    f"❌ Hunyuan3D Space not accessible.\n"
                    f"The tencent/Hunyuan3D-2.1 Space may be down or moved.\n"
                    f"Check: https://huggingface.co/spaces/tencent/Hunyuan3D-2.1"
                ) from e
            else:
                raise RuntimeError(
                    f"❌ Hunyuan3D generation failed: {e}\n"
                    f"Check the Hunyuan3D Space status and try again."
                ) from e