Update app.py

app.py

@@ -27,18 +27,24 @@ logger = logging.getLogger(__name__)
 # 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",
+        "torch torchvision torchaudio", # Let pip handle the specific wheels
+        "facexlib", # Dependency for GFPGAN
+        "basicsr",  # Dependency for RealESRGAN/GFPGAN
         "gfpgan",
-        "realesrgan"
+        "realesrgan",
+        "huggingface_hub" # Ensure hf_hub_download is available
     ]
+
+    # Use a standard index-url or let pip find the best one
+    # Forcing CPU might prevent GPU usage if available
+    # Let's try without forcing CPU first, Hugging Face Spaces often handles this.
+    # If you specifically need CPU only, you might re-add --index-url https://download.pytorch.org/whl/cpu
     
     for package in packages_to_install:
         try:
@@ -46,17 +52,34 @@ def install_dependencies():
             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")
 
 # Try to install dependencies on startup
 try:
     install_dependencies()
+    # Import libraries AFTER installation
+    import cv2
+    import torch
+    import numpy as np
+    from PIL import Image, ImageEnhance
+    from huggingface_hub import hf_hub_download
+    try:
+        from realesrgan import RealESRGAN
+    except ImportError:
+        logger.warning("RealESRGAN import failed after installation attempt.")
+        RealESRGAN = None # Set to None if import fails
+    try:
+        from gfpgan import GFPGANer
+    except ImportError:
+         logger.warning("GFPGANer import failed after installation attempt.")
+         GFPGANer = None # Set to None if import fails
+
     time.sleep(2)  # Give some time for packages to settle
 except Exception as e:
-    logger.error(f"Failed to install dependencies: {str(e)}")
+    logger.error(f"Failed to install dependencies or import libraries: {str(e)}")
 
-# Check for GPU or CPU
+# Check for GPU or CPU AFTER torch is potentially installed
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 logger.info(f"Using device: {device}")
 
@@ -79,7 +102,7 @@ MODEL_OPTIONS = {
         "filename": "GFPGANv1.4.pth",
         "type": "face",
         "method": "gfpgan",
-        "scale": 1
+        "scale": 1 # GFPGAN is primarily for face restoration, upscaling is secondary/handled by bg_upsampler
     },
     "HDR Enhancement": {
         "type": "hdr",
@@ -94,134 +117,156 @@ 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:
-        # OpenCV based models (always available as fallback)
+        # OpenCV based models (always available as fallback if opencv-contrib is installed)
         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
-            )
-            
-            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
-        
+            try:
+                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
+                )
+
+                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
+            except Exception as e:
+                 logger.error(f"Error loading OpenCV SR model: {str(e)}")
+                 # Fallback to None if OpenCV SR fails
+                 return None, f"Failed to load OpenCV SR model: {str(e)}"
+
+
         # Real-ESRGAN models
         elif config["method"] == "realesrgan":
+            if RealESRGAN is None:
+                 logger.warning("RealESRGAN class not found, falling back to OpenCV SR.")
+                 return load_model("OpenCV Super Resolution") # Fallback
+
             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
                 )
-                
+
+                # Initialize RealESRGAN with the correct device
                 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")
-        
+            except Exception as e:
+                logger.error(f"Error loading Real-ESRGAN model: {str(e)}")
+                logger.warning("Falling back to OpenCV Super Resolution")
+                return load_model("OpenCV Super Resolution") # Fallback
+
         # GFPGAN for face enhancement
         elif config["method"] == "gfpgan":
+            if GFPGANer is None:
+                 logger.warning("GFPGANer class not found, falling back to OpenCV SR.")
+                 return load_model("OpenCV Super Resolution") # Fallback
+
             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
                 )
-                
+
+                # GFPGANer initialization
+                # Note: If you want background upsampling with GFPGAN, you need to initialize bg_upsampler
+                # e.g., bg_upsampler=RealESRGANer(model_path='...', model_name='RealESRGAN_x4plus.pth', ...)
+                # For simplicity and focusing on face, bg_upsampler=None is used here.
                 face_enhancer = GFPGANer(
                     model_path=model_path,
-                    upscale=config["scale"],
-                    arch='clean',
+                    upscale=config["scale"], # This upscale might be ignored if paste_back is True and no bg_upsampler
+                    arch='clean', # Use 'clean' arch for GFPGANv1.4
                     channel_multiplier=2,
-                    bg_upsampler=None
+                    bg_upsampler=None # No background upsampling
                 )
-                
+
                 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")
-        
+            except Exception as e:
+                logger.error(f"Error loading GFPGAN model: {str(e)}")
+                logger.warning("Falling back to OpenCV Super Resolution")
+                return load_model("OpenCV Super Resolution") # Fallback
+
         # 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:
-            raise ValueError(f"Unknown model method: {config['method']}")
-    
+            return None, f"Unknown model method: {config['method']}"
+
     except Exception as e:
-        logger.error(f"Error loading model {model_name}: {str(e)}")
+        logger.error(f"Unexpected error during model loading for {model_name}: {str(e)}")
         import traceback
         traceback.print_exc()
-        
-        # Always provide a fallback method
+
+        # Always provide a fallback method if the desired one completely fails
         if model_name != "OpenCV Super Resolution":
-            logger.info("Falling back to OpenCV Super Resolution")
+            logger.info("Critical error loading model, falling back to OpenCV Super Resolution")
             return load_model("OpenCV Super Resolution")
         else:
-            return None, f"Failed to load model: {str(e)}"
+            # If OpenCV SR itself fails, something is fundamentally wrong
+            return None, f"Failed to load any model, including fallback: {str(e)}"
+
 
 # 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:
+        # Assume it's already a numpy array (e.g., from Gradio internal handling)
         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 for display
@@ -229,14 +274,16 @@ 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)
-    
+        # Ensure the range is correct before casting
+        img_bgr = np.clip(img_bgr, 0, 255)
+        img_bgr = img_bgr.astype(np.uint8)
+
     # Convert BGR to RGB for PIL
     img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
-    
+
     return Image.fromarray(img_rgb)
 
 # HDR enhancement function
@@ -244,152 +291,199 @@ 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
+
+    # Convert to float32 for processing, range [0, 1]
     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)
-    
-    # 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
-    
-    # Blend original and enhanced L channel
-    l = l * (1 - strength) + l_enhanced * strength
-    
-    # Merge channels
-    img_lab_enhanced = cv2.merge([l, a, b])
-    img_rgb_enhanced = cv2.cvtColor(img_lab_enhanced, cv2.COLOR_LAB2RGB)
-    
-    # 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
+
+    # --- Exposure Fusion based approach (more robust) ---
+    try:
+        # Estimate camera response function (merge_mertens is more robust)
+        merge_mertens = cv2.createMergeMertens(contrast_weight=1.0, saturation_weight=1.0, exposure_weight=0.0)
+        # You'd ideally need multiple exposures for true HDR merge.
+        # Simulating this by generating slightly adjusted exposures might not be ideal.
+        # Let's use a simpler single-image tone mapping or CLAHE on different channels.
+        
+        # Using CLAHE on L channel (from LAB) and potentially V channel (from HSV)
+        img_lab = cv2.cvtColor(img_float, cv2.COLOR_RGB2LAB)
+        l, a, b = cv2.split(img_lab)
+
+        # Apply CLAHE to L channel
+        # ClipLimit proportional to strength
+        clahe_l = cv2.createCLAHE(clipLimit=max(1.0, 5.0 * strength), tileGridSize=(8, 8))
+        # CLAHE works on uint8, so scale L channel
+        l_uint8 = np.clip(l * 255.0, 0, 255).astype(np.uint8)
+        l_enhanced_uint8 = clahe_l.apply(l_uint8)
+        l_enhanced = l_enhanced_uint8.astype(np.float32) / 255.0
+
+        # Blend original and enhanced L channel based on strength
+        l_final = l * (1 - strength) + l_enhanced * strength
+
+        # Merge LAB and convert back to RGB
+        img_lab_enhanced = cv2.merge([l_final, a, b])
+        img_rgb_enhanced = cv2.cvtColor(img_lab_enhanced, cv2.COLOR_LAB2RGB)
+
+        # --- Additional Enhancements (optional, based on strength) ---
+        # Vibrance/Saturation adjustment (HSV)
+        img_hsv = cv2.cvtColor(img_rgb_enhanced, cv2.COLOR_RGB2HSV)
+        h, s, v = cv2.split(img_hsv)
+
+        # Increase saturation, more for less saturated pixels
+        saturation_factor = 0.4 * strength # Adjust factor as needed
+        s_enhanced = np.clip(s + (s * saturation_factor * (1 - s)), 0, 1)
+
+        # Slight brightness adjustment
+        brightness_factor = 0.1 * strength
+        v_enhanced = np.clip(v + (v * brightness_factor), 0, 1)
+
+
+        # Merge HSV and convert back to RGB
+        img_rgb_enhanced_hsv = cv2.cvtColor(cv2.merge([h, s_enhanced, v_enhanced]), cv2.COLOR_HSV2RGB)
+
+        # --- Subtle Detail Enhancement (Unsharp Masking effect) ---
+        # Convert back to uint8 for blurring
+        img_uint8_detail = (np.clip(img_rgb_enhanced_hsv, 0, 1) * 255).astype(np.uint8)
+        blur = cv2.GaussianBlur(img_uint8_detail, (0, 0), 5) # Kernel size 5, sigma automatically calculated
+        # Convert blur back to float for calculation
+        blur_float = blur.astype(np.float32) / 255.0
+
+        detail = img_rgb_enhanced_hsv - blur_float
+        # Add detail back, scaled by strength
+        img_final_float = np.clip(img_rgb_enhanced_hsv + detail * (0.8 * strength), 0, 1)
+
+        # Convert back to BGR (uint8) for output
+        img_bgr_enhanced = (img_final_float * 255).astype(np.uint8)
+        img_bgr_enhanced = cv2.cvtColor(img_bgr_enhanced, cv2.COLOR_RGB2BGR)
+
+        return img_bgr_enhanced
+
+    except Exception as e:
+        logger.error(f"Error during HDR enhancement: {str(e)}")
+        # Return original image if enhancement fails
+        return img_bgr
+
 
 # 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:
         # Load model
-        model, model_type = load_model(model_name)
-        if isinstance(model_type, str) and model_type.startswith("Failed"):
-            return model_type, None
-        
+        model, model_info = load_model(model_name)
+        if isinstance(model_info, str) and model_info.startswith("Failed"):
+             # If loading fails, model is None, info is the error message
+             return model_info, None
+
+        model_type = model_info # model_info now holds the model type string
+
         # 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)
+            logger.info(f"Applying denoising with strength {denoise}")
+            # Adjust h and hColor based on denoise slider
+            # Recommended range for h is 10 for color images (adjust based on noise level)
+            h_val = int(denoise * 20 + 10) # Map 0-1 slider to approx 10-30 h value
             img_bgr = cv2.fastNlMeansDenoisingColored(
-                img_bgr, None, 
-                h=strength_value, 
-                hColor=strength_value, 
-                templateWindowSize=7, 
+                img_bgr, None,
+                h=h_val,
+                hColor=h_val,
+                templateWindowSize=7,
                 searchWindowSize=21
             )
-        
+
+        output_bgr = img_bgr # Initialize output with potentially denoised image
+
         # Process based on model type
         if model_type == "upscale":
+            if model is None:
+                 return f"Upscaling model '{model_name}' is not loaded or available.", None
             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:
-                # Default to OpenCV upscaling
-                sr = cv2.dnn_superres.DnnSuperResImpl_create()
-                sr.upsample(img_bgr)
-        
+                # Real-ESRGAN model object has a 'predict' method
+                output_bgr = model.predict(img_bgr)
+
+            # No else needed, as load_model should handle fallbacks
+
         elif model_type == "face":
+            if model is None:
+                 return f"Face enhancement model '{model_name}' is not loaded or available.", None
             logger.info(f"Enhancing face with {model_name}")
-            
+
             if model_name == "GFPGAN (Face Enhancement)":
+                # GFPGAN model object has an 'enhance' method
                 try:
                     # GFPGAN returns (cropped_faces, restored_faces, restored_img)
+                    # restored_img is the pasted-back result
                     _, _, output_bgr = model.enhance(
-                        img_bgr, 
-                        has_aligned=False, 
-                        only_center_face=False, 
+                        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:
-                # Default upscaling for face if specific model fails
-                sr = cv2.dnn_superres.DnnSuperResImpl_create()
-                output_bgr = sr.upsample(img_bgr)
-        
+                    logger.error(f"Error enhancing face with GFPGAN: {str(e)}")
+                    # If GFPGAN fails, don't just return, try basic upscaling or original
+                    # For now, let's just log and return original or denoised image
+                    output_bgr = img_bgr # Keep the denoised (or original) image
+                    return f"Error applying GFPGAN: {str(e)}. Returning base image.", postprocess_image(output_bgr)
+
         elif model_type == "hdr":
-            logger.info("Applying HDR enhancement")
-            # Custom HDR enhancement
+            # HDR enhancement doesn't use an external model object, it's a function call
+            logger.info(f"Applying HDR enhancement with strength {strength}")
             output_bgr = enhance_hdr(img_bgr, strength=strength)
-        
+
         else:
-            return f"Unknown model type: {model_type}", None
-        
-        # Apply sharpening if requested
+            # Should not happen if MODEL_OPTIONS is correct
+            return f"Unknown model type for processing: {model_type}", None
+
+
+        # Apply sharpening if requested (apply to the output of the main process)
         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)
-        
+            logger.info(f"Applying sharpening with strength {sharpen}")
+            # Simple unsharp mask effect
+            kernel = np.array([
+                [0, -1, 0],
+                [-1, 5, -1],
+                [0, -1, 0]
+            ], np.float32)
+            # We can adjust the strength by blending original and sharpened, or using a kernel with varying center weight
+            # A simpler approach is blending:
+            sharpened_img = cv2.filter2D(output_bgr, -1, kernel)
+            # Blend original output and sharpened output
+            output_bgr = cv2.addWeighted(output_bgr, 1.0 - sharpen, sharpened_img, sharpen, 0)
+
+
         # Post-process and return image
         enhanced_image = postprocess_image(output_bgr)
-        
+
         return "Image enhanced successfully!", enhanced_image
-    
+
     except Exception as e:
-        logger.error(f"Error processing image: {str(e)}")
+        logger.error(f"An error occurred during image processing: {str(e)}")
         import traceback
         traceback.print_exc()
-        return f"Error: {str(e)}", None
+        # Attempt to return original image on error
+        if image is not None:
+             try:
+                 original_img_pil = Image.fromarray(cv2.cvtColor(preprocess_image(image), cv2.COLOR_BGR2RGB))
+                 return f"Processing failed: {str(e)}. Returning original image.", original_img_pil
+             except Exception as post_e:
+                 logger.error(f"Failed to return original image after error: {str(post_e)}")
+                 return f"Processing failed: {str(e)}. Could not return image.", None
+        else:
+             return f"Processing failed: {str(e)}. No image provided.", None
+
 
 # Gradio interface
 with gr.Blocks(title="Image Upscale & Enhancement - By FebryEnsz") as demo:
@@ -397,93 +491,105 @@ with gr.Blocks(title="Image Upscale & Enhancement - By FebryEnsz") as demo:
         """
         # 🖼️ Image Upscale & Enhancement
         ### By FebryEnsz
-        
+
         Upload an image and enhance it with AI-powered upscaling and enhancement.
-        
+
         **Features:**
-        - Super-resolution upscaling (4x)
-        - Face enhancement for portraits
+        - Super-resolution upscaling (4x) using Real-ESRGAN or OpenCV
+        - Face enhancement for portraits using GFPGAN
         - HDR enhancement for better contrast and details
+        - Additional Denoise and Sharpen options
         """
     )
-    
+
     with gr.Row():
         with gr.Column(scale=1):
-            image_input = gr.Image(label="Upload Image", type="pil")
+            image_input = gr.Image(label="Upload Image", type="pil", image_mode="RGB") # Explicitly request RGB
             
-            with gr.Box():
+            # Changed gr.Box() to gr.Group()
+            with gr.Group(): # Replaced gr.Box()
                 gr.Markdown("### Enhancement Options")
                 model_choice = gr.Dropdown(
                     choices=list(MODEL_OPTIONS.keys()),
                     label="Model Selection",
-                    value="OpenCV Super Resolution"
+                    value="OpenCV Super Resolution",
+                    allow_flagging="never" # Optional: disable flagging
                 )
-                
+
                 with gr.Accordion("Advanced Settings", open=False):
+                    # Keep strength_slider visible but update label based on model
                     strength_slider = gr.Slider(
                         minimum=0.1,
                         maximum=1.0,
-                        step=0.1,
-                        label="Enhancement Strength",
+                        step=0.05, # Added more steps for finer control
+                        label="Enhancement Strength", # Default label
                         value=0.8,
+                        visible=True # Ensure it's visible
                     )
-                    
+
                     denoise_slider = gr.Slider(
                         minimum=0.0,
                         maximum=1.0,
-                        step=0.1,
-                        label="Noise Reduction",
+                        step=0.05, # Added more steps
+                        label="Noise Reduction Strength",
                         value=0.0,
                     )
-                    
+
                     sharpen_slider = gr.Slider(
                         minimum=0.0,
                         maximum=1.0,
-                        step=0.1,
-                        label="Sharpening",
+                        step=0.05, # Added more steps
+                        label="Sharpening Strength",
                         value=0.0,
                     )
-            
+
             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")
-    
+            output_image = gr.Image(label="Enhanced Image", type="pil") # Specify type="pil" consistently
+
     # Handle model change to update UI
+    # This function only needs to update the label of the strength slider
     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")
+            return gr.update(label="HDR Intensity")
         elif model_type == "face":
-            return gr.update(visible=True, label="Enhancement Strength")
+             return gr.update(label="Face Enhancement Strength")
+        elif model_type == "upscale":
+             return gr.update(label="Enhancement Strength") # Keep a generic label for upscale
         else:
-            return gr.update(visible=True, label="Enhancement Strength")
-            
+             return gr.update(label="Enhancement Strength") # Default
+
     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]
+        outputs=[output_text, output_image],
+        api_name="enhance" # Optional: give it an API name
     )
-    
+
     # 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
-        
+        - 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.
+        - Sharpening can add detail but may also increase noise if applied too strongly.
+
         ---
-        Version 2.0 | Running on: """ + ("GPU 🚀" if torch.cuda.is_available() else "CPU ⚙️")
+        Version 2.1 | Running on: """ + (f"GPU 🚀 ({torch.cuda.get_device_name(0)})" if torch.cuda.is_available() else "CPU ⚙️") + """
+        """
     )
 
 # Launch the app
 if __name__ == "__main__":
-    demo.launch()
+    # Use share=True for a temporary public link (useful for debugging, but not needed for Spaces)
+    # Use enable_queue=True for better handling of concurrent requests on Spaces
+    demo.launch(enable_queue=True)