Create app.py

app.py

@@ -0,0 +1,310 @@
+# app.py
+# Image Upscale and Enhancement with Multiple Models
+# By FebryEnsz
+# SDK: Gradio
+# Hosted on Hugging Face Spaces
+
+import gradio as gr
+import torch
+from PIL import Image
+import numpy as np
+from huggingface_hub import hf_hub_download
+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)
+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.")
+    BASESR_AVAILABLE = False
+
+# RealESRGAN import
+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
+
+
+# --- 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}
+
+# --- 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}")
+
+# Function to load the selected model
+def load_model(model_name):
+    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:
+        return None, f"Error: Model '{model_name}' not supported 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.")
+            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
+            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.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
+            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 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 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
+                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:
+            raise ValueError(f"Configuration missing for model: {model_name}")
+
+        # Cache the loaded model
+        cached_model = model
+        cached_model_name = 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."
+    except Exception as e:
+        print(f"Error loading model {model_name}: {e}")
+        # Clear cache on error
+        cached_model = None
+        cached_model_name = None
+        return None, f"Error loading 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)
+    # 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)
+    return Image.fromarray(img)
+
+# Main processing function
+def enhance_image(image: Image.Image, model_name: str):
+    if image is None:
+        return "Please upload an image.", None
+    
+    status_message = f"Processing image with {model_name}..."
+    
+    # Load the selected model and its type
+    model, model_info = load_model(model_name)
+    
+    if model is None:
+        # model_info contains the error message if loading failed
+        return model_info, None
+    
+    model_type = model_info # model_info is the type string ('upscale' or 'face')
+
+    try:
+        # Preprocess the image (PIL RGB -> OpenCV BGR)
+        img_np_bgr = preprocess_image(image)
+        
+        # Process based on model type and specific model implementation
+        if model_type == "upscale":
+            print(f"Applying {model_name} upscaling...")
+            if model_name == "Real-ESRGAN-x4":
+                 # RealESRGAN works with uint8 BGR numpy directly
+                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
+                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
+                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}")
+                
+            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 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
+
+
+            else:
+                 raise ValueError(f"Unknown face enhancement model: {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)
+        
+        return status_message, enhanced_image
+    
+    except ImportError as ie:
+         return f"Error processing image: Missing dependency - {ie}", None
+    except Exception as e:
+        print(f"Error during processing: {e}")
+        import traceback
+        traceback.print_exc() # Print full traceback for debugging
+        return f"Error processing image: {str(e)}", None
+
+# Gradio interface
+with gr.Blocks(title="Image Upscale & Enhancement - By FebryEnsz") as demo:
+    gr.Markdown(
+        """
+        # 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).
+        
+        **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.
+        """
+    )
+    
+    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)
+                 enhance_button = gr.Button("Enhance Image", interactive=False)
+                 print("No models are available because dependencies are missing.")
+            else:
+                model_choice = gr.Dropdown(
+                    choices=available_models,
+                    label="Select Model",
+                    value=available_models[0] # Default to the first available model
+                )
+                # Removed scale_slider as models are fixed scale (x4)
+                enhance_button = gr.Button("Enhance Image")
+            
+        with gr.Column():
+            output_text = gr.Textbox(label="Status", max_lines=2)
+            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]
+        )
+
+# 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
+
+    demo.launch()