app_sdxl_backup.py

"""
ColorCraft SDXL - LineArt Storage Efficient Edition
==================================================

Storage-efficient LineArt approach:
- On-demand model loading with cache management
- Storage monitoring and cleanup
- LineArt preprocessor output display
- Single model approach to minimize storage usage
"""

import os
import gc
import shutil
import gradio as gr
import torch
import numpy as np
from PIL import Image

# Environment setup
os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1")

# Simple compatibility fix at the very beginning
try:
    import huggingface_hub
    if not hasattr(huggingface_hub, 'cached_download'):
        from huggingface_hub import hf_hub_download
        huggingface_hub.cached_download = hf_hub_download
        print("✅ Applied cached_download compatibility shim")
except Exception as e:
    print(f"⚠️ Compatibility shim error: {e}")

# Now import diffusers after the fix
try:
    from diffusers import (
        StableDiffusionXLControlNetPipeline, 
        StableDiffusionXLAdapterPipeline,
        ControlNetModel,
        T2IAdapter,
        AutoencoderKL,
        DDIMScheduler
    )
    print("✅ Diffusers imported successfully")
except ImportError as e:
    print(f"❌ Diffusers import error: {e}")

# Try ControlNet preprocessor imports
try:
    from controlnet_aux import LineartDetector
    print("✅ LineartDetector imported successfully")
except ImportError as e:
    print(f"❌ LineartDetector import error: {e}")
    LineartDetector = None

# Try depth preprocessor (using correct import)
try:
    from controlnet_aux import MidasDetector
    print("✅ MidasDetector (depth) imported successfully")
    DepthEstimator = MidasDetector  # Alias for compatibility
except ImportError as e:
    print(f"❌ MidasDetector import error: {e}")
    DepthEstimator = None

# Fallback depth model
try:
    from transformers import DPTImageProcessor, DPTForDepthEstimation
    print("✅ DPT depth model imported successfully")
except ImportError as e:
    print(f"❌ DPT import error: {e}")
    DPTImageProcessor = None
    DPTForDepthEstimation = None

import warnings
warnings.filterwarnings("ignore")

# =============================================================================
# STORAGE-EFFICIENT CONFIGURATION
# =============================================================================

SDXL_MODEL_ID = "stabilityai/stable-diffusion-xl-base-1.0"
VAE_MODEL_ID = "madebyollin/sdxl-vae-fp16-fix"

# LineArt model options (verified SDXL compatible)
LINEART_MODELS = {
    "standard": {
        "type": "controlnet",
        "id": "ShermanG/ControlNet-Standard-Lineart-for-SDXL"
    },
    "anime": {
        "type": "controlnet", 
        "id": "r3gm/controlnet-lineart-anime-sdxl-fp16"
    },
    "tencent": {
        "type": "t2i_adapter",
        "id": "TencentARC/t2i-adapter-lineart-sdxl-1.0"
    }
}

DEPTH_MODEL_ID = "diffusers/controlnet-depth-sdxl-1.0"

# Cache directories with storage management
CACHE_ROOT = "/data" if os.path.exists("/data") else "."
HF_CACHE_DIR = f"{CACHE_ROOT}/hf_cache"
os.makedirs(HF_CACHE_DIR, exist_ok=True)

# =============================================================================
# STORAGE MANAGEMENT
# =============================================================================

def get_storage_info():
    """Get current storage usage"""
    try:
        if os.path.exists("/data"):
            total, used, free = shutil.disk_usage("/data")
            return {
                "total_gb": total / (1024**3),
                "used_gb": used / (1024**3), 
                "free_gb": free / (1024**3),
                "usage_percent": (used / total) * 100
            }
        else:
            return {
                "total_gb": 0,
                "used_gb": 0,
                "free_gb": 0,
                "usage_percent": 0
            }
    except Exception as e:
        print(f"Storage info error: {e}")
        return {"total_gb": 0, "used_gb": 0, "free_gb": 0, "usage_percent": 0}

def format_storage_status():
    """Format storage status for display"""
    info = get_storage_info()
    if info["total_gb"] > 0:
        return f"Storage: {info['free_gb']:.1f}GB free of {info['total_gb']:.1f}GB total ({info['usage_percent']:.1f}% used)"
    else:
        return "Ephemeral storage (no persistent cache)"

def clear_cache_if_needed(min_free_gb=2.0):
    """Clear cache if storage is low"""
    info = get_storage_info()
    if info["free_gb"] < min_free_gb:
        print(f"⚠️ Low storage: {info['free_gb']:.1f}GB free, clearing cache...")
        
        # Clear HuggingFace cache
        if os.path.exists(HF_CACHE_DIR):
            try:
                shutil.rmtree(HF_CACHE_DIR)
                os.makedirs(HF_CACHE_DIR, exist_ok=True)
                print("✅ HF cache cleared")
            except Exception as e:
                print(f"Cache clear error: {e}")
        
        # Force garbage collection
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
            
            return True
            return False
    
# =============================================================================
# STORAGE-EFFICIENT LINEART PIPELINE
# =============================================================================

class StorageEfficientLineArtPipeline:
    def __init__(self):
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.dtype = torch.float16 if torch.cuda.is_available() else torch.float32
        self.pipeline = None
        self.lineart_detector = None
        self.depth_estimator = None
        self.dpt_processor = None
        self.dpt_model = None
        self.last_storage_check = None
        print(f"🚀 Storage-Efficient LineArt initialized on {self.device}")
    
    def check_storage_before_load(self):
        """Check storage before loading models"""
        info = get_storage_info()
        print(f"📊 Storage check: {info['free_gb']:.1f}GB free, {info['usage_percent']:.1f}% used")
        
        # Clear cache if very low on space
        if info['free_gb'] < 1.0:
            print("⚠️ Critical storage! Clearing all caches...")
            clear_cache_if_needed(min_free_gb=0.5)
            self.clear_pipeline()
            return False
        elif info['free_gb'] < 3.0:
            print("⚠️ Low storage! May need to clear cache...")
            return True
        else:
            print("✅ Storage OK for model loading")
            return True
    
    def clear_pipeline(self):
        """Clear current pipeline to free memory"""
        if self.pipeline is not None:
            print("🗑️ Clearing pipeline to free storage...")
            del self.pipeline
            self.pipeline = None
        
        if self.lineart_detector is not None:
            print("🗑️ Clearing LineArt detector...")
            del self.lineart_detector
            self.lineart_detector = None
        
        if self.depth_estimator is not None:
            print("🗑️ Clearing depth estimator...")
            del self.depth_estimator
            self.depth_estimator = None
            
        if self.dpt_model is not None:
            print("🗑️ Clearing DPT model...")
            del self.dpt_model
            self.dpt_model = None
            
        if self.dpt_processor is not None:
            del self.dpt_processor
            self.dpt_processor = None
        
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
    
    def load_lineart_detector(self):
        """Load LineArt detector with storage check and fallback"""
        if self.lineart_detector is not None:
            return self.lineart_detector
            
        if LineartDetector is None:
            print("❌ LineartDetector not available")
            return None
        
        # Check storage before loading
        if not self.check_storage_before_load():
            return None
            
        # Try multiple approaches for LineArt detector
        try:
            print("🎨 Loading LineArt Detector...")
            
            # Try 1: Default from_pretrained
            try:
                self.lineart_detector = LineartDetector.from_pretrained(
                    "lllyasviel/Annotators",
                    cache_dir=HF_CACHE_DIR
                )
                print("✅ LineArt Detector loaded (method 1)!")
                return self.lineart_detector
            except Exception as e1:
                print(f"⚠️ Method 1 failed: {e1}")
            
            # Try 2: Direct instantiation
            try:
                self.lineart_detector = LineartDetector()
                print("✅ LineArt Detector loaded (method 2)!")
                return self.lineart_detector
            except Exception as e2:
                print(f"⚠️ Method 2 failed: {e2}")
            
            # Try 3: Alternative approach without cache_dir
            try:
                self.lineart_detector = LineartDetector.from_pretrained("lllyasviel/Annotators")
                print("✅ LineArt Detector loaded (method 3)!")
                return self.lineart_detector
            except Exception as e3:
                print(f"⚠️ Method 3 failed: {e3}")
                
            print("❌ All LineArt Detector loading methods failed")
            return None
            
        except Exception as e:
            print(f"❌ LineArt Detector loading failed: {str(e)}")
            return None
    
    def load_depth_preprocessor(self):
        """Load depth preprocessor with fallback options"""
        if self.depth_estimator is not None:
            return self.depth_estimator
            
        # Check storage before loading
        if not self.check_storage_before_load():
            return None
            
        try:
            print("🏔️ Loading Depth Estimator...")
            
            # Try ControlNet-aux MidasDetector first
            if DepthEstimator is not None:
                try:
                    self.depth_estimator = DepthEstimator.from_pretrained(
                        "lllyasviel/Annotators",
                        cache_dir=HF_CACHE_DIR
                    )
                    print("✅ MidasDetector (depth) loaded!")
                    return self.depth_estimator
                except Exception as e1:
                    print(f"⚠️ MidasDetector failed: {e1}")
                    
                # Try direct instantiation if from_pretrained fails
                try:
                    self.depth_estimator = DepthEstimator()
                    print("✅ MidasDetector (depth) loaded with direct instantiation!")
                    return self.depth_estimator
                except Exception as e2:
                    print(f"⚠️ MidasDetector direct instantiation failed: {e2}")
            
            # Fallback to DPT model
            if DPTImageProcessor is not None and DPTForDepthEstimation is not None:
                try:
                    self.dpt_processor = DPTImageProcessor.from_pretrained(
                        "Intel/dpt-hybrid-midas",
                        cache_dir=HF_CACHE_DIR
                    )
                    self.dpt_model = DPTForDepthEstimation.from_pretrained(
                        "Intel/dpt-hybrid-midas",
                        torch_dtype=self.dtype,
                        cache_dir=HF_CACHE_DIR
                    ).to(self.device)
                    print("✅ DPT depth model loaded as fallback!")
                    return "dpt_fallback"
                except Exception as e2:
                    print(f"⚠️ DPT fallback failed: {e2}")
                    
            print("❌ All depth preprocessor loading methods failed")
            return None
            
        except Exception as e:
            print(f"❌ Depth preprocessor loading failed: {str(e)}")
            return None
    
    def process_depth_map(self, input_image):
        """Process image to depth map"""
        try:
            if self.depth_estimator is not None:
                # Use ControlNet-aux MidasDetector
                depth_image = self.depth_estimator(input_image)
                print("✅ Depth processed with MidasDetector")
                return depth_image
                
            elif self.dpt_model is not None and self.dpt_processor is not None:
                # Use DPT fallback
                inputs = self.dpt_processor(images=input_image, return_tensors="pt").to(self.device)
                
                with torch.no_grad():
                    outputs = self.dpt_model(**inputs)
                    predicted_depth = outputs.predicted_depth
                
                # Convert to PIL Image
                prediction = torch.nn.functional.interpolate(
                    predicted_depth.unsqueeze(1),
                    size=input_image.size[::-1],
                    mode="bicubic",
                    align_corners=False,
                )
                output = prediction.squeeze().cpu().numpy()
                formatted = (output * 255 / np.max(output)).astype("uint8")
                depth_image = Image.fromarray(formatted).convert("RGB")
                print("✅ Depth processed with DPT fallback")
                return depth_image
            else:
                print("❌ No depth preprocessor available")
                return None
                
        except Exception as e:
            print(f"❌ Depth processing failed: {str(e)}")
            return None
    
    def apply_pure_bw_post_processing(self, image):
        """Apply edge detection and thresholding for pure black and white"""
        try:
            import cv2
            
            # Convert PIL to numpy
            img_array = np.array(image.convert('RGB'))
            
            # Convert to grayscale
            gray = cv2.cvtColor(img_array, cv2.COLOR_RGB2GRAY)
            
            # Apply Gaussian blur to reduce noise
            blurred = cv2.GaussianBlur(gray, (3, 3), 0)
            
            # Apply Canny edge detection
            edges = cv2.Canny(blurred, 50, 150)
            
            # Apply morphological operations to clean up lines
            kernel = np.ones((2, 2), np.uint8)
            edges = cv2.morphologyEx(edges, cv2.MORPH_CLOSE, kernel)
            
            # Apply binary thresholding for pure black and white
            _, binary = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
            
            # Combine edges with binary threshold
            combined = cv2.bitwise_or(edges, binary)
            
            # Invert so lines are black on white background
            inverted = cv2.bitwise_not(combined)
            
            # Convert back to PIL
            result_image = Image.fromarray(inverted).convert('RGB')
            print("✅ Pure B&W post-processing applied")
            
            return result_image
            
        except Exception as e:
            print(f"⚠️ Post-processing failed: {e}, returning original")
            return image
    
    def load_pipeline(self, lineart_model_key="standard"):
        """Load pipeline with storage management - supports both ControlNet and T2I-Adapter"""
        # Create a unique pipeline key that includes the model selection
        pipeline_key = f"pipeline_{lineart_model_key}"
        
        # Check if we have this specific pipeline cached
        if hasattr(self, pipeline_key) and getattr(self, pipeline_key) is not None:
            return getattr(self, pipeline_key)
        
        # Check storage before loading
        if not self.check_storage_before_load():
            return None
            
        try:
            model_config = LINEART_MODELS.get(lineart_model_key, LINEART_MODELS["standard"])
            model_type = model_config["type"]
            model_id = model_config["id"]
            
            print(f"🎨 Loading {model_type.upper()} model: {model_id}")
            
            # Load VAE first
            vae = AutoencoderKL.from_pretrained(
                VAE_MODEL_ID,
                torch_dtype=self.dtype,
                cache_dir=HF_CACHE_DIR
            )
            
            if model_type == "controlnet":
                # Load ControlNet model
                controlnet = ControlNetModel.from_pretrained(
                    model_id,
                    torch_dtype=self.dtype,
                    cache_dir=HF_CACHE_DIR,
                    variant="fp16" if self.dtype == torch.float16 else None
                )
                
                # Create ControlNet pipeline
                pipeline = StableDiffusionXLControlNetPipeline.from_pretrained(
                    SDXL_MODEL_ID,
                    controlnet=controlnet,
                    vae=vae,
                    torch_dtype=self.dtype,
                    cache_dir=HF_CACHE_DIR,
                    use_safetensors=True,
                    variant="fp16" if self.dtype == torch.float16 else None
                )
                
            elif model_type == "t2i_adapter":
                # Load T2I-Adapter model
                adapter = T2IAdapter.from_pretrained(
                    model_id,
                    torch_dtype=self.dtype,
                    cache_dir=HF_CACHE_DIR,
                    variant="fp16" if self.dtype == torch.float16 else None
                )
                
                # Create T2I-Adapter pipeline
                pipeline = StableDiffusionXLAdapterPipeline.from_pretrained(
                    SDXL_MODEL_ID,
                    adapter=adapter,
                    vae=vae,
                    torch_dtype=self.dtype,
                    cache_dir=HF_CACHE_DIR,
                    use_safetensors=True,
                    variant="fp16" if self.dtype == torch.float16 else None
                )
            
            else:
                raise ValueError(f"Unknown model type: {model_type}")
            
            # Check storage after model loading
            info = get_storage_info()
            print(f"📊 After model load: {info['free_gb']:.1f}GB free")
            
            if info['free_gb'] < 1.0:
                print("❌ Not enough space for full pipeline!")
                return None
            
            # Basic optimizations
            pipeline.enable_model_cpu_offload()
            
            # Cache this specific pipeline
            setattr(self, pipeline_key, pipeline)
            
            # Final storage check
            info = get_storage_info()
            print(f"📊 After pipeline load: {info['free_gb']:.1f}GB free")
            
            print(f"✅ {model_type.upper()} pipeline loaded successfully for {lineart_model_key}!")
            return pipeline
            
        except Exception as e:
            print(f"❌ Pipeline loading failed: {str(e)}")
            # Try to clear cache and return None
            clear_cache_if_needed(min_free_gb=0.5)
            return None
    
    def generate_with_lineart(self, input_image, prompt="coloring book page, black and white line art", 
                             num_steps=20, guidance_scale=7.5, controlnet_scale=1.2, 
                             style_preset="clean", use_negative_prompt=True, output_resolution=768,
                             use_depth_preprocessor=False, lineart_model="standard", apply_post_processing=False):
        """Generate with LineArt and advanced controls for coloring book optimization"""
        try:
            # Load detector first
            detector = self.load_lineart_detector()
            
            # Process input image to LineArt
            input_image = Image.fromarray(input_image).convert("RGB")
            
            if detector is None:
                print("⚠️ LineArt Detector failed, using simple edge detection fallback...")
                # Simple fallback: convert to grayscale and apply basic edge detection
                import cv2
                gray = cv2.cvtColor(np.array(input_image), cv2.COLOR_RGB2GRAY)
                edges = cv2.Canny(gray, 50, 150)
                # Convert back to 3-channel for consistency
                lineart_image = Image.fromarray(cv2.cvtColor(edges, cv2.COLOR_GRAY2RGB))
                print("✅ Fallback edge detection complete")
            else:
                print("🎨 Processing image with LineArt detector...")
                lineart_image = detector(input_image)
                print("✅ LineArt processing complete")
            
            # Load pipeline with selected LineArt model
            pipeline = self.load_pipeline(lineart_model)
            if pipeline is None:
                return None, lineart_image, None, "❌ Pipeline failed to load (storage full?)"
            
            # Check storage before generation
            info = get_storage_info()
            if info['free_gb'] < 0.5:
                return None, lineart_image, None, f"❌ Insufficient storage for generation: {info['free_gb']:.1f}GB free"
            
            # Build optimized prompts based on style preset
            style_prompts = {
                "ultra_clean": "pure black and white line art, coloring book page, simple outlines, minimal details, vector art style, clean lines, no shading, no colors, stark contrast",
                "clean": "black and white line art, coloring book page, clear outlines, simple design, minimal shading, clean vector style",
                "detailed": "detailed black and white line art, coloring book page, intricate outlines, fine details, precise lines",
                "bold": "bold black and white line art, coloring book page, thick outlines, strong lines, high contrast",
                "simple": "simple black and white line art, coloring book page, basic outlines, easy to color"
            }
            
            # Get style-specific prompt additions
            style_addition = style_prompts.get(style_preset, style_prompts["clean"])
            final_prompt = f"{style_addition}, {prompt}" if prompt.strip() else style_addition
            
            # Optimized negative prompts for coloring book style
            negative_prompts = {
                "ultra_clean": "color, colors, colored, shading, shadows, gradients, realistic, photographic, photography, 3d render, painting, watercolor, oil painting, detailed textures, complex lighting, depth, perspective, artistic style, sketch",
                "clean": "color, colors, shading, gradients, realistic, photographic, complex details, artistic style",
                "detailed": "color, colors, heavy shading, realistic, photographic",
                "bold": "color, colors, fine details, realistic, photographic",
                "simple": "color, colors, complex details, shading, realistic, photographic"
            }
            
            final_negative = negative_prompts.get(style_preset, negative_prompts["clean"]) if use_negative_prompt else ""
            
            print(f"🎨 Generating with style: {style_preset}")
            print(f"📝 Prompt: {final_prompt[:100]}...")
            print(f"🚫 Negative: {final_negative[:50]}...")
            
                        # Calculate output dimensions
            aspect_ratio = input_image.width / input_image.height
            if aspect_ratio > 1:  # Landscape
                output_width = min(output_resolution, 1536)  # Max 1536 for memory
                output_height = int(output_width / aspect_ratio)
            else:  # Portrait or square
                output_height = min(output_resolution, 1536)
                output_width = int(output_height * aspect_ratio)
            
            # Ensure dimensions are multiples of 8 (required by SDXL)
            output_width = (output_width // 8) * 8
            output_height = (output_height // 8) * 8
            
            print(f"🖼️ Output resolution: {output_width}x{output_height}")
            
            # Get model type to determine parameter names
            model_config = LINEART_MODELS.get(lineart_model, LINEART_MODELS["standard"])
            model_type = model_config["type"]
            
            # Generate with type-specific parameters
            if model_type == "controlnet":
                result = pipeline(
                    prompt=final_prompt,
                    negative_prompt=final_negative,
                    image=lineart_image,
                    num_inference_steps=num_steps,
                    guidance_scale=guidance_scale,
                    controlnet_conditioning_scale=controlnet_scale,
                    width=output_width,
                    height=output_height
                )
            elif model_type == "t2i_adapter":
                result = pipeline(
                    prompt=final_prompt,
                    negative_prompt=final_negative,
                    image=lineart_image,
                    num_inference_steps=num_steps,
                    guidance_scale=guidance_scale,
                    adapter_conditioning_scale=controlnet_scale,  # Different parameter name for T2I-Adapter
                    width=output_width,
                    height=output_height
                )
            else:
                raise ValueError(f"Unknown model type: {model_type}")
            
            # Get the generated image
            generated_image = result.images[0]
            
            # Apply post-processing if enabled
            if apply_post_processing:
                generated_image = self.apply_pure_bw_post_processing(generated_image)
            
            # Process depth if enabled
            depth_image = None
            if use_depth_preprocessor:
                depth_estimator = self.load_depth_preprocessor()
                if depth_estimator:
                    depth_image = self.process_depth_map(input_image)
                    if depth_image:
                        print("✅ Depth map generated for display")
            
            post_processing_text = " + Post-Processing" if apply_post_processing else ""
            print("✅ Generation complete")
            return generated_image, lineart_image, depth_image, f"✅ Generated with {style_preset} style, {num_steps} steps{post_processing_text}!"
            
        except Exception as e:
            error_msg = f"❌ Generation failed: {str(e)}"
            print(error_msg)
            
            # If it's a storage error, try to clear cache
            if "out of memory" in str(e).lower() or "space" in str(e).lower():
                clear_cache_if_needed(min_free_gb=0.5)
                self.clear_pipeline()
                error_msg += " (Cleared cache due to storage issue)"
            
            return None, None, None, error_msg

# =============================================================================
# GRADIO INTERFACE WITH STORAGE MONITORING
# =============================================================================

def create_storage_efficient_interface():
    """Create storage-efficient interface with monitoring"""
    
    # Initialize pipeline
    pipeline_manager = StorageEfficientLineArtPipeline()
    
    def generate_with_monitoring(input_image, custom_prompt, style_preset, num_steps, 
                               guidance_scale, controlnet_scale, use_negative_prompt,
                               output_resolution, use_depth_preprocessor, lineart_model, apply_post_processing):
        if input_image is None:
            return None, None, None, "Please upload an image", format_storage_status()
        
        # Use custom prompt or default
        prompt = custom_prompt.strip() if custom_prompt.strip() else ""
        
        # Generate with all user controls
        result_image, lineart_preview, depth_preview, status = pipeline_manager.generate_with_lineart(
            input_image, 
            prompt, 
            num_steps=num_steps,
            guidance_scale=guidance_scale,
            controlnet_scale=controlnet_scale,
            style_preset=style_preset,
            use_negative_prompt=use_negative_prompt,
            output_resolution=output_resolution,
            use_depth_preprocessor=use_depth_preprocessor,
            lineart_model=lineart_model,
            apply_post_processing=apply_post_processing
        )
        
        # Update storage status
        storage_status = format_storage_status()
        
        return result_image, lineart_preview, depth_preview, status, storage_status
    
    def manual_clear_cache():
        pipeline_manager.clear_pipeline()
        cleared = clear_cache_if_needed(min_free_gb=0.0)  # Force clear
        status = "✅ Cache cleared!" if cleared else "✅ Cache clear attempted"
        storage_status = format_storage_status()
        return status, storage_status
    
    with gr.Blocks(title="ColorCraft SDXL - Advanced LineArt Controls") as demo:
        
        gr.HTML("""
        <h1>🎨 ColorCraft SDXL - Advanced LineArt Controls</h1>
        <p>Fine-tuned controls for clean black & white coloring books • LineArt preprocessor • Storage monitoring</p>
        """)
        
        with gr.Row():
            with gr.Column(scale=1):
                gr.HTML("<h3>📷 Input & Controls</h3>")
                
                input_image = gr.Image(
                    label="Upload Image",
                    type="numpy",
                    height=300
                )
                
                # Style Preset - Most Important Control
                style_preset = gr.Dropdown(
                    label="🎨 Style Preset",
                    choices=["ultra_clean", "clean", "detailed", "bold", "simple"],
                    value="ultra_clean",
                    info="Ultra Clean = Pure B&W, Clean = Minimal shading"
                )
                
                # LineArt Model Selection
                lineart_model = gr.Dropdown(
                    label="🖌️ LineArt Model",
                    choices=[
                        ("Standard ControlNet", "standard"),
                        ("Anime ControlNet (r3gm)", "anime"), 
                        ("TencentARC T2I-Adapter", "tencent")
                    ],
                    value="standard",
                    info="Test different LineArt approaches: ControlNet vs T2I-Adapter architectures"
                )
                
                custom_prompt = gr.Textbox(
                    label="Custom Prompt (optional)",
                    placeholder="Additional prompt (style preset will be added automatically)",
                    lines=2,
                    value=""
                )
                
                # Fine-tuning Controls
                with gr.Accordion("⚙️ Advanced Settings", open=True):
                    with gr.Row():
                        num_steps = gr.Slider(
                            label="Inference Steps",
                            minimum=10,
                            maximum=50,
                            value=25,
                            step=5,
                            info="More steps = higher quality, slower"
                        )
                        
                        guidance_scale = gr.Slider(
                            label="Guidance Scale",
                            minimum=3.0,
                            maximum=15.0,
                            value=9.0,
                            step=0.5,
                            info="Higher = follows prompt more strictly"
                        )
                    
                    with gr.Row():
                        controlnet_scale = gr.Slider(
                            label="ControlNet Strength",
                            minimum=0.5,
                            maximum=2.0,
                            value=1.4,
                            step=0.1,
                            info="Higher = follows input lines more closely"
                        )
                        
                        use_negative_prompt = gr.Checkbox(
                            label="Use Negative Prompt",
                            value=True,
                            info="Helps remove colors and shading"
                        )
                    
                    with gr.Row():
                        output_resolution = gr.Slider(
                            label="Output Resolution",
                            minimum=512,
                            maximum=1536,
                            value=1024,
                            step=64,
                            info="Higher = better quality, slower generation"
                        )
                        
                        use_depth_preprocessor = gr.Checkbox(
                            label="Add Depth Preprocessing",
                            value=False,
                            info="Experimental: Add structural depth guidance"
                        )
                    
                    with gr.Row():
                        apply_post_processing = gr.Checkbox(
                            label="Pure B&W Post-Processing",
                            value=False,
                            info="Apply edge detection + thresholding for pure black/white"
                        )
                
                generate_btn = gr.Button(
                    "🎨 Generate LineArt",
                    variant="primary",
                    size="lg"
                )
            
            with gr.Column(scale=2):
                gr.HTML("<h3>🎨 Generated Results</h3>")
                
                output_image = gr.Image(
                    label="Generated Coloring Book",
                    height=300
                )
                
                with gr.Row():
                    lineart_preview = gr.Image(
                        label="LineArt Preprocessor Output",
                        height=200
                    )
                    
                    depth_preview = gr.Image(
                        label="Depth Preprocessor Output",
                        height=200
                    )
                    
                    with gr.Column():
                        status_output = gr.Textbox(
                            label="Generation Status",
                            value="Ready to generate...",
                            lines=3
                        )
                        
                        storage_status = gr.Textbox(
                            label="Storage Status",
                            value=format_storage_status(),
                            lines=2
                        )
        
        # Storage Management
        with gr.Accordion("💾 Storage Management", open=True):
            with gr.Row():
                clear_cache_btn = gr.Button("🗑️ Clear Cache & Free Storage")
                refresh_storage_btn = gr.Button("🔄 Refresh Storage Status")
                
            gr.HTML("""
            <div style="padding: 10px; background-color: #f0f0f0; border-radius: 5px;">
                <strong>Storage Tips:</strong><br/>
                • Models require ~8-12GB total<br/>
                • Generation needs ~2GB free space<br/>
                • Clear cache if generation fails<br/>
                • Click refresh to update storage status
            </div>
            """)
        
        # Add style preset guide
        with gr.Accordion("📖 Style Guide", open=False):
            gr.HTML("""
            <div style="padding: 10px; background-color: #f8f9fa; border-radius: 5px;">
                <h4>🎨 Style Presets Explained:</h4>
                <ul>
                    <li><strong>Ultra Clean:</strong> Pure black & white, minimal details, vector-like</li>
                    <li><strong>Clean:</strong> Simple outlines, minimal shading, easy to color</li>
                    <li><strong>Detailed:</strong> More intricate lines and fine details</li>
                    <li><strong>Bold:</strong> Thick outlines, high contrast</li>
                    <li><strong>Simple:</strong> Basic outlines, perfect for beginners</li>
                </ul>
                <h4>⚙️ Settings Tips:</h4>
                <ul>
                    <li><strong>Steps:</strong> 15-20 for speed, 25-35 for quality</li>
                    <li><strong>Guidance:</strong> 7-9 for natural, 10-12 for strict prompt following</li>
                    <li><strong>ControlNet:</strong> 1.0-1.2 balanced, 1.4+ for strict line following</li>
                    <li><strong>Resolution:</strong> 768 for speed, 1024+ for detail, 1536 max quality</li>
                    <li><strong>Depth Preprocessing:</strong> Experimental structural guidance (slower)</li>
                </ul>
            </div>
            """)
        
        # Event handlers
        generate_btn.click(
            generate_with_monitoring,
            inputs=[input_image, custom_prompt, style_preset, num_steps, 
                   guidance_scale, controlnet_scale, use_negative_prompt,
                   output_resolution, use_depth_preprocessor, lineart_model, apply_post_processing],
            outputs=[output_image, lineart_preview, depth_preview, status_output, storage_status]
        )
        
        clear_cache_btn.click(
            manual_clear_cache,
            outputs=[status_output, storage_status]
        )
        
        refresh_storage_btn.click(
            lambda: format_storage_status(),
            outputs=storage_status
        )
        
        # Note: Auto-update removed due to Gradio compatibility issues
        # Storage status updates manually via buttons
    
    return demo

# =============================================================================
# MAIN EXECUTION
# =============================================================================

if __name__ == "__main__":
    print("🚀 Starting ColorCraft SDXL - Storage Efficient LineArt")
    print(f"💾 Cache directory: {HF_CACHE_DIR}")
    
    # Initial storage check
    info = get_storage_info()
    print(f"📊 Initial storage: {info['free_gb']:.1f}GB free of {info['total_gb']:.1f}GB total")
    
    # Test imports
    print("\n🔍 Testing imports...")
    try:
        print(f"✅ PyTorch: {torch.__version__}")
        print(f"✅ Device: {torch.cuda.get_device_name() if torch.cuda.is_available() else 'CPU'}")
        
        import diffusers
        print(f"✅ Diffusers: {diffusers.__version__}")
        
        import transformers
        print(f"✅ Transformers: {transformers.__version__}")
        
        import huggingface_hub
        print(f"✅ HuggingFace Hub: {huggingface_hub.__version__}")
        
        if LineartDetector:
            print("✅ ControlNet-Aux: Available")
        else:
            print("❌ ControlNet-Aux: Not available")
            
    except Exception as e:
        print(f"❌ Import test failed: {e}")
    
    # Create and launch interface
    demo = create_storage_efficient_interface()
    demo.queue(max_size=3)  # Smaller queue to save memory
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        show_error=True
    )