Update app.py

app.py

@@ -1,6 +1,7 @@
 # app.py
-# Updated: 2025-04-05 18:49:54 PM IST (Ludhiana, Punjab, India)
-# Fix: Simplified CodeFormer() initialization call to resolve TypeError.
+# Updated: 2025-04-05 18:56:50 PM IST (Ludhiana, Punjab, India)
+# Fix: Removed .to(device) call after CodeFormer() initialization due to AttributeError.
+# Added check for .eval() method.
 
 import gradio as gr
 import torch
@@ -37,34 +38,44 @@ except Exception as e:
 
 # --- Initialize Model (only if import succeeded) ---
 if 'CodeFormer' in globals(): # Check if import was successful
-    print("Attempting to initialize CodeFormer model with simplified call...")
+    print("Attempting to initialize CodeFormer model (simplified init, no .to(device))...")
     try:
-        # Use CPU explicitly
+        # Define the target device (though we won't call .to() on the model object)
         device = torch.device("cpu")
-        print(f"Using device: {device}")
+        print(f"Target device check: {device}")
 
         # *** MODIFIED LINE BELOW ***
-        # Call constructor WITHOUT detailed arguments like dim_embd, etc.
-        # Assume the package uses defaults and handles pretrained weights.
-        codeformer_net = CodeFormer().to(device)
-        # If the above line fails, another common pattern for packages is:
-        # codeformer_net = CodeFormer(pretrained=True).to(device) # Uncomment and try this if the first fails
+        # Initialize WITHOUT .to(device)
+        codeformer_net = CodeFormer()
+        print("CodeFormer() object created.")
+
+        # Check if the object has an eval method (common for models)
+        # If it doesn't, it might not be a standard torch.nn.Module
+        try:
+            codeformer_net.eval()
+            print("codeformer_net.eval() called successfully.")
+        except AttributeError:
+            print("Note: codeformer_net object does not have an eval() method. Assuming it's not needed or handled internally.")
+        except Exception as e_eval:
+             print(f"Warning: Calling codeformer_net.eval() failed: {e_eval}")
+
 
-        codeformer_net.eval()
         is_initialized = True
-        print("CodeFormer model initialized successfully (using simplified init).")
+        print("CodeFormer model object initialization appears successful (removed .to(device)).")
 
     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()}"
+        # This error might occur if CodeFormer() tries to load weights internally and fails
+        init_error_message = f"Error: Could not find CodeFormer model weights during init. The package might have failed to download them automatically. Check logs. Details: {e}\n{traceback.format_exc()}"
         print(init_error_message)
         codeformer_net = None # Ensure model is None if init fails
     except TypeError as e:
          # Catch TypeErrors specifically from the init call
-         init_error_message = f"TypeError during CodeFormer init. The constructor signature might be different in this package version (e.g., does it need 'pretrained=True'?). Details: {e}\n{traceback.format_exc()}"
+         init_error_message = f"TypeError during CodeFormer init. The constructor signature might be wrong. Details: {e}\n{traceback.format_exc()}"
          print(init_error_message)
          codeformer_net = None
     except Exception as e:
-        init_error_message = f"Error initializing CodeFormer model: {e}\n{traceback.format_exc()}"
+        # Catch any other initialization errors
+        init_error_message = f"Error initializing CodeFormer model object: {e}\n{traceback.format_exc()}"
         print(init_error_message)
         codeformer_net = None # Ensure model is None if init fails
 else:
@@ -115,9 +126,8 @@ def enhance_image(input_img, fidelity_weight, background_enhance, face_upsample)
         # Ensure parameters match the CodeFormer API you have installed.
         # `w` corresponds to fidelity_weight.
         print("Starting CodeFormer enhancement...")
-        with torch.no_grad():
-            # Ensure enhance method signature matches the package API
-            # Assuming enhance(image, w, face_upsample, bg_upsampler, adain=True) is correct
+        # *** Assuming enhance method exists and works on CPU by default ***
+        with torch.no_grad(): # Keep torch.no_grad() as internal ops might use Torch
             output_bgr, _, _ = codeformer_net.enhance(
                 img_bgr,
                 w=fidelity_weight,
@@ -146,7 +156,7 @@ def enhance_image(input_img, fidelity_weight, background_enhance, face_upsample)
 # --- Gradio Interface Definition ---
 title = "CodeFormer Image Enhancement (CPU Demo)"
 # Dynamically update description based on model load status
-status_message = 'Model Loaded Successfully.' if is_initialized else f'Model Load FAILED: {init_error_message}'
+status_message = 'Model initialization appears successful.' if is_initialized else f'Model Load FAILED: {init_error_message}'
 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).
@@ -206,3 +216,4 @@ if __name__ == "__main__":
         print(f"FATAL: Failed to launch Gradio interface: {e}\n{traceback.format_exc()}")
         # Exit if launch fails. Logs should show the error.
         exit(1)
+