Update app.py

app.py

@@ -6,311 +6,484 @@
 
 import gradio as gr
 import torch
-from PIL import Image
 import numpy as np
-from huggingface_hub import hf_hub_download
+from PIL import Image, ImageEnhance
 import cv2
-import os # Import os for path handling
-
-# --- Dependency Imports (Need to be installed via pip or manual clone) ---
-# BasicSR related imports (for SwinIR, EDSR, CodeFormer utilities)
-# Wrap imports in try/except to handle missing libraries
-BASESR_AVAILABLE = False
-try:
-    from basicsr.archs.swinir_arch import SwinIR as SwinIR_Arch
-    from basicsr.archs.edsr_arch import EDSR as EDSR_Arch
-    from basicsr.utils import img2tensor, tensor2img
-    BASESR_AVAILABLE = True
-except ImportError:
-    print("Warning: basicsr not found. SwinIR, EDSR, and CodeFormer (using basicsr utils) will not be available.")
-
-# RealESRGAN import
-REALESRGAN_AVAILABLE = False
-try:
-    from realesrgan import RealESRGAN
-    REALESRGAN_AVAILABLE = True
-except ImportError:
-    print("Warning: realesrgan not found. Real-ESRGAN-x4 will not be available.")
-
-# CodeFormer import (Often requires manual setup or specific installation)
-# We assume it's importable if basicsr is available AND the CodeFormer library itself
-# was somehow installed (e.g., via cloning and manual setup).
-# Given the previous error, direct pip install from git often fails.
-# We'll primarily rely on basicsr utilities, but a proper CodeFormer instance
-# might still require its dedicated installation.
-CODEFORMER_AVAILABLE = False
-if BASESR_AVAILABLE: # CodeFormer often depends on basicsr utilities
-     try:
-         # Attempting a common import path if CodeFormer is installed separately
-         # This might need adjustment based on your CodeFormer install method
-         from CodeFormer import CodeFormer # Adjust import based on your CodeFormer install path
-         CODEFORMER_AVAILABLE = True
-     except ImportError:
-         print("Warning: CodeFormer library not directly importable. CodeFormer model might not work correctly.")
-         # If basicsr is available, we might still list the model but it might fail later if CodeFormer class isn't there
-         pass # Allow BASESR_AVAILABLE to potentially enable the config entry
-
+import os
+import sys
+import subprocess
+import time
+from huggingface_hub import hf_hub_download
 
-# --- Model Configuration ---
-# Dictionary of available models and their configuration
-# format: "UI Name": {"repo_id": "hf_repo_id", "filename": "weight_filename", "type": "upscale" or "face", ...}
-MODEL_CONFIGS = {}
+# Create cache directory for models
+CACHE_DIR = os.path.join(os.path.expanduser("~"), ".cache", "image_enhancer")
+os.makedirs(CACHE_DIR, exist_ok=True)
 
-if REALESRGAN_AVAILABLE:
-     MODEL_CONFIGS["Real-ESRGAN-x4"] = {"repo_id": "RealESRGAN/RealESRGAN_x4plus", "filename": "RealESRGAN_x4plus.pth", "type": "upscale", "scale": 4}
+# Set up logging
+import logging
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
 
-if BASESR_AVAILABLE:
-     MODEL_CONFIGS["SwinIR-4x"] = {"repo_id": "SwinIR/SwinIR-Large", "filename": "SwinIR_4x.pth", "type": "upscale", "scale": 4}
-     MODEL_CONFIGS["EDSR-x4"] = {"repo_id": "EDSR/edsr_x4", "filename": "edsr_x4.pth", "type": "upscale", "scale": 4}
-     # Add CodeFormer config only if basicsr is available, and potentially CODEFORMER_AVAILABLE is True
-     # Even if CODEFORMER_AVAILABLE is False, listing it might rely on basicsr providing necessary components (less likely)
-     # Given installation issues, let's only add it if the library is actually importable.
-     if CODEFORMER_AVAILABLE:
-         MODEL_CONFIGS["CodeFormer (Face Enhancement)"] = {"repo_id": "CodeFormer/codeformer", "filename": "codeformer.pth", "type": "face"}
-     else:
-         print("CodeFormer (Face Enhancement) model will not be listed due to import issues.")
+# Install required packages at runtime for Hugging Face Spaces
+def install_dependencies():
+    logger.info("Checking and installing dependencies...")
+    
+    packages_to_install = [
+        "opencv-python",
+        "opencv-contrib-python",  # For dnn_superres module
+        "numpy",
+        "pillow",
+        "torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu",
+        "facexlib",
+        "basicsr",
+        "gfpgan",
+        "realesrgan"
+    ]
+    
+    for package in packages_to_install:
+        try:
+            logger.info(f"Installing {package}")
+            subprocess.check_call([sys.executable, "-m", "pip", "install", package])
+        except Exception as e:
+            logger.warning(f"Error installing {package}: {str(e)}")
+    
+    logger.info("Dependencies installation complete")
 
-# No need to filter anymore, configs are only added if available
+# Try to install dependencies on startup
+try:
+    install_dependencies()
+    time.sleep(2)  # Give some time for packages to settle
+except Exception as e:
+    logger.error(f"Failed to install dependencies: {str(e)}")
 
-# --- Model Loading Cache ---
-cached_model = {}
-cached_model_name = None
+# Check for GPU or CPU
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-print(f"Using device: {device}") # This shows which device PyTorch detects
-
-# Function to load the selected model
-def load_model(model_name: str):
-    global cached_model, cached_model_name
-
-    if model_name == cached_model_name and cached_model is not None:
-        print(f"Using cached model: {model_name}")
-        return cached_model, MODEL_CONFIGS[model_name]['type']
-
-    print(f"Loading model: {model_name}")
-    config = MODEL_CONFIGS.get(model_name)
-    if config is None:
-        # This case should ideally not happen if UI choices are filtered,
-        # but good for safety.
-        return None, f"Error: Model '{model_name}' not configured or dependencies missing."
+logger.info(f"Using device: {device}")
 
+# Dictionary of available models and their configuration
+MODEL_OPTIONS = {
+    "OpenCV Super Resolution": {
+        "type": "upscale",
+        "method": "opencv",
+        "scale": 4
+    },
+    "Real-ESRGAN-x4": {
+        "repo_id": "xinntao/Real-ESRGAN",
+        "filename": "RealESRGAN_x4plus.pth",
+        "type": "upscale",
+        "method": "realesrgan",
+        "scale": 4
+    },
+    "GFPGAN (Face Enhancement)": {
+        "repo_id": "TencentARC/GFPGAN",
+        "filename": "GFPGANv1.4.pth",
+        "type": "face",
+        "method": "gfpgan",
+        "scale": 1
+    },
+    "HDR Enhancement": {
+        "type": "hdr",
+        "method": "custom",
+        "scale": 1
+    }
+}
+
+# Cache for loaded models
+model_cache = {}
+
+# Function to load the selected model with robust fallbacks
+def load_model(model_name):
+    global model_cache
+    
+    # Return cached model if available
+    if model_name in model_cache:
+        logger.info(f"Using cached model: {model_name}")
+        return model_cache[model_name]
+    
+    logger.info(f"Loading model: {model_name}")
+    config = MODEL_OPTIONS.get(model_name)
+    if not config:
+        return None, f"Model {model_name} not found in configuration"
+    
+    model_type = config["type"]
+    
     try:
-        model_type = config['type']
-        model_path = hf_hub_download(repo_id=config['repo_id'], filename=config['filename'])
-
-        if model_name == "Real-ESRGAN-x4":
-            if not REALESRGAN_AVAILABLE: raise ImportError("realesrgan was not imported correctly.")
-            model = RealESRGAN(device, scale=config['scale'])
-            model.load_weights(model_path)
-        
-        elif model_name == "SwinIR-4x":
-            if not BASESR_AVAILABLE: raise ImportError("basicsr was not imported correctly.")
-            model = SwinIR_Arch(
-                upscale=config['scale'], in_chans=3, img_size=64, window_size=8,
-                compress_ratio= -1, dilate_basis=-1, res_range=-1, attn_type='linear'
+        # OpenCV based models (always available as fallback)
+        if config["method"] == "opencv":
+            logger.info("Loading OpenCV Super Resolution model")
+            sr = cv2.dnn_superres.DnnSuperResImpl_create()
+            
+            # Use EDSR as default model
+            model_path = hf_hub_download(
+                repo_id="eugenesiow/edsr",
+                filename="EDSR_x4.pb",
+                cache_dir=CACHE_DIR
             )
-            pretrained_dict = torch.load(model_path, map_location=device)
-            model.load_state_dict(pretrained_dict, strict=True)
-            model.eval()
-            model.to(device)
-        
-        elif model_name == "EDSR-x4":
-            if not BASESR_AVAILABLE: raise ImportError("basicsr was not imported correctly.")
-            model = EDSR_Arch(num_feat=64, num_block=16, upscale=config['scale'])
-            pretrained_dict = torch.load(model_path, map_location=device)
-            model.load_state_dict(pretrained_dict, strict=True)
-            model.eval()
-            model.to(device)
             
-        elif model_name == "CodeFormer (Face Enhancement)":
-             if not CODEFORMER_AVAILABLE:
-                  # This check is redundant if config is only added when available,
-                  # but good practice.
-                  raise ImportError("CodeFormer library was not imported correctly.")
-
-             # Ensure model_path is correct (downloaded via hf_hub_download)
-             # CodeFormer loading often needs specific handling for checkpoints (params_ema)
-             # This part is sensitive to the exact CodeFormer version/structure
-             # Assuming a similar loading pattern to basicsr models:
-             if CODEFORMER_AVAILABLE:
-                model = CodeFormer(dim_embd=512, codebook_size=1024, n_head=8, n_layer=9) # Basic instantiation
-                pretrained_dict = torch.load(model_path, map_location=device)['params_ema']
-                keys = list(pretrained_dict.keys())
-                if keys and keys[0].startswith('module.'):
-                    pretrained_dict = {k[7:]: v for k, v in pretrained_dict.items()}
-                model.load_state_dict(pretrained_dict, strict=True)
-                model.eval()
-                model.to(device)
-             else:
-                 # Fallback check, should not be reached if config is filtered
-                 raise ImportError("CodeFormer library not available.")
-
-
+            sr.readModel(model_path)
+            sr.setModel("edsr", 4)
+            
+            # Set backend to cuda if available
+            if torch.cuda.is_available():
+                sr.setPreferableBackend(cv2.dnn.DNN_BACKEND_CUDA)
+                sr.setPreferableTarget(cv2.dnn.DNN_TARGET_CUDA)
+            
+            model_cache[model_name] = (sr, model_type)
+            return sr, model_type
+        
+        # Real-ESRGAN models
+        elif config["method"] == "realesrgan":
+            try:
+                from realesrgan import RealESRGAN
+                logger.info("Loading Real-ESRGAN model")
+                
+                model_path = hf_hub_download(
+                    repo_id=config["repo_id"],
+                    filename=config["filename"],
+                    cache_dir=CACHE_DIR
+                )
+                
+                model = RealESRGAN(device, scale=config["scale"])
+                model.load_weights(model_path)
+                
+                model_cache[model_name] = (model, model_type)
+                return model, model_type
+            except ImportError:
+                logger.warning("RealESRGAN not available, falling back to OpenCV")
+                return load_model("OpenCV Super Resolution")
+        
+        # GFPGAN for face enhancement
+        elif config["method"] == "gfpgan":
+            try:
+                from gfpgan import GFPGANer
+                logger.info("Loading GFPGAN model")
+                
+                model_path = hf_hub_download(
+                    repo_id=config["repo_id"],
+                    filename=config["filename"],
+                    cache_dir=CACHE_DIR
+                )
+                
+                face_enhancer = GFPGANer(
+                    model_path=model_path,
+                    upscale=config["scale"],
+                    arch='clean',
+                    channel_multiplier=2,
+                    bg_upsampler=None
+                )
+                
+                model_cache[model_name] = (face_enhancer, model_type)
+                return face_enhancer, model_type
+            except ImportError:
+                logger.warning("GFPGAN not available, falling back to OpenCV")
+                return load_model("OpenCV Super Resolution")
+        
+        # HDR Enhancement (custom implementation)
+        elif config["method"] == "custom":
+            # No model to load for custom HDR
+            model_cache[model_name] = (None, model_type)
+            return None, model_type
+        
         else:
-            # This should not be reached with filtered configs
-            raise ValueError(f"Configuration missing or invalid for model: {model_name}")
-
-        # Cache the loaded model
-        cached_model = model
-        cached_model_name = model_name
-
-        return model, model_type
-    
-    except ImportError as ie:
-        # This catches errors if the library was *somehow* listed as available
-        # but then failed on a deeper import within load_model
-        print(f"Dependency check failed during load for {model_name}: {ie}")
-        # Clear cache on error
-        cached_model = None
-        cached_model_name = None
-        return None, f"Error: Dependency not fully available - {ie}. Model cannot be loaded."
+            raise ValueError(f"Unknown model method: {config['method']}")
+    
     except Exception as e:
-        print(f"Error loading model {model_name}: {e}")
+        logger.error(f"Error loading model {model_name}: {str(e)}")
         import traceback
-        traceback.print_exc() # Print full traceback for debugging
-        # Clear cache on error
-        cached_model = None
-        cached_model_name = None
-        return None, f"Error loading model: {str(e)}"
+        traceback.print_exc()
+        
+        # Always provide a fallback method
+        if model_name != "OpenCV Super Resolution":
+            logger.info("Falling back to OpenCV Super Resolution")
+            return load_model("OpenCV Super Resolution")
+        else:
+            return None, f"Failed to load model: {str(e)}"
 
-# Function to preprocess image (PIL RGB to OpenCV BGR numpy)
-def preprocess_image(image: Image.Image) -> np.ndarray:
-    img = np.array(image)
-    img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
-    return img
+# Function to preprocess image for processing
+def preprocess_image(image):
+    """Convert PIL image to numpy array for processing"""
+    if image is None:
+        return None
+    
+    if isinstance(image, Image.Image):
+        # Convert PIL image to numpy array
+        img = np.array(image)
+    else:
+        img = image
+    
+    # Handle grayscale images by converting to RGB
+    if len(img.shape) == 2:
+        img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)
+    
+    # Handle RGBA images by removing alpha channel
+    if img.shape[2] == 4:
+        img = img[:, :, :3]
+    
+    # Convert RGB to BGR for OpenCV processing
+    img_bgr = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
+    
+    return img_bgr
 
-# Function to postprocess image (OpenCV BGR numpy to PIL RGB)
-def postprocess_image(img: np.ndarray) -> Image.Image:
-    if img.dtype != np.uint8:
-         img = np.clip(img, 0, 255).astype(np.uint8)
-    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
-    return Image.fromarray(img)
+# Function to postprocess image for display
+def postprocess_image(img_bgr):
+    """Convert processed BGR image back to RGB PIL image"""
+    if img_bgr is None:
+        return None
+    
+    # Ensure image is uint8
+    if img_bgr.dtype != np.uint8:
+        img_bgr = np.clip(img_bgr, 0, 255).astype(np.uint8)
+    
+    # Convert BGR to RGB for PIL
+    img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
+    
+    return Image.fromarray(img_rgb)
 
-# Main processing function
-def enhance_image(image: Image.Image, model_name: str):
-    if image is None:
-        # Return tuple of (message, image)
-        return "Please upload an image.", None
+# HDR enhancement function
+def enhance_hdr(img_bgr, strength=1.0):
+    """Custom HDR enhancement using OpenCV"""
+    # Convert BGR to RGB
+    img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
+    
+    # Convert to float32 for processing
+    img_float = img_rgb.astype(np.float32) / 255.0
+    
+    # Convert to LAB color space for better contrast enhancement
+    img_lab = cv2.cvtColor(img_float, cv2.COLOR_RGB2LAB)
+    l, a, b = cv2.split(img_lab)
     
-    status_message = f"Processing image with {model_name}..."
+    # Apply CLAHE (Contrast Limited Adaptive Histogram Equalization)
+    clahe = cv2.createCLAHE(clipLimit=3.0 * strength, tileGridSize=(8, 8))
+    l_enhanced = clahe.apply(np.clip(l * 255, 0, 255).astype(np.uint8)) / 255.0
     
-    # Load the selected model and its type
-    model, model_info = load_model(model_name)
+    # Blend original and enhanced L channel
+    l = l * (1 - strength) + l_enhanced * strength
     
-    if model is None:
-        # model_info contains the error message from load_model
-        return model_info, None
+    # Merge channels
+    img_lab_enhanced = cv2.merge([l, a, b])
+    img_rgb_enhanced = cv2.cvtColor(img_lab_enhanced, cv2.COLOR_LAB2RGB)
     
-    model_type = model_info # model_info is the type string ('upscale' or 'face')
+    # Add vibrance (increase saturation of low-saturation areas)
+    hsv = cv2.cvtColor(img_rgb_enhanced, cv2.COLOR_RGB2HSV)
+    h, s, v = cv2.split(hsv)
+    
+    # Increase saturation adaptively (more for lower saturation, less for already saturated pixels)
+    saturation_factor = 0.3 * strength
+    s_enhanced = np.clip(s * (1 + saturation_factor * (1 - s)), 0, 1)
+    
+    # Increase brightness slightly
+    v_enhanced = np.clip(v * (1 + 0.1 * strength), 0, 1)
+    
+    # Merge HSV channels and convert back to RGB
+    hsv_enhanced = cv2.merge([h, s_enhanced, v_enhanced])
+    img_enhanced = cv2.cvtColor(hsv_enhanced, cv2.COLOR_HSV2RGB)
+    
+    # Apply subtle detail enhancement
+    kernel_size = 5
+    blur = cv2.GaussianBlur(img_enhanced, (kernel_size, kernel_size), 0)
+    detail = img_enhanced - blur
+    img_enhanced = np.clip(img_enhanced + detail * (0.5 * strength), 0, 1)
+    
+    # Convert back to BGR for output
+    img_enhanced = (img_enhanced * 255).astype(np.uint8)
+    img_bgr_enhanced = cv2.cvtColor(img_enhanced, cv2.COLOR_RGB2BGR)
+    
+    return img_bgr_enhanced
 
+# Main image enhancement function
+def enhance_image(image, model_name, strength=1.0, denoise=0.0, sharpen=0.0):
+    """Enhance image using selected model with additional processing options"""
+    if image is None:
+        return "Please upload an image.", None
+    
     try:
-        # Preprocess the image (PIL RGB -> OpenCV BGR)
-        img_np_bgr = preprocess_image(image)
+        # Load model
+        model, model_type = load_model(model_name)
+        if isinstance(model_type, str) and model_type.startswith("Failed"):
+            return model_type, None
         
-        # Process based on model type and specific model implementation
+        # Preprocess image
+        img_bgr = preprocess_image(image)
+        if img_bgr is None:
+            return "Failed to process image", None
+        
+        # Apply denoising if requested
+        if denoise > 0:
+            strength_value = int(denoise * 10)
+            img_bgr = cv2.fastNlMeansDenoisingColored(
+                img_bgr, None, 
+                h=strength_value, 
+                hColor=strength_value, 
+                templateWindowSize=7, 
+                searchWindowSize=21
+            )
+        
+        # Process based on model type
         if model_type == "upscale":
-            print(f"Applying {model_name} upscaling...")
-            if model_name == "Real-ESRGAN-x4":
-                if not REALESRGAN_AVAILABLE: raise ImportError("Real-ESRGAN library not available.")
-                output_np_bgr = model.predict(img_np_bgr)
-            elif model_name in ["SwinIR-4x", "EDSR-x4"]:
-                if not BASESR_AVAILABLE: raise ImportError(f"basicsr library not available for {model_name}")
-                # HWC BGR uint8 -> CHW RGB float (0-1) -> send to device
-                img_tensor = img2tensor(img_np_bgr.astype(np.float32) / 255., bgr2rgb=True, float32=True).unsqueeze(0).to(device)
-
-                with torch.no_grad():
-                    output_tensor = model(img_tensor)
-
-                # CHW RGB float (0-1) -> HWC RGB uint8 -> Convert to BGR
-                output_np_rgb = tensor2img(output_tensor, rgb2bgr=False, min_max=(0, 1))
-                output_np_bgr = cv2.cvtColor(output_np_rgb, cv2.COLOR_RGB2BGR)
-                
+            logger.info(f"Upscaling image with {model_name}")
+            
+            if model_name == "OpenCV Super Resolution":
+                # OpenCV super resolution
+                output_bgr = model.upsample(img_bgr)
+            
+            elif model_name == "Real-ESRGAN-x4":
+                # Real-ESRGAN upscaling
+                try:
+                    output_bgr = model.predict(img_bgr)
+                except Exception as e:
+                    logger.error(f"Error with Real-ESRGAN: {str(e)}")
+                    # Fall back to OpenCV
+                    fallback_model, _ = load_model("OpenCV Super Resolution")
+                    output_bgr = fallback_model.upsample(img_bgr)
+            
             else:
-                raise ValueError(f"Unknown upscale model type configuration: {model_name}")
-                
-            status_message = f"Image upscaled successfully with {model_name}!"
-
+                # Default to OpenCV upscaling
+                sr = cv2.dnn_superres.DnnSuperResImpl_create()
+                sr.upsample(img_bgr)
+        
         elif model_type == "face":
-            print(f"Applying {model_name} face enhancement...")
-            if model_name == "CodeFormer (Face Enhancement)":
-                if not CODEFORMER_AVAILABLE: raise ImportError("CodeFormer library not available.")
-                # Ensure model is on correct device
-                if next(model.parameters()).device != device:
-                    model.to(device)
-                
-                # Call the enhance method (adjust parameters w, adain as needed)
-                # CodeFormer enhance typically returns a tuple, first element is image
-                output_np_bgr = model.enhance(img_np_bgr, w=0.5, adain=True)[0] # Example call
+            logger.info(f"Enhancing face with {model_name}")
+            
+            if model_name == "GFPGAN (Face Enhancement)":
+                try:
+                    # GFPGAN returns (cropped_faces, restored_faces, restored_img)
+                    _, _, output_bgr = model.enhance(
+                        img_bgr, 
+                        has_aligned=False, 
+                        only_center_face=False, 
+                        paste_back=True
+                    )
+                except Exception as e:
+                    logger.error(f"Error with GFPGAN: {str(e)}")
+                    # Fall back to basic upscaling
+                    fallback_model, _ = load_model("OpenCV Super Resolution")
+                    output_bgr = fallback_model.upsample(img_bgr)
             else:
-                 raise ValueError(f"Unknown face enhancement model type configuration: {model_name}")
-
-            status_message = f"Face enhancement applied successfully with {model_name}!"
-
-        # Postprocess the output image (OpenCV BGR -> PIL RGB)
-        enhanced_image = postprocess_image(output_np_bgr)
+                # Default upscaling for face if specific model fails
+                sr = cv2.dnn_superres.DnnSuperResImpl_create()
+                output_bgr = sr.upsample(img_bgr)
+        
+        elif model_type == "hdr":
+            logger.info("Applying HDR enhancement")
+            # Custom HDR enhancement
+            output_bgr = enhance_hdr(img_bgr, strength=strength)
         
-        return status_message, enhanced_image
+        else:
+            return f"Unknown model type: {model_type}", None
+        
+        # Apply sharpening if requested
+        if sharpen > 0:
+            sharpen_kernel = np.array([
+                [-1, -1, -1],
+                [-1, 9 + sharpen * 2, -1],
+                [-1, -1, -1]
+            ])
+            output_bgr = cv2.filter2D(output_bgr, -1, sharpen_kernel)
+        
+        # Post-process and return image
+        enhanced_image = postprocess_image(output_bgr)
+        
+        return "Image enhanced successfully!", enhanced_image
     
-    except ImportError as ie:
-         # This catches errors if the library was imported initially but
-         # failed later when its functions/classes were called.
-         print(f"Error processing image due to missing dependency call: {ie}")
-         return f"Error processing image: Required library function not found - {ie}", None
     except Exception as e:
-        print(f"Error during processing: {e}")
+        logger.error(f"Error processing image: {str(e)}")
         import traceback
-        traceback.print_exc() # Print full traceback for debugging
-        return f"Error processing image: {str(e)}", None
+        traceback.print_exc()
+        return f"Error: {str(e)}", None
 
 # Gradio interface
 with gr.Blocks(title="Image Upscale & Enhancement - By FebryEnsz") as demo:
     gr.Markdown(
         """
-        # Image Upscale & Enhancement
-        **By FebryEnsz**
+        # 🖼️ Image Upscale & Enhancement
+        ### By FebryEnsz
         
-        Upload an image and select a model to enhance it. Choose from multiple models for upscaling (to make it 'HD' or higher resolution) or face enhancement (to improve facial details and focus).
+        Upload an image and enhance it with AI-powered upscaling and enhancement.
         
-        **Note:** This app requires specific Python libraries (`torch`, `basicsr`, `realesrgan`, `CodeFormer`) to be installed for all models to be available. If a model option is missing, its required library was not successfully installed or found during startup. Check the Space build logs for installation errors.
+        **Features:**
+        - Super-resolution upscaling (4x)
+        - Face enhancement for portraits
+        - HDR enhancement for better contrast and details
         """
     )
     
-    # Filter available choices based on loaded configs
-    available_models = list(MODEL_CONFIGS.keys())
-    
     with gr.Row():
-        with gr.Column():
+        with gr.Column(scale=1):
             image_input = gr.Image(label="Upload Image", type="pil")
             
-            if not available_models:
-                 model_choice = gr.Textbox(label="Select Model", value="No models available. Check build logs for dependency errors.", interactive=False)
-                 enhance_button = gr.Button("Enhance Image", interactive=False)
-                 print("No models are available because dependencies are missing.")
-            else:
+            with gr.Box():
+                gr.Markdown("### Enhancement Options")
                 model_choice = gr.Dropdown(
-                    choices=available_models,
-                    label="Select Model",
-                    value=available_models[0] # Default to the first available model
+                    choices=list(MODEL_OPTIONS.keys()),
+                    label="Model Selection",
+                    value="OpenCV Super Resolution"
                 )
-                # Removed scale_slider
-                enhance_button = gr.Button("Enhance Image")
+                
+                with gr.Accordion("Advanced Settings", open=False):
+                    strength_slider = gr.Slider(
+                        minimum=0.1,
+                        maximum=1.0,
+                        step=0.1,
+                        label="Enhancement Strength",
+                        value=0.8,
+                    )
+                    
+                    denoise_slider = gr.Slider(
+                        minimum=0.0,
+                        maximum=1.0,
+                        step=0.1,
+                        label="Noise Reduction",
+                        value=0.0,
+                    )
+                    
+                    sharpen_slider = gr.Slider(
+                        minimum=0.0,
+                        maximum=1.0,
+                        step=0.1,
+                        label="Sharpening",
+                        value=0.0,
+                    )
             
-        with gr.Column():
-            output_text = gr.Textbox(label="Status", max_lines=3)
+            enhance_button = gr.Button("✨ Enhance Image", variant="primary")
+            
+        with gr.Column(scale=1):
+            output_text = gr.Textbox(label="Status")
             output_image = gr.Image(label="Enhanced Image")
     
-    # Connect the button to the processing function
-    if available_models: # Only connect if models are available
-        enhance_button.click(
-            fn=enhance_image,
-            inputs=[image_input, model_choice],
-            outputs=[output_text, output_image]
-        )
+    # Handle model change to update UI
+    def on_model_change(model_name):
+        model_config = MODEL_OPTIONS.get(model_name, {})
+        model_type = model_config.get("type", "")
+        
+        # Update UI based on model type
+        if model_type == "hdr":
+            return gr.update(visible=True, label="HDR Intensity")
+        elif model_type == "face":
+            return gr.update(visible=True, label="Enhancement Strength")
+        else:
+            return gr.update(visible=True, label="Enhancement Strength")
+            
+    model_choice.change(on_model_change, inputs=[model_choice], outputs=[strength_slider])
+    
+    # Connect button to function
+    enhance_button.click(
+        fn=enhance_image,
+        inputs=[image_input, model_choice, strength_slider, denoise_slider, sharpen_slider],
+        outputs=[output_text, output_image]
+    )
+    
+    # Footer information
+    gr.Markdown(
+        """
+        ### Tips
+        - For best results with face enhancement, ensure faces are clearly visible
+        - HDR enhancement works best with images that have both bright and dark areas
+        - For noisy images, try increasing the noise reduction slider
+        
+        ---
+        Version 2.0 | Running on: """ + ("GPU 🚀" if torch.cuda.is_available() else "CPU ⚙️")
+    )
 
-# Launch the Gradio app
+# Launch the app
 if __name__ == "__main__":
-    # Set torch backend for potentially better performance on some systems
-    # Removed MPS fallback for simplicity unless specifically needed and tested
-    # if torch.backends.mps.is_available():
-    #     os.environ['PYTORCH_ENABLE_MPS_FALLBACK'] = '1'
-
     demo.launch()