Update app.py

app.py

@@ -39,14 +39,14 @@ class PipelineManager:
         self.nsfw_detector_loaded = False
         
     def clear_memory(self):
-        """Aggressive memory cleanup"""
+        """Aggressive memory cleanup to free up GPU/CPU memory."""
         if torch.cuda.is_available():
             torch.cuda.empty_cache()
             torch.cuda.synchronize()
         gc.collect()
     
     def load_nsfw_detector(self) -> bool:
-        """Load NSFW detection model"""
+        """Load NSFW detection model (CLIP) with error handling."""
         if self.nsfw_detector_loaded:
             return True
             
@@ -70,12 +70,13 @@ class PipelineManager:
     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.
+        Falls back to prompt-based detection if CLIP model fails or is not loaded.
         """
         try:
             # Load NSFW detector if not already loaded
             if not self.nsfw_detector_loaded:
                 if not self.load_nsfw_detector():
+                    # If NSFW detector cannot be loaded, fall back to prompt-based
                     return self._fallback_nsfw_detection(prompt)
             
             # CLIP-based NSFW detection
@@ -132,11 +133,12 @@ class PipelineManager:
                 return is_nsfw_result, confidence
                 
         except Exception as e:
-            logger.error(f"NSFW detection error: {e}")
+            logger.error(f"NSFW detection error (CLIP model failed): {e}")
+            # Fallback to prompt-based detection if CLIP model encounters an error
             return self._fallback_nsfw_detection(prompt)
     
     def _fallback_nsfw_detection(self, prompt: str = "") -> Tuple[bool, float]:
-        """Fallback NSFW detection based on prompt analysis"""
+        """Fallback NSFW detection based on prompt keyword analysis."""
         nsfw_keywords = [
             'nude', 'naked', 'nsfw', 'explicit', 'sexual', 'erotic', 'porn',
             'adult', 'xxx', 'sex', 'breast', 'nipple', 'genital', 'provocative'
@@ -148,13 +150,15 @@ class PipelineManager:
                 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)
+        # Random chance for demonstration (consider removing 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)
         
         return False, random.uniform(0.1, 0.3)
-        """Load models with enhanced error handling and memory optimization"""
+    
+    def load_models(self) -> bool:
+        """Load Stable Diffusion XL models (txt2img and img2img) with enhanced error handling and memory optimization."""
         if self.model_loaded:
             return True
             
@@ -176,14 +180,14 @@ class PipelineManager:
                 torch_dtype=DTYPE,
                 use_safetensors=True,
                 variant="fp16" if DEVICE == "cuda" else None,
-                safety_checker=None,  # Disable for faster loading
+                safety_checker=None,  # Disable for faster loading, using custom NSFW check
                 requires_safety_checker=False
             )
             
-            # Memory optimizations
+            # Apply memory optimizations to txt2img pipeline
             self._optimize_pipeline(self.txt2img_pipe)
             
-            # Create img2img pipeline sharing components
+            # Create img2img pipeline sharing components to save memory
             self.img2img_pipe = StableDiffusionXLImg2ImgPipeline(
                 vae=self.txt2img_pipe.vae,
                 text_encoder=self.txt2img_pipe.text_encoder,
@@ -192,10 +196,11 @@ class PipelineManager:
                 tokenizer_2=self.txt2img_pipe.tokenizer_2,
                 unet=self.txt2img_pipe.unet,
                 scheduler=self.txt2img_pipe.scheduler,
-                safety_checker=None,
+                safety_checker=None, # Disable for faster loading, using custom NSFW check
                 requires_safety_checker=False
             )
             
+            # Apply memory optimizations to img2img pipeline
             self._optimize_pipeline(self.img2img_pipe)
             
             self.model_loaded = True
@@ -208,22 +213,23 @@ class PipelineManager:
             return False
     
     def _optimize_pipeline(self, pipeline):
-        """Apply memory optimizations to pipeline"""
+        """Apply memory optimizations to a given diffusion pipeline."""
         pipeline.enable_attention_slicing()
         pipeline.enable_vae_slicing()
         
         if DEVICE == "cuda":
-            # Use sequential CPU offloading for better memory management
+            # Use sequential CPU offloading for better memory management on GPU
             pipeline.enable_sequential_cpu_offload()
-            # Enable memory efficient attention if available
+            # Enable memory efficient attention if xformers is available
             try:
                 pipeline.enable_xformers_memory_efficient_attention()
-            except:
+            except Exception: # Catch any error if xformers is not installed/configured
                 logger.info("xformers not available, using default attention")
         else:
+            # Move pipeline to CPU if CUDA is not available
             pipeline = pipeline.to(DEVICE)
 
-# Global pipeline manager
+# Global pipeline manager instance
 pipe_manager = PipelineManager()
 
 # Enhanced prompt templates
@@ -246,11 +252,13 @@ EXAMPLE_PROMPTS = [
 ]
 
 def enhance_prompt(prompt: str, add_quality: bool = True) -> str:
-    """Smart prompt enhancement"""
+    """
+    Enhances the given prompt with quality tags unless they are already present.
+    """
     if not prompt.strip():
         return ""
     
-    # Don't add quality tags if they're already present
+    # Don't add quality tags if they're already present in the prompt (case-insensitive)
     if any(tag in prompt.lower() for tag in ["score_", "masterpiece", "best quality"]):
         return prompt
     
@@ -259,22 +267,28 @@ def enhance_prompt(prompt: str, add_quality: bool = True) -> str:
     return prompt
 
 def validate_and_fix_dimensions(width: int, height: int) -> Tuple[int, int]:
-    """Ensure SDXL-compatible dimensions with better aspect ratio handling"""
+    """
+    Ensures SDXL-compatible dimensions (multiples of 64) and reasonable aspect ratios.
+    Clamps dimensions between 512 and 1024.
+    """
     # Round to nearest multiple of 64
     width = max(512, min(1024, ((width + 31) // 64) * 64))
     height = max(512, min(1024, ((height + 31) // 64) * 64))
     
     # Ensure reasonable aspect ratios (prevent extremely wide/tall images)
     aspect_ratio = width / height
-    if aspect_ratio > 2.0:  # Too wide
+    if aspect_ratio > 2.0:  # Too wide, adjust height
         height = width // 2
-    elif aspect_ratio < 0.5:  # Too tall
+    elif aspect_ratio < 0.5:  # Too tall, adjust width
         width = height // 2
     
     return width, height
 
 def create_metadata_png(image: Image.Image, params: Dict[str, Any]) -> str:
-    """Create PNG with embedded metadata"""
+    """
+    Creates a temporary PNG file with embedded metadata from the generation parameters.
+    Returns the path to the created PNG file.
+    """
     temp_path = tempfile.mktemp(suffix=".png", prefix="cyberrealistic_")
     
     meta = PngImagePlugin.PngInfo()
@@ -282,7 +296,7 @@ def create_metadata_png(image: Image.Image, params: Dict[str, Any]) -> str:
         if value is not None:
             meta.add_text(key, str(value))
     
-    # Add generation timestamp
+    # Add generation timestamp and model info
     meta.add_text("Generated", datetime.now().strftime("%Y-%m-%d %H:%M:%S UTC"))
     meta.add_text("Model", f"{MODEL_REPO}/{MODEL_FILENAME}")
     
@@ -290,7 +304,9 @@ def create_metadata_png(image: Image.Image, params: Dict[str, Any]) -> str:
     return temp_path
 
 def format_generation_info(params: Dict[str, Any], generation_time: float) -> str:
-    """Format generation information display"""
+    """
+    Formats the generation information into a human-readable string for display.
+    """
     info_lines = [
         f"✅ Generated in {generation_time:.1f}s",
         f"📐 Resolution: {params.get('width', 'N/A')}×{params.get('height', 'N/A')}",
@@ -305,20 +321,23 @@ def format_generation_info(params: Dict[str, Any], generation_time: float) -> st
     
     return "\n".join(info_lines)
 
-@spaces.GPU(duration=120)  # Increased duration for model loading
+@spaces.GPU(duration=120)  # Increased duration for model loading and generation
 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"""
+    """
+    Handles text-to-image generation, including parameter processing, model inference,
+    NSFW detection, and metadata creation.
+    """
     
     if not prompt.strip():
-        return None, None, "❌ Please enter a prompt"
+        return None, None, "❌ Please enter a prompt."
     
-    # Lazy load models
+    # Lazy load models if not already loaded
     if not pipe_manager.load_models():
         return None, None, "❌ Failed to load model. Please try again."
     
     try:
-        pipe_manager.clear_memory()
+        pipe_manager.clear_memory() # Clear memory before generation
         
         # Process parameters
         width, height = validate_and_fix_dimensions(width, height)
@@ -328,12 +347,12 @@ def generate_txt2img(prompt: str, negative_prompt: str, steps: int, guidance_sca
         enhanced_prompt = enhance_prompt(prompt, add_quality)
         generator = torch.Generator(device=DEVICE).manual_seed(seed)
         
-        # Generation parameters
+        # Generation parameters dictionary
         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
+            "num_inference_steps": min(max(steps, 10), 50),  # Clamp steps to a reasonable range
+            "guidance_scale": max(1.0, min(guidance_scale, 20.0)),  # Clamp guidance scale
             "width": width,
             "height": height,
             "generator": generator,
@@ -348,11 +367,11 @@ def generate_txt2img(prompt: str, negative_prompt: str, steps: int, guidance_sca
         
         generation_time = time.time() - start_time
         
-        # NSFW Detection
+        # Perform NSFW Detection on the generated image
         is_nsfw_result, nsfw_confidence = pipe_manager.is_nsfw(result.images[0], enhanced_prompt)
         
         if is_nsfw_result:
-            # Create a blurred/censored version or return error
+            # If NSFW, blur the image and return a warning message
             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."
             
@@ -376,7 +395,7 @@ def generate_txt2img(prompt: str, negative_prompt: str, steps: int, guidance_sca
             
             return blurred_image, png_path, info_text
         
-        # Prepare metadata
+        # If not NSFW, prepare metadata and save the original image
         metadata = {
             "prompt": enhanced_prompt,
             "negative_prompt": negative_prompt or DEFAULT_NEGATIVE,
@@ -402,38 +421,42 @@ def generate_txt2img(prompt: str, negative_prompt: str, steps: int, guidance_sca
         logger.error(f"Generation error: {e}")
         return None, None, f"❌ Generation failed: {str(e)}"
     finally:
-        pipe_manager.clear_memory()
+        pipe_manager.clear_memory() # Ensure memory is cleared even if an error occurs
 
-@spaces.GPU(duration=120)
+@spaces.GPU(duration=120) # Increased duration for model loading and generation
 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"""
+    """
+    Handles image-to-image generation, including image preprocessing, parameter processing,
+    model inference, NSFW detection, and metadata creation.
+    """
     
     if input_image is None:
-        return None, None, "❌ Please upload an input image"
+        return None, None, "❌ Please upload an input image."
     
     if not prompt.strip():
-        return None, None, "❌ Please enter a prompt"
+        return None, None, "❌ Please enter a prompt."
     
+    # Lazy load models if not already loaded
     if not pipe_manager.load_models():
         return None, None, "❌ Failed to load model. Please try again."
     
     try:
-        pipe_manager.clear_memory()
+        pipe_manager.clear_memory() # Clear memory before generation
         
-        # Process input image
+        # Process input image: convert to RGB if necessary
         if input_image.mode != 'RGB':
             input_image = input_image.convert('RGB')
         
-        # Smart resizing maintaining aspect ratio
+        # Smart resizing maintaining aspect ratio to fit within max_dimension
         original_size = input_image.size
         max_dimension = 1024
         
         if max(original_size) > max_dimension:
             input_image.thumbnail((max_dimension, max_dimension), Image.Resampling.LANCZOS)
         
-        # Ensure SDXL compatible dimensions
+        # Ensure SDXL compatible dimensions (multiples of 64)
         w, h = validate_and_fix_dimensions(*input_image.size)
         input_image = input_image.resize((w, h), Image.Resampling.LANCZOS)
         
@@ -444,14 +467,14 @@ def generate_img2img(input_image: Image.Image, prompt: str, negative_prompt: str
         enhanced_prompt = enhance_prompt(prompt, add_quality)
         generator = torch.Generator(device=DEVICE).manual_seed(seed)
         
-        # Generation parameters
+        # Generation parameters dictionary
         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)),
+            "num_inference_steps": min(max(steps, 10), 50), # Clamp steps
+            "guidance_scale": max(1.0, min(guidance_scale, 20.0)), # Clamp guidance scale
+            "strength": max(0.1, min(strength, 1.0)), # Clamp strength
             "generator": generator,
             "output_type": "pil"
         }
@@ -464,11 +487,11 @@ def generate_img2img(input_image: Image.Image, prompt: str, negative_prompt: str
         
         generation_time = time.time() - start_time
         
-        # NSFW Detection
+        # Perform NSFW Detection on the transformed image
         is_nsfw_result, nsfw_confidence = pipe_manager.is_nsfw(result.images[0], enhanced_prompt)
         
         if is_nsfw_result:
-            # Create blurred version for inappropriate content
+            # If NSFW, blur the image and return a warning message
             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."
             
@@ -492,7 +515,7 @@ def generate_img2img(input_image: Image.Image, prompt: str, negative_prompt: str
             
             return blurred_image, png_path, info_text
         
-        # Prepare metadata
+        # If not NSFW, prepare metadata and save the original image
         metadata = {
             "prompt": enhanced_prompt,
             "negative_prompt": negative_prompt or DEFAULT_NEGATIVE,
@@ -518,15 +541,18 @@ def generate_img2img(input_image: Image.Image, prompt: str, negative_prompt: str
         logger.error(f"Generation error: {e}")
         return None, None, f"❌ Generation failed: {str(e)}"
     finally:
-        pipe_manager.clear_memory()
+        pipe_manager.clear_memory() # Ensure memory is cleared even if an error occurs
 
 def get_random_prompt():
-    """Get a random example prompt"""
+    """Returns a random example prompt from a predefined list."""
     return random.choice(EXAMPLE_PROMPTS)
 
 # Enhanced Gradio interface
 def create_interface():
-    """Create the Gradio interface"""
+    """
+    Creates and returns the Gradio Blocks interface for the CyberRealistic Pony Generator.
+    This includes tabs for Text-to-Image and Image-to-Image, along with controls and outputs.
+    """
     
     with gr.Blocks(
         title="CyberRealistic Pony - SDXL Generator",
@@ -729,42 +755,17 @@ def create_interface():
     
     return demo
 
-# Initialize and launch
+# Initialize and launch the Gradio application
 if __name__ == "__main__":
     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.queue(max_size=20)  # Enable queuing for better user experience
     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
+        share=False  # Set to True if you want a public link (e.g., for Hugging Face Spaces)
     )
-        
-        # 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])
-    
-    return demo
-
-# Initialize and launch
-if __name__ == "__main__":
-    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
-    )