Upload 2 files

app.py

@@ -237,7 +237,11 @@ class RealESRGANStrategy(UpscalerStrategy):
         dtype = torch.float16 if Config.DEVICE == 'cuda' else torch.bfloat16
         
         try:
-            with torch.autocast(device_type=Config.DEVICE, dtype=dtype):
+            # Explicitly disable autocast on CPU for RealESRGAN to avoid "PythonFallbackKernel" errors
+            # This seems to be a regression in recent PyTorch versions on CPU with some ops
+            context = torch.autocast(device_type=Config.DEVICE, dtype=dtype) if Config.DEVICE != 'cpu' else torch.no_grad()
+            
+            with context:
                 if tile_config['tile'] > 0:
                      output_tensor = manual_tile_upscale(
                         self.model, 
@@ -282,6 +286,8 @@ class Swin2SRStrategy(UpscalerStrategy):
                 logger.info(f"{self.name} loaded successfully.")
             except Exception as e:
                 logger.error(f"Failed to load Swin2SR: {e}")
+                # Swin2SR loading failure is often due to transformers version mismatch or device issues
+                # We re-raise to let the UI handle it, but log the specific error
                 raise
 
     def upscale(self, image: Image.Image, **kwargs) -> Image.Image:
@@ -296,7 +302,11 @@ class Swin2SRStrategy(UpscalerStrategy):
 
         inputs = self.processor(images=image, return_tensors="pt").to(Config.DEVICE)
         
-        with torch.no_grad():
+        # Swin2SR on CPU can be finicky with autocast/tracing. 
+        # Explicitly disable autocast for Swin2SR on CPU to avoid "PythonFallbackKernel" errors
+        context = torch.no_grad()
+        
+        with context:
             outputs = self.model(**inputs)
         
         output = outputs.reconstruction.data.squeeze().float().cpu().clamp_(0, 1).numpy()
@@ -400,8 +410,7 @@ class SpanStrategy(UpscalerStrategy):
                             self.model = torch.compile(self.model)
                             logger.info("✓ torch.compile enabled (default mode)")
                         self.compiled = True
-                    except Exception as e:
-                        logger.warning(f"⚠ torch.compile failed: {e}")
+                    except Exception:
                         self.compiled = True
 
                 logger.info(f"{self.name} loaded successfully.")
@@ -423,10 +432,15 @@ class SpanStrategy(UpscalerStrategy):
         h, w = img_np.shape[:2]
         tile_config = select_tile_config(h, w)
         
-        dtype = torch.float16 if Config.DEVICE == 'cuda' else torch.bfloat16
+        # Disable AMP for SPAN on CPU to avoid "UntypedStorage" weakref errors in inductor
+        # SPAN architecture seems sensitive to autocast + compile on CPU
+        dtype = torch.float32 if Config.DEVICE == 'cpu' else torch.float16
         
         try:
-            with torch.autocast(device_type=Config.DEVICE, dtype=dtype):
+            # Only use autocast if not CPU or if explicitly desired
+            context = torch.autocast(device_type=Config.DEVICE, dtype=dtype) if Config.DEVICE != 'cpu' else torch.no_grad()
+            
+            with context:
                 if tile_config['tile'] > 0:
                      output_tensor = manual_tile_upscale(
                         self.model,