core/pipeline.py

"""Main asset generation pipeline orchestration."""

import time
from pathlib import Path
from typing import Optional

from core.config import QUALITY_PRESETS
from core.types import GenerationResult, AssetMetadata
from generators import FluxGenerator, HunyuanGenerator
from processors import BlenderProcessor, AssetValidator
from utils import CacheManager, SecurityManager


class AssetPipeline:
    """Orchestrates the complete asset generation pipeline."""
    
    def __init__(self):
        self.flux = FluxGenerator()
        self.hunyuan = HunyuanGenerator()  # API client with better error handling
        self.blender = BlenderProcessor()
        self.validator = AssetValidator()
        self.cache = CacheManager()
        self.security = SecurityManager()
        
        self.output_dir = Path("outputs")
        self.temp_dir = Path("temp")
        self.script_dir = Path("scripts")
        
        # Create directories
        self.output_dir.mkdir(exist_ok=True)
        self.temp_dir.mkdir(exist_ok=True)
    
    def generate(
        self,
        prompt: str,
        quality: str = "High",
        progress_callback: Optional[callable] = None
    ) -> GenerationResult:
        """
        Generate 3D asset from text prompt.
        
        Args:
            prompt: Text description of asset
            quality: Quality preset (Fast/Balanced/High/Ultra)
            progress_callback: Optional callback for progress updates
        
        Returns:
            GenerationResult with GLB path and metadata
        """
        start_time = time.time()
        
        def update_progress(value: float, desc: str):
            if progress_callback:
                progress_callback(value, desc=desc)
            print(f"[Pipeline] {desc} ({value*100:.0f}%)")
        
        try:
            # Step 1: Security checks
            update_progress(0.0, "Validating input...")
            prompt = self.security.sanitize_prompt(prompt)
            self.security.check_rate_limit()
            
            # Step 2: Check cache
            update_progress(0.05, "Checking cache...")
            if cached_path := self.cache.get_cached_result(prompt, quality):
                elapsed = time.time() - start_time
                
                preset = QUALITY_PRESETS[quality]
                metadata = AssetMetadata(
                    prompt=prompt,
                    quality=quality,
                    flux_steps=preset.flux_steps,
                    hunyuan_steps=preset.hunyuan_steps,
                    file_size_mb=cached_path.stat().st_size / 1e6,
                    generation_time_s=elapsed,
                    optimized=True,
                    rigged=False
                )
                
                return GenerationResult(
                    glb_path=cached_path,
                    status_message="✨ Loaded from cache (saved GPU quota!)",
                    metadata=metadata,
                    cached=True
                )
            
            # Step 3: Get quality preset
            preset = QUALITY_PRESETS.get(quality, QUALITY_PRESETS["High"])
            
            # Step 4: Generate 2D image (FLUX)
            update_progress(0.1, f"Generating 2D image (FLUX {preset.flux_steps} steps)...")
            image_path = self.flux.generate(prompt, preset, self.temp_dir)
            
            # Step 5: Generate 3D model (Hunyuan3D)
            update_progress(0.5, f"Converting to 3D (Hunyuan3D {preset.hunyuan_steps} steps)...")
            glb_path = self.hunyuan.generate(image_path, preset, self.temp_dir)
            
            # Step 6: Validate GLB
            update_progress(0.8, "Validating 3D model...")
            is_valid, validation_msg = self.validator.validate_glb(glb_path)
            if not is_valid:
                raise ValueError(f"GLB validation failed: {validation_msg}")
            
            # Step 7: Blender optimization (if available)
            update_progress(0.85, "Optimizing for game engine...")
            
            raw_path = self.output_dir / f"asset_raw_{int(time.time())}.glb"
            import shutil
            shutil.copy(glb_path, raw_path)
            
            optimized = False
            optimization_msg = ""
            
            if self.blender.is_available():
                optimized_path = self.output_dir / f"asset_optimized_{int(time.time())}.glb"
                script_path = self.script_dir / "blender_optimize.py"
                
                success, message = self.blender.optimize(raw_path, optimized_path, script_path)
                
                if success:
                    final_path = optimized_path
                    optimized = True
                    optimization_msg = f"\n✅ {message}"
                else:
                    final_path = raw_path
                    optimization_msg = f"\n⚠️ Optimization skipped: {message}"
            else:
                final_path = raw_path
                optimization_msg = "\n⚠️ Blender not available, using raw output"
            
            # Step 8: Save to cache
            update_progress(0.95, "Saving to cache...")
            self.cache.save_to_cache(prompt, quality, final_path)
            
            # Step 9: Cleanup temp files
            if image_path.exists():
                image_path.unlink()
            
            # Step 10: Create result
            elapsed = time.time() - start_time
            
            metadata = AssetMetadata(
                prompt=prompt,
                quality=quality,
                flux_steps=preset.flux_steps,
                hunyuan_steps=preset.hunyuan_steps,
                file_size_mb=final_path.stat().st_size / 1e6,
                generation_time_s=elapsed,
                optimized=optimized,
                rigged=False
            )
            
            status_msg = f"✨ Generated in {elapsed:.1f}s"
            status_msg += f"\n📊 {preset.flux_steps} FLUX steps, {preset.hunyuan_steps} Hunyuan3D steps"
            status_msg += optimization_msg
            
            update_progress(1.0, "Complete!")
            
            return GenerationResult(
                glb_path=final_path,
                status_message=status_msg,
                metadata=metadata,
                cached=False
            )
            
        except Exception as e:
            print(f"[Pipeline] Error: {e}")
            raise