Update app.py

app.py

@@ -146,6 +146,7 @@ opt.colors = [
     # '#92C6EC',
     # '#FECAC0',
 ]
+opt.excluded_keys = ['inpainting_mode', 'is_runing', 'active_palettes', 'current_palette', 'model']
 
 
 ### Event handlers
@@ -294,7 +295,7 @@ def import_state(state, json_text):
     current_palette = state.current_palette
     # active_palettes = state.active_palettes
     state_dict = json.loads(json_text)
-    for k in ('inpainting_mode', 'is_runing', 'active_palettes', 'current_palette'):
+    for k in opt.excluded_keys:
         if k in state_dict:
             del state_dict[k]
     state = argparse.Namespace(**state_dict)
@@ -362,15 +363,15 @@ def register(state, drawpad):
     # prompts, negative_prompts = preprocess_prompts(
     #     prompts, negative_prompts, style_name=state.style_name, quality_name=state.quality_name)
 
-    model.update_background(
+    state.model.update_background(
         background.convert('RGB'),
         prompt=None,
         negative_prompt=None,
     )
-    state.prompts[0] = model.background.prompt
-    state.neg_prompts[0] = model.background.negative_prompt
+    state.prompts[0] = state.model.background.prompt
+    state.neg_prompts[0] = state.model.background.negative_prompt
 
-    model.update_layers(
+    state.model.update_layers(
         prompts=prompts,
         negative_prompts=negative_prompts,
         masks=masks.to(device),
@@ -384,10 +385,10 @@ def register(state, drawpad):
 
 @spaces.GPU(duration=120)
 def run(state, drawpad):
-    model.device = torch.device('cuda')
-    model.reset_seed(model.generator, opt.seed)
-    model.reset_latent()
-    model.prepare()
+    state.model.device = torch.device('cuda')
+    state.model.reset_seed(state.model.generator, opt.seed)
+    state.model.reset_latent()
+    state.model.prepare()
 
     state = register(state, drawpad)
     state.is_running = True
@@ -438,7 +439,7 @@ def draw(state, drawpad):
     #     mask_stds = [0] + [state.mask_stds[v] for v in has_masks]
 
     for i in range(len(has_masks)):
-        model.update_single_layer(
+        state.model.update_single_layer(
             idx=i,
             mask=masks[i],
             mask_strength=mask_strengths[i],
@@ -516,10 +517,6 @@ css = f"""
     width: 100%;
     aspect-ratio: {opt.width} / {opt.height};
 }}
-
-.layer-wrap {{
-    display: none;
-}}
 """
 
 for i in range(opt.max_palettes + 1):
@@ -604,6 +601,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 variables (one-hot).
         state.active_palettes = 5
@@ -1142,7 +1140,7 @@ async () => {{
     #     api_name='quality_select',
     # )
 
-    iface.btn_export_state.click(lambda x: vars(x), state, iface.json_state_export)
+    iface.btn_export_state.click(lambda x: {k: v for k, v in vars(x).items() if k not in opt.excluded_keys}, state, iface.json_state_export)
     iface.btn_import_state.click(import_state, [state, iface.tbox_state_import], [
         state,
         *iface.btn_semantics,