Update app.py

app.py

@@ -27,6 +27,8 @@ print('Use device:', device)
 
 net = torch.jit.load(f'weights/pkp-v1.{device}.jit.pt')
 
+model_net = torch.load(f'weights/colorizer.pt')
+
 
 def resize_original(img: Image.Image):
     if img is None:
@@ -53,7 +55,7 @@ def resize_original(img: Image.Image):
     return gr.Image.update(value=guide_img.convert('RGBA')), guide_img.convert('RGBA')
 
 
-def colorize(img: Dict[str, Image.Image], guide_img: Image.Image, seed: int, hint_mode: str):
+def colorize(img: Dict[str Image.Image], guide_img: Image.Image, seed: int, hint_mode: str):
     if not isinstance(img, dict):
         return gr.update(visible=True)
 
@@ -84,6 +86,23 @@ def colorize(img: Dict[str, Image.Image], guide_img: Image.Image, seed: int, hin
     return Image.fromarray(out).convert('RGB')
 
 
+def colorize2(img: Image.Image, model_option: str):
+    if not isinstance(img, dict):
+        return gr.update(visible=True)
+
+    if hint_mode == "Model 1":
+        model_int = 0
+    elif hint_mode == "Model 2":
+        model_int = 0
+
+    with torch.inference_mode():
+        out2 = model(input)
+        out = sample[0].cpu().numpy().transpose([1,2,0])
+        out = np.uint8(((out + 1) / 2 * 255).clip(0,255))
+    
+    return Image.fromarray(out).convert('RGB')
+
+
 with gr.Blocks() as demo:
     gr.Markdown('''<center><h1>Image Colorization With Hint</h1></center>
 <h2>Colorize your images/sketches with hint points.</h2>
@@ -110,6 +129,25 @@ with gr.Blocks() as demo:
             btn = gr.Button("Run")
         with gr.Column():
             output = gr.Image(type="pil", label="Output", interactive=False)
+    with gr.Row():
+        with gr.Column():
+            inp2 = gr.Image(
+                source="upload", 
+                type="pil", 
+                label="Sketch",
+                interactive=True
+            )
+            inp_store2 = gr.Image(
+                type="pil", 
+                interactive=False
+            )
+            inp_store2.visible = False
+        with gr.Column():
+            # seed = gr.Slider(1, 2**32, step=1, label="Seed", interactive=True, randomize=True)
+            model_option = gr.Radio(["Model 1", "Model 2"], value="Model 1", label="Model 2")
+            btn2 = gr.Button("Run Colorization")
+        with gr.Column():
+            output2 = gr.Image(type="pil", label="Output2", interactive=False)
     gr.Markdown('''
 Upon uploading an image, kindly give color hints at specific points, and then run the model. Average inference time is about 52 seconds.<br />
 ''')
@@ -124,6 +162,11 @@ Upon uploading an image, kindly give color hints at specific points, and then ru
         [inp, inp_store, seed, hint_mode],
         output
     )
+    btn2.click(
+        colorize2, 
+        [inp, model_option],
+        output2
+    )
 
 if __name__ == "__main__":
     demo.launch()