import gradio as gr import matplotlib.pyplot as plt from PIL import 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): image_file = Image.fromarray(np.array(image_file)[:, ::-1, :], "RGB") image_file = normalize_test(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()