Update app.py

app.py

@@ -1,5 +1,6 @@
 # app.py
-# Updated: 2025-04-05 18:43:16 IST (Ludhiana, Punjab, India)
+# Updated: 2025-04-05 18:49:54 PM IST (Ludhiana, Punjab, India)
+# Fix: Simplified CodeFormer() initialization call to resolve TypeError.
 
 import gradio as gr
 import torch
@@ -11,7 +12,7 @@ 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')}")
+print(f"Current Time (IST): {time.strftime('%Y-%m-%d %H:%M:%S')}") # Using system time from HF Space
 
 # Suppress specific warnings or all warnings if needed
 warnings.filterwarnings("ignore")
@@ -36,39 +37,32 @@ 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...")
+    print("Attempting to initialize CodeFormer model with simplified call...")
     try:
         # 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.
+        # *** 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
 
         codeformer_net.eval()
         is_initialized = True
-        print("CodeFormer model initialized successfully.")
+        print("CodeFormer model initialized successfully (using simplified init).")
 
     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 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()}"
+         print(init_error_message)
+         codeformer_net = None
     except Exception as e:
         init_error_message = f"Error initializing CodeFormer model: {e}\n{traceback.format_exc()}"
         print(init_error_message)
@@ -122,6 +116,8 @@ def enhance_image(input_img, fidelity_weight, background_enhance, face_upsample)
         # `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
             output_bgr, _, _ = codeformer_net.enhance(
                 img_bgr,
                 w=fidelity_weight,
@@ -149,11 +145,13 @@ 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}'
 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.
+**Status:** {status_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>"
 
@@ -165,7 +163,7 @@ if os.path.exists("examples"):
         ["examples/face2.png", 0.5, True, True],
         ["examples/bg1.png", 0.8, True, False],
     ]
-    print("Example files found.")
+    print("Example files folder found.")
 else:
     print("Note: 'examples' folder not found. Gradio examples will be empty.")
 
@@ -194,7 +192,6 @@ try:
 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
 
 
@@ -209,4 +206,3 @@ 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)
-