Update app.py

app.py

@@ -1,23 +1,241 @@
-import tensorflow as tf 
-import pathlib
 import gradio as gr
 import matplotlib.pyplot as plt 
-from huggingface_hub import from_pretrained_keras
-import numpy as np
 
-# Normalizing the images to [-1, 1]
-def normalize_test(input_image):
-  input_image = tf.cast(input_image, tf.float32)
-  input_image = (input_image / 127.5) - 1
-  return input_image
-  
-def resize(input_image, height, width):
-  input_image = tf.image.resize(input_image, [height, width],
-                                method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
-  return input_image
-  
+import torch.nn as nn
+import torch.nn.functional as F
+import torch
+
+from pathlib import Path
+from re import TEMPLATE
+from typing import Optional, Union
+import os
+
+from huggingface_hub import PyTorchModelHubMixin, HfApi, HfFolder, Repository
+
+TEMPLATE_MODEL_CARD_PATH = "dummy"
+
+class HugGANModelHubMixin(PyTorchModelHubMixin):
+    """A mixin to push PyTorch Models to the Hugging Face Hub. This
+    mixin was adapted from the PyTorchModelHubMixin to also push a template
+    README.md for the HugGAN sprint.
+    """
+
+    def push_to_hub(
+        self,
+        repo_path_or_name: Optional[str] = None,
+        repo_url: Optional[str] = None,
+        commit_message: Optional[str] = "Add model",
+        organization: Optional[str] = None,
+        private: Optional[bool] = None,
+        api_endpoint: Optional[str] = None,
+        use_auth_token: Optional[Union[bool, str]] = None,
+        git_user: Optional[str] = None,
+        git_email: Optional[str] = None,
+        config: Optional[dict] = None,
+        skip_lfs_files: bool = False,
+        default_model_card: Optional[str] = TEMPLATE_MODEL_CARD_PATH
+    ) -> str:
+        """
+        Upload model checkpoint or tokenizer files to the Hub while
+        synchronizing a local clone of the repo in `repo_path_or_name`.
+        Parameters:
+            repo_path_or_name (`str`, *optional*):
+                Can either be a repository name for your model or tokenizer in
+                the Hub or a path to a local folder (in which case the
+                repository will have the name of that local folder). If not
+                specified, will default to the name given by `repo_url` and a
+                local directory with that name will be created.
+            repo_url (`str`, *optional*):
+                Specify this in case you want to push to an existing repository
+                in the hub. If unspecified, a new repository will be created in
+                your namespace (unless you specify an `organization`) with
+                `repo_name`.
+            commit_message (`str`, *optional*):
+                Message to commit while pushing. Will default to `"add config"`,
+                `"add tokenizer"` or `"add model"` depending on the type of the
+                class.
+            organization (`str`, *optional*):
+                Organization in which you want to push your model or tokenizer
+                (you must be a member of this organization).
+            private (`bool`, *optional*):
+                Whether the repository created should be private.
+            api_endpoint (`str`, *optional*):
+                The API endpoint to use when pushing the model to the hub.
+            use_auth_token (`bool` or `str`, *optional*):
+                The token to use as HTTP bearer authorization for remote files.
+                If `True`, will use the token generated when running
+                `transformers-cli login` (stored in `~/.huggingface`). Will
+                default to `True` if `repo_url` is not specified.
+            git_user (`str`, *optional*):
+                will override the `git config user.name` for committing and
+                pushing files to the hub.
+            git_email (`str`, *optional*):
+                will override the `git config user.email` for committing and
+                pushing files to the hub.
+            config (`dict`, *optional*):
+                Configuration object to be saved alongside the model weights.
+            default_model_card (`str`, *optional*):
+                Path to a markdown file to use as your default model card.
+        Returns:
+            The url of the commit of your model in the given repository.
+        """
+
+        if repo_path_or_name is None and repo_url is None:
+            raise ValueError(
+                "You need to specify a `repo_path_or_name` or a `repo_url`."
+            )
+
+        if use_auth_token is None and repo_url is None:
+            token = HfFolder.get_token()
+            if token is None:
+                raise ValueError(
+                    "You must login to the Hugging Face hub on this computer by typing `huggingface-cli login` and "
+                    "entering your credentials to use `use_auth_token=True`. Alternatively, you can pass your own "
+                    "token as the `use_auth_token` argument."
+                )
+        elif isinstance(use_auth_token, str):
+            token = use_auth_token
+        else:
+            token = None
+
+        if repo_path_or_name is None:
+            repo_path_or_name = repo_url.split("/")[-1]
+
+        # If no URL is passed and there's no path to a directory containing files, create a repo
+        if repo_url is None and not os.path.exists(repo_path_or_name):
+            repo_id = Path(repo_path_or_name).name
+            if organization:
+                repo_id = f"{organization}/{repo_id}"
+            repo_url = HfApi(endpoint=api_endpoint).create_repo(
+                repo_id=repo_id,
+                token=token,
+                private=private,
+                repo_type=None,
+                exist_ok=True,
+            )
+
+        repo = Repository(
+            repo_path_or_name,
+            clone_from=repo_url,
+            use_auth_token=use_auth_token,
+            git_user=git_user,
+            git_email=git_email,
+            skip_lfs_files=skip_lfs_files
+        )
+        repo.git_pull(rebase=True)
+
+        # Save the files in the cloned repo
+        self.save_pretrained(repo_path_or_name, config=config)
+
+        model_card_path = Path(repo_path_or_name) / 'README.md'
+        if not model_card_path.exists():
+            model_card_path.write_text(TEMPLATE_MODEL_CARD_PATH.read_text())
+
+        # Commit and push!
+        repo.git_add()
+        repo.git_commit(commit_message)
+        return repo.git_push()
+
+
+def weights_init_normal(m):
+    classname = m.__class__.__name__
+    if classname.find("Conv") != -1:
+        torch.nn.init.normal_(m.weight.data, 0.0, 0.02)
+    elif classname.find("BatchNorm2d") != -1:
+        torch.nn.init.normal_(m.weight.data, 1.0, 0.02)
+        torch.nn.init.constant_(m.bias.data, 0.0)
+
+
+##############################
+#           U-NET
+##############################
+
+
+class UNetDown(nn.Module):
+    def __init__(self, in_size, out_size, normalize=True, dropout=0.0):
+        super(UNetDown, self).__init__()
+        layers = [nn.Conv2d(in_size, out_size, 4, 2, 1, bias=False)]
+        if normalize:
+            layers.append(nn.InstanceNorm2d(out_size))
+        layers.append(nn.LeakyReLU(0.2))
+        if dropout:
+            layers.append(nn.Dropout(dropout))
+        self.model = nn.Sequential(*layers)
+
+    def forward(self, x):
+        return self.model(x)
+
+
+class UNetUp(nn.Module):
+    def __init__(self, in_size, out_size, dropout=0.0):
+        super(UNetUp, self).__init__()
+        layers = [
+            nn.ConvTranspose2d(in_size, out_size, 4, 2, 1, bias=False),
+            nn.InstanceNorm2d(out_size),
+            nn.ReLU(inplace=True),
+        ]
+        if dropout:
+            layers.append(nn.Dropout(dropout))
+
+        self.model = nn.Sequential(*layers)
+
+    def forward(self, x, skip_input):
+        x = self.model(x)
+        x = torch.cat((x, skip_input), 1)
+
+        return x
+
+
+class GeneratorUNet(nn.Module, HugGANModelHubMixin):
+    def __init__(self, in_channels=3, out_channels=3):
+        super(GeneratorUNet, self).__init__()
+
+        self.down1 = UNetDown(in_channels, 64, normalize=False)
+        self.down2 = UNetDown(64, 128)
+        self.down3 = UNetDown(128, 256)
+        self.down4 = UNetDown(256, 512, dropout=0.5)
+        self.down5 = UNetDown(512, 512, dropout=0.5)
+        self.down6 = UNetDown(512, 512, dropout=0.5)
+        self.down7 = UNetDown(512, 512, dropout=0.5)
+        self.down8 = UNetDown(512, 512, normalize=False, dropout=0.5)
+
+        self.up1 = UNetUp(512, 512, dropout=0.5)
+        self.up2 = UNetUp(1024, 512, dropout=0.5)
+        self.up3 = UNetUp(1024, 512, dropout=0.5)
+        self.up4 = UNetUp(1024, 512, dropout=0.5)
+        self.up5 = UNetUp(1024, 256)
+        self.up6 = UNetUp(512, 128)
+        self.up7 = UNetUp(256, 64)
+
+        self.final = nn.Sequential(
+            nn.Upsample(scale_factor=2),
+            nn.ZeroPad2d((1, 0, 1, 0)),
+            nn.Conv2d(128, out_channels, 4, padding=1),
+            nn.Tanh(),
+        )
+
+    def forward(self, x):
+        # U-Net generator with skip connections from encoder to decoder
+        d1 = self.down1(x)
+        d2 = self.down2(d1)
+        d3 = self.down3(d2)
+        d4 = self.down4(d3)
+        d5 = self.down5(d4)
+        d6 = self.down6(d5)
+        d7 = self.down7(d6)
+        d8 = self.down8(d7)
+        u1 = self.up1(d8, d7)
+        u2 = self.up2(u1, d6)
+        u3 = self.up3(u2, d5)
+        u4 = self.up4(u3, d4)
+        u5 = self.up5(u4, d3)
+        u6 = self.up6(u5, d2)
+        u7 = self.up7(u6, d1)
+
+        return self.final(u7)
+
 def load_image_infer(image_file):
-  input_image = resize(image_file, 256, 256)
+  imageA = Image.fromarray(np.array(imageA)[:, ::-1, :], "RGB")
   input_image = normalize_test(input_image)
 
   return input_image
@@ -27,20 +245,19 @@ def generate_images(test_input):
   prediction = generator(np.expand_dims(test_input, axis=0), training=True)
   fig = plt.figure(figsize=(128, 128))
   title = ['Predicted Image']
-
+  
   plt.title('Predicted Image')
   # Getting the pixel values in the [0, 1] range to plot.
   plt.imshow(prediction[0,:,:,:] * 0.5 + 0.5)
   plt.axis('off')
   return fig
-  
-  
-generator = from_pretrained_keras("keras-io/pix2pix-generator")
 
+generator = GeneratorUNet()
+generator.from_pretrained("huggan/pix2pix-edge2shoes")
 
 img = gr.inputs.Image(shape=(256,256))
 plot = gr.outputs.Image(type="plot")
 
-description = "Conditional GAN model that translates image-to-image."
+description = "Pix2pix model that translates image-to-image."
 gr.Interface(generate_images, inputs = img, outputs = plot, 
-title = "Pix2Pix Shoes Reconstructor", description = description, examples = [["./img.png"]]).launch()
+title = "Pix2Pix Shoes Reconstructor", description = description).launch()