added new model and black

app.py

@@ -1,5 +1,6 @@
 # Hack for spaces
 import os
+
 os.system("pip uninstall -y gradio")
 os.system("pip install -r requirements.txt")
 
@@ -19,7 +20,7 @@ import albumentations as A
 import albumentations.pytorch as al_pytorch
 import torchvision
 from pl_bolts.models.gans import Pix2Pix
-
+from pl_bolts.models.gans.pix2pix.components import PatchGAN
 
 
 """ Class """
@@ -58,6 +59,24 @@ class OverpoweredPix2Pix(Pix2Pix):
         )
 
 
+class PatchGanChanged(OverpoweredPix2Pix):
+    def __init__(self, in_channels, out_channels):
+        super(PatchGanChanged, self).__init__(
+            in_channels=in_channels, out_channels=out_channels
+        )
+        self.patch_gan = self.get_dense_PatchGAN(self.patch_gan)
+
+    @staticmethod
+    def get_dense_PatchGAN(disc: PatchGAN) -> PatchGAN:
+        """Add final layer to gan"""
+        disc.final = torch.nn.Sequential(
+            disc.final,
+            torch.nn.Flatten(),
+            torch.nn.Linear(16 * 16, 1),
+        )
+        return disc
+
+
 """ Load the model """
 # train_64_val_16_patchgan_1val_plbolts_model_chkpt = "model/lightning_bolts_model/modified_path_gan.ckpt"
 train_64_val_16_plbolts_model_chkpt = (
@@ -66,6 +85,7 @@ train_64_val_16_plbolts_model_chkpt = (
 train_16_val_1_plbolts_model_chkpt = (
     "model/lightning_bolts_model/epoch=99-step=89000.ckpt"
 )
+modified_patchgan_chkpt = "model/lightning_bolts_model/modified_patchgan.ckpt"
 # model_checkpoint_path = "model/pix2pix_lightning_model/version_0/checkpoints/epoch=199-step=355600.ckpt"
 # model_checkpoint_path = "model/pix2pix_lightning_model/gen.pth"
 
@@ -81,6 +101,10 @@ train_16_val_1_plbolts_model = OverpoweredPix2Pix.load_from_checkpoint(
 )
 train_16_val_1_plbolts_model.eval()
 
+#
+modified_patchgan_model = PatchGanChanged.load_from_checkpoint(modified_patchgan_chkpt)
+modified_patchgan_model.eval()
+
 
 def predict(img: Image, type_of_model: str):
     """Create predictions"""
@@ -104,7 +128,7 @@ def predict(img: Image, type_of_model: str):
     elif type_of_model == "train batch size 64, val batch size 16":
         model = train_64_val_16_plbolts_model
     else:
-        raise Exception("NOT YET SUPPORTED")
+        model = modified_patchgan_model
 
     with torch.no_grad():
         result = model.gen(inference_img)
@@ -120,6 +144,13 @@ def predict2(img: Image):
     return predict(img=img, type_of_model="train batch size 64, val batch size 16")
 
 
+def predict3(img: Image):
+    return predict(
+        img=img,
+        type_of_model="train batch size 64, val batch size 16, patch gan has 1 output score instead of 16*16",
+    )
+
+
 model_input = gr.inputs.Radio(
     [
         "train batch size 16, val batch size 1",
@@ -177,6 +208,20 @@ with gr.Blocks() as demo:
                     outputs=image_output2,
                     fn=predict2,
                 )
+        with gr.TabItem("Single Value Discriminator"):
+            with gr.Row():
+                image_input3 = gr.inputs.Image(type="pil")
+                image_output3 = gr.outputs.Image(
+                    type="pil",
+                )
+            colour_3 = gr.Button("Colour it!")
+            with gr.Row():
+                gr.Examples(
+                    examples=img_examples,
+                    inputs=image_input3,
+                    outputs=image_output3,
+                    fn=predict3,
+                )
 
     colour_1.click(
         fn=predict1,
@@ -188,6 +233,11 @@ with gr.Blocks() as demo:
         inputs=image_input2,
         outputs=image_output2,
     )
+    colour_3.click(
+        fn=predict3,
+        inputs=image_input3,
+        outputs=image_output3,
+    )
 
 demo.title = "Colour your sketches!"
 demo.launch()

model/lightning_bolts_model/modified_patchgan.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7b6b85940399eb68eca7a62b603cd62ac2bc813bbec70a16df83842da73dd14a
+size 686280151