File size: 6,415 Bytes
0e805d4
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
"""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()
        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