Update app.py

app.py

@@ -9,10 +9,16 @@ import uuid
 import torch
 from PIL import Image
 
-# Force CPU usage
+# Force CPU usage and disable CUDA completely
 torch.backends.cudnn.enabled = False
+os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
 torch.cuda.is_available = lambda: False
-os.environ['CUDA_VISIBLE_DEVICES'] = ''
+torch.cuda.device_count = lambda: 0
+torch.cuda.get_device_name = lambda x: 'cpu'
+torch.cuda.current_device = lambda: 0
+torch.cuda.set_device = lambda x: None
+torch.Tensor.cuda = lambda self, device=None: self
+torch.nn.Module.cuda = lambda self, device=None: self
 
 demo_path = Path(__file__).resolve().parent
 root_path = demo_path
@@ -67,6 +73,16 @@ example_previews = [
     [thumbnails[7], 'Prompt: sunglasses'],
 ]
 
+# Monkey patch any remaining CUDA calls in the models
+original_to = torch.Tensor.to
+def patched_to(self, *args, **kwargs):
+    if len(args) > 0 and isinstance(args[0], str) and args[0] == 'cuda':
+        return original_to(self, 'cpu')
+    if 'device' in kwargs and kwargs['device'] == 'cuda':
+        kwargs['device'] = 'cpu'
+    return original_to(self, *args, **kwargs)
+torch.Tensor.to = patched_to
+
 # Load models with CPU only
 models.pre_download_inpainting_models()
 inpainting_models = OrderedDict([
@@ -74,6 +90,32 @@ inpainting_models = OrderedDict([
     ("Stable-Inpainting 2.0", 'sd2_inp'),
     ("Stable-Inpainting 1.5", 'sd15_inp')
 ])
+
+# Patch model loading to ensure CPU usage
+original_load_model = models.load_inpainting_model
+def patched_load_model(*args, **kwargs):
+    kwargs['device'] = 'cpu'
+    model = original_load_model(*args, **kwargs)
+    model.to('cpu')
+    return model
+models.load_inpainting_model = patched_load_model
+
+original_sr_load_model = models.sd2_sr.load_model
+def patched_sr_load_model(*args, **kwargs):
+    kwargs['device'] = 'cpu'
+    model = original_sr_load_model(*args, **kwargs)
+    model.to('cpu')
+    return model
+models.sd2_sr.load_model = patched_sr_load_model
+
+original_sam_load_model = models.sam.load_model
+def patched_sam_load_model(*args, **kwargs):
+    kwargs['device'] = 'cpu'
+    model = original_sam_load_model(*args, **kwargs)
+    model.to('cpu')
+    return model
+models.sam.load_model = patched_sam_load_model
+
 sr_model = models.sd2_sr.load_model(device='cpu')
 sam_predictor = models.sam.load_model(device='cpu')
 
@@ -140,7 +182,6 @@ def inpainting_run(model_name, use_rasg, use_painta, prompt, imageMask,
     hr_image, seed, eta, negative_prompt, positive_prompt, ddim_steps,
     guidance_scale=7.5, batch_size=1, session_id=''
 ):
-    # Remove torch.cuda.empty_cache() since we're not using GPU
     set_model_from_name(model_name)
 
     method = ['default']
@@ -206,8 +247,6 @@ def upscale_run(
     if len(gallery) == 0:
         return Image.open(root_path / '__assets__/demo/sr_info.png')
 
-    # Remove torch.cuda.empty_cache() since we're not using GPU
-
     seed = int(seed)
     img_index = int(img_index)