Update app.py

app.py

@@ -14,6 +14,8 @@ 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
@@ -21,52 +23,63 @@ try:
     BASESR_AVAILABLE = True
 except ImportError:
     print("Warning: basicsr not found. SwinIR, EDSR, and CodeFormer (using basicsr utils) will not be available.")
-    BASESR_AVAILABLE = False
 
 # 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.")
-    REALESRGAN_AVAILABLE = False
 
-# CodeFormer import (This assumes CodeFormer is installed and importable,
-# or integrated into basicsr's structure) - often requires manual setup.
-# We will use basicsr's utilities for CodeFormer if available, and try a direct import if possible.
-try:
-    # Attempting a common import path if CodeFormer is installed separately
-    from CodeFormer import CodeFormer # Adjust import based on your CodeFormer install
-    CODEFORMER_AVAILABLE = True
-except ImportError:
-    print("Warning: CodeFormer not found. CodeFormer (Face Enhancement) will not be available.")
-    CODEFORMER_AVAILABLE = False
+# 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
 
 
 # --- 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 = {
-    "Real-ESRGAN-x4": {"repo_id": "RealESRGAN/RealESRGAN_x4plus", "filename": "RealESRGAN_x4plus.pth", "type": "upscale", "scale": 4} if REALESRGAN_AVAILABLE else None,
-    "SwinIR-4x": {"repo_id": "SwinIR/SwinIR-Large", "filename": "SwinIR_4x.pth", "type": "upscale", "scale": 4} if BASESR_AVAILABLE else None,
-    "EDSR-x4": {"repo_id": "EDSR/edsr_x4", "filename": "edsr_x4.pth", "type": "upscale", "scale": 4} if BASESR_AVAILABLE else None,
-    # Note: CodeFormer often requires its own setup. Assuming basicsr utils might help,
-    # but its core logic is in the CodeFormer library.
-    "CodeFormer (Face Enhancement)": {"repo_id": "CodeFormer/codeformer", "filename": "codeformer.pth", "type": "face"} if CODEFORMER_AVAILABLE or BASESR_AVAILABLE else None, # Use CodeFormer if installed, otherwise rely on basicsr utilities being present
-}
-
-# Filter out unavailable models
-MODEL_CONFIGS = {k: v for k, v in MODEL_CONFIGS.items() if v is not None}
+# format: "UI Name": {"repo_id": "hf_repo_id", "filename": "weight_filename", "type": "upscale" or "face", ...}
+MODEL_CONFIGS = {}
+
+if REALESRGAN_AVAILABLE:
+     MODEL_CONFIGS["Real-ESRGAN-x4"] = {"repo_id": "RealESRGAN/RealESRGAN_x4plus", "filename": "RealESRGAN_x4plus.pth", "type": "upscale", "scale": 4}
+
+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.")
+
+# No need to filter anymore, configs are only added if available
 
 # --- Model Loading Cache ---
-# Use a simple cache to avoid reloading the same model multiple times
 cached_model = {}
 cached_model_name = None
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-print(f"Using device: {device}")
+print(f"Using device: {device}") # This shows which device PyTorch detects
 
 # Function to load the selected model
-def load_model(model_name):
+def load_model(model_name: str):
     global cached_model, cached_model_name
 
     if model_name == cached_model_name and cached_model is not None:
@@ -76,35 +89,32 @@ def load_model(model_name):
     print(f"Loading model: {model_name}")
     config = MODEL_CONFIGS.get(model_name)
     if config is None:
-        return None, f"Error: Model '{model_name}' not supported or dependencies missing."
+        # 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."
 
     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 not installed.")
+            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 not installed.")
-            # SwinIR requires specific initialization parameters
-            # These match the SwinIR_4x.pth model from the repo
+            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'
             )
-            # Load weights, handling potential key mismatches if necessary
             pretrained_dict = torch.load(model_path, map_location=device)
-            model.load_state_dict(pretrained_dict, strict=True) # strict=False if keys might mismatch
-            model.eval() # Set to evaluation mode
+            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 not installed.")
-            # EDSR architecture needs scale, num_feat, num_block
-            # Assuming typical values for EDSR_x4 from the repo
+            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)
@@ -112,33 +122,32 @@ def load_model(model_name):
             model.to(device)
             
         elif model_name == "CodeFormer (Face Enhancement)":
-             if not (CODEFORMER_AVAILABLE or BASESR_AVAILABLE): raise ImportError("CodeFormer or basicsr not installed.")
-             # CodeFormer loading is more complex, often requiring instantiation with specific args
-             # and potentially related models (like GFPGAN for background).
-             # For simplicity here, we assume a basic CodeFormer instance can be created.
-             # This part might need adjustment based on your CodeFormer installation.
+             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:
-                # This is a simplified instantiation; a real CodeFormer usage might need more args
-                model = CodeFormer(dim_embd=512, codebook_size=1024, n_head=8, n_layer=9)
-                pretrained_dict = torch.load(model_path, map_location=device)['params_ema'] # CodeFormer often saves params_ema
-                # Need to handle potential DataParallel prefix if saved from DP
+                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)
-             elif BASESR_AVAILABLE:
-                 # Fallback: If CodeFormer library isn't directly importable but basicsr is,
-                 # we *cannot* instantiate the CodeFormer model itself unless basicsr provides it.
-                 # This option is likely only possible if CodeFormer is installed *within* a basicsr environment
-                 # or if basicsr provides the architecture. Given the complexity, let's just raise an error
-                 # if CODEFORMER_AVAILABLE is False.
-                 raise ImportError("CodeFormer library not found. BasicSR utilities alone are not enough to instantiate CodeFormer.")
+             else:
+                 # Fallback check, should not be reached if config is filtered
+                 raise ImportError("CodeFormer library not available.")
 
 
         else:
-            raise ValueError(f"Configuration missing for model: {model_name}")
+            # 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
@@ -147,10 +156,17 @@ def load_model(model_name):
         return model, model_type
     
     except ImportError as ie:
-        print(f"Dependency missing for {model_name}: {ie}")
-        return None, f"Error: Missing dependency - {ie}. Please ensure model libraries are installed."
+        # 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."
     except Exception as e:
         print(f"Error loading model {model_name}: {e}")
+        import traceback
+        traceback.print_exc() # Print full traceback for debugging
         # Clear cache on error
         cached_model = None
         cached_model_name = None
@@ -159,13 +175,11 @@ def load_model(model_name):
 # Function to preprocess image (PIL RGB to OpenCV BGR numpy)
 def preprocess_image(image: Image.Image) -> np.ndarray:
     img = np.array(image)
-    # OpenCV uses BGR, PIL uses RGB. Need conversion.
     img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
     return img
 
 # Function to postprocess image (OpenCV BGR numpy to PIL RGB)
 def postprocess_image(img: np.ndarray) -> Image.Image:
-    # Ensure image is in the correct range and type before converting
     if img.dtype != np.uint8:
          img = np.clip(img, 0, 255).astype(np.uint8)
     img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
@@ -174,6 +188,7 @@ def postprocess_image(img: np.ndarray) -> Image.Image:
 # 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
     
     status_message = f"Processing image with {model_name}..."
@@ -182,7 +197,7 @@ def enhance_image(image: Image.Image, model_name: str):
     model, model_info = load_model(model_name)
     
     if model is None:
-        # model_info contains the error message if loading failed
+        # model_info contains the error message from load_model
         return model_info, None
     
     model_type = model_info # model_info is the type string ('upscale' or 'face')
@@ -195,51 +210,38 @@ def enhance_image(image: Image.Image, model_name: str):
         if model_type == "upscale":
             print(f"Applying {model_name} upscaling...")
             if model_name == "Real-ESRGAN-x4":
-                 # RealESRGAN works with uint8 BGR numpy directly
+                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 is required for {model_name}")
-                # These models often work with float tensors (0-1 range)
-                # Using basicsr utils: HWC BGR uint8 -> CHW RGB float (0-1) -> send to device
+                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)
 
-                # Using basicsr utils: CHW RGB float (0-1) -> HWC RGB uint8 -> Convert to BGR for postprocessing
+                # 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)
                 
             else:
-                raise ValueError(f"Unknown upscale model: {model_name}")
+                raise ValueError(f"Unknown upscale model type configuration: {model_name}")
                 
             status_message = f"Image upscaled successfully with {model_name}!"
 
         elif model_type == "face":
             print(f"Applying {model_name} face enhancement...")
             if model_name == "CodeFormer (Face Enhancement)":
-                if not (CODEFORMER_AVAILABLE or BASESR_AVAILABLE):
-                    raise ImportError(f"CodeFormer or basicsr is required for {model_name}")
-                # CodeFormer's enhance method typically expects uint8 BGR numpy
-                # It might return multiple outputs, the first is usually the enhanced image
-                # Example: CodeFormer's inference script might return (restored_img, bboxes)
-                # We assume the image is the first element.
-                # Note: CodeFormer often needs additional setup/parameters for GFPGAN, etc.
-                # This is a simplified call.
-                # Ensure model is on correct device before call
+                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)
                 
-                # A minimal CodeFormer enhancement might look like this, but the actual API
-                # depends on the CodeFormer library version/structure you're using.
-                # The original CodeFormer repo's inference takes numpy BGR.
-                # This is a *placeholder* call assuming such a method exists and works like this:
-                output_np_bgr = model.enhance(img_np_bgr, w=0.5, adain=True)[0] # w and adain are common params
-
-
+                # 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
             else:
-                 raise ValueError(f"Unknown face enhancement model: {model_name}")
+                 raise ValueError(f"Unknown face enhancement model type configuration: {model_name}")
 
             status_message = f"Face enhancement applied successfully with {model_name}!"
 
@@ -249,7 +251,10 @@ def enhance_image(image: Image.Image, model_name: str):
         return status_message, enhanced_image
     
     except ImportError as ie:
-         return f"Error processing image: Missing dependency - {ie}", None
+         # 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}")
         import traceback
@@ -265,19 +270,19 @@ with gr.Blocks(title="Image Upscale & Enhancement - By FebryEnsz") as demo:
         
         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).
         
-        **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 is not installed or found.
+        **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.
         """
     )
     
+    # Filter available choices based on loaded configs
+    available_models = list(MODEL_CONFIGS.keys())
+    
     with gr.Row():
         with gr.Column():
             image_input = gr.Image(label="Upload Image", type="pil")
             
-            # Filter available choices based on loaded configs
-            available_models = list(MODEL_CONFIGS.keys())
-            
             if not available_models:
-                 model_choice = gr.Textbox(label="Select Model", value="No models available. Check dependencies.", interactive=False)
+                 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:
@@ -286,11 +291,11 @@ with gr.Blocks(title="Image Upscale & Enhancement - By FebryEnsz") as demo:
                     label="Select Model",
                     value=available_models[0] # Default to the first available model
                 )
-                # Removed scale_slider as models are fixed scale (x4)
+                # Removed scale_slider
                 enhance_button = gr.Button("Enhance Image")
             
         with gr.Column():
-            output_text = gr.Textbox(label="Status", max_lines=2)
+            output_text = gr.Textbox(label="Status", max_lines=3)
             output_image = gr.Image(label="Enhanced Image")
     
     # Connect the button to the processing function
@@ -304,7 +309,8 @@ with gr.Blocks(title="Image Upscale & Enhancement - By FebryEnsz") as demo:
 # Launch the Gradio app
 if __name__ == "__main__":
     # Set torch backend for potentially better performance on some systems
-    if torch.backends.mps.is_available():
-        os.environ['PYTORCH_ENABLE_MPS_FALLBACK'] = '1' # Optional: enable fallback for MPS
+    # 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()