FIX: Import spaces before torch to prevent CUDA initialization error

app.py

@@ -14,1217 +14,877 @@ Author: AI Agent Framework Specialist
 Version: 2.0.0 Production
 """
 
-from spaces import GPU
+# IMPORT SPACES FIRST - Before any CUDA-related imports
+try:
+    from spaces import GPU
+    SPACES_AVAILABLE = True
+except ImportError:
+    SPACES_AVAILABLE = False
+    print("Warning: spaces package not available, GPU acceleration disabled")
+
 import os
 import sys
 import time
 import asyncio
-import hashlib
 import logging
-import traceback
-import warnings
-from typing import Optional, Tuple, Dict, Any, Union, List
-from contextlib import asynccontextmanager
+import hashlib
+import gc
+import psutil
+import threading
+from datetime import datetime, timedelta
+from typing import Optional, Dict, Any, List, Tuple
 from dataclasses import dataclass
-from enum import Enum
+from pathlib import Path
+import json
 
-# Third-party imports
-import gradio as gr
+# Now import CUDA-related packages
 import torch
 import numpy as np
 from PIL import Image
-import psutil
-from functools import lru_cache
-from datetime import datetime, timedelta
-
-# Diffusers and model imports
-from diffusers import DiffusionPipeline, StableDiffusionImg2ImgPipeline
-from diffusers.utils import logging as diffusers_logging
-from spaces import GPU
+import gradio as gr
 
-# Suppress noisy warnings
-warnings.filterwarnings("ignore", category=UserWarning)
-diffusers_logging.set_verbosity_error()
+# Import diffusers after spaces
+try:
+    from diffusers import DiffusionPipeline
+    from diffusers.utils import logging as diffusers_logging
+    DIFFUSERS_AVAILABLE = True
+except ImportError:
+    DIFFUSERS_AVAILABLE = False
+    print("Warning: diffusers not properly installed")
 
-# Configure logging early
+# Configure logging
 logging.basicConfig(
     level=logging.INFO,
     format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
     handlers=[
-        logging.StreamHandler(sys.stdout),
-        logging.FileHandler('z_image_turbo.log', mode='a')
+        logging.FileHandler('z_image_turbo.log'),
+        logging.StreamHandler(sys.stdout)
     ]
 )
 logger = logging.getLogger(__name__)
 
+# Suppress verbose logging
+diffusers_logging.set_verbosity_error()
+
+# Constants
+MODEL_ID = "Tongyi-MAI/Z-Image-Turbo"
+DEFAULT_ASPECT_RATIO = "1:1"
+STYLE_PRESETS = [
+    "None",
+    "Cinematic",
+    "Photographic",
+    "Anime",
+    "Oil Painting",
+    "Watercolor",
+    "Cyberpunk",
+    "Fantasy Art",
+    "3D Render",
+    "Vintage"
+]
+ASPECT_RATIOS = {
+    "1:1": (512, 512),
+    "16:9": (768, 432),
+    "9:16": (432, 768),
+    "4:3": (576, 432),
+    "3:4": (432, 576),
+    "3:2": (612, 408),
+    "2:3": (408, 612)
+}
 
-class ErrorCode(Enum):
-    """Enumeration of application error codes"""
-    SUCCESS = 0
-    MODEL_LOAD_ERROR = 1001
-    GENERATION_ERROR = 1002
-    TRANSFORM_ERROR = 1003
-    INVALID_INPUT = 2001
-    RESOURCE_ERROR = 3001
-    NETWORK_ERROR = 4001
-    CACHE_ERROR = 5001
-    UNKNOWN_ERROR = 9999
+# Custom CSS for better UI
+CUSTOM_CSS = """
+.footer {
+    text-align: center;
+    margin-top: 20px;
+    padding: 10px;
+    background: linear-gradient(45deg, #667eea 0%, #764ba2 100%);
+    border-radius: 10px;
+    color: white;
+}
 
+.generate-btn {
+    background: linear-gradient(45deg, #f093fb 0%, #f5576c 100%) !important;
+    border: none !important;
+    font-weight: bold !important;
+}
+
+.main-container {
+    max-width: 1200px;
+    margin: 0 auto;
+}
+
+.system-monitor {
+    font-family: 'Courier New', monospace;
+    background: #f8f9fa;
+    padding: 10px;
+    border-radius: 5px;
+    margin: 10px 0;
+}
+
+.error-message {
+    background: #fee;
+    border: 1px solid #fcc;
+    padding: 10px;
+    border-radius: 5px;
+    color: #c00;
+}
+
+.success-message {
+    background: #efe;
+    border: 1px solid #cfc;
+    padding: 10px;
+    border-radius: 5px;
+    color: #080;
+}
+"""
 
 @dataclass
 class GenerationResult:
-    """Data class for generation results"""
+    """Result of image generation with metadata"""
     success: bool
     image: Optional[Image.Image] = None
-    seed: int = 0
-    message: str = ""
-    error_code: ErrorCode = ErrorCode.SUCCESS
-    generation_time: float = 0.0
-    metadata: Optional[Dict[str, Any]] = None
+    error: Optional[str] = None
+    error_code: Optional[int] = None
+    generation_time: Optional[float] = None
+    cache_hit: bool = False
+    optimization_used: List[str] = None
+
+class CacheManager:
+    """Intelligent cache with LRU eviction and TTL"""
+
+    def __init__(self, max_size: int = 50, ttl_minutes: int = 30):
+        self.max_size = max_size
+        self.ttl = timedelta(minutes=ttl_minutes)
+        self.cache: Dict[str, Tuple[Image.Image, datetime]] = {}
+        self.access_times: Dict[str, datetime] = {}
+        self.lock = threading.Lock()
+
+    def _generate_key(self, prompt: str, negative_prompt: str, style: str,
+                     aspect_ratio: str, guidance_scale: float,
+                     inference_steps: int, seed: int) -> str:
+        """Generate cache key from parameters"""
+        key_data = f"{prompt}|{negative_prompt}|{style}|{aspect_ratio}|{guidance_scale}|{inference_steps}|{seed}"
+        return hashlib.md5(key_data.encode()).hexdigest()
+
+    def get(self, prompt: str, negative_prompt: str, style: str,
+            aspect_ratio: str, guidance_scale: float,
+            inference_steps: int, seed: int) -> Optional[Image.Image]:
+        """Get cached image if available and not expired"""
+        key = self._generate_key(prompt, negative_prompt, style, aspect_ratio,
+                                guidance_scale, inference_steps, seed)
+
+        with self.lock:
+            if key in self.cache:
+                image, timestamp = self.cache[key]
+                if datetime.now() - timestamp < self.ttl:
+                    self.access_times[key] = datetime.now()
+                    logger.info(f"Cache hit for key: {key[:8]}...")
+                    return image.copy()
+                else:
+                    # Expired entry
+                    del self.cache[key]
+                    if key in self.access_times:
+                        del self.access_times[key]
+        return None
 
+    def put(self, prompt: str, negative_prompt: str, style: str,
+            aspect_ratio: str, guidance_scale: float,
+            inference_steps: int, seed: int, image: Image.Image):
+        """Cache image with LRU eviction"""
+        key = self._generate_key(prompt, negative_prompt, style, aspect_ratio,
+                                guidance_scale, inference_steps, seed)
+
+        with self.lock:
+            # Evict if necessary
+            if len(self.cache) >= self.max_size and key not in self.cache:
+                # Find least recently used
+                lru_key = min(self.access_times.keys(),
+                            key=lambda k: self.access_times[k])
+                del self.cache[lru_key]
+                del self.access_times[lru_key]
+                logger.info(f"Evicted LRU entry: {lru_key[:8]}...")
+
+            self.cache[key] = (image.copy(), datetime.now())
+            self.access_times[key] = datetime.now()
+            logger.info(f"Cached new image: {key[:8]}...")
+
+    def clear(self):
+        """Clear all cache entries"""
+        with self.lock:
+            self.cache.clear()
+            self.access_times.clear()
+        logger.info("Cache cleared")
+
+    def get_stats(self) -> Dict[str, Any]:
+        """Get cache statistics"""
+        with self.lock:
+            return {
+                "size": len(self.cache),
+                "max_size": self.max_size,
+                "usage_percent": (len(self.cache) / self.max_size) * 100
+            }
 
 class SystemMonitor:
-    """Monitor system resources and performance"""
+    """Monitor system resources and performance metrics"""
 
     def __init__(self):
-        self.start_time = time.time()
+        self.start_time = datetime.now()
         self.generation_count = 0
+        self.success_count = 0
         self.error_count = 0
-        self.cache_hits = 0
-        self.cache_misses = 0
+        self.total_generation_time = 0
+        self.lock = threading.Lock()
+
+    def log_generation(self, success: bool, generation_time: float):
+        """Log generation metrics"""
+        with self.lock:
+            self.generation_count += 1
+            self.total_generation_time += generation_time
+            if success:
+                self.success_count += 1
+            else:
+                self.error_count += 1
 
-    def get_system_info(self) -> Dict[str, Any]:
-        """Get current system resource information"""
-        try:
-            memory = psutil.virtual_memory()
-            gpu_memory = self._get_gpu_memory()
+    def get_stats(self) -> Dict[str, Any]:
+        """Get comprehensive system stats"""
+        with self.lock:
+            uptime = datetime.now() - self.start_time
+            avg_gen_time = (self.total_generation_time / self.generation_count
+                          if self.generation_count > 0 else 0)
 
-            return {
-                "uptime_seconds": time.time() - self.start_time,
-                "memory_used_gb": memory.used / (1024**3),
-                "memory_percent": memory.percent,
-                "gpu_memory_used_gb": gpu_memory,
+            stats = {
+                # System resources
                 "cpu_percent": psutil.cpu_percent(interval=0.1),
-                "active_generations": self.generation_count,
+                "memory_percent": psutil.virtual_memory().percent,
+                "disk_percent": psutil.disk_usage('/').percent,
+
+                # GPU info if available
+                "gpu_available": torch.cuda.is_available(),
+                "gpu_memory": None,
+
+                # Performance metrics
+                "uptime_seconds": uptime.total_seconds(),
+                "uptime_str": str(uptime).split('.')[0],
+                "generation_count": self.generation_count,
+                "success_count": self.success_count,
                 "error_count": self.error_count,
-                "cache_hit_rate": self.get_cache_hit_rate()
+                "success_rate": (self.success_count / self.generation_count * 100
+                               if self.generation_count > 0 else 0),
+                "avg_generation_time": round(avg_gen_time, 2),
+                "generations_per_minute": (self.generation_count / uptime.total_seconds() * 60
+                                         if uptime.total_seconds() > 0 else 0)
             }
-        except Exception as e:
-            logger.error(f"Error getting system info: {e}")
-            return {}
 
-    def _get_gpu_memory(self) -> float:
-        """Get GPU memory usage in GB"""
-        try:
+            # Add GPU stats if available
             if torch.cuda.is_available():
-                return torch.cuda.memory_allocated() / (1024**3)
-        except:
-            pass
-        return 0.0
-
-    def get_cache_hit_rate(self) -> float:
-        """Calculate cache hit rate percentage"""
-        total = self.cache_hits + self.cache_misses
-        return (self.cache_hits / total * 100) if total > 0 else 0.0
-
-    def record_generation(self, success: bool):
-        """Record a generation attempt"""
-        self.generation_count += 1
-        if not success:
-            self.error_count += 1
-
-    def record_cache_hit(self):
-        """Record a cache hit"""
-        self.cache_hits += 1
-
-    def record_cache_miss(self):
-        """Record a cache miss"""
-        self.cache_misses += 1
+                stats["gpu_memory"] = {
+                    "allocated": torch.cuda.memory_allocated() / 1024**3,
+                    "cached": torch.cuda.memory_reserved() / 1024**3,
+                    "total": torch.cuda.get_device_properties(0).total_memory / 1024**3
+                }
 
+            return stats
 
 class ModelManager:
-    """Manages model loading, optimization, and resource handling"""
+    """Handle model loading and optimization"""
 
     def __init__(self):
-        self.model_name = "Tongyi-MAI/Z-Image-Turbo"
-        self.pipe_t2i = None
-        self.pipe_i2i = None
-        self.is_loaded = False
+        self.pipeline = None
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.dtype = self._get_optimal_dtype()
         self.optimizations_applied = []
-        self._load_lock = asyncio.Lock()
-
-    async def load_models(self) -> bool:
-        """Load models with proper error handling and fallbacks"""
-        async with self._load_lock:
-            if self.is_loaded:
-                return True
-
-            try:
-                logger.info(f"Loading model: {self.model_name}")
-                start_time = time.time()
-
-                # Determine optimal dtype based on hardware
-                dtype = self._get_optimal_dtype()
-
-                # Load base pipeline
-                self.pipe_t2i = DiffusionPipeline.from_pretrained(
-                    self.model_name,
-                    torch_dtype=dtype,
-                    use_safetensors=True,
-                    variant=self._get_variant(dtype),
-                    low_cpu_mem_usage=True
-                )
-
-                # Create img2img pipeline
-                self.pipe_i2i = StableDiffusionImg2ImgPipeline(
-                    vae=self.pipe_t2i.vae,
-                    text_encoder=self.pipe_t2i.text_encoder,
-                    tokenizer=self.pipe_t2i.tokenizer,
-                    unet=self.pipe_t2i.unet,
-                    scheduler=self.pipe_t2i.scheduler,
-                    safety_checker=None,
-                    feature_extractor=None,
-                    requires_safety_checker=False
-                )
-
-                # Apply optimizations
-                await self._apply_optimizations()
-
-                load_time = time.time() - start_time
-                logger.info(f"Models loaded successfully in {load_time:.2f}s")
-                logger.info(f"Applied optimizations: {', '.join(self.optimizations_applied)}")
-
-                self.is_loaded = True
-                return True
-
-            except Exception as e:
-                logger.error(f"Failed to load models: {e}")
-                logger.error(traceback.format_exc())
-                return False
 
     def _get_optimal_dtype(self) -> torch.dtype:
-        """Determine optimal data type based on hardware"""
-        try:
-            # Check for bfloat16 support (better for newer GPUs)
-            if torch.cuda.is_available() and torch.cuda.is_bf16_supported():
-                logger.info("Using bfloat16 for optimal performance")
+        """Determine optimal data type for the hardware"""
+        if torch.cuda.is_available():
+            # Check GPU capabilities
+            gpu_props = torch.cuda.get_device_properties(0)
+            if gpu_props.major >= 8:  # Ampere and newer
                 return torch.bfloat16
-            # Fall back to float16 for compatibility
-            elif torch.cuda.is_available():
-                logger.info("Using float16 for CUDA")
-                return torch.float16
-            # CPU fallback
             else:
-                logger.info("Using float32 for CPU")
-                return torch.float32
-        except:
-            logger.warning("Could not detect optimal dtype, using float32")
-            return torch.float32
-
-    def _get_variant(self, dtype: torch.dtype) -> Optional[str]:
-        """Get model variant based on dtype"""
-        return "fp16" if dtype == torch.float16 else None
-
-    async def _apply_optimizations(self):
-        """Apply performance optimizations with proper fallbacks"""
-
-        # 1. Try xformers optimization
-        if self._try_enable_xformers():
-            self.optimizations_applied.append("xformers")
-
-        # 2. Try model CPU offloading for memory efficiency
-        if self._try_enable_cpu_offload():
-            self.optimizations_applied.append("cpu_offload")
-
-        # 3. Try PyTorch 2.0+ compilation
-        if self._try_enable_torch_compile():
-            self.optimizations_applied.append("torch_compile")
-
-        # 4. Enable VAE slicing for memory efficiency
-        self._enable_vae_slicing()
+                return torch.float16
+        return torch.float32
 
-        # 5. Clear CUDA cache
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
+    def _try_enable_torch_compile(self) -> bool:
+        """Try to enable torch.compile for better performance"""
+        try:
+            if hasattr(torch, 'compile') and torch.__version__ >= "2.0":
+                logger.info("PyTorch 2.0+ detected, enabling compilation")
+                return True
+            else:
+                logger.info("PyTorch < 2.0 detected, compilation not available")
+                return False
+        except Exception as e:
+            logger.warning(f"Could not enable torch.compile: {e}")
+            return False
 
     def _try_enable_xformers(self) -> bool:
-        """Try to enable xformers with fallback"""
+        """Try to enable xformers for memory efficiency"""
         try:
+            import xformers
             import xformers.ops
-            self.pipe_t2i.enable_xformers_memory_efficient_attention()
-            self.pipe_i2i.enable_xformers_memory_efficient_attention()
-            logger.info("✓ Enabled xformers memory efficient attention")
+            logger.info("xformers is available and will be used")
             return True
         except ImportError:
-            logger.info("⚠ xformers not available, using default attention")
+            logger.info("xformers not available, using standard attention")
             return False
         except Exception as e:
-            logger.warning(f"⚠ Could not enable xformers: {e}")
+            logger.warning(f"Could not enable xformers: {e}")
             return False
 
-    def _try_enable_cpu_offload(self) -> bool:
-        """Try to enable CPU offloading with fallback"""
-        try:
-            # Only enable if we have limited GPU memory
-            if torch.cuda.is_available():
-                gpu_memory = torch.cuda.get_device_properties(0).total_memory
-                if gpu_memory < 8 * 1024**3:  # Less than 8GB
-                    self.pipe_t2i.enable_sequential_cpu_offload()
-                    self.pipe_i2i.enable_sequential_cpu_offload()
-                    logger.info("✓ Enabled sequential CPU offloading")
-                    return True
-        except Exception as e:
-            logger.warning(f"⚠ Could not enable CPU offload: {e}")
-        return False
-
-    def _try_enable_torch_compile(self) -> bool:
-        """Try to enable torch.compile with version check and fallback"""
-        try:
-            # Check PyTorch version
-            torch_version = torch.__version__.split('+')[0]
-            major, minor = map(int, torch_version.split('.')[:2])
-
-            if major >= 2:
-                logger.info("PyTorch 2.0+ detected, attempting compilation...")
-                self.pipe_t2i.unet = torch.compile(
-                    self.pipe_t2i.unet,
-                    mode="reduce-overhead",
-                    fullgraph=False  # More compatible
-                )
-                self.pipe_i2i.unet = torch.compile(
-                    self.pipe_i2i.unet,
-                    mode="reduce-overhead",
-                    fullgraph=False
-                )
-                logger.info("✓ Successfully compiled UNet with torch.compile")
-                return True
-            else:
-                logger.info(f"⚠ PyTorch {torch_version} < 2.0, compilation not available")
-        except Exception as e:
-            logger.warning(f"⚠ Could not compile UNet: {e}")
-        return False
+    def load_model(self):
+        """Load and optimize the model"""
+        if self.pipeline is not None:
+            return True
 
-    def _enable_vae_slicing(self):
-        """Enable VAE slicing for memory efficiency"""
         try:
-            self.pipe_t2i.vae.enable_slicing()
-            self.pipe_i2i.vae.enable_slicing()
-            logger.info("✓ Enabled VAE slicing")
-        except Exception as e:
-            logger.warning(f"⚠ Could not enable VAE slicing: {e}")
-
-
-class CacheManager:
-    """Manages caching for generated images and analyses"""
-
-    def __init__(self, max_size: int = 100):
-        self.max_size = max_size
-        self.image_cache: Dict[str, Tuple[Image.Image, datetime]] = {}
-        self.analysis_cache: Dict[str, Tuple[str, datetime]] = {}
-        self.cache_ttl = timedelta(hours=24)
-
-    def get_cache_key(self, *args) -> str:
-        """Generate consistent cache key"""
-        key_str = "|".join(str(arg) for arg in args)
-        return hashlib.sha256(key_str.encode()).hexdigest()[:16]
-
-    def get_cached_image(self, cache_key: str) -> Optional[Image.Image]:
-        """Get cached image if valid"""
-        if cache_key in self.image_cache:
-            image, timestamp = self.image_cache[cache_key]
-            if datetime.now() - timestamp < self.cache_ttl:
-                return image
-            else:
-                del self.image_cache[cache_key]
-        return None
+            logger.info(f"Loading model {MODEL_ID} on {self.device} with dtype {self.dtype}")
+
+            # Load pipeline
+            self.pipeline = DiffusionPipeline.from_pretrained(
+                MODEL_ID,
+                torch_dtype=self.dtype,
+                use_safetensors=True,
+                variant=None  # Remove variant to avoid fp16 issues
+            )
 
-    def cache_image(self, cache_key: str, image: Image.Image):
-        """Cache an image with LRU eviction"""
-        # Remove oldest if at capacity
-        if len(self.image_cache) >= self.max_size:
-            oldest_key = min(self.image_cache.keys(),
-                           key=lambda k: self.image_cache[k][1])
-            del self.image_cache[oldest_key]
-
-        self.image_cache[cache_key] = (image, datetime.now())
-
-    def get_cached_analysis(self, cache_key: str) -> Optional[str]:
-        """Get cached analysis if valid"""
-        if cache_key in self.analysis_cache:
-            analysis, timestamp = self.analysis_cache[cache_key]
-            if datetime.now() - timestamp < self.cache_ttl:
-                return analysis
+            # Move to device
+            self.pipeline = self.pipeline.to(self.device)
+
+            # Enable optimizations
+            if self.device == "cuda":
+                # Try xformers
+                if self._try_enable_xformers():
+                    self.pipeline.enable_xformers_memory_efficient_attention()
+                    self.optimizations_applied.append("xformers")
+
+                # Try CPU offloading for memory efficiency
+                if torch.cuda.get_device_properties(0).total_memory < 8 * 1024**3:  # < 8GB
+                    logger.info("Low GPU memory detected, enabling CPU offloading")
+                    self.pipeline.enable_sequential_cpu_offload()
+                    self.optimizations_applied.append("cpu_offload")
+                else:
+                    self.optimizations_applied.append("gpu_only")
+
+                # Try torch.compile
+                if self._try_enable_torch_compile():
+                    # Compile the UNet for better performance
+                    try:
+                        self.pipeline.unet = torch.compile(self.pipeline.unet, mode="reduce-overhead")
+                        self.optimizations_applied.append("torch_compile")
+                    except Exception as e:
+                        logger.warning(f"Could not compile UNet: {e}")
             else:
-                del self.analysis_cache[cache_key]
-        return None
-
-    def cache_analysis(self, cache_key: str, analysis: str):
-        """Cache an analysis"""
-        if len(self.analysis_cache) >= self.max_size:
-            oldest_key = min(self.analysis_cache.keys(),
-                           key=lambda k: self.analysis_cache[k][1])
-            del self.analysis_cache[oldest_key]
+                self.optimizations_applied.append("cpu_only")
 
-        self.analysis_cache[cache_key] = (analysis, datetime.now())
+            # Enable VAE slicing for memory efficiency
+            self.pipeline.enable_vae_slicing()
+            self.optimizations_applied.append("vae_slicing")
 
-    def clear_expired(self):
-        """Clear expired cache entries"""
-        now = datetime.now()
-        expired_images = [k for k, (_, t) in self.image_cache.items()
-                         if now - t >= self.cache_ttl]
-        for k in expired_images:
-            del self.image_cache[k]
-
-        expired_analyses = [k for k, (_, t) in self.analysis_cache.items()
-                           if now - t >= self.cache_ttl]
-        for k in expired_analyses:
-            del self.analysis_cache[k]
+            logger.info(f"Model loaded successfully with optimizations: {', '.join(self.optimizations_applied)}")
+            return True
 
+        except Exception as e:
+            logger.error(f"Failed to load model: {e}")
+            self.pipeline = None
+            return False
 
 class ImageProcessor:
-    """Handles image generation and transformation with error handling"""
+    """Process image generation and transformation"""
 
-    def __init__(self, model_manager: ModelManager, cache_manager: CacheManager):
+    def __init__(self, model_manager: ModelManager, cache_manager: CacheManager,
+                 system_monitor: SystemMonitor):
         self.model_manager = model_manager
         self.cache_manager = cache_manager
-        self.style_suffixes = {
-            "None": "",
-            "Photorealistic": ", photorealistic, ultra detailed, 8k, professional photography",
-            "Cinematic": ", cinematic lighting, movie scene, dramatic atmosphere, film grain",
-            "Anime": ", anime style, vibrant colors, cel shaded, studio ghibli inspired",
-            "Digital Art": ", digital art, detailed illustration, concept art",
-            "Oil Painting": ", oil painting, classical art, rich textures",
-            "Watercolor": ", watercolor painting, soft edges, artistic",
-            "3D Render": ", 3D render, octane render, detailed 3D",
-            "Fantasy": ", fantasy art, magical, ethereal atmosphere",
-            "Sci-Fi": ", sci-fi art, futuristic, high-tech"
-        }
-
-    @GPU(duration=120)
-    async def generate_image(
-        self,
-        prompt: str,
-        style: str = "None",
-        ratio: str = "1:1 Square (1024x1024)",
-        steps: int = 9,
-        seed: int = 42,
-        randomize: bool = True,
-        guidance_scale: float = 0.0
-    ) -> GenerationResult:
+        self.system_monitor = system_monitor
+
+    def generate_image(self, prompt: str, negative_prompt: str = "", style: str = "None",
+                      aspect_ratio: str = "1:1", guidance_scale: float = 7.5,
+                      inference_steps: int = 4, seed: int = -1,
+                      use_cache: bool = True) -> GenerationResult:
         """Generate image with comprehensive error handling"""
-        result = GenerationResult(success=False)
         start_time = time.time()
 
-        try:
-            # Validate inputs
-            if not prompt or not prompt.strip():
-                result.error_code = ErrorCode.INVALID_INPUT
-                result.message = "Prompt cannot be empty"
-                return result
-
-            # Ensure models are loaded
-            if not await self.model_manager.load_models():
-                result.error_code = ErrorCode.MODEL_LOAD_ERROR
-                result.message = "Failed to load models"
+        # Check cache first
+        if use_cache:
+            cached_image = self.cache_manager.get(
+                prompt, negative_prompt, style, aspect_ratio,
+                guidance_scale, inference_steps, seed
+            )
+            if cached_image:
+                result = GenerationResult(
+                    success=True,
+                    image=cached_image,
+                    generation_time=time.time() - start_time,
+                    cache_hit=True
+                )
+                self.system_monitor.log_generation(True, result.generation_time)
                 return result
 
-            # Parse dimensions
-            width, height = self._parse_aspect_ratio(ratio)
-
-            # Handle seed
-            if randomize:
-                seed = torch.randint(0, 2**32 - 1, (1,)).item()
+        # Validate model
+        if not self.model_manager.load_model():
+            return GenerationResult(
+                success=False,
+                error="Model not loaded",
+                error_code=1001,
+                generation_time=time.time() - start_time
+            )
 
-            generator = torch.Generator().manual_seed(seed)
+        try:
+            # Set seed
+            if seed != -1:
+                torch.manual_seed(seed)
+                if torch.cuda.is_available():
+                    torch.cuda.manual_seed_all(seed)
+
+            # Get dimensions
+            width, height = ASPECT_RATIOS.get(aspect_ratio, ASPECT_RATIOS[DEFAULT_ASPECT_RATIO])
+
+            # Apply style to prompt
+            if style != "None":
+                style_prompt = f"{prompt}, {style.lower()} style"
+            else:
+                style_prompt = prompt
 
-            # Enhance prompt with style
-            enhanced_prompt = prompt + self.style_suffixes.get(style, "")
+            # Generate
+            logger.info(f"Generating image: {style_prompt[:50]}...")
+
+            # Adjust parameters for quality/speed balance
+            if inference_steps < 4:
+                inference_steps = 4  # Minimum for quality
+            elif inference_steps > 50:
+                inference_steps = 50  # Maximum for efficiency
+
+            with torch.cuda.amp.autocast() if self.model_manager.device == "cuda" else torch.no_grad():
+                result_image = self.model_manager.pipeline(
+                    prompt=style_prompt,
+                    negative_prompt=negative_prompt,
+                    num_inference_steps=inference_steps,
+                    guidance_scale=guidance_scale,
+                    width=width,
+                    height=height,
+                    num_images_per_prompt=1,
+                    generator=torch.Generator(device=self.model_manager.device).manual_seed(seed) if seed != -1 else None
+                ).images[0]
 
-            # Optimize steps based on complexity
-            optimized_steps = self._optimize_steps(enhanced_prompt, steps)
+            # Cache the result
+            if use_cache:
+                self.cache_manager.put(
+                    prompt, negative_prompt, style, aspect_ratio,
+                    guidance_scale, inference_steps, seed, result_image
+                )
 
-            # Generate
-            logger.info(f"Generating: {enhanced_prompt[:50]}... | {width}x{height} | {optimized_steps} steps")
-
-            output = self.model_manager.pipe_t2i(
-                prompt=enhanced_prompt,
-                width=width,
-                height=height,
-                num_inference_steps=optimized_steps,
-                guidance_scale=guidance_scale,
-                generator=generator,
-                output_type="pil"
+            generation_time = time.time() - start_time
+            result = GenerationResult(
+                success=True,
+                image=result_image,
+                generation_time=generation_time,
+                cache_hit=False,
+                optimization_used=self.model_manager.optimizations_applied.copy()
             )
 
-            # Success
-            result.success = True
-            result.image = output.images[0]
-            result.seed = seed
-            result.message = "Generated successfully"
-            result.generation_time = time.time() - start_time
-
-            logger.info(f"Generated in {result.generation_time:.2f}s")
+            self.system_monitor.log_generation(True, generation_time)
+            logger.info(f"Image generated successfully in {generation_time:.2f}s")
+            return result
 
         except torch.cuda.OutOfMemoryError:
-            result.error_code = ErrorCode.RESOURCE_ERROR
-            result.message = "GPU out of memory. Try smaller image or restart space."
-            logger.error("GPU OOM during generation")
-        except Exception as e:
-            result.error_code = ErrorCode.GENERATION_ERROR
-            result.message = f"Generation failed: {str(e)}"
-            logger.error(f"Generation error: {e}")
-            logger.error(traceback.format_exc())
-
-        return result
-
-    @GPU(duration=120)
-    async def transform_image(
-        self,
-        input_image: Image.Image,
-        prompt: str,
-        style: str = "None",
-        strength: float = 0.8,
-        steps: int = 9,
-        seed: int = 42,
-        randomize: bool = True,
-        guidance_scale: float = 0.0
-    ) -> GenerationResult:
-        """Transform image with comprehensive error handling"""
-        result = GenerationResult(success=False)
-        start_time = time.time()
-
-        try:
-            # Validate inputs
-            if input_image is None:
-                result.error_code = ErrorCode.INVALID_INPUT
-                result.message = "Please upload an image"
-                return result
-
-            if not prompt or not prompt.strip():
-                result.error_code = ErrorCode.INVALID_INPUT
-                result.message = "Prompt cannot be empty"
-                return result
-
-            # Ensure models are loaded
-            if not await self.model_manager.load_models():
-                result.error_code = ErrorCode.MODEL_LOAD_ERROR
-                result.message = "Failed to load models"
-                return result
+            logger.error("CUDA out of memory")
+            # Clear cache and try to free memory
+            self.cache_manager.clear()
+            gc.collect()
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+
+            self.system_monitor.log_generation(False, time.time() - start_time)
+            return GenerationResult(
+                success=False,
+                error="GPU out of memory. Try smaller image size or restart space.",
+                error_code=3001,
+                generation_time=time.time() - start_time
+            )
 
-            # Preprocess image
-            processed_image = self._preprocess_image(input_image)
+        except Exception as e:
+            logger.error(f"Generation failed: {e}")
+            self.system_monitor.log_generation(False, time.time() - start_time)
+            return GenerationResult(
+                success=False,
+                error=str(e),
+                error_code=1002,
+                generation_time=time.time() - start_time
+            )
 
-            # Handle seed
-            if randomize:
-                seed = torch.randint(0, 2**32 - 1, (1,)).item()
+    def transform_image(self, source_image: Image.Image, prompt: str,
+                       negative_prompt: str = "", strength: float = 0.7,
+                       guidance_scale: float = 7.5, inference_steps: int = 4,
+                       seed: int = -1) -> GenerationResult:
+        """Transform existing image"""
+        start_time = time.time()
 
-            generator = torch.Generator().manual_seed(seed)
+        # Validate model
+        if not self.model_manager.load_model():
+            return GenerationResult(
+                success=False,
+                error="Model not loaded",
+                error_code=1001,
+                generation_time=time.time() - start_time
+            )
 
-            # Enhance prompt
-            enhanced_prompt = prompt + self.style_suffixes.get(style, "")
+        try:
+            # Set seed
+            if seed != -1:
+                torch.manual_seed(seed)
+                if torch.cuda.is_available():
+                    torch.cuda.manual_seed_all(seed)
 
-            # Optimize steps based on strength
-            effective_steps = max(4, int(steps * strength)) if strength > 0 else steps
+            # Prepare image
+            source_image = source_image.convert("RGB")
 
             # Transform
-            logger.info(f"Transforming: {enhanced_prompt[:50]}... | strength={strength}")
-
-            output = self.model_manager.pipe_i2i(
-                prompt=enhanced_prompt,
-                image=processed_image,
-                strength=strength,
-                num_inference_steps=effective_steps,
-                guidance_scale=guidance_scale,
-                generator=generator,
-                output_type="pil"
+            logger.info(f"Transforming image with prompt: {prompt[:50]}...")
+
+            with torch.cuda.amp.autocast() if self.model_manager.device == "cuda" else torch.no_grad():
+                result_image = self.model_manager.pipeline(
+                    image=source_image,
+                    prompt=prompt,
+                    negative_prompt=negative_prompt,
+                    strength=strength,
+                    num_inference_steps=inference_steps,
+                    guidance_scale=guidance_scale,
+                    generator=torch.Generator(device=self.model_manager.device).manual_seed(seed) if seed != -1 else None
+                ).images[0]
+
+            generation_time = time.time() - start_time
+            result = GenerationResult(
+                success=True,
+                image=result_image,
+                generation_time=generation_time,
+                optimization_used=self.model_manager.optimizations_applied.copy()
             )
 
-            # Success
-            result.success = True
-            result.image = output.images[0]
-            result.seed = seed
-            result.message = "Transformed successfully"
-            result.generation_time = time.time() - start_time
-
-            logger.info(f"Transformed in {result.generation_time:.2f}s")
+            self.system_monitor.log_generation(True, generation_time)
+            logger.info(f"Image transformed successfully in {generation_time:.2f}s")
+            return result
 
-        except torch.cuda.OutOfMemoryError:
-            result.error_code = ErrorCode.RESOURCE_ERROR
-            result.message = "GPU out of memory. Try smaller image or restart space."
-            logger.error("GPU OOM during transform")
         except Exception as e:
-            result.error_code = ErrorCode.TRANSFORM_ERROR
-            result.message = f"Transform failed: {str(e)}"
-            logger.error(f"Transform error: {e}")
-            logger.error(traceback.format_exc())
-
-        return result
-
-    def _parse_aspect_ratio(self, ratio: str) -> Tuple[int, int]:
-        """Parse aspect ratio string to dimensions"""
-        ratios = {
-            "1:1": (1024, 1024),
-            "16:9": (1344, 768),
-            "9:16": (768, 1344),
-            "4:3": (1152, 896),
-            "3:4": (896, 1152)
-        }
-
-        # Extract ratio from string
-        for key, (w, h) in ratios.items():
-            if key in ratio:
-                return w, h
-
-        # Default to 1:1
-        return 1024, 1024
-
-    def _optimize_steps(self, prompt: str, base_steps: int) -> int:
-        """Optimize step count based on prompt complexity"""
-        # Calculate complexity score
-        words = len(prompt.split())
-        commas = prompt.count(',')
-        periods = prompt.count('.')
-
-        complexity = words + (commas * 2) + (periods * 2)
-
-        # Adjust steps
-        if complexity < 10:
-            return max(4, base_steps - 2)
-        elif complexity > 30:
-            return min(16, base_steps + 2)
-        else:
-            return base_steps
-
-    def _preprocess_image(self, image: Image.Image) -> Image.Image:
-        """Preprocess image for img2img pipeline"""
-        # Convert to RGB
-        if image.mode != "RGB":
-            image = image.convert("RGB")
-
-        # Resize to standard dimensions (maintain aspect ratio)
-        w, h = image.size
-
-        # Calculate new dimensions (multiple of 16)
-        max_size = 1024
-        aspect_ratio = w / h
-
-        if w > h:
-            new_w = min(max_size, w)
-            new_h = int(new_w / aspect_ratio)
-        else:
-            new_h = min(max_size, h)
-            new_w = int(new_h * aspect_ratio)
-
-        # Round to nearest multiple of 16
-        new_w = (new_w // 16) * 16
-        new_h = (new_h // 16) * 16
-
-        # Ensure minimum dimensions
-        new_w = max(512, new_w)
-        new_h = max(512, new_h)
-
-        return image.resize((new_w, new_h), Image.LANCZOS)
-
+            logger.error(f"Transform failed: {e}")
+            self.system_monitor.log_generation(False, time.time() - start_time)
+            return GenerationResult(
+                success=False,
+                error=str(e),
+                error_code=1003,
+                generation_time=time.time() - start_time
+            )
 
-# Initialize global components
-system_monitor = SystemMonitor()
+# Create global instances
 model_manager = ModelManager()
-cache_manager = CacheManager(max_size=100)
-image_processor = ImageProcessor(model_manager, cache_manager)
-
-# UI Constants
-STYLES = ["None", "Photorealistic", "Cinematic", "Anime", "Digital Art",
-          "Oil Painting", "Watercolor", "3D Render", "Fantasy", "Sci-Fi"]
-
-RATIOS = [
-    "1:1 Square (1024x1024)",
-    "16:9 Landscape (1344x768)",
-    "9:16 Portrait (768x1344)",
-    "4:3 Standard (1152x896)"
-]
-
-# CSS for enhanced UI
-CSS = """
-:root {
-    --primary: #3b82f6;
-    --primary-dark: #2563eb;
-    --secondary: #10b981;
-    --background: #f8fafc;
-    --surface: #ffffff;
-    --error: #ef4444;
-    --warning: #f59e0b;
-    --success: #22c55e;
-    --border-radius: 12px;
-    --shadow: 0 4px 6px -1px rgb(0 0 0 / 0.1);
-}
-
-/* Main container */
-.gradio-container {
-    font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif;
-    background: var(--background);
-}
-
-/* Headers */
-h1 {
-    color: #1e293b;
-    font-weight: 700;
-    font-size: 2.5rem;
-    margin-bottom: 0.5rem;
-}
-
-h2 {
-    color: #334155;
-    font-weight: 600;
-    font-size: 1.5rem;
-    margin-top: 1.5rem;
-}
-
-/* Buttons */
-.gradio-button {
-    border-radius: var(--border-radius);
-    font-weight: 600;
-    transition: all 0.2s ease;
-}
-
-.gradio-button.primary {
-    background: var(--primary);
-    border: none;
-}
-
-.gradio-button.primary:hover {
-    background: var(--primary-dark);
-    transform: translateY(-1px);
-    box-shadow: var(--shadow);
-}
-
-/* Cards */
-.border {
-    border: 1px solid #e2e8f0 !important;
-    border-radius: var(--border-radius) !important;
-    background: var(--surface);
-}
-
-/* Status indicators */
-.status-success {
-    color: var(--success);
-    font-weight: 600;
-}
-
-.status-error {
-    color: var(--error);
-    font-weight: 600;
-}
-
-.status-warning {
-    color: var(--warning);
-    font-weight: 600;
-}
-
-/* Performance metrics */
-.metric-card {
-    background: var(--surface);
-    padding: 1rem;
-    border-radius: var(--border-radius);
-    box-shadow: var(--shadow);
-}
-
-.metric-value {
-    font-size: 2rem;
-    font-weight: 700;
-    color: var(--primary);
-}
-
-.metric-label {
-    color: #64748b;
-    font-size: 0.875rem;
-    margin-top: 0.25rem;
-}
-
-/* Animations */
-@keyframes pulse {
-    0%, 100% { opacity: 1; }
-    50% { opacity: 0.5; }
-}
-
-.loading {
-    animation: pulse 2s cubic-bezier(0.4, 0, 0.6, 1) infinite;
-}
-
-/* Responsive design */
-@media (max-width: 768px) {
-    .gradio-row {
-        flex-direction: column !important;
-    }
-}
-"""
-
-
-async def handle_generation(
-    prompt: str,
-    style: str,
-    ratio: str,
-    steps: int,
-    seed: int,
-    randomize: bool,
-    guidance_scale: float
-) -> Tuple[Optional[Image.Image], int, str]:
-    """Handle image generation with caching"""
-    try:
-        # Check cache first
-        cache_key = cache_manager.get_cache_key(prompt, style, ratio, steps, seed)
-        cached_image = cache_manager.get_cached_image(cache_key)
-
-        if cached_image:
-            system_monitor.record_cache_hit()
-            logger.info("Returning cached generation")
-            return cached_image, seed, "✅ Retrieved from cache"
-
-        system_monitor.record_cache_miss()
-
-        # Generate new image
-        result = await image_processor.generate_image(
-            prompt=prompt,
-            style=style,
-            ratio=ratio,
-            steps=steps,
-            seed=seed,
-            randomize=randomize,
-            guidance_scale=guidance_scale
-        )
-
-        if result.success:
-            # Cache the result
-            cache_manager.cache_image(cache_key, result.image)
-            system_monitor.record_generation(True)
-            return result.image, result.seed, f"✅ {result.message} ({result.generation_time:.1f}s)"
-        else:
-            system_monitor.record_generation(False)
-            return None, seed, f"❌ {result.message}"
-
-    except Exception as e:
-        system_monitor.record_generation(False)
-        logger.error(f"Generation handler error: {e}")
-        return None, seed, f"❌ Unexpected error: {str(e)}"
-
-
-async def handle_transform(
-    input_image: Image.Image,
-    prompt: str,
-    style: str,
-    strength: float,
-    steps: int,
-    seed: int,
-    randomize: bool,
-    guidance_scale: float
-) -> Tuple[Optional[Image.Image], int, str]:
-    """Handle image transformation"""
-    try:
-        result = await image_processor.transform_image(
-            input_image=input_image,
-            prompt=prompt,
-            style=style,
-            strength=strength,
-            steps=steps,
-            seed=seed,
-            randomize=randomize,
-            guidance_scale=guidance_scale
-        )
-
-        if result.success:
-            system_monitor.record_generation(True)
-            return result.image, result.seed, f"✅ {result.message} ({result.generation_time:.1f}s)"
+cache_manager = CacheManager()
+system_monitor = SystemMonitor()
+image_processor = ImageProcessor(model_manager, cache_manager, system_monitor)
+
+# Preload model if possible
+model_manager.load_model()
+
+def format_system_stats(stats: Dict[str, Any]) -> str:
+    """Format system stats for display"""
+    html = "<div class='system-monitor'>"
+    html += "<h4>🖥️ System Resources</h4>"
+    html += f"<strong>CPU:</strong> {stats['cpu_percent']:.1f}%<br>"
+    html += f"<strong>Memory:</strong> {stats['memory_percent']:.1f}%<br>"
+    html += f"<strong>Disk:</strong> {stats['disk_percent']:.1f}%<br>"
+
+    if stats['gpu_available']:
+        gpu_mem = stats.get('gpu_memory', {})
+        if gpu_mem:
+            html += f"<strong>GPU Memory:</strong> {gpu_mem.get('allocated', 0):.1f}GB / {gpu_mem.get('total', 0):.1f}GB<br>"
         else:
-            system_monitor.record_generation(False)
-            return None, seed, f"❌ {result.message}"
+            html += "<strong>GPU:</strong> Available<br>"
+    else:
+        html += "<strong>GPU:</strong> Not Available<br>"
+
+    html += "<h4>📊 Performance Metrics</h4>"
+    html += f"<strong>Uptime:</strong> {stats['uptime_str']}<br>"
+    html += f"<strong>Generations:</strong> {stats['generation_count']}<br>"
+    html += f"<strong>Success Rate:</strong> {stats['success_rate']:.1f}%<br>"
+    html += f"<strong>Avg Time:</strong> {stats['avg_generation_time']}s<br>"
+    html += f"<strong>Gen/Min:</strong> {stats['generations_per_minute']:.1f}<br>"
+
+    cache_stats = cache_manager.get_stats()
+    html += "<h4>💾 Cache</h4>"
+    html += f"<strong>Size:</strong> {cache_stats['size']}/{cache_stats['max_size']}<br>"
+    html += f"<strong>Usage:</strong> {cache_stats['usage_percent']:.1f}%<br>"
+
+    if model_manager.optimizations_applied:
+        html += "<h4>⚡ Active Optimizations</h4>"
+        html += f"{' • '.join(model_manager.optimizations_applied)}<br>"
+
+    html += "</div>"
+    return html
+
+def generate_image_wrapper(prompt: str, negative_prompt: str, style: str,
+                          aspect_ratio: str, guidance_scale: float,
+                          inference_steps: int, seed: int, progress=gr.Progress()) -> Tuple[Optional[Image.Image], str]:
+    """Wrapper for image generation with progress tracking"""
+    progress(0.1, desc="Preparing generation...")
+
+    result = image_processor.generate_image(
+        prompt=prompt,
+        negative_prompt=negative_prompt,
+        style=style,
+        aspect_ratio=aspect_ratio,
+        guidance_scale=guidance_scale,
+        inference_steps=int(inference_steps),
+        seed=int(seed)
+    )
 
-    except Exception as e:
-        system_monitor.record_generation(False)
-        logger.error(f"Transform handler error: {e}")
-        return None, seed, f"❌ Unexpected error: {str(e)}"
+    progress(0.9, desc="Finalizing...")
+
+    if result.success:
+        message = f"✅ Generated in {result.generation_time:.2f}s"
+        if result.cache_hit:
+            message += " (from cache)"
+        return result.image, message
+    else:
+        error_msg = f"❌ Error {result.error_code}: {result.error}"
+        return None, error_msg
+
+def transform_image_wrapper(source_image: Image.Image, prompt: str,
+                           negative_prompt: str, strength: float,
+                           guidance_scale: float, inference_steps: int,
+                           seed: int, progress=gr.Progress()) -> Tuple[Optional[Image.Image], str]:
+    """Wrapper for image transformation with progress tracking"""
+    if source_image is None:
+        return None, "❌ Please upload an image"
+
+    progress(0.1, desc="Preparing transformation...")
+
+    result = image_processor.transform_image(
+        source_image=source_image,
+        prompt=prompt,
+        negative_prompt=negative_prompt,
+        strength=strength,
+        guidance_scale=guidance_scale,
+        inference_steps=int(inference_steps),
+        seed=int(seed)
+    )
 
+    progress(0.9, desc="Finalizing...")
 
-def create_interface() -> gr.Blocks:
-    """Create the Gradio interface"""
+    if result.success:
+        return result.image, f"✅ Transformed in {result.generation_time:.2f}s"
+    else:
+        return None, f"❌ Error {result.error_code}: {result.error}"
 
+# Build Gradio interface
+def build_interface():
+    """Build the Gradio interface"""
     with gr.Blocks(
-        title="Z Image Turbo - Production",
-        theme=gr.themes.Soft(),
-        css=CSS
+        title="Z Image Turbo",
+        css=CUSTOM_CSS,
+        theme=gr.themes.Soft()
     ) as demo:
-
-        # Header
-        gr.HTML("""
-        <div style="text-align: center; padding: 2rem 0;">
-            <h1 style="margin: 0;">⚡ Z Image Turbo</h1>
-            <p style="color: #64748b; font-size: 1.1rem; margin-top: 0.5rem;">
-                Production-Ready Image Generation with Advanced Optimizations
-            </p>
-        </div>
-        """)
+        gr.Markdown("# 🎨 Z Image Turbo")
+        gr.Markdown("High-performance image generation and transformation")
 
         with gr.Tabs():
             # Generation Tab
-            with gr.Tab("🎨 Generate", elem_id="generate-tab"):
+            with gr.TabItem("✨ Generate"):
                 with gr.Row():
-                    with gr.Column(scale=3):
-                        gen_prompt = gr.Textbox(
+                    with gr.Column(scale=2):
+                        prompt_input = gr.Textbox(
                             label="Prompt",
                             placeholder="Describe the image you want to generate...",
-                            lines=3,
-                            max_lines=5
+                            lines=3
                         )
-
-                        gen_style = gr.Dropdown(
-                            choices=STYLES,
-                            value="None",
-                            label="Style",
-                            info="Apply a style to your generation"
-                        )
-
-                        gen_ratio = gr.Dropdown(
-                            choices=RATIOS,
-                            value="1:1 Square (1024x1024)",
-                            label="Aspect Ratio"
+                        negative_prompt_input = gr.Textbox(
+                            label="Negative Prompt",
+                            placeholder="What you don't want in the image...",
+                            lines=2,
+                            value=""
                         )
 
                         with gr.Row():
-                            gen_steps = gr.Slider(
-                                minimum=4,
-                                maximum=16,
-                                value=9,
-                                step=1,
-                                label="Inference Steps",
-                                info="More steps = better quality but slower"
+                            style_dropdown = gr.Dropdown(
+                                label="Style",
+                                choices=STYLE_PRESETS,
+                                value="None"
                             )
-
-                            gen_guidance = gr.Slider(
-                                minimum=0.0,
-                                maximum=20.0,
-                                value=0.0,
-                                step=0.5,
-                                label="Guidance Scale",
-                                info="Higher = more prompt adherence"
+                            aspect_ratio_dropdown = gr.Dropdown(
+                                label="Aspect Ratio",
+                                choices=list(ASPECT_RATIOS.keys()),
+                                value=DEFAULT_ASPECT_RATIO
                             )
 
                         with gr.Row():
-                            gen_seed = gr.Number(
-                                label="Seed",
-                                value=42,
-                                precision=0,
-                                info="Fixed seed for reproducible results"
+                            guidance_scale = gr.Slider(
+                                label="Guidance Scale",
+                                minimum=1.0,
+                                maximum=20.0,
+                                value=7.5,
+                                step=0.5
                             )
-                            gen_randomize = gr.Checkbox(
-                                label="Randomize Seed",
-                                value=True,
-                                info="Generate with random seed"
+                            inference_steps = gr.Slider(
+                                label="Inference Steps",
+                                minimum=1,
+                                maximum=50,
+                                value=4,
+                                step=1
                             )
 
-                        gen_btn = gr.Button(
+                        seed_input = gr.Number(
+                            label="Seed (-1 for random)",
+                            value=-1,
+                            precision=0
+                        )
+
+                        generate_btn = gr.Button(
                             "🚀 Generate",
                             variant="primary",
-                            size="lg",
-                            elem_classes=["generate-button"]
+                            elem_classes=["generate-btn"]
                         )
 
-                    with gr.Column(scale=2):
-                        gen_output = gr.Image(
+                    with gr.Column(scale=1):
+                        output_image = gr.Image(
                             label="Generated Image",
-                            type="pil",
-                            format="png",
-                            interactive=False,
-                            show_share_button=True,
-                            show_download_button=True,
-                            elem_classes=["output-image"]
+                            type="pil"
                         )
-
-                        gen_status = gr.Textbox(
-                            label="Status",
-                            interactive=False,
-                            max_lines=2,
-                            elem_classes=["status-text"]
+                        output_info = gr.Textbox(
+                            label="Information",
+                            interactive=False
                         )
 
-                        gen_seed_out = gr.Number(
-                            label="Seed Used",
-                            interactive=False,
-                            precision=0
-                        )
-
-                # Event handler
-                gen_btn.click(
-                    fn=lambda *args: asyncio.run(handle_generation(*args)),
+                generate_btn.click(
+                    fn=generate_image_wrapper,
                     inputs=[
-                        gen_prompt, gen_style, gen_ratio,
-                        gen_steps, gen_seed, gen_randomize, gen_guidance
+                        prompt_input,
+                        negative_prompt_input,
+                        style_dropdown,
+                        aspect_ratio_dropdown,
+                        guidance_scale,
+                        inference_steps,
+                        seed_input
                     ],
-                    outputs=[gen_output, gen_seed_out, gen_status]
+                    outputs=[output_image, output_info]
                 )
 
             # Transform Tab
-            with gr.Tab("🖼️ Transform", elem_id="transform-tab"):
-                gr.Markdown("""
-                ### Transform an existing image with AI
-                Upload an image and provide a prompt to guide the transformation.
-                """)
-
+            with gr.TabItem("🔄 Transform"):
                 with gr.Row():
-                    with gr.Column(scale=3):
-                        trans_input = gr.Image(
-                            label="Input Image",
-                            type="pil",
-                            sources=["upload", "webcam"]
+                    with gr.Column(scale=2):
+                        source_image_input = gr.Image(
+                            label="Source Image",
+                            type="pil"
                         )
-
-                        trans_prompt = gr.Textbox(
+                        transform_prompt = gr.Textbox(
                             label="Transform Prompt",
                             placeholder="Describe how to transform the image...",
-                            lines=2
+                            lines=3
                         )
-
-                        trans_style = gr.Dropdown(
-                            choices=STYLES,
-                            value="None",
-                            label="Style"
+                        transform_negative_prompt = gr.Textbox(
+                            label="Negative Prompt",
+                            placeholder="What to avoid in the transformation...",
+                            lines=2,
+                            value=""
                         )
 
                         with gr.Row():
-                            trans_strength = gr.Slider(
+                            transform_strength = gr.Slider(
+                                label="Transform Strength",
                                 minimum=0.0,
                                 maximum=1.0,
-                                value=0.8,
-                                step=0.1,
-                                label="Transformation Strength",
-                                info="Higher = more changes"
+                                value=0.7,
+                                step=0.1
                             )
-
-                            trans_steps = gr.Slider(
-                                minimum=4,
-                                maximum=16,
-                                value=9,
-                                step=1,
-                                label="Inference Steps"
+                            transform_guidance = gr.Slider(
+                                label="Guidance Scale",
+                                minimum=1.0,
+                                maximum=20.0,
+                                value=7.5,
+                                step=0.5
                             )
 
-                        with gr.Row():
-                            trans_seed = gr.Number(
-                                label="Seed",
-                                value=42,
-                                precision=0
-                            )
-                            trans_randomize = gr.Checkbox(
-                                label="Randomize Seed",
-                                value=True
-                            )
+                        transform_steps = gr.Slider(
+                            label="Inference Steps",
+                            minimum=1,
+                            maximum=50,
+                            value=4,
+                            step=1
+                        )
 
-                        trans_btn = gr.Button(
-                            "✨ Transform",
-                            variant="primary",
-                            size="lg"
+                        transform_seed = gr.Number(
+                            label="Seed (-1 for random)",
+                            value=-1,
+                            precision=0
                         )
 
-                    with gr.Column(scale=2):
-                        trans_output = gr.Image(
-                            label="Transformed Image",
-                            type="pil",
-                            format="png",
-                            interactive=False,
-                            show_share_button=True
+                        transform_btn = gr.Button(
+                            "🔄 Transform",
+                            variant="primary"
                         )
 
-                        trans_status = gr.Textbox(
-                            label="Status",
-                            interactive=False,
-                            max_lines=2
+                    with gr.Column(scale=1):
+                        transformed_image = gr.Image(
+                            label="Transformed Image",
+                            type="pil"
+                        )
+                        transform_info = gr.Textbox(
+                            label="Information",
+                            interactive=False
                         )
 
-                # Event handler
-                trans_btn.click(
-                    fn=lambda *args: asyncio.run(handle_transform(*args)),
+                transform_btn.click(
+                    fn=transform_image_wrapper,
                     inputs=[
-                        trans_input, trans_prompt, trans_style,
-                        trans_strength, trans_steps, trans_seed,
-                        trans_randomize, gen_guidance
+                        source_image_input,
+                        transform_prompt,
+                        transform_negative_prompt,
+                        transform_strength,
+                        transform_guidance,
+                        transform_steps,
+                        transform_seed
                     ],
-                    outputs=[trans_output, trans_seed_out, trans_status]
+                    outputs=[transformed_image, transform_info]
                 )
 
             # System Monitor Tab
-            with gr.Tab("📊 System Monitor", elem_id="monitor-tab"):
-                gr.Markdown("""
-                ### System Performance Metrics
-                Real-time monitoring of system resources and application performance.
-                """)
-
-                # Resource metrics
-                with gr.Row():
-                    with gr.Column():
-                        gr.Markdown("#### 🖥️ System Resources")
-
-                        with gr.Row():
-                            mem_usage = gr.Number(
-                                label="Memory Usage (GB)",
-                                precision=2,
-                                elem_classes=["metric-value"]
-                            )
-                            cpu_usage = gr.Number(
-                                label="CPU Usage (%)",
-                                precision=1,
-                                elem_classes=["metric-value"]
-                            )
-
-                        gpu_mem = gr.Number(
-                            label="GPU Memory (GB)",
-                            precision=2,
-                            elem_classes=["metric-value"]
-                        )
-
-                    with gr.Column():
-                        gr.Markdown("#### 📈 Application Metrics")
-
-                        with gr.Row():
-                            uptime = gr.Number(
-                                label="Uptime (seconds)",
-                                precision=0,
-                                elem_classes=["metric-value"]
-                            )
-                            generations = gr.Number(
-                                label="Total Generations",
-                                precision=0,
-                                elem_classes=["metric-value"]
-                            )
-
-                        cache_rate = gr.Number(
-                            label="Cache Hit Rate (%)",
-                            precision=1,
-                            elem_classes=["metric-value"]
-                        )
-
-                # Optimization status
-                with gr.Row():
-                    opt_status = gr.JSON(
-                        label="Optimization Status",
-                        value={
-                            "model_loaded": False,
-                            "optimizations": [],
-                            "last_update": datetime.now().isoformat()
-                        }
-                    )
-
-                # Refresh button
-                refresh_btn = gr.Button("🔄 Refresh", size="sm")
-
-                # Refresh handler
-                def refresh_metrics():
-                    """Refresh all metrics"""
-                    info = system_monitor.get_system_info()
-
-                    return (
-                        info.get("memory_used_gb", 0),
-                        info.get("cpu_percent", 0),
-                        info.get("gpu_memory_used_gb", 0),
-                        info.get("uptime_seconds", 0),
-                        info.get("active_generations", 0),
-                        info.get("cache_hit_rate", 0),
-                        {
-                            "model_loaded": model_manager.is_loaded,
-                            "optimizations": model_manager.optimizations_applied,
-                            "last_update": datetime.now().isoformat()
-                        }
-                    )
+            with gr.TabItem("📊 System Monitor"):
+                system_stats = gr.HTML(
+                    value=format_system_stats(system_monitor.get_stats()),
+                    label="System Statistics"
+                )
+                refresh_btn = gr.Button("🔄 Refresh")
+                clear_cache_btn = gr.Button("🗑️ Clear Cache")
 
                 refresh_btn.click(
-                    fn=refresh_metrics,
-                    outputs=[
-                        mem_usage, cpu_usage, gpu_mem,
-                        uptime, generations, cache_rate, opt_status
-                    ]
+                    fn=lambda: format_system_stats(system_monitor.get_stats()),
+                    outputs=[system_stats]
                 )
 
-                # Auto-refresh every 5 seconds
-                demo.load(
-                    fn=refresh_metrics,
-                    outputs=[
-                        mem_usage, cpu_usage, gpu_mem,
-                        uptime, generations, cache_rate, opt_status
-                    ],
-                    every=5
+                clear_cache_btn.click(
+                    fn=lambda: (cache_manager.clear(), format_system_stats(system_monitor.get_stats()))[1],
+                    outputs=[system_stats]
                 )
 
-            # About Tab
-            with gr.Tab("ℹ️ About", elem_id="about-tab"):
-                gr.Markdown("""
-                # Z Image Turbo - Production Edition
-
-                ## Features
-                - ✅ **PyTorch 2.0+ Compilation** with graceful fallback
-                - ✅ **xformers Optimization** with CPU fallback
-                - ✅ **Memory Management** with CPU offloading
-                - ✅ **Caching System** for improved performance
-                - ✅ **Comprehensive Error Handling**
-                - ✅ **Real-time Monitoring**
-                - ✅ **Production-Ready Architecture**
-
-                ## Model
-                - **Base Model**: [Tongyi-MAI/Z-Image-Turbo](https://huggingface.co/Tongyi-MAI/Z-Image-Turbo)
-                - **Architecture**: DiT-based diffusion model
-                - **Optimized for**: Fast generation with high quality
-
-                ## System Requirements
-                - GPU with at least 6GB VRAM recommended
-                - PyTorch 2.0+ for optimal performance
-                - Optional: xformers for memory efficiency
-
-                ## Changelog
-                ### v2.0.0 Production
-                - Added comprehensive error handling
-                - Implemented PyTorch compilation with fallback
-                - Added xformers optimization with CPU fallback
-                - Integrated caching system
-                - Added real-time monitoring
-                - Improved resource management
-
-                ---
-                Created with ❤️ by AI Agent Framework Specialist
-                """)
+        # Footer
+        gr.HTML("""
+        <div class="footer">
+            <p>🚀 Z Image Turbo - Production Edition v2.0.0</p>
+            <p>Created with ❤️ by AI Agent Framework Specialist</p>
+        </div>
+        """)
 
     return demo
 
-
-# Health check endpoint
-async def health_check() -> Dict[str, Any]:
-    """Application health check"""
-    return {
-        "status": "healthy" if model_manager.is_loaded else "loading",
-        "model_loaded": model_manager.is_loaded,
-        "optimizations": model_manager.optimizations_applied,
-        "uptime": time.time() - system_monitor.start_time,
-        "generation_count": system_monitor.generation_count,
-        "error_count": system_monitor.error_count,
-        "cache_hit_rate": system_monitor.get_cache_hit_rate()
-    }
-
-
-# Main application entry
+# Create and launch the demo
 if __name__ == "__main__":
-    logger.info("Starting Z Image Turbo - Production Edition")
+    logger.info("Starting Z Image Turbo application...")
 
-    # Create interface
-    demo = create_interface()
-
-    # Configure for Hugging Face Spaces
-    demo.queue(
-        api_open=False,
-        max_size=20,
-        default_concurrency_limit=1
-    )
+    demo = build_interface()
 
     # Launch with optimizations
     demo.launch(
         share=False,
         show_error=True,
-        show_tips=True,
         max_threads=40,
-        prevent_thread_lock=False
-    )
-
-    logger.info("Application launched successfully")
+        prevent_thread_lock=False,
+        enable_queue=True
+    )