Update app.py

app.py

@@ -33,13 +33,27 @@ def auto_download_if_needed(weight_path):
         os.system(f"wget https://github.com/Kiteretsu77/APISR/releases/download/{version_path}")
         os.system(f"mv {filename} pretrained")
 
+def load_model_with_device(loader_func, weight_path, scale, device):
+    # First load the state dict
+    state_dict = torch.load(weight_path, map_location=device)
+    
+    # Initialize the model
+    generator = loader_func(weight_path, scale=scale)
+    
+    # Load the state dict and move to device
+    if hasattr(generator, 'load_state_dict'):
+        generator.load_state_dict(state_dict)
+    generator = generator.to(device)
+    
+    return generator
+
 def inference(img_path, model_name):
     try:
         # Determine device - use GPU if available, otherwise CPU
         device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
         weight_dtype = torch.float32
         
-        # Load the model with appropriate device mapping
+        # Model configurations
         model_configs = {
             "4xGRL": ("pretrained/4x_APISR_GRL_GAN_generator.pth", load_grl, 4),
             "4xRRDB": ("pretrained/4x_APISR_RRDB_GAN_generator.pth", load_rrdb, 4),
@@ -53,13 +67,19 @@ def inference(img_path, model_name):
         weight_path, loader_func, scale = model_configs[model_name]
         auto_download_if_needed(weight_path)
         
-        # Load model with explicit device mapping
-        generator = loader_func(
-            weight_path, 
-            scale=scale,
-            map_location=device
-        )
-        generator = generator.to(device=device, dtype=weight_dtype)
+        # Load model and move to appropriate device
+        try:
+            generator = load_model_with_device(loader_func, weight_path, scale, device)
+        except RuntimeError as e:
+            if "out of memory" in str(e):
+                # If we run out of CUDA memory, try loading on CPU instead
+                device = torch.device('cpu')
+                generator = load_model_with_device(loader_func, weight_path, scale, device)
+            else:
+                raise e
+
+        generator = generator.to(dtype=weight_dtype)
+        generator.eval()  # Set to evaluation mode
 
         print(f"Processing {img_path} on {device}")
         print(f"Current time: {datetime.datetime.now(pytz.timezone('US/Eastern'))}")