Update app.py

app.py

@@ -1,271 +1,23 @@
 import gradio as gr
 from torchvision.transforms import Compose, Resize, ToTensor, Normalize
-import matplotlib.pyplot as plt 
 from PIL import Image
-import numpy as np
-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):
-    # Configure dataloaders
-    transform = Compose([
-                Resize((args.image_size, args.image_size), Image.BICUBIC),
-                ToTensor(),
-                Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
-            ])
-    image_file = Image.fromarray(np.array(image_file)[:, ::-1, :], "RGB")
-    image_file = transform(image_file)
-
-    return image_file
-  
-def generate_images(test_input):
-  test_input = load_image_infer(test_input)
-  prediction = generator(test_input).data
-  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 = GeneratorUNet()
-generator.from_pretrained("huggan/pix2pix-edge2shoes")
-
-img = gr.inputs.Image(shape=(256,256))
-plot = gr.outputs.Image(type="plot")
-
-description = "Pix2pix model that translates image-to-image."
-gr.Interface(generate_images, inputs = img, outputs = plot, 
-title = "Pix2Pix Shoes Reconstructor", description = description).launch()
+from torchvision.utils import save_image
+
+from huggan.pytorch.pix2pix.modeling_pix2pix import GeneratorUNet
+
+transform = Compose(
+    [
+        Resize((256, 256), Image.BICUBIC),
+        ToTensor(),
+        Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
+    ]
+)
+model = GeneratorUNet.from_pretrained('huggan/pix2pix-cityscapes')
+
+def predict_fn(img):
+    inp = transform(img).unsqueeze(0)
+    out = model(inp)
+    save_image(out, 'out.png', normalize=True)
+    return 'out.png'
+
+gr.Interface(predict_fn, inputs=gr.inputs.Image(type='pil'), outputs='image', examples=[['sample.jpeg']]).launch()