Update app.py

app.py

@@ -1,3 +1,6 @@
+# app.py
+# Updated: 2025-04-05 18:43:16 IST (Ludhiana, Punjab, India)
+
 import gradio as gr
 import torch
 import cv2
@@ -5,81 +8,77 @@ import numpy as np
 import os
 import time
 import warnings
+import traceback # Import traceback for detailed error printing
+
+print("--- Script Start ---")
+print(f"Current Time (IST): {time.strftime('%Y-%m-%d %H:%M:%S')}")
 
 # Suppress specific warnings or all warnings if needed
 warnings.filterwarnings("ignore")
 
-# Try importing CodeFormer, handle potential import errors
+# --- Globals ---
+codeformer_net = None
+is_initialized = False
+init_error_message = ""
+
+# --- Attempt to Import CodeFormer ---
 try:
     from codeformer import CodeFormer
-except ImportError:
-    print("Error: CodeFormer not found. Make sure it's installed correctly (check requirements.txt).")
-    # Optionally, try adding the repo path if cloned (more complex setup)
-    # sys.path.append('CodeFormer') # If you cloned the repo into a folder named CodeFormer
-    # from basicsr.utils.registry import ARCH_REGISTRY
-    raise
-
-print("Imports successful.")
-
-# --- Configuration ---
-# Automatically select CPU
-device = torch.device("cpu")
-print(f"Using device: {device}")
-
-# Initialize CodeFormer - Model weights will be downloaded automatically on first run
-# Ensure you have internet access in the Space for the download.
-print("Initializing CodeFormer model...")
-try:
-    # Adjust model path if needed, but pretrained=True should handle downloads
-    # Check the documentation for the 'codeformer' package if this fails.
-    # Common parameters: bg_upsampler='realesrgan', face_upsample=True
-    codeformer_net = CodeFormer(
-        dim_embd=512,
-        codebook_size=1024,
-        n_head=8,
-        n_layers=9,
-        connect_list=['32', '64', '128', '256']
-    ).to(device)
-    
-    # Load the pre-trained model weights
-    # Adjust the path based on how the package stores weights or if downloaded manually
-    # This path assumes the standard download location used by `load_state_dict_from_url`
-    # It might differ based on the specific 'codeformer' pip package version.
-    # If this fails, check where the package downloads/expects the .pth file.
-    model_path = 'weights/CodeFormer/codeformer.pth' # Default path often used
-    
-    # Check if the default path exists, otherwise rely on package's internal loading if possible
-    # A robust package might have a load_pretrained() method. Check its usage.
-    # This explicit loading might be needed if the package is minimal.
-    # Let's assume the package handles loading implicitly or requires a different call.
-    # Simpler approach: Rely on package potentially loading during init or a specific method.
-    # If the above CodeFormer() init doesn't load weights, check package docs.
-    # For now, let's assume the package *might* need explicit loading IF NOT BUILT-IN:
-    
-    # Placeholder checkpoint loading - adjust based on actual package behavior
-    # This might be automatically handled by the package; if the app fails here,
-    # investigate how the specific `codeformer` pip package loads weights.
+    print("Successfully imported CodeFormer.")
+except ImportError as e:
+    init_error_message = f"Error: CodeFormer package not found or import failed. Check requirements.txt installation. Details: {e}"
+    print(init_error_message)
+    # Keep script running so Gradio might load with an error message
+except Exception as e:
+    init_error_message = f"An unexpected error occurred during import: {e}\n{traceback.format_exc()}"
+    print(init_error_message)
+
+
+# --- Initialize Model (only if import succeeded) ---
+if 'CodeFormer' in globals(): # Check if import was successful
+    print("Attempting to initialize CodeFormer model...")
     try:
-        # Example: Load weights/CodeFormer/codeformer.pth
-        # This path needs to be correct relative to where HF downloads/caches it, or package internal path
-        # It's often complex to pinpoint the exact cache location in HF Spaces
-        # A safer bet is often using a model hub integration if available, or ensuring the package handles it well.
-        # For now, we'll *assume* the package loads weights correctly or fails gracefully if not found
-        # checkpoint = torch.load(model_path)['params_ema']
-        # codeformer_net.load_state_dict(checkpoint)
-        print("Model weights assumed to be loaded by package or implicitly.") # Placeholder message
-    except FileNotFoundError:
-         print(f"Warning: Pretrained weights not found at default path '{model_path}'. Relying on package's internal loading mechanism if available.")
+        # Use CPU explicitly
+        device = torch.device("cpu")
+        print(f"Using device: {device}")
+
+        # Initialize CodeFormer.
+        # Relying on the package to handle pretrained weight download/loading.
+        # Common parameters: bg_upsampler='realesrgan', face_upsample=True
+        # These details might depend on the exact version of the 'codeformer' pip package.
+        codeformer_net = CodeFormer(
+            dim_embd=512,
+            codebook_size=1024,
+            n_head=8,
+            n_layers=9,
+            connect_list=['32', '64', '128', '256']
+            # Add other necessary parameters based on the package's specific API if needed
+        ).to(device)
+
+        # Load pretrained weights - The package *should* ideally handle this,
+        # either during init or via a specific method. If this implicit loading fails,
+        # you might need to investigate the specific package's API for loading weights.
+        # We are assuming here the class instantiation handles it or makes weights available.
+        # Let's skip manual torch.load unless proven necessary by failure.
+
+        codeformer_net.eval()
+        is_initialized = True
+        print("CodeFormer model initialized successfully.")
+
+    except FileNotFoundError as e:
+        init_error_message = f"Error: Could not find CodeFormer model weights. The package might have failed to download them automatically, or they are expected at a specific path. Check package docs and Space logs. Details: {e}\n{traceback.format_exc()}"
+        print(init_error_message)
+        codeformer_net = None # Ensure model is None if init fails
     except Exception as e:
-         print(f"Error loading weights explicitly: {e}. Relying on package's internal loading.")
+        init_error_message = f"Error initializing CodeFormer model: {e}\n{traceback.format_exc()}"
+        print(init_error_message)
+        codeformer_net = None # Ensure model is None if init fails
+else:
+    # If import failed, ensure message reflects that
+    if not init_error_message: # Safety net if import error wasn't captured somehow
+       init_error_message = "CodeFormer could not be imported. Cannot initialize model."
+    print("Skipping model initialization due to import failure.")
 
-    codeformer_net.eval()
-    print("CodeFormer model initialized successfully.")
-except Exception as e:
-    print(f"Error initializing CodeFormer model: {e}")
-    # Provide helpful error message in the UI if initialization fails
-    gr.Error(f"Failed to load CodeFormer model. Check logs. Error: {e}")
-    codeformer_net = None # Set to None to prevent processing attempts
 
 # --- Processing Function ---
 def enhance_image(input_img, fidelity_weight, background_enhance, face_upsample):
@@ -91,39 +90,50 @@ def enhance_image(input_img, fidelity_weight, background_enhance, face_upsample)
         background_enhance (bool): Whether to enhance background using RealESRGAN.
         face_upsample (bool): Whether to further upsample restored faces.
     Returns:
-        np.ndarray: Enhanced image (RGB format).
-        str: Processing time message.
+        np.ndarray: Enhanced image (RGB format) or None on error.
+        str: Status or processing time message.
     """
-    if codeformer_net is None:
-        return None, "Error: CodeFormer model not loaded."
-        
+    print("--- enhance_image function called ---") # Log function entry
+
+    if not is_initialized or codeformer_net is None:
+        error_msg = f"ERROR: CodeFormer model is not available. Initialization failed. Check logs for details. Message: {init_error_message}"
+        print(error_msg)
+        # Return None for the image and the error message for the status textbox
+        return None, error_msg
+
     if input_img is None:
+        print("Error: No input image provided.")
         return None, "Error: No input image provided."
 
     print(f"Processing image with fidelity: {fidelity_weight}, bg_enhance: {background_enhance}, face_upsample: {face_upsample}")
     start_time = time.time()
 
     try:
-        # Gradio provides RGB, CodeFormer often expects BGR internally via OpenCV
+        # 1. Convert RGB (from Gradio) to BGR (often expected by OpenCV/CodeFormer backend)
         img_bgr = cv2.cvtColor(input_img, cv2.COLOR_RGB2BGR)
+        print("Input image converted from RGB to BGR.")
+
+        # 2. Select background upsampler based on checkbox
+        bg_upsampler = 'realesrgan' if background_enhance else None
+        print(f"Background upsampler selected: {bg_upsampler}")
 
-        # Enhance the image - Use the correct method from the CodeFormer package
-        # The method might be called 'enhance', 'process', 'restore', etc.
-        # Check the package documentation for the exact API.
-        # Assuming a method like `codeformer_net.enhance(...)` or similar exists:
-        # The exact parameters (like `w`, `adain`) depend on the CodeFormer implementation.
-        # `w` typically corresponds to fidelity_weight.
+        # 3. Run CodeFormer enhancement
+        # Ensure parameters match the CodeFormer API you have installed.
+        # `w` corresponds to fidelity_weight.
+        print("Starting CodeFormer enhancement...")
         with torch.no_grad():
             output_bgr, _, _ = codeformer_net.enhance(
                 img_bgr,
                 w=fidelity_weight,
-                adain=True, # Adain usually enabled
+                adain=True, # AdaIN is commonly used
                 face_upsample=face_upsample,
-                bg_upsampler='realesrgan' if background_enhance else None # Use bg_upsampler if requested
+                bg_upsampler=bg_upsampler
             )
+        print("CodeFormer enhancement finished.")
 
-        # Convert back to RGB for Gradio display
+        # 4. Convert BGR output back to RGB for Gradio display
         output_rgb = cv2.cvtColor(output_bgr, cv2.COLOR_BGR2RGB)
+        print("Output image converted from BGR to RGB.")
 
         end_time = time.time()
         processing_time = end_time - start_time
@@ -132,45 +142,71 @@ def enhance_image(input_img, fidelity_weight, background_enhance, face_upsample)
         return output_rgb, time_msg
 
     except Exception as e:
-        print(f"Error during enhancement: {e}")
-        import traceback
-        traceback.print_exc()
-        return None, f"Error during processing: {e}"
+        # Catch errors during the enhancement process
+        error_details = f"Error during enhancement: {e}\n{traceback.format_exc()}"
+        print(error_details)
+        return None, error_details # Return error message to the status textbox
 
-# --- Gradio Interface ---
+# --- Gradio Interface Definition ---
 title = "CodeFormer Image Enhancement (CPU Demo)"
-description = """
+description = f"""
 Upload an image to enhance its quality, particularly for faces, using CodeFormer.
 **Note:** This demo runs on a free Hugging Face CPU. Processing will be **SLOW** (expect seconds to minutes per image).
 Adjust the fidelity weight (0 = max quality enhancement, 1 = closer to original). Optionally enhance background and upsample faces.
+**Status:** {'Model Loaded Successfully.' if is_initialized else f'Model Load FAILED: {init_error_message}'} Check Logs for details if failed.
 """
 article = "<p style='text-align: center'>CodeFormer CPU Demo | <a href='https://github.com/sczhou/CodeFormer' target='_blank'>Official Repo</a></p>"
 
-iface = gr.Interface(
-    fn=enhance_image,
-    inputs=[
-        gr.Image(label="Upload Image", type="numpy"),
-        gr.Slider(minimum=0.0, maximum=1.0, step=0.05, value=0.7, label="Fidelity Weight (0 = Max Quality, 1 = Max Fidelity)"),
-        gr.Checkbox(label="Enhance Background (Uses RealESRGAN)", value=True),
-        gr.Checkbox(label="Upsample Restored Faces", value=True)
-    ],
-    outputs=[
-        gr.Image(label="Enhanced Image", type="numpy"),
-        gr.Textbox(label="Processing Time")
-    ],
-    title=title,
-    description=description,
-    article=article,
-    examples=[
-        ["examples/face1.png", 0.7, True, True], # Add example files to an 'examples' folder in your Space
+# Define examples (Make sure the 'examples' folder and files exist in your Space repo)
+example_list = []
+if os.path.exists("examples"):
+    example_list = [
+        ["examples/face1.png", 0.7, True, True],
         ["examples/face2.png", 0.5, True, True],
         ["examples/bg1.png", 0.8, True, False],
-    ],
-    allow_flagging="never" # Can change to "manual" or "auto" if needed
-)
+    ]
+    print("Example files found.")
+else:
+    print("Note: 'examples' folder not found. Gradio examples will be empty.")
+
+
+print("Defining Gradio interface...")
+try:
+    iface = gr.Interface(
+        fn=enhance_image,
+        inputs=[
+            gr.Image(label="Upload Image", type="numpy"),
+            gr.Slider(minimum=0.0, maximum=1.0, step=0.05, value=0.7, label="Fidelity Weight (0 = Max Quality, 1 = Max Fidelity)"),
+            gr.Checkbox(label="Enhance Background (Uses RealESRGAN)", value=True),
+            gr.Checkbox(label="Upsample Restored Faces", value=True)
+        ],
+        outputs=[
+            gr.Image(label="Enhanced Image", type="numpy"),
+            gr.Textbox(label="Status / Processing Time") # Output for status messages and time
+        ],
+        title=title,
+        description=description,
+        article=article,
+        examples=example_list,
+        allow_flagging="never"
+    )
+    print("Gradio interface defined successfully.")
+except Exception as e:
+    print(f"FATAL: Failed to define Gradio interface: {e}\n{traceback.format_exc()}")
+    # If interface definition fails, we can't launch.
+    # Raising an error might provide more info in logs, or just print and exit.
+    raise RuntimeError("Could not create Gradio Interface.") from e
+
 
 # --- Launch the App ---
 if __name__ == "__main__":
-    iface.launch()
-    print("Gradio app launched.")
+    print("Attempting to launch Gradio app...")
+    try:
+        iface.launch()
+        print("Gradio app launched successfully.")
+        print("--- Script End (App Running) ---")
+    except Exception as e:
+        print(f"FATAL: Failed to launch Gradio interface: {e}\n{traceback.format_exc()}")
+        # Exit if launch fails. Logs should show the error.
+        exit(1)