Update app.py

app.py

@@ -9,6 +9,8 @@ import uuid
 import torch
 from PIL import Image
 
+import spaces
+
 demo_path = Path(__file__).resolve().parent
 root_path = demo_path
 sys.path.append(str(root_path))
@@ -70,12 +72,20 @@ inpainting_models = OrderedDict([
     ("Stable-Inpainting 2.0", 'sd2_inp'),
     ("Stable-Inpainting 1.5", 'sd15_inp')
 ])
-sr_model = models.sd2_sr.load_model(device='cuda:1')
-sam_predictor = models.sam.load_model(device='cuda:0')
+sr_model = None
+sam_predictor = None
 
 inp_model_name = list(inpainting_models.keys())[0]
-inp_model = models.load_inpainting_model(
-    inpainting_models[inp_model_name], device='cuda:0', cache=True)
+inp_model = None
+
+
+@spaces.GPU(duration=120)
+def load_models():
+    global sr_model, sam_predictor, inp_model
+    sr_model = models.sd2_sr.load_model(device='cuda')
+    sam_predictor = models.sam.load_model(device='cuda')
+    inp_model = models.load_inpainting_model(
+        inpainting_models[inp_model_name], device='cuda', cache=True)
 
 
 def set_model_from_name(new_inp_model_name):
@@ -84,7 +94,7 @@ def set_model_from_name(new_inp_model_name):
     if new_inp_model_name != inp_model_name:
         print (f"Activating Inpaintng Model: {new_inp_model_name}")
         inp_model = models.load_inpainting_model(
-            inpainting_models[new_inp_model_name], device='cuda:0', cache=True)
+            inpainting_models[new_inp_model_name], device='cuda', cache=True)
         inp_model_name = new_inp_model_name
 
 
@@ -132,6 +142,7 @@ def recover_user_session(session_id):
     return hr_image, hr_mask, gallery, prompt
 
 
+@spaces.GPU(duration=120)
 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=''
@@ -193,6 +204,7 @@ def inpainting_run(model_name, use_rasg, use_painta, prompt, imageMask,
     return blended_images, session_id
 
 
+@spaces.GPU(duration=120)
 def upscale_run(
     ddim_steps, seed, use_sam_mask, session_id, img_index,
     negative_prompt='', positive_prompt='high resolution professional photo'
@@ -375,5 +387,6 @@ with gr.Blocks(css=demo_path / 'style.css') as demo:
         _js="function(a, b, c, d, e){ return [a, b, c, d, selected_gallery_index()] }",
     )
 
+load_models()
 demo.queue(max_size=20)
 demo.launch(share=True, allowed_paths=[str(TMP_DIR)])