Spaces:
Runtime error
Runtime error
| 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() |