Update app.py

app.py

@@ -6,11 +6,12 @@ from diffusers import (
     StableDiffusionXLControlNetPipeline,
     ControlNetModel,
     AutoencoderKL,
-    DPMSolverMultistepScheduler
+    DPMSolverMultistepScheduler,
+    EulerAncestralDiscreteScheduler
 )
 from diffusers.models.attention_processor import AttnProcessor2_0
 from insightface.app import FaceAnalysis
-from PIL import Image
+from PIL import Image, ImageEnhance, ImageFilter
 import numpy as np
 import cv2
 from transformers import pipeline as transformers_pipeline
@@ -29,6 +30,12 @@ class RetroArtConverter:
     def __init__(self):
         self.device = device
         self.dtype = dtype
+        self.models_loaded = {
+            'custom_checkpoint': False,
+            'custom_vae': False,
+            'lora': False,
+            'instantid': False
+        }
         
         # Initialize face analysis for InstantID
         print("Loading face analysis model...")
@@ -43,7 +50,7 @@ class RetroArtConverter:
             self.face_detection_enabled = True
         except Exception as e:
             print(f"⚠️ Face detection not available: {e}")
-            print("Continuing without face detection (will still work fine)")
+            print("Continuing without face detection")
             self.face_app = None
             self.face_detection_enabled = False
         
@@ -64,9 +71,10 @@ class RetroArtConverter:
             ).to(self.device)
             print("✓ InstantID ControlNet loaded successfully")
             self.instantid_enabled = True
+            self.models_loaded['instantid'] = True
         except Exception as e:
             print(f"⚠️ InstantID ControlNet not available: {e}")
-            print("Running without InstantID (identity may not be preserved)")
+            print("Running without InstantID")
             self.controlnet_instantid = None
             self.instantid_enabled = False
         
@@ -83,13 +91,15 @@ class RetroArtConverter:
                 torch_dtype=self.dtype
             ).to(self.device)
             print("✓ Custom VAE loaded successfully")
+            self.models_loaded['custom_vae'] = True
         except Exception as e:
-            print(f"Warning: Could not load custom VAE: {e}")
-            print("Using default SDXL VAE")
+            print(f"⚠️ Could not load custom VAE: {e}")
+            print("Using high-quality SDXL VAE instead")
             self.vae = AutoencoderKL.from_pretrained(
                 "madebyollin/sdxl-vae-fp16-fix",
                 torch_dtype=self.dtype
             ).to(self.device)
+            self.models_loaded['custom_vae'] = False
         
         # Load depth estimator for preprocessing
         print("Loading depth estimator...")
@@ -123,9 +133,10 @@ class RetroArtConverter:
                 use_safetensors=True
             ).to(self.device)
             print("✓ Custom checkpoint loaded successfully")
+            self.models_loaded['custom_checkpoint'] = True
         except Exception as e:
-            print(f"Warning: Could not load custom checkpoint: {e}")
-            print("Using default SDXL")
+            print(f"⚠️ Could not load custom checkpoint: {e}")
+            print("Using default SDXL base model")
             self.pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
                 "stabilityai/stable-diffusion-xl-base-1.0",
                 controlnet=controlnets,
@@ -133,8 +144,9 @@ class RetroArtConverter:
                 torch_dtype=self.dtype,
                 use_safetensors=True
             ).to(self.device)
+            self.models_loaded['custom_checkpoint'] = False
         
-        # Load LORA from HuggingFace Hub (requires PEFT)
+        # Load LORA from HuggingFace Hub
         print("Loading LORA (retroart) from HuggingFace Hub...")
         try:
             lora_path = hf_hub_download(
@@ -144,18 +156,24 @@ class RetroArtConverter:
             )
             self.pipe.load_lora_weights(lora_path)
             print("✓ LORA loaded successfully")
+            self.models_loaded['lora'] = True
         except Exception as e:
-            print(f"Warning: Could not load LORA: {e}")
+            print(f"⚠️ Could not load LORA: {e}")
             print("Running without LORA")
+            self.models_loaded['lora'] = False
         
-        # Optimize pipeline
-        self.pipe.scheduler = DPMSolverMultistepScheduler.from_config(
+        # Use EulerAncestral scheduler for better quality
+        self.pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(
             self.pipe.scheduler.config
         )
         
-        self.pipe.enable_vae_slicing()
+        # Disable VAE slicing for better quality (use only if you have VRAM issues)
+        # self.pipe.enable_vae_slicing()
+        
+        # Enable attention slicing for memory efficiency
         self.pipe.unet.set_attn_processor(AttnProcessor2_0())
         
+        # Try to enable xformers if available
         if self.device == "cuda":
             try:
                 self.pipe.enable_xformers_memory_efficient_attention()
@@ -167,16 +185,46 @@ class RetroArtConverter:
         self.using_multiple_controlnets = isinstance(controlnets, list)
         print(f"Pipeline initialized with {'multiple' if self.using_multiple_controlnets else 'single'} ControlNet(s)")
         
+        print("\n=== MODEL STATUS ===")
+        for model, loaded in self.models_loaded.items():
+            status = "✓ LOADED" if loaded else "✗ FALLBACK"
+            print(f"{model}: {status}")
+        print("===================\n")
+        
         print("Model initialization complete!")
     
-    def get_depth_map(self, image):
-        """Generate depth map from input image"""
+    def enhance_image_quality(self, image):
+        """Enhance input image quality before processing"""
+        # Sharpen slightly
+        enhancer = ImageEnhance.Sharpness(image)
+        image = enhancer.enhance(1.2)
+        
+        # Enhance contrast slightly
+        enhancer = ImageEnhance.Contrast(image)
+        image = enhancer.enhance(1.1)
+        
+        return image
+    
+    def get_depth_map(self, image, enhance=True):
+        """Generate depth map from input image with quality improvements"""
+        # Enhance image before depth estimation if needed
+        if enhance:
+            image = self.enhance_image_quality(image)
+        
         depth = self.depth_estimator(image)
         depth_image = depth['depth']
         
         depth_array = np.array(depth_image)
-        depth_normalized = (depth_array - depth_array.min()) / (depth_array.max() - depth_array.min()) * 255
+        
+        # Better normalization with histogram stretching
+        depth_min, depth_max = np.percentile(depth_array, [2, 98])
+        depth_normalized = np.clip((depth_array - depth_min) / (depth_max - depth_min + 1e-8), 0, 1) * 255
         depth_normalized = depth_normalized.astype(np.uint8)
+        
+        # Apply slight gaussian blur to reduce noise
+        depth_normalized = cv2.GaussianBlur(depth_normalized, (3, 3), 0)
+        
+        # Convert to 3-channel image
         depth_colored = cv2.cvtColor(depth_normalized, cv2.COLOR_GRAY2RGB)
         
         return Image.fromarray(depth_colored)
@@ -200,19 +248,6 @@ class RetroArtConverter:
             print(f"Face embedding extraction error: {e}")
             return None
     
-    def prepare_face_image(self, image, face_bbox):
-        """Prepare face image for InstantID ControlNet"""
-        x1, y1, x2, y2 = map(int, face_bbox)
-        # Add some padding
-        padding = 20
-        x1 = max(0, x1 - padding)
-        y1 = max(0, y1 - padding)
-        x2 = min(image.width, x2 + padding)
-        y2 = min(image.height, y2 + padding)
-        
-        face_image = image.crop((x1, y1, x2, y2))
-        return face_image
-    
     def calculate_target_size(self, original_width, original_height, max_dimension=1024):
         """Calculate target size maintaining aspect ratio"""
         aspect_ratio = original_width / original_height
@@ -235,14 +270,15 @@ class RetroArtConverter:
         input_image,
         prompt="retro pixel art game, 16-bit style, vibrant colors",
         negative_prompt="blurry, low quality, modern, photorealistic, 3d render",
-        num_inference_steps=30,
+        num_inference_steps=40,  # Increased for better quality
         guidance_scale=7.5,
-        controlnet_conditioning_scale=0.8,
+        controlnet_conditioning_scale=0.6,  # Reduced for less depth influence
         lora_scale=0.85,
-        identity_preservation=0.8,  # NEW PARAMETER
-        image_scale=0.2  # NEW PARAMETER for InstantID strength
+        identity_preservation=0.8,
+        image_scale=0.2,
+        enhance_quality=True  # New parameter
     ):
-        """Main generation function"""
+        """Main generation function with quality improvements"""
         
         # Resize image maintaining aspect ratio
         original_width, original_height = input_image.size
@@ -250,14 +286,19 @@ class RetroArtConverter:
         
         print(f"Resizing from {original_width}x{original_height} to {target_width}x{target_height}")
         
+        # Use LANCZOS for high-quality resizing
         resized_image = input_image.resize((target_width, target_height), Image.LANCZOS)
         
-        # Generate depth map
+        # Optionally enhance image quality
+        if enhance_quality:
+            resized_image = self.enhance_image_quality(resized_image)
+        
+        # Generate depth map with quality enhancements
         print("Generating depth map...")
-        depth_image = self.get_depth_map(resized_image)
+        depth_image = self.get_depth_map(resized_image, enhance=enhance_quality)
         depth_image = depth_image.resize((target_width, target_height), Image.LANCZOS)
         
-        # Determine if we're using multiple ControlNets based on initialization
+        # Determine if we're using multiple ControlNets
         using_multiple_controlnets = self.using_multiple_controlnets
         
         # Extract face embeddings if InstantID is enabled
@@ -280,16 +321,20 @@ class RetroArtConverter:
                 prompt = f"portrait, detailed face, facial features, {prompt}"
         
         # Set LORA scale
-        if hasattr(self.pipe, 'set_adapters'):
+        if hasattr(self.pipe, 'set_adapters') and self.models_loaded['lora']:
             try:
                 self.pipe.set_adapters(["retroart"], adapter_weights=[lora_scale])
-            except:
-                print("Could not set LORA adapters, continuing without")
+                print(f"LORA scale set to: {lora_scale}")
+            except Exception as e:
+                print(f"Could not set LORA adapters: {e}")
+        
+        # Enhanced negative prompt for better quality
+        enhanced_negative_prompt = f"{negative_prompt}, worst quality, low quality, normal quality, lowres, watermark, signature, text, jpeg artifacts, noise, grainy"
         
         # Prepare pipeline kwargs
         pipe_kwargs = {
             "prompt": prompt,
-            "negative_prompt": negative_prompt,
+            "negative_prompt": enhanced_negative_prompt,
             "num_inference_steps": num_inference_steps,
             "guidance_scale": guidance_scale,
             "width": target_width,
@@ -299,7 +344,6 @@ class RetroArtConverter:
         
         # Add control images and scales based on ControlNet configuration
         if using_multiple_controlnets and has_detected_faces:
-            # Multiple ControlNets: depth + InstantID
             print("Using multiple ControlNets (Depth + InstantID)")
             control_images = [depth_image, resized_image]
             conditioning_scales = [controlnet_conditioning_scale, image_scale]
@@ -312,23 +356,21 @@ class RetroArtConverter:
                 pipe_kwargs["cross_attention_kwargs"] = {"ip_adapter_image_embeds": [face_embeddings]}
         
         elif using_multiple_controlnets and not has_detected_faces:
-            # Multiple ControlNets initialized but no faces detected
-            # Pass images for both controlnets but with zero weight for InstantID
             print("Multiple ControlNets available but no faces detected, using depth only")
-            control_images = [depth_image, depth_image]  # Use depth for both
-            conditioning_scales = [controlnet_conditioning_scale, 0.0]  # Zero weight for InstantID
+            control_images = [depth_image, depth_image]
+            conditioning_scales = [controlnet_conditioning_scale, 0.0]
             
             pipe_kwargs["image"] = control_images
             pipe_kwargs["controlnet_conditioning_scale"] = conditioning_scales
         
         else:
-            # Single ControlNet (depth only)
             print("Using single ControlNet (Depth only)")
             pipe_kwargs["image"] = depth_image
             pipe_kwargs["controlnet_conditioning_scale"] = controlnet_conditioning_scale
         
         # Generate image
         print("Generating retro art...")
+        print(f"Steps: {num_inference_steps}, Guidance: {guidance_scale}")
         result = self.pipe(**pipe_kwargs)
         
         return result.images[0]
@@ -347,8 +389,9 @@ def process_image(
     guidance_scale,
     controlnet_scale,
     lora_scale,
-    identity_preservation,  # NEW
-    image_scale  # NEW
+    identity_preservation,
+    image_scale,
+    enhance_quality
 ):
     if image is None:
         return None
@@ -362,8 +405,9 @@ def process_image(
             guidance_scale=guidance_scale,
             controlnet_conditioning_scale=controlnet_scale,
             lora_scale=lora_scale,
-            identity_preservation=identity_preservation,  # NEW
-            image_scale=image_scale  # NEW
+            identity_preservation=identity_preservation,
+            image_scale=image_scale,
+            enhance_quality=enhance_quality
         )
         return result
     except Exception as e:
@@ -373,59 +417,75 @@ def process_image(
         raise gr.Error(f"Generation failed: {str(e)}")
 
 # Create Gradio interface
-with gr.Blocks(title="RetroArt Converter") as demo:
+with gr.Blocks(title="RetroArt Converter", theme=gr.themes.Soft()) as demo:
     gr.Markdown("""
-    # 🎮 RetroArt Converter
+    # 🎮 RetroArt Converter - Quality Enhanced
     
-    Convert any image into retro game art style!
+    Convert any image into retro game art style with improved quality!
     
     **Features:**
-    - Custom SDXL checkpoint (Horizon)
+    - High-quality depth estimation and preprocessing
+    - Enhanced prompts for better results
+    - Custom SDXL checkpoint (Horizon) 
     - Pixelate VAE for authentic retro look
     - RetroArt LORA for style enhancement
-    - Face preservation with InstantID (if available)
-    - Depth-aware generation with ControlNet
+    - Face preservation with InstantID
     """)
     
+    # Model status display
+    if converter.models_loaded:
+        status_text = "**Loaded Models:**\n"
+        status_text += f"- Custom Checkpoint: {'✓' if converter.models_loaded['custom_checkpoint'] else '✗ (using SDXL base)'}\n"
+        status_text += f"- Custom VAE: {'✓' if converter.models_loaded['custom_vae'] else '✗ (using default VAE)'}\n"
+        status_text += f"- LORA: {'✓' if converter.models_loaded['lora'] else '✗ (disabled)'}\n"
+        status_text += f"- InstantID: {'✓' if converter.models_loaded['instantid'] else '✗ (disabled)'}\n"
+        gr.Markdown(status_text)
+    
     with gr.Row():
         with gr.Column():
             input_image = gr.Image(label="Input Image", type="pil")
             
             prompt = gr.Textbox(
                 label="Prompt",
-                value="retro pixel art game, 16-bit style, vibrant colors, detailed",
+                value="masterpiece, best quality, retro pixel art game, 16-bit style, vibrant colors, highly detailed",
                 lines=3
             )
             
             negative_prompt = gr.Textbox(
                 label="Negative Prompt",
-                value="blurry, low quality, modern, photorealistic, 3d render, ugly, distorted",
+                value="blurry, low quality, modern, photorealistic, 3d render, ugly, distorted, deformed",
                 lines=2
             )
             
-            with gr.Accordion("Advanced Settings", open=False):
+            enhance_quality = gr.Checkbox(
+                label="Enable Quality Enhancement",
+                value=True,
+                info="Sharpen and enhance input image before processing"
+            )
+            
+            with gr.Accordion("Quality Settings", open=True):
                 steps = gr.Slider(
                     minimum=20,
-                    maximum=50,
-                    value=30,
-                    step=1,
-                    label="Inference Steps"
+                    maximum=70,
+                    value=40,
+                    step=5,
+                    label="Inference Steps (more = better quality but slower)"
                 )
                 
                 guidance_scale = gr.Slider(
-                    minimum=1,
+                    minimum=3,
                     maximum=15,
                     value=7.5,
                     step=0.5,
-                    label="Guidance Scale"
+                    label="Guidance Scale (how closely to follow prompt)"
                 )
                 
                 controlnet_scale = gr.Slider(
                     minimum=0,
-                    maximum=2,
-                    value=0.8,
-                    step=0.1,
-                    label="ControlNet Depth Scale"
+                    maximum=1.5,
+                    value=0.6,
+                    step=0.05,
+                    label="ControlNet Depth Scale (lower = more creative)"
                 )
                 
                 lora_scale = gr.Slider(
@@ -435,14 +495,14 @@ with gr.Blocks(title="RetroArt Converter") as demo:
                     step=0.05,
                     label="RetroArt LORA Scale"
                 )
-                
-                # NEW PARAMETERS
+            
+            with gr.Accordion("Identity Settings (for portraits)", open=False):
                 identity_preservation = gr.Slider(
                     minimum=0,
                     maximum=1.5,
                     value=0.8,
                     step=0.1,
-                    label="Identity Preservation (InstantID strength)"
+                    label="Identity Preservation"
                 )
                 
                 image_scale = gr.Slider(
@@ -453,16 +513,34 @@ with gr.Blocks(title="RetroArt Converter") as demo:
                     label="InstantID Image Scale"
                 )
             
-            generate_btn = gr.Button("🎨 Generate Retro Art", variant="primary")
+            generate_btn = gr.Button("🎨 Generate Retro Art", variant="primary", size="lg")
         
         with gr.Column():
             output_image = gr.Image(label="Retro Art Output")
+            
+            gr.Markdown("""
+            ### Tips for Best Quality:
+            1. **Use high-resolution input images** (at least 512x512)
+            2. **Increase inference steps** to 50-60 for maximum quality
+            3. **Lower ControlNet scale** (0.5-0.6) for more stylization
+            4. **Adjust guidance scale:** 7-9 for balanced results
+            5. **Enable quality enhancement** for sharper inputs
+            6. Try different prompts with quality keywords: "masterpiece, best quality, highly detailed"
+            """)
     
     gr.Examples(
         examples=[
-            ["example_portrait.jpg", "retro pixel art portrait, 16-bit game character", "blurry, modern", 30, 7.5, 0.8, 0.85, 0.8, 0.2],
+            [
+                "example_portrait.jpg", 
+                "masterpiece, best quality, retro pixel art portrait, 16-bit game character, vibrant colors",
+                "blurry, modern, low quality",
+                40, 7.5, 0.6, 0.85, 0.8, 0.2, True
+            ],
+        ],
+        inputs=[
+            input_image, prompt, negative_prompt, steps, guidance_scale, 
+            controlnet_scale, lora_scale, identity_preservation, image_scale, enhance_quality
         ],
-        inputs=[input_image, prompt, negative_prompt, steps, guidance_scale, controlnet_scale, lora_scale, identity_preservation, image_scale],
         outputs=[output_image],
         fn=process_image,
         cache_examples=False
@@ -470,7 +548,10 @@ with gr.Blocks(title="RetroArt Converter") as demo:
     
     generate_btn.click(
         fn=process_image,
-        inputs=[input_image, prompt, negative_prompt, steps, guidance_scale, controlnet_scale, lora_scale, identity_preservation, image_scale],
+        inputs=[
+            input_image, prompt, negative_prompt, steps, guidance_scale, 
+            controlnet_scale, lora_scale, identity_preservation, image_scale, enhance_quality
+        ],
         outputs=[output_image]
     )