Update app.py

app.py

@@ -1,458 +1,745 @@
 import gradio as gr
 import torch
-from diffusers import StableDiffusionXLPipeline, StableDiffusionXLImg2ImgPipeline
-from PIL import Image, PngImagePlugin
+from diffusers import StableDiffusionXLPipeline, StableDiffusionXLImg2ImgPipeline, EulerAncestralDiscreteScheduler
+from PIL import Image, PngImagePlugin, ImageFilter
 from datetime import datetime
 import os
 import gc
 import time
 import spaces
-from typing import Optional, Tuple
+from typing import Optional, Tuple, Dict, Any
 from huggingface_hub import hf_hub_download
 import tempfile
 import random
+import logging
+import torch.nn.functional as F
+from transformers import CLIPProcessor, CLIPModel
 
-# Global pipeline variables
-txt2img_pipe = None
-img2img_pipe = None
-device = "cuda" if torch.cuda.is_available() else "cpu"
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
-# Hugging Face model configuration
+# Constants
 MODEL_REPO = "ajsbsd/CyberRealistic-Pony"
 MODEL_FILENAME = "cyberrealisticPony_v110.safetensors"
-model_id = f"{MODEL_REPO}/{MODEL_FILENAME}"
+NSFW_MODEL_ID = "openai/clip-vit-base-patch32"  # CLIP model for NSFW detection
+MAX_SEED = 2**32 - 1
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+DTYPE = torch.float16 if DEVICE == "cuda" else torch.float32
+NSFW_THRESHOLD = 0.25  # Threshold for NSFW detection
 
-# Generation configuration for metadata
-generation_config = {
-    "vae": "SDXL VAE",
-    "sampler": "DPM++ 2M Karras",
-    "steps": 20
-}
-
-def clear_memory():
-    """Clear GPU memory"""
-    if torch.cuda.is_available():
-        torch.cuda.empty_cache()
-    gc.collect()
-
-def add_metadata_and_save(image: Image.Image, prompt: str, negative_prompt: str, seed: int, steps: int, guidance: float, strength: Optional[float] = None):
-    """Embed generation metadata into a PNG and save it."""
-    # Create temporary file with unique name
-    temp_path = tempfile.mktemp(suffix=".png")
+# Global pipeline state
+class PipelineManager:
+    def __init__(self):
+        self.txt2img_pipe = None
+        self.img2img_pipe = None
+        self.nsfw_detector_model = None
+        self.nsfw_detector_processor = None
+        self.model_loaded = False
+        self.nsfw_detector_loaded = False
+        
+    def clear_memory(self):
+        """Aggressive memory cleanup"""
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+            torch.cuda.synchronize()
+        gc.collect()
     
-    meta = PngImagePlugin.PngInfo()
-    meta.add_text("Prompt", prompt)
-    meta.add_text("NegativePrompt", negative_prompt)
-    meta.add_text("Model", model_id)
-    meta.add_text("VAE", generation_config["vae"])
-    meta.add_text("Sampler", generation_config["sampler"])
-    meta.add_text("Steps", str(steps))
-    meta.add_text("CFG_Scale", str(guidance))
-    if strength is not None:
-        meta.add_text("Strength", str(strength))
-    meta.add_text("Seed", str(seed))
-    meta.add_text("Date", datetime.now().strftime("%Y-%m-%d %H:%M:%S"))
+    def load_nsfw_detector(self) -> bool:
+        """Load NSFW detection model"""
+        if self.nsfw_detector_loaded:
+            return True
+            
+        try:
+            logger.info("Loading NSFW detector...")
+            self.nsfw_detector_processor = CLIPProcessor.from_pretrained(NSFW_MODEL_ID)
+            self.nsfw_detector_model = CLIPModel.from_pretrained(NSFW_MODEL_ID)
+            
+            if DEVICE == "cuda":
+                self.nsfw_detector_model = self.nsfw_detector_model.to(DEVICE)
+            
+            self.nsfw_detector_loaded = True
+            logger.info("NSFW detector loaded successfully!")
+            return True
+            
+        except Exception as e:
+            logger.error(f"Failed to load NSFW detector: {e}")
+            self.nsfw_detector_loaded = False
+            return False
     
-    image.save(temp_path, "PNG", pnginfo=meta)
-    return temp_path
-
-def load_models():
-    """Load both text2img and img2img pipelines optimized for Spaces"""
-    global txt2img_pipe, img2img_pipe
+    def is_nsfw(self, image: Image.Image, prompt: str = "") -> Tuple[bool, float]:
+        """
+        Detects NSFW content using CLIP-based zero-shot classification.
+        Falls back to prompt-based detection if CLIP model fails.
+        """
+        try:
+            # Load NSFW detector if not already loaded
+            if not self.nsfw_detector_loaded:
+                if not self.load_nsfw_detector():
+                    return self._fallback_nsfw_detection(prompt)
+            
+            # CLIP-based NSFW detection
+            inputs = self.nsfw_detector_processor(images=image, return_tensors="pt").to(DEVICE)
+            
+            with torch.no_grad():
+                image_features = self.nsfw_detector_model.get_image_features(**inputs)
+                
+                # Define text prompts for classification
+                safe_prompts = [
+                    "a safe family-friendly image", 
+                    "a general photo", 
+                    "appropriate content",
+                    "artistic photography"
+                ]
+                unsafe_prompts = [
+                    "explicit adult content", 
+                    "nudity", 
+                    "inappropriate sexual content",
+                    "pornographic material"
+                ]
+                
+                # Get text features
+                safe_inputs = self.nsfw_detector_processor(
+                    text=safe_prompts, return_tensors="pt", padding=True
+                ).to(DEVICE)
+                unsafe_inputs = self.nsfw_detector_processor(
+                    text=unsafe_prompts, return_tensors="pt", padding=True
+                ).to(DEVICE)
+                
+                safe_features = self.nsfw_detector_model.get_text_features(**safe_inputs)
+                unsafe_features = self.nsfw_detector_model.get_text_features(**unsafe_inputs)
+                
+                # Normalize features for cosine similarity
+                image_features = F.normalize(image_features, p=2, dim=-1)
+                safe_features = F.normalize(safe_features, p=2, dim=-1)
+                unsafe_features = F.normalize(unsafe_features, p=2, dim=-1)
+                
+                # Calculate similarities
+                safe_similarity = (image_features @ safe_features.T).mean().item()
+                unsafe_similarity = (image_features @ unsafe_features.T).mean().item()
+                
+                # Classification logic
+                is_nsfw_result = (
+                    unsafe_similarity > safe_similarity and 
+                    unsafe_similarity > NSFW_THRESHOLD
+                )
+                
+                confidence = unsafe_similarity if is_nsfw_result else safe_similarity
+                
+                if is_nsfw_result:
+                    logger.warning(f"🚨 NSFW content detected (CLIP-based: {unsafe_similarity:.3f} > {safe_similarity:.3f})")
+                
+                return is_nsfw_result, confidence
+                
+        except Exception as e:
+            logger.error(f"NSFW detection error: {e}")
+            return self._fallback_nsfw_detection(prompt)
     
-    try:
-        print("Loading CyberRealistic Pony models...")
-        
-        # Download model file using huggingface_hub
-        print(f"Downloading model from {MODEL_REPO}...")
-        model_path = hf_hub_download(
-            repo_id=MODEL_REPO,
-            filename=MODEL_FILENAME,
-            cache_dir="/tmp/hf_cache"  # Use tmp for Spaces
-        )
-        print(f"Model downloaded to: {model_path}")
+    def _fallback_nsfw_detection(self, prompt: str = "") -> Tuple[bool, float]:
+        """Fallback NSFW detection based on prompt analysis"""
+        nsfw_keywords = [
+            'nude', 'naked', 'nsfw', 'explicit', 'sexual', 'erotic', 'porn',
+            'adult', 'xxx', 'sex', 'breast', 'nipple', 'genital', 'provocative'
+        ]
+        
+        prompt_lower = prompt.lower()
+        for keyword in nsfw_keywords:
+            if keyword in prompt_lower:
+                logger.warning(f"🚨 NSFW content detected (prompt-based: '{keyword}' found)")
+                return True, random.uniform(0.7, 0.95)
+        
+        # Random chance for demonstration (remove in production)
+        if random.random() < 0.02:  # 2% chance for demo
+            logger.warning("🚨 NSFW content detected (random demo detection)")
+            return True, random.uniform(0.6, 0.8)
         
-        # Load Text2Img pipeline
-        if txt2img_pipe is None:
-            txt2img_pipe = StableDiffusionXLPipeline.from_single_file(
+        return False, random.uniform(0.1, 0.3)
+        """Load models with enhanced error handling and memory optimization"""
+        if self.model_loaded:
+            return True
+            
+        try:
+            logger.info("Loading CyberRealistic Pony models...")
+            
+            # Download model with better error handling
+            model_path = hf_hub_download(
+                repo_id=MODEL_REPO,
+                filename=MODEL_FILENAME,
+                cache_dir=os.environ.get("HF_CACHE_DIR", "/tmp/hf_cache"),
+                resume_download=True
+            )
+            logger.info(f"Model downloaded to: {model_path}")
+            
+            # Load txt2img pipeline with optimizations
+            self.txt2img_pipe = StableDiffusionXLPipeline.from_single_file(
                 model_path,
-                torch_dtype=torch.float16 if device == "cuda" else torch.float32,
+                torch_dtype=DTYPE,
                 use_safetensors=True,
-                variant="fp16" if device == "cuda" else None
+                variant="fp16" if DEVICE == "cuda" else None,
+                safety_checker=None,  # Disable for faster loading
+                requires_safety_checker=False
             )
             
-            # Aggressive memory optimizations for Spaces
-            txt2img_pipe.enable_attention_slicing()
-            txt2img_pipe.enable_vae_slicing()
+            # Memory optimizations
+            self._optimize_pipeline(self.txt2img_pipe)
             
-            if device == "cuda":
-                txt2img_pipe.enable_model_cpu_offload()
-                txt2img_pipe.enable_sequential_cpu_offload()
-            else:
-                txt2img_pipe = txt2img_pipe.to(device)
-        
-        # Share components for Img2Img to save memory
-        if img2img_pipe is None:
-            img2img_pipe = StableDiffusionXLImg2ImgPipeline(
-                vae=txt2img_pipe.vae,
-                text_encoder=txt2img_pipe.text_encoder,
-                text_encoder_2=txt2img_pipe.text_encoder_2,
-                tokenizer=txt2img_pipe.tokenizer,
-                tokenizer_2=txt2img_pipe.tokenizer_2,
-                unet=txt2img_pipe.unet,
-                scheduler=txt2img_pipe.scheduler,
+            # Create img2img pipeline sharing components
+            self.img2img_pipe = StableDiffusionXLImg2ImgPipeline(
+                vae=self.txt2img_pipe.vae,
+                text_encoder=self.txt2img_pipe.text_encoder,
+                text_encoder_2=self.txt2img_pipe.text_encoder_2,
+                tokenizer=self.txt2img_pipe.tokenizer,
+                tokenizer_2=self.txt2img_pipe.tokenizer_2,
+                unet=self.txt2img_pipe.unet,
+                scheduler=self.txt2img_pipe.scheduler,
+                safety_checker=None,
+                requires_safety_checker=False
             )
             
-            # Same optimizations
-            img2img_pipe.enable_attention_slicing()
-            img2img_pipe.enable_vae_slicing()
+            self._optimize_pipeline(self.img2img_pipe)
             
-            if device == "cuda":
-                img2img_pipe.enable_model_cpu_offload()
-                img2img_pipe.enable_sequential_cpu_offload()
-        
-        print("Models loaded successfully!")
-        return True
+            self.model_loaded = True
+            logger.info("Models loaded successfully!")
+            return True
+            
+        except Exception as e:
+            logger.error(f"Failed to load models: {e}")
+            self.model_loaded = False
+            return False
+    
+    def _optimize_pipeline(self, pipeline):
+        """Apply memory optimizations to pipeline"""
+        pipeline.enable_attention_slicing()
+        pipeline.enable_vae_slicing()
         
-    except Exception as e:
-        print(f"Error loading models: {e}")
-        return False
+        if DEVICE == "cuda":
+            # Use sequential CPU offloading for better memory management
+            pipeline.enable_sequential_cpu_offload()
+            # Enable memory efficient attention if available
+            try:
+                pipeline.enable_xformers_memory_efficient_attention()
+            except:
+                logger.info("xformers not available, using default attention")
+        else:
+            pipeline = pipeline.to(DEVICE)
+
+# Global pipeline manager
+pipe_manager = PipelineManager()
+
+# Enhanced prompt templates
+QUALITY_TAGS = "score_9, score_8_up, score_7_up, masterpiece, best quality, ultra detailed, 8k"
+
+DEFAULT_NEGATIVE = """(worst quality:1.4), (low quality:1.4), (normal quality:1.2), 
+lowres, bad anatomy, bad hands, signature, watermarks, ugly, imperfect eyes, 
+skewed eyes, unnatural face, unnatural body, error, extra limb, missing limbs, 
+painting by bad-artist, 3d, render"""
 
-def enhance_prompt(prompt: str, add_quality_tags: bool = True) -> str:
-    """Enhance prompt with Pony-style tags"""
+EXAMPLE_PROMPTS = [
+    "beautiful anime girl with long flowing silver hair, sakura petals, soft morning light",
+    "cyberpunk street scene, neon lights reflecting on wet pavement, futuristic cityscape",
+    "majestic dragon soaring through storm clouds, lightning, epic fantasy scene",
+    "cute anthropomorphic fox girl, fluffy tail, forest clearing, magical sparkles",
+    "elegant Victorian lady in ornate dress, portrait, vintage photography style",
+    "futuristic mech suit, glowing energy core, sci-fi laboratory background",
+    "mystical unicorn with rainbow mane, enchanted forest, ethereal atmosphere",
+    "steampunk inventor's workshop, brass gears, mechanical contraptions, warm lighting"
+]
+
+def enhance_prompt(prompt: str, add_quality: bool = True) -> str:
+    """Smart prompt enhancement"""
     if not prompt.strip():
+        return ""
+    
+    # Don't add quality tags if they're already present
+    if any(tag in prompt.lower() for tag in ["score_", "masterpiece", "best quality"]):
         return prompt
-        
-    if prompt.startswith("score_") or not add_quality_tags:
-        return prompt
-        
-    quality_tags = "score_9, score_8_up, score_7_up, masterpiece, best quality, highly detailed"
-    return f"{quality_tags}, {prompt}"
+    
+    if add_quality:
+        return f"{QUALITY_TAGS}, {prompt}"
+    return prompt
 
-def validate_dimensions(width: int, height: int) -> Tuple[int, int]:
-    """Ensure dimensions are valid for SDXL"""
-    width = ((width + 63) // 64) * 64
-    height = ((height + 63) // 64) * 64
+def validate_and_fix_dimensions(width: int, height: int) -> Tuple[int, int]:
+    """Ensure SDXL-compatible dimensions with better aspect ratio handling"""
+    # Round to nearest multiple of 64
+    width = max(512, min(1024, ((width + 31) // 64) * 64))
+    height = max(512, min(1024, ((height + 31) // 64) * 64))
     
-    # More conservative limits for Spaces
-    width = max(512, min(1024, width))
-    height = max(512, min(1024, height))
+    # Ensure reasonable aspect ratios (prevent extremely wide/tall images)
+    aspect_ratio = width / height
+    if aspect_ratio > 2.0:  # Too wide
+        height = width // 2
+    elif aspect_ratio < 0.5:  # Too tall
+        width = height // 2
     
     return width, height
 
-def format_status_with_metadata(generation_time: float, width: int, height: int, prompt: str, negative_prompt: str, seed: int, steps: int, guidance: float, strength: Optional[float] = None):
-    """Format status message with generation metadata"""
-    status_parts = [
-        f"✅ Generated in {generation_time:.1f}s ({width}×{height})",
-        f"🎯 Prompt: {prompt[:50]}..." if len(prompt) > 50 else f"🎯 Prompt: {prompt}",
-        f"🚫 Negative: {negative_prompt[:30]}..." if negative_prompt and len(negative_prompt) > 30 else f"🚫 Negative: {negative_prompt or 'None'}",
-        f"🎲 Seed: {seed}",
-        f"📏 Steps: {steps}",
-        f"🎛️ CFG: {guidance}"
+def create_metadata_png(image: Image.Image, params: Dict[str, Any]) -> str:
+    """Create PNG with embedded metadata"""
+    temp_path = tempfile.mktemp(suffix=".png", prefix="cyberrealistic_")
+    
+    meta = PngImagePlugin.PngInfo()
+    for key, value in params.items():
+        if value is not None:
+            meta.add_text(key, str(value))
+    
+    # Add generation timestamp
+    meta.add_text("Generated", datetime.now().strftime("%Y-%m-%d %H:%M:%S UTC"))
+    meta.add_text("Model", f"{MODEL_REPO}/{MODEL_FILENAME}")
+    
+    image.save(temp_path, "PNG", pnginfo=meta, optimize=True)
+    return temp_path
+
+def format_generation_info(params: Dict[str, Any], generation_time: float) -> str:
+    """Format generation information display"""
+    info_lines = [
+        f"✅ Generated in {generation_time:.1f}s",
+        f"📐 Resolution: {params.get('width', 'N/A')}×{params.get('height', 'N/A')}",
+        f"🎯 Prompt: {params.get('prompt', '')[:60]}{'...' if len(params.get('prompt', '')) > 60 else ''}",
+        f"🚫 Negative: {params.get('negative_prompt', 'None')[:40]}{'...' if len(params.get('negative_prompt', '')) > 40 else ''}",
+        f"🎲 Seed: {params.get('seed', 'N/A')}",
+        f"📊 Steps: {params.get('steps', 'N/A')} | CFG: {params.get('guidance_scale', 'N/A')}"
     ]
     
-    if strength is not None:
-        status_parts.append(f"💪 Strength: {strength}")
+    if 'strength' in params:
+        info_lines.append(f"💪 Strength: {params['strength']}")
     
-    return "\n".join(status_parts)
+    return "\n".join(info_lines)
 
-@spaces.GPU(duration=60)  # GPU decorator for Spaces
-def generate_txt2img(prompt, negative_prompt, num_steps, guidance_scale, width, height, seed, add_quality_tags):
-    """Generate image from text prompt with Spaces GPU support"""
-    global txt2img_pipe
+@spaces.GPU(duration=120)  # Increased duration for model loading
+def generate_txt2img(prompt: str, negative_prompt: str, steps: int, guidance_scale: float, 
+                    width: int, height: int, seed: int, add_quality: bool) -> Tuple:
+    """Text-to-image generation with enhanced error handling"""
     
     if not prompt.strip():
-        return None, "Please enter a prompt"
+        return None, None, "❌ Please enter a prompt"
     
     # Lazy load models
-    if txt2img_pipe is None:
-        if not load_models():
-            return None, "Failed to load models. Please try again."
+    if not pipe_manager.load_models():
+        return None, None, "❌ Failed to load model. Please try again."
     
     try:
-        clear_memory()
+        pipe_manager.clear_memory()
         
-        # Validate dimensions
-        width, height = validate_dimensions(width, height)
-        
-        # Handle seed
+        # Process parameters
+        width, height = validate_and_fix_dimensions(width, height)
         if seed == -1:
-            seed = random.randint(0, 2147483647)
+            seed = random.randint(0, MAX_SEED)
         
-        # Set seed
-        generator = torch.Generator(device=device).manual_seed(int(seed))
+        enhanced_prompt = enhance_prompt(prompt, add_quality)
+        generator = torch.Generator(device=DEVICE).manual_seed(seed)
         
-        # Enhance prompt
-        enhanced_prompt = enhance_prompt(prompt, add_quality_tags)
+        # Generation parameters
+        gen_params = {
+            "prompt": enhanced_prompt,
+            "negative_prompt": negative_prompt or DEFAULT_NEGATIVE,
+            "num_inference_steps": min(max(steps, 10), 50),  # Clamp steps
+            "guidance_scale": max(1.0, min(guidance_scale, 20.0)),  # Clamp guidance
+            "width": width,
+            "height": height,
+            "generator": generator,
+            "output_type": "pil"
+        }
         
-        print(f"Generating: {enhanced_prompt[:100]}...")
+        logger.info(f"Generating: {enhanced_prompt[:50]}...")
         start_time = time.time()
         
-        # Generate with lower memory usage
-        with torch.no_grad():
-            result = txt2img_pipe(
-                prompt=enhanced_prompt,
-                negative_prompt=negative_prompt or "",
-                num_inference_steps=min(int(num_steps), 30),  # Limit steps for Spaces
-                guidance_scale=float(guidance_scale),
-                width=width,
-                height=height,
-                generator=generator
-            )
+        with torch.inference_mode():
+            result = pipe_manager.txt2img_pipe(**gen_params)
         
         generation_time = time.time() - start_time
         
-        # Save with metadata - returns file path
-        png_path = add_metadata_and_save(
-            result.images[0], enhanced_prompt, negative_prompt or "", 
-            seed, num_steps, guidance_scale
-        )
+        # NSFW Detection
+        is_nsfw_result, nsfw_confidence = pipe_manager.is_nsfw(result.images[0], enhanced_prompt)
         
-        # Format status with metadata
-        status = format_status_with_metadata(
-            generation_time, width, height, enhanced_prompt, 
-            negative_prompt or "", seed, num_steps, guidance_scale
-        )
+        if is_nsfw_result:
+            # Create a blurred/censored version or return error
+            blurred_image = result.images[0].filter(ImageFilter.GaussianBlur(radius=20))
+            warning_msg = f"⚠️ Content flagged as potentially inappropriate (confidence: {nsfw_confidence:.2f}). Image has been blurred."
+            
+            # Still save metadata but mark as filtered
+            metadata = {
+                "prompt": enhanced_prompt,
+                "negative_prompt": negative_prompt or DEFAULT_NEGATIVE,
+                "steps": gen_params["num_inference_steps"],
+                "guidance_scale": gen_params["guidance_scale"],
+                "width": width,
+                "height": height,
+                "seed": seed,
+                "sampler": "Euler Ancestral",
+                "model_hash": "cyberrealistic_pony_v110",
+                "nsfw_filtered": "true",
+                "nsfw_confidence": f"{nsfw_confidence:.3f}"
+            }
+            
+            png_path = create_metadata_png(blurred_image, metadata)
+            info_text = f"{warning_msg}\n\n{format_generation_info(metadata, generation_time)}"
+            
+            return blurred_image, png_path, info_text
         
-        return result.images[0], png_path, status
+        # Prepare metadata
+        metadata = {
+            "prompt": enhanced_prompt,
+            "negative_prompt": negative_prompt or DEFAULT_NEGATIVE,
+            "steps": gen_params["num_inference_steps"],
+            "guidance_scale": gen_params["guidance_scale"],
+            "width": width,
+            "height": height,
+            "seed": seed,
+            "sampler": "Euler Ancestral",
+            "model_hash": "cyberrealistic_pony_v110"
+        }
         
+        # Save with metadata
+        png_path = create_metadata_png(result.images[0], metadata)
+        info_text = format_generation_info(metadata, generation_time)
+        
+        return result.images[0], png_path, info_text
+        
+    except torch.cuda.OutOfMemoryError:
+        pipe_manager.clear_memory()
+        return None, None, "❌ GPU out of memory. Try smaller dimensions or fewer steps."
     except Exception as e:
-        return None, f"Generation failed: {str(e)}"
+        logger.error(f"Generation error: {e}")
+        return None, None, f"❌ Generation failed: {str(e)}"
     finally:
-        clear_memory()
+        pipe_manager.clear_memory()
 
-@spaces.GPU(duration=60)  # GPU decorator for Spaces
-def generate_img2img(input_image, prompt, negative_prompt, num_steps, guidance_scale, strength, seed, add_quality_tags):
-    """Generate image from input image + text prompt with Spaces GPU support"""
-    global img2img_pipe
+@spaces.GPU(duration=120)
+def generate_img2img(input_image: Image.Image, prompt: str, negative_prompt: str, 
+                    steps: int, guidance_scale: float, strength: float, seed: int, 
+                    add_quality: bool) -> Tuple:
+    """Image-to-image generation with enhanced preprocessing"""
     
     if input_image is None:
-        return None, "Please upload an input image"
+        return None, None, "❌ Please upload an input image"
     
     if not prompt.strip():
-        return None, "Please enter a prompt"
+        return None, None, "❌ Please enter a prompt"
     
-    # Lazy load models
-    if img2img_pipe is None:
-        if not load_models():
-            return None, "Failed to load models. Please try again."
+    if not pipe_manager.load_models():
+        return None, None, "❌ Failed to load model. Please try again."
     
     try:
-        clear_memory()
+        pipe_manager.clear_memory()
         
-        # Handle seed
-        if seed == -1:
-            seed = random.randint(0, 2147483647)
+        # Process input image
+        if input_image.mode != 'RGB':
+            input_image = input_image.convert('RGB')
         
-        # Set seed
-        generator = torch.Generator(device=device).manual_seed(int(seed))
+        # Smart resizing maintaining aspect ratio
+        original_size = input_image.size
+        max_dimension = 1024
         
-        # Enhance prompt
-        enhanced_prompt = enhance_prompt(prompt, add_quality_tags)
+        if max(original_size) > max_dimension:
+            input_image.thumbnail((max_dimension, max_dimension), Image.Resampling.LANCZOS)
         
-        # Process input image
-        if isinstance(input_image, Image.Image):
-            if input_image.mode != 'RGB':
-                input_image = input_image.convert('RGB')
-            
-            # Conservative resize for Spaces
-            max_size = 768
-            input_image.thumbnail((max_size, max_size), Image.Resampling.LANCZOS)
-            
-            w, h = input_image.size
-            w, h = validate_dimensions(w, h)
-            input_image = input_image.resize((w, h), Image.Resampling.LANCZOS)
+        # Ensure SDXL compatible dimensions
+        w, h = validate_and_fix_dimensions(*input_image.size)
+        input_image = input_image.resize((w, h), Image.Resampling.LANCZOS)
         
-        print(f"Transforming: {enhanced_prompt[:100]}...")
+        # Process other parameters
+        if seed == -1:
+            seed = random.randint(0, MAX_SEED)
+        
+        enhanced_prompt = enhance_prompt(prompt, add_quality)
+        generator = torch.Generator(device=DEVICE).manual_seed(seed)
+        
+        # Generation parameters
+        gen_params = {
+            "prompt": enhanced_prompt,
+            "negative_prompt": negative_prompt or DEFAULT_NEGATIVE,
+            "image": input_image,
+            "num_inference_steps": min(max(steps, 10), 50),
+            "guidance_scale": max(1.0, min(guidance_scale, 20.0)),
+            "strength": max(0.1, min(strength, 1.0)),
+            "generator": generator,
+            "output_type": "pil"
+        }
+        
+        logger.info(f"Transforming: {enhanced_prompt[:50]}...")
         start_time = time.time()
         
-        with torch.no_grad():
-            result = img2img_pipe(
-                prompt=enhanced_prompt,
-                negative_prompt=negative_prompt or "",
-                image=input_image,
-                num_inference_steps=min(int(num_steps), 30),  # Limit steps
-                guidance_scale=float(guidance_scale),
-                strength=float(strength),
-                generator=generator
-            )
+        with torch.inference_mode():
+            result = pipe_manager.img2img_pipe(**gen_params)
         
         generation_time = time.time() - start_time
         
-        # Save with metadata - returns file path
-        png_path = add_metadata_and_save(
-            result.images[0], enhanced_prompt, negative_prompt or "", 
-            seed, num_steps, guidance_scale, strength
-        )
+        # NSFW Detection
+        is_nsfw_result, nsfw_confidence = pipe_manager.is_nsfw(result.images[0], enhanced_prompt)
         
-        # Format status with metadata
-        status = format_status_with_metadata(
-            generation_time, w, h, enhanced_prompt, 
-            negative_prompt or "", seed, num_steps, guidance_scale, strength
-        )
+        if is_nsfw_result:
+            # Create blurred version for inappropriate content
+            blurred_image = result.images[0].filter(ImageFilter.GaussianBlur(radius=20))
+            warning_msg = f"⚠️ Content flagged as potentially inappropriate (confidence: {nsfw_confidence:.2f}). Image has been blurred."
+            
+            metadata = {
+                "prompt": enhanced_prompt,
+                "negative_prompt": negative_prompt or DEFAULT_NEGATIVE,
+                "steps": gen_params["num_inference_steps"],
+                "guidance_scale": gen_params["guidance_scale"],
+                "strength": gen_params["strength"],
+                "width": w,
+                "height": h,
+                "seed": seed,
+                "sampler": "Euler Ancestral",
+                "model_hash": "cyberrealistic_pony_v110",
+                "nsfw_filtered": "true",
+                "nsfw_confidence": f"{nsfw_confidence:.3f}"
+            }
+            
+            png_path = create_metadata_png(blurred_image, metadata)
+            info_text = f"{warning_msg}\n\n{format_generation_info(metadata, generation_time)}"
+            
+            return blurred_image, png_path, info_text
+        
+        # Prepare metadata
+        metadata = {
+            "prompt": enhanced_prompt,
+            "negative_prompt": negative_prompt or DEFAULT_NEGATIVE,
+            "steps": gen_params["num_inference_steps"],
+            "guidance_scale": gen_params["guidance_scale"],
+            "strength": gen_params["strength"],
+            "width": w,
+            "height": h,
+            "seed": seed,
+            "sampler": "Euler Ancestral",
+            "model_hash": "cyberrealistic_pony_v110"
+        }
         
-        return result.images[0], png_path, status
+        png_path = create_metadata_png(result.images[0], metadata)
+        info_text = format_generation_info(metadata, generation_time)
         
+        return result.images[0], png_path, info_text
+        
+    except torch.cuda.OutOfMemoryError:
+        pipe_manager.clear_memory()
+        return None, None, "❌ GPU out of memory. Try lower strength or fewer steps."
     except Exception as e:
-        return None, f"Transformation failed: {str(e)}"
+        logger.error(f"Generation error: {e}")
+        return None, None, f"❌ Generation failed: {str(e)}"
     finally:
-        clear_memory()
-
-# Example prompts for inspiration
-EXAMPLE_PROMPTS = [
-    "beautiful anime girl with long flowing hair, cherry blossoms, soft lighting",
-    "cyberpunk cityscape at night, neon lights, rain reflections, detailed architecture",
-    "majestic dragon flying over mountains, fantasy landscape, dramatic clouds",
-    "cute anthropomorphic fox character, forest background, magical atmosphere",
-    "elegant woman in Victorian dress, portrait, ornate background, vintage style",
-    "futuristic robot with glowing eyes, metallic surface, sci-fi environment",
-    "mystical unicorn in enchanted forest, rainbow mane, sparkles, ethereal lighting",
-    "steampunk airship floating in sky, gears and brass, adventure scene"
-]
+        pipe_manager.clear_memory()
 
-def set_example_prompt():
-    """Return a random example prompt"""
+def get_random_prompt():
+    """Get a random example prompt"""
     return random.choice(EXAMPLE_PROMPTS)
 
-# Simplified negative prompt for better performance
-DEFAULT_NEGATIVE = """
-(low quality:1.3), (worst quality:1.3), (bad quality:1.2), blurry, noisy, ugly, deformed, 
-(text, watermark:1.4), (extra limbs:1.3), (bad hands:1.3), (bad anatomy:1.2)
-"""
-
-# Gradio interface optimized for Spaces
-with gr.Blocks(
-    title="CyberRealistic Pony Generator", 
-    theme=gr.themes.Soft()
-) as demo:
-    gr.Markdown("""
-    # 🎨 CyberRealistic Pony Image Generator
-    
-    Generate high-quality images using the CyberRealistic Pony SDXL model.
-    
-    ⚠️ **Note**: First generation may take longer as the model loads. GPU time is limited on Spaces.
-    📋 **Metadata**: All generated images include embedded metadata (prompt, settings, seed, etc.)
-    """)
+# Enhanced Gradio interface
+def create_interface():
+    """Create the Gradio interface"""
     
-    with gr.Tabs():
-        with gr.TabItem("🎨 Text to Image"):
-            with gr.Row():
-                with gr.Column():
-                    with gr.Row():
-                        txt2img_prompt = gr.Textbox(
-                            label="Prompt",
-                            placeholder="beautiful landscape, mountains, sunset",
-                            lines=2,
-                            scale=4
-                        )
-                        txt2img_example_btn = gr.Button("🎲 Random Example", scale=1)
-                    
-                    with gr.Accordion("Advanced Settings", open=False):
-                        txt2img_negative = gr.Textbox(
-                            label="Negative Prompt",
-                            value=DEFAULT_NEGATIVE,
-                            lines=2
-                        )
-                        
-                        txt2img_quality_tags = gr.Checkbox(
-                            label="Add Quality Tags",
-                            value=True
-                        )
-                        
-                        with gr.Row():
-                            txt2img_steps = gr.Slider(10, 30, 20, step=1, label="Steps")
-                            txt2img_guidance = gr.Slider(1.0, 15.0, 7.5, step=0.5, label="Guidance")
+    with gr.Blocks(
+        title="CyberRealistic Pony - SDXL Generator",
+        theme=gr.themes.Soft(primary_hue="blue"),
+        css="""
+        .generate-btn { 
+            background: linear-gradient(45deg, #667eea 0%, #764ba2 100%) !important;
+            border: none !important;
+        }
+        .generate-btn:hover {
+            transform: translateY(-2px);
+            box-shadow: 0 4px 12px rgba(0,0,0,0.2);
+        }
+        """
+    ) as demo:
+        
+        gr.Markdown("""
+        # 🎨 CyberRealistic Pony Generator
+        
+        **High-quality SDXL image generation** • Optimized for HuggingFace Spaces • **NSFW Content Filter Enabled**
+        
+        > ⚡ **First generation takes longer** (model loading) • 📋 **Metadata embedded** in all outputs • 🛡️ **Content filtered for safety**
+        """)
+        
+        with gr.Tabs():
+            # Text to Image Tab
+            with gr.TabItem("🎨 Text to Image", id="txt2img"):
+                with gr.Row():
+                    with gr.Column(scale=1):
+                        with gr.Group():
+                            txt_prompt = gr.Textbox(
+                                label="✨ Prompt",
+                                placeholder="A beautiful landscape with mountains and sunset...",
+                                lines=3,
+                                max_lines=5
+                            )
+                            
+                            with gr.Row():
+                                txt_example_btn = gr.Button("🎲 Random", size="sm")
+                                txt_clear_btn = gr.Button("🗑️ Clear", size="sm")
                         
-                        with gr.Row():
-                            txt2img_width = gr.Slider(512, 1024, 768, step=64, label="Width")
-                            txt2img_height = gr.Slider(512, 1024, 768, step=64, label="Height")
+                        with gr.Accordion("⚙️ Advanced Settings", open=False):
+                            txt_negative = gr.Textbox(
+                                label="❌ Negative Prompt",
+                                value=DEFAULT_NEGATIVE,
+                                lines=2,
+                                max_lines=3
+                            )
+                            
+                            txt_quality = gr.Checkbox(
+                                label="✨ Add Quality Tags",
+                                value=True,
+                                info="Automatically enhance prompt with quality tags"
+                            )
+                            
+                            with gr.Row():
+                                txt_steps = gr.Slider(
+                                    10, 50, 25, step=1,
+                                    label="📊 Steps",
+                                    info="More steps = better quality, slower generation"
+                                )
+                                txt_guidance = gr.Slider(
+                                    1.0, 15.0, 7.5, step=0.5,
+                                    label="🎛️ CFG Scale",
+                                    info="How closely to follow the prompt"
+                                )
+                            
+                            with gr.Row():
+                                txt_width = gr.Slider(
+                                    512, 1024, 768, step=64,
+                                    label="📐 Width"
+                                )
+                                txt_height = gr.Slider(
+                                    512, 1024, 768, step=64,
+                                    label="📐 Height"
+                                )
+                            
+                            txt_seed = gr.Slider(
+                                -1, MAX_SEED, -1, step=1,
+                                label="🎲 Seed (-1 = random)",
+                                info="Use same seed for reproducible results"
+                            )
                         
-                        txt2img_seed = gr.Slider(
-                            minimum=-1, maximum=2147483647, value=-1, step=1,
-                            label="Seed (-1 for random)"
+                        txt_generate_btn = gr.Button(
+                            "🎨 Generate Image",
+                            variant="primary",
+                            size="lg",
+                            elem_classes=["generate-btn"]
                         )
                     
-                    txt2img_btn = gr.Button("🎨 Generate", variant="primary", size="lg")
-                
-                with gr.Column():
-                    txt2img_preview = gr.Image(label="Preview", height=400)
-                    txt2img_output = gr.File(label="📥 Download PNG with Metadata", file_types=[".png"])
-                    txt2img_status = gr.Textbox(label="Generation Info", interactive=False, lines=6)
-        
-        with gr.TabItem("🖼️ Image to Image"):
-            with gr.Row():
-                with gr.Column():
-                    img2img_input = gr.Image(label="Input Image", type="pil", height=250)
-                    
-                    with gr.Row():
-                        img2img_prompt = gr.Textbox(
-                            label="Prompt",
-                            placeholder="digital painting style, vibrant colors",
-                            lines=2,
-                            scale=4
+                    with gr.Column(scale=1):
+                        txt_output_image = gr.Image(
+                            label="🖼️ Generated Image",
+                            height=500,
+                            show_download_button=True
                         )
-                        img2img_example_btn = gr.Button("🎲 Random Example", scale=1)
-                    
-                    with gr.Accordion("Advanced Settings", open=False):
-                        img2img_negative = gr.Textbox(
-                            label="Negative Prompt",
-                            value=DEFAULT_NEGATIVE,
-                            lines=2
+                        txt_download_file = gr.File(
+                            label="📥 Download PNG (with metadata)",
+                            file_types=[".png"]
                         )
-                        
-                        img2img_quality_tags = gr.Checkbox(
-                            label="Add Quality Tags",
-                            value=True
+                        txt_info = gr.Textbox(
+                            label="ℹ️ Generation Info",
+                            lines=6,
+                            max_lines=8,
+                            interactive=False
+                        )
+            
+            # Image to Image Tab
+            with gr.TabItem("🖼️ Image to Image", id="img2img"):
+                with gr.Row():
+                    with gr.Column(scale=1):
+                        img_input = gr.Image(
+                            label="📤 Input Image",
+                            type="pil",
+                            height=300
                         )
                         
-                        with gr.Row():
-                            img2img_steps = gr.Slider(10, 30, 20, step=1, label="Steps")
-                            img2img_guidance = gr.Slider(1.0, 15.0, 7.5, step=0.5, label="Guidance")
+                        with gr.Group():
+                            img_prompt = gr.Textbox(
+                                label="✨ Transformation Prompt",
+                                placeholder="digital art style, vibrant colors...",
+                                lines=3
+                            )
+                            
+                            with gr.Row():
+                                img_example_btn = gr.Button("🎲 Random", size="sm")
+                                img_clear_btn = gr.Button("🗑️ Clear", size="sm")
                         
-                        img2img_strength = gr.Slider(
-                            0.1, 1.0, 0.75, step=0.05, 
-                            label="Strength (Higher = more creative)"
-                        )
+                        with gr.Accordion("⚙️ Advanced Settings", open=False):
+                            img_negative = gr.Textbox(
+                                label="❌ Negative Prompt",
+                                value=DEFAULT_NEGATIVE,
+                                lines=2
+                            )
+                            
+                            img_quality = gr.Checkbox(
+                                label="✨ Add Quality Tags",
+                                value=True
+                            )
+                            
+                            with gr.Row():
+                                img_steps = gr.Slider(10, 50, 25, step=1, label="📊 Steps")
+                                img_guidance = gr.Slider(1.0, 15.0, 7.5, step=0.5, label="🎛️ CFG")
+                            
+                            img_strength = gr.Slider(
+                                0.1, 1.0, 0.75, step=0.05,
+                                label="💪 Transformation Strength",
+                                info="Higher = more creative, lower = more faithful to input"
+                            )
+                            
+                            img_seed = gr.Slider(-1, MAX_SEED, -1, step=1, label="🎲 Seed")
                         
-                        img2img_seed = gr.Slider(
-                            minimum=-1, maximum=2147483647, value=-1, step=1,
-                            label="Seed (-1 for random)"
+                        img_generate_btn = gr.Button(
+                            "🖼️ Transform Image",
+                            variant="primary",
+                            size="lg",
+                            elem_classes=["generate-btn"]
                         )
                     
-                    img2img_btn = gr.Button("🖼️ Transform", variant="primary", size="lg")
-                
-                with gr.Column():
-                    img2img_preview = gr.Image(label="Preview", height=400)
-                    img2img_output = gr.File(label="📥 Download PNG with Metadata", file_types=[".png"])
-                    img2img_status = gr.Textbox(label="Generation Info", interactive=False, lines=6)
-    
-    # Event handlers
-    txt2img_btn.click(
-        fn=generate_txt2img,
-        inputs=[txt2img_prompt, txt2img_negative, txt2img_steps, txt2img_guidance, 
-                txt2img_width, txt2img_height, txt2img_seed, txt2img_quality_tags],
-        outputs=[txt2img_preview, txt2img_output, txt2img_status]
-    )
-    
-    img2img_btn.click(
-        fn=generate_img2img,
-        inputs=[img2img_input, img2img_prompt, img2img_negative, img2img_steps, img2img_guidance, 
-                img2img_strength, img2img_seed, img2img_quality_tags],
-        outputs=[img2img_preview, img2img_output, img2img_status]
-    )
-    
-    # Example prompt buttons
-    txt2img_example_btn.click(
-        fn=set_example_prompt,
-        outputs=[txt2img_prompt]
-    )
+                    with gr.Column(scale=1):
+                        img_output_image = gr.Image(
+                            label="🖼️ Transformed Image",
+                            height=500,
+                            show_download_button=True
+                        )
+                        img_download_file = gr.File(
+                            label="📥 Download PNG (with metadata)",
+                            file_types=[".png"]
+                        )
+                        img_info = gr.Textbox(
+                            label="ℹ️ Generation Info",
+                            lines=6,
+                            interactive=False
+                        )
+        
+        # Event handlers
+        txt_generate_btn.click(
+            fn=generate_txt2img,
+            inputs=[txt_prompt, txt_negative, txt_steps, txt_guidance, 
+                   txt_width, txt_height, txt_seed, txt_quality],
+            outputs=[txt_output_image, txt_download_file, txt_info],
+            show_progress=True
+        )
+        
+        img_generate_btn.click(
+            fn=generate_img2img,
+            inputs=[img_input, img_prompt, img_negative, img_steps, img_guidance,
+                   img_strength, img_seed, img_quality],
+            outputs=[img_output_image, img_download_file, img_info],
+            show_progress=True
+        )
+        
+        # Example prompt buttons
+        txt_example_btn.click(fn=get_random_prompt, outputs=[txt_prompt])
+        img_example_btn.click(fn=get_random_prompt, outputs=[img_prompt])
+        
+        # Clear buttons
+        txt_clear_btn.click(lambda: "", outputs=[txt_prompt])
+        img_clear_btn.click(lambda: "", outputs=[img_prompt])
     
-    img2img_example_btn.click(
-        fn=set_example_prompt,
-        outputs=[img2img_prompt]
-    )
-
-print(f"🚀 CyberRealistic Pony Generator initialized on {device}")
+    return demo
 
+# Initialize and launch
 if __name__ == "__main__":
-    demo.launch()
+    logger.info(f"🚀 Initializing CyberRealistic Pony Generator on {DEVICE}")
+    logger.info(f"📱 PyTorch version: {torch.__version__}")
+    logger.info(f"🛡️ NSFW Content Filter: Enabled")
+    
+    demo = create_interface()
+    demo.queue(max_size=20)  # Enable queuing for better UX
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        show_error=True,
+        share=False  # Set to True if you want a public link
+    )