Update app.py

app.py

@@ -383,7 +383,7 @@ def register(state, drawpad, model):
     return state
 
 
-@spaces.GPU(duration=120)
+@spaces.GPU(duration=(opt.prep_time + opt.run_time + 5))
 def run(state, drawpad):
     model.device = torch.device('cuda')
     model.reset_seed(model.generator, opt.seed)
@@ -412,10 +412,8 @@ def show_element():
     return gr.update(visible=True)
 
 
+@spaces.GPU
 def draw(state, drawpad):
-    if not hasattr(state, 'model') or state.model is None:
-        print('[WARNING] Model is not registered, update ignored.')
-        return
     if not state.is_running:
         print('[WARNING] Streaming is currently off, update ignored.')
         return
@@ -444,7 +442,7 @@ def draw(state, drawpad):
     #     mask_stds = [0] + [state.mask_stds[v] for v in has_masks]
 
     for i in range(len(has_masks)):
-        state.model.update_single_layer(
+        model.update_single_layer(
             idx=i,
             mask=masks[i],
             mask_strength=mask_strengths[i],