Support concurrent per-call model choice (Before using a global model)

app.py

@@ -1,6 +1,7 @@
 import gradio as gr
 import os
 import cv2
+from regex import D
 import torch
 import numpy as np
 import argparse
@@ -28,10 +29,12 @@ actor.load_state_dict(torch.load(actor_path))
 actor = actor.to(device).eval()
 Decoder = Decoder.to(device).eval()
 
+decoders = {"Default": Decoder}
+actors = {"Default": actor}
 
-def decode(x, canvas): # b * (10 + 3)
+def decode(x, canvas, decoder = Decoder): # b * (10 + 3)
     x = x.view(-1, 10 + 3)
-    stroke = 1 - Decoder(x[:, :10])
+    stroke = 1 - decoder(x[:, :10])
     stroke = stroke.view(-1, width, width, 1)
     color_stroke = stroke * x[:, -3:].view(-1, 1, 1, 3)
     stroke = stroke.permute(0, 3, 1, 2)
@@ -98,7 +101,9 @@ def save_img(res, imgid, origin_shape, output_name, divide=False):
 
 
 
-def paint_img(img, max_step = 40):
+def paint_img(img, max_step = 40, model_choices = "Default"):
+    Decoder = decoders[model_choices]
+    actor = actors[model_choices]
     max_step = int(max_step)
     # imgid = 0
     # output_name = os.path.join('output', str(len(os.listdir('output'))) if os.path.exists('output') else '0')
@@ -130,7 +135,7 @@ def paint_img(img, max_step = 40):
         for i in range(max_step):
             stepnum = T * i / max_step
             actions = actor(torch.cat([canvas, img, stepnum, coord], 1))
-            canvas, res = decode(actions, canvas)
+            canvas, res = decode(actions, canvas, Decoder)
             for j in range(5):
                 # save_img(res[j], imgid)
                 # imgid += 1
@@ -152,7 +157,7 @@ def paint_img(img, max_step = 40):
             for i in range(max_step):
                 stepnum = T * i / max_step
                 actions = actor(torch.cat([canvas, patch_img, stepnum, coord], 1))
-                canvas, res = decode(actions, canvas)
+                canvas, res = decode(actions, canvas, Decoder)
                 # print('divided canvas step {}, L2Loss = {}'.format(i, ((canvas - patch_img) ** 2).mean()))
                 for j in range(5):
                     # save_img(res[j], imgid, True)
@@ -168,8 +173,8 @@ def paint_img(img, max_step = 40):
         yield output
 
 
-def change_model(choice: str):
-    global Decoder, actor
+def load_model_if_needed(choice: str):
+    # global Decoder, actor
     if choice == "Default":
         actor_path = 'ckpts/actor.pkl'
         renderer_path = 'ckpts/renderer.pkl'
@@ -182,11 +187,19 @@ def change_model(choice: str):
     else:
         actor_path = 'ckpts/actor_notrans.pkl'
         renderer_path = 'ckpts/bezierwotrans.pkl'                
-
-    Decoder.load_state_dict(torch.load(renderer_path, map_location= "cpu"))
-    actor.load_state_dict(torch.load(actor_path, map_location= "cpu"))
-    actor = actor.to(device).eval()
-    Decoder = Decoder.to(device).eval()
+    if choice not in decoders:
+        Decoder = FCN()
+        Decoder.load_state_dict(torch.load(renderer_path, map_location= "cpu"))
+        Decoder = Decoder.to(device).eval()
+        decoders[choice] = Decoder
+    if choice not in actors:
+        actor = ResNet(9, 18, 65) # action_bundle = 5, 65 = 5 * 13
+        actor.load_state_dict(torch.load(actor_path, map_location= "cpu"))
+        actor = actor.to(device).eval()
+        actors[choice] = actor
+    
+    
+    
 
 from typing import Generator
 def wrapper(func):
@@ -233,9 +246,9 @@ with gr.Blocks() as demo:
             output.render()
 
 
-    dropdown.select(change_model, dropdown)
+    dropdown.select(load_model_if_needed, dropdown)
     click_event = translate_btn.click(lambda x: gr.Button(value="Cancel", variant="stop") if x == "Paint" else gr.Button(value="Paint", variant="primary"), translate_btn, translate_btn)\
-                 .then(wrapper(paint_img), inputs=[translate_btn, input_image, step], outputs=output, trigger_mode = 'multiple')\
+                 .then(wrapper(paint_img), inputs=[translate_btn, input_image, step, dropdown], outputs=output, trigger_mode = 'multiple')\
                  .then(lambda x: gr.Button(value="Paint", variant="primary"), translate_btn, translate_btn)
     clr_btn.click(None, None, cancels=[click_event])
     examples = gr.Examples(examples=examples,