Update app.py

app.py

@@ -320,7 +320,7 @@ def import_state(state, json_text):
 ### Main worker
 
 
-def register(state, drawpad):
+def register(state, drawpad, model):
     seed_everything(state.seed if state.seed >=0 else np.random.randint(2147483647))
     print('Generate!')
 
@@ -362,15 +362,15 @@ def register(state, drawpad):
     # prompts, negative_prompts = preprocess_prompts(
     #     prompts, negative_prompts, style_name=state.style_name, quality_name=state.quality_name)
 
-    state.model.update_background(
+    model.update_background(
         background.convert('RGB'),
         prompt=None,
         negative_prompt=None,
     )
-    state.prompts[0] = state.model.background.prompt
-    state.neg_prompts[0] = state.model.background.negative_prompt
+    state.prompts[0] = model.background.prompt
+    state.neg_prompts[0] = model.background.negative_prompt
 
-    state.model.update_layers(
+    model.update_layers(
         prompts=prompts,
         negative_prompts=negative_prompts,
         masks=masks.to(device),
@@ -384,23 +384,23 @@ def register(state, drawpad):
 
 @spaces.GPU(duration=120)
 def run(state, drawpad):
-    state.model = model
-    state.model.device = torch.device('cuda')
-    state.model.reset_seed(state.model.generator, opt.seed)
-    state.model.reset_latent()
-    state.model.prepare()
+    model.device = torch.device('cuda')
+    model.reset_seed(model.generator, opt.seed)
+    model.reset_latent()
+    model.prepare()
 
-    state = register(state, drawpad)
+    state = register(state, drawpad, model)
     state.is_running = True
 
     tic = time.time()
     while True:
-        yield [state, state.model()]
+        yield [state, model()]
         toc = time.time()
         tdelta = toc - tic
         if tdelta > opt.run_time:
             state.is_running = False
-            return [state, state.model()]
+            state.model = None
+            return [state, model()]
 
 
 def hide_element():
@@ -412,7 +412,11 @@ def show_element():
 
 
 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
 
     user_input = np.asarray(drawpad['layers'][0]) # (H, W, 4)
@@ -601,7 +605,7 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css, head=head) as demo:
         state.model_id = opt.model
         state.style_name = '(None)'
         state.quality_name = 'Standard v3.1'
-        state.model = model
+        state.model = None
 
         # State variables (one-hot).
         state.active_palettes = 5