Fix ZeroGPU compatibility: load model on CPU at startup, move to CUDA inside @spaces.GPU

app.py

@@ -19,6 +19,7 @@ import gradio as gr
 import numpy as np
 import shutil
 import os
+import spaces
 
 
 # Download checkpoint from HF Hub at startup
@@ -42,7 +43,9 @@ for og, edit in EXAMPLE_PAIRS:
 magic_fixup = MagicFixup(model_path=checkpoint_path)
 
 
+@spaces.GPU
 def sample(original_image, coarse_edit):
+    magic_fixup.model.cuda()
     to_tensor = torchvision.transforms.ToTensor()
     with autocast("cuda"):
         w, h = coarse_edit.size

run_magicfu.py

@@ -105,7 +105,6 @@ def load_model_from_config(config, ckpt, verbose=False):
             print("unexpected keys:")
             print(u)
 
-    model.cuda()
     model.eval()
     return model
 
@@ -124,7 +123,7 @@ def get_model(config_path, ckpt_path):
         print(u)
 
     
-    model = model.to(device)
+    model = model.cpu()
     return model
 
 def get_grid(size):
@@ -193,10 +192,10 @@ class MagicFixup:
                         inpaint_image = image_tensor#*mask_tensor
 
                         test_model_kwargs={}
-                        test_model_kwargs['inpaint_mask']=mask_tensor.to(device)
-                        test_model_kwargs['inpaint_image']=inpaint_image.to(device)
-                        clean_ref_tensor = clean_ref_tensor.to(device)
-                        ref_tensor=ref_tensor.to(device)
+                        test_model_kwargs['inpaint_mask']=mask_tensor.to(self.model.device)
+                        test_model_kwargs['inpaint_image']=inpaint_image.to(self.model.device)
+                        clean_ref_tensor = clean_ref_tensor.to(self.model.device)
+                        ref_tensor=ref_tensor.to(self.model.device)
                         uc = None
                         if scale != 1.0:
                             uc = self.model.learnable_vector