update to latest version of gradio

.pre-commit-config.yaml

@@ -1,60 +0,0 @@
-repos:
-  - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v4.6.0
-    hooks:
-      - id: check-executables-have-shebangs
-      - id: check-json
-      - id: check-merge-conflict
-      - id: check-shebang-scripts-are-executable
-      - id: check-toml
-      - id: check-yaml
-      - id: end-of-file-fixer
-      - id: mixed-line-ending
-        args: ["--fix=lf"]
-      - id: requirements-txt-fixer
-      - id: trailing-whitespace
-  - repo: https://github.com/myint/docformatter
-    rev: v1.7.5
-    hooks:
-      - id: docformatter
-        args: ["--in-place"]
-  - repo: https://github.com/pycqa/isort
-    rev: 5.13.2
-    hooks:
-      - id: isort
-        args: ["--profile", "black"]
-  - repo: https://github.com/pre-commit/mirrors-mypy
-    rev: v1.10.0
-    hooks:
-      - id: mypy
-        args: ["--ignore-missing-imports"]
-        additional_dependencies:
-          [
-            "types-python-slugify",
-            "types-requests",
-            "types-PyYAML",
-            "types-pytz",
-          ]
-  - repo: https://github.com/psf/black
-    rev: 24.4.2
-    hooks:
-      - id: black
-        language_version: python3.10
-        args: ["--line-length", "119"]
-  - repo: https://github.com/kynan/nbstripout
-    rev: 0.7.1
-    hooks:
-      - id: nbstripout
-        args:
-          [
-            "--extra-keys",
-            "metadata.interpreter metadata.kernelspec cell.metadata.pycharm",
-          ]
-  - repo: https://github.com/nbQA-dev/nbQA
-    rev: 1.8.5
-    hooks:
-      - id: nbqa-black
-      - id: nbqa-pyupgrade
-        args: ["--py37-plus"]
-      - id: nbqa-isort
-        args: ["--float-to-top"]

.vscode/settings.json

@@ -1,30 +0,0 @@
-{
-    "editor.formatOnSave": true,
-    "files.insertFinalNewline": false,
-    "[python]": {
-        "editor.defaultFormatter": "ms-python.black-formatter",
-        "editor.formatOnType": true,
-        "editor.codeActionsOnSave": {
-            "source.organizeImports": "explicit"
-        }
-    },
-    "[jupyter]": {
-        "files.insertFinalNewline": false
-    },
-    "black-formatter.args": [
-        "--line-length=119"
-    ],
-    "isort.args": ["--profile", "black"],
-    "flake8.args": [
-        "--max-line-length=119"
-    ],
-    "ruff.lint.args": [
-        "--line-length=119"
-    ],
-    "notebook.output.scrolling": true,
-    "notebook.formatOnCellExecution": true,
-    "notebook.formatOnSave.enabled": true,
-        "notebook.codeActionsOnSave": {
-            "source.organizeImports": "explicit"
-        }
-}

README.md

@@ -4,12 +4,9 @@ emoji: 🐨
 colorFrom: indigo
 colorTo: red
 sdk: gradio
-sdk_version: 4.36.1
+sdk_version: 3.0.5
 app_file: app.py
 pinned: false
-suggested_hardware: t4-small
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces#reference
-
-https://arxiv.org/abs/2202.12211

app.py

@@ -2,105 +2,155 @@
 
 from __future__ import annotations
 
-import pathlib
+import argparse
+import functools
+import os
+import pickle
+import sys
 
 import gradio as gr
 import numpy as np
-
-from model import Model
-
-DESCRIPTION = "# [Self-Distilled StyleGAN](https://github.com/self-distilled-stylegan/self-distilled-internet-photos)"
-
-
-def get_sample_image_url(name: str) -> str:
-    sample_image_dir = "https://huggingface.co/spaces/hysts/Self-Distilled-StyleGAN/resolve/main/samples"
-    return f"{sample_image_dir}/{name}.jpg"
-
-
-def get_sample_image_markdown(name: str) -> str:
-    url = get_sample_image_url(name)
-    size = name.split("_")[1]
-    truncation_type = "_".join(name.split("_")[2:])
-    return f"""
-    - size: {size}x{size}
-    - seed: 0-99
-    - truncation: 0.7
-    - truncation type: {truncation_type}
-    ![sample images]({url})"""
-
-
-def get_cluster_center_image_url(model_name: str) -> str:
-    cluster_center_image_dir = (
-        "https://huggingface.co/spaces/hysts/Self-Distilled-StyleGAN/resolve/main/cluster_center_images"
-    )
-    return f"{cluster_center_image_dir}/{model_name}.jpg"
-
-
-def get_cluster_center_image_markdown(model_name: str) -> str:
-    url = get_cluster_center_image_url(model_name)
-    return f"![cluster center images]({url})"
-
-
-model = Model()
-
-with gr.Blocks(css="style.css") as demo:
-    gr.Markdown(DESCRIPTION)
-
-    with gr.Tabs():
-        with gr.TabItem("App"):
-            with gr.Row():
-                with gr.Column():
-                    with gr.Group():
-                        model_name = gr.Dropdown(label="Model", choices=model.MODEL_NAMES, value=model.MODEL_NAMES[0])
-                        seed = gr.Slider(label="Seed", minimum=0, maximum=np.iinfo(np.uint32).max, step=1, value=0)
-                        psi = gr.Slider(label="Truncation psi", minimum=0, maximum=2, step=0.05, value=0.7)
-                        truncation_type = gr.Dropdown(
-                            label="Truncation Type", choices=model.TRUNCATION_TYPES, value=model.TRUNCATION_TYPES[0]
-                        )
-                        run_button = gr.Button("Run")
-                with gr.Column():
-                    result = gr.Image(label="Result", elem_id="result")
-
-        with gr.TabItem("Sample Images"):
-            with gr.Row():
-                paths = sorted(pathlib.Path("samples").glob("*"))
-                names = [path.stem for path in paths]
-                model_name2 = gr.Dropdown(label="Type", choices=names, value="dogs_1024_multimodal_lpips")
-            with gr.Row():
-                text = get_sample_image_markdown(model_name2.value)
-                sample_images = gr.Markdown(text)
-
-        with gr.TabItem("Cluster Center Images"):
-            with gr.Row():
-                model_name3 = gr.Dropdown(label="Model", choices=model.MODEL_NAMES, value=model.MODEL_NAMES[0])
-            with gr.Row():
-                text = get_cluster_center_image_markdown(model_name3.value)
-                cluster_center_images = gr.Markdown(value=text)
-
-    model_name.change(
-        fn=model.set_model,
-        inputs=model_name,
-    )
-    run_button.click(
-        fn=model.set_model_and_generate_image,
-        inputs=[
-            model_name,
-            seed,
-            psi,
-            truncation_type,
+import torch
+import torch.nn as nn
+from huggingface_hub import hf_hub_download
+
+sys.path.insert(0, 'stylegan3')
+
+TITLE = 'Self-Distilled StyleGAN'
+DESCRIPTION = '''This is an unofficial demo for models provided in https://github.com/self-distilled-stylegan/self-distilled-internet-photos.
+
+Expected execution time on Hugging Face Spaces: 2s
+'''
+SAMPLE_IMAGE_DIR = 'https://huggingface.co/spaces/hysts/Self-Distilled-StyleGAN/resolve/main/samples'
+ARTICLE = f'''## Generated images
+- truncation: 0.7
+### Dogs
+- size: 1024x1024
+- seed: 0-99
+![Dogs]({SAMPLE_IMAGE_DIR}/dogs.jpg)
+### Elephants
+- size: 512x512
+- seed: 0-99
+![Elephants]({SAMPLE_IMAGE_DIR}/elephants.jpg)
+### Horses
+- size: 256x256
+- seed: 0-99
+![Horses]({SAMPLE_IMAGE_DIR}/horses.jpg)
+### Bicycles
+- size: 256x256
+- seed: 0-99
+![Bicycles]({SAMPLE_IMAGE_DIR}/bicycles.jpg)
+### Lions
+- size: 512x512
+- seed: 0-99
+![Lions]({SAMPLE_IMAGE_DIR}/lions.jpg)
+### Giraffes
+- size: 512x512
+- seed: 0-99
+![Giraffes]({SAMPLE_IMAGE_DIR}/giraffes.jpg)
+### Parrots
+- size: 512x512
+- seed: 0-99
+![Parrots]({SAMPLE_IMAGE_DIR}/parrots.jpg)
+
+<center><img src="https://visitor-badge.glitch.me/badge?page_id=hysts.self-distilled-stylegan" alt="visitor badge"/></center>
+'''
+
+TOKEN = os.environ['TOKEN']
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--device', type=str, default='cpu')
+    parser.add_argument('--theme', type=str)
+    parser.add_argument('--live', action='store_true')
+    parser.add_argument('--share', action='store_true')
+    parser.add_argument('--port', type=int)
+    parser.add_argument('--disable-queue',
+                        dest='enable_queue',
+                        action='store_false')
+    parser.add_argument('--allow-flagging', type=str, default='never')
+    return parser.parse_args()
+
+
+def generate_z(z_dim: int, seed: int, device: torch.device) -> torch.Tensor:
+    return torch.from_numpy(np.random.RandomState(seed).randn(
+        1, z_dim)).to(device).float()
+
+
+@torch.inference_mode()
+def generate_image(model_name: str, seed: int, truncation_psi: float,
+                   model_dict: dict[str, nn.Module],
+                   device: torch.device) -> np.ndarray:
+    model = model_dict[model_name]
+    seed = int(np.clip(seed, 0, np.iinfo(np.uint32).max))
+
+    z = generate_z(model.z_dim, seed, device)
+    label = torch.zeros([1, model.c_dim], device=device)
+
+    out = model(z, label, truncation_psi=truncation_psi)
+    out = (out.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255).to(torch.uint8)
+    return out[0].cpu().numpy()
+
+
+def load_model(model_name: str, device: torch.device) -> nn.Module:
+    path = hf_hub_download('hysts/Self-Distilled-StyleGAN',
+                           f'models/{model_name}_pytorch.pkl',
+                           use_auth_token=TOKEN)
+    with open(path, 'rb') as f:
+        model = pickle.load(f)['G_ema']
+    model.eval()
+    model.to(device)
+    with torch.inference_mode():
+        z = torch.zeros((1, model.z_dim)).to(device)
+        label = torch.zeros([1, model.c_dim], device=device)
+        model(z, label)
+    return model
+
+
+def main():
+    args = parse_args()
+    device = torch.device(args.device)
+
+    model_names = [
+        'dogs_1024',
+        'elephants_512',
+        'horses_256',
+        'bicycles_256',
+        'lions_512',
+        'giraffes_512',
+        'parrots_512',
+    ]
+
+    model_dict = {name: load_model(name, device) for name in model_names}
+
+    func = functools.partial(generate_image,
+                             model_dict=model_dict,
+                             device=device)
+    func = functools.update_wrapper(func, generate_image)
+
+    gr.Interface(
+        func,
+        [
+            gr.inputs.Radio(
+                model_names, type='value', default='dogs_1024', label='Model'),
+            gr.inputs.Number(default=0, label='Seed'),
+            gr.inputs.Slider(
+                0, 2, step=0.05, default=0.7, label='Truncation psi'),
         ],
-        outputs=result,
-    )
-    model_name2.change(
-        fn=get_sample_image_markdown,
-        inputs=model_name2,
-        outputs=sample_images,
-    )
-    model_name3.change(
-        fn=get_cluster_center_image_markdown,
-        inputs=model_name3,
-        outputs=cluster_center_images,
+        gr.outputs.Image(type='numpy', label='Output'),
+        title=TITLE,
+        description=DESCRIPTION,
+        article=ARTICLE,
+        theme=args.theme,
+        allow_flagging=args.allow_flagging,
+        live=args.live,
+    ).launch(
+        enable_queue=args.enable_queue,
+        server_port=args.port,
+        share=args.share,
     )
 
-if __name__ == "__main__":
-    demo.queue(max_size=10).launch()
+
+if __name__ == '__main__':
+    main()

model.py

@@ -1,192 +0,0 @@
-from __future__ import annotations
-
-import pathlib
-import pickle
-import sys
-
-import lpips
-import numpy as np
-import torch
-import torch.nn as nn
-from huggingface_hub import hf_hub_download
-
-current_dir = pathlib.Path(__file__).parent
-submodule_dir = current_dir / "stylegan3"
-sys.path.insert(0, submodule_dir.as_posix())
-
-
-class LPIPS(lpips.LPIPS):
-    @staticmethod
-    def preprocess(image: np.ndarray) -> torch.Tensor:
-        data = torch.from_numpy(image).float() / 255
-        data = data * 2 - 1
-        return data.permute(2, 0, 1).unsqueeze(0)
-
-    @torch.inference_mode()
-    def compute_features(self, data: torch.Tensor) -> list[torch.Tensor]:
-        data = self.scaling_layer(data)
-        data = self.net(data)
-        return [lpips.normalize_tensor(x) for x in data]
-
-    @torch.inference_mode()
-    def compute_distance(self, features0: list[torch.Tensor], features1: list[torch.Tensor]) -> float:
-        res = 0
-        for lin, x0, x1 in zip(self.lins, features0, features1):
-            d = (x0 - x1) ** 2
-            y = lin(d)
-            y = lpips.lpips.spatial_average(y)
-            res += y.item()
-        return res
-
-
-class Model:
-
-    MODEL_NAMES = [
-        "dogs_1024",
-        "elephants_512",
-        "horses_256",
-        "bicycles_256",
-        "lions_512",
-        "giraffes_512",
-        "parrots_512",
-    ]
-    TRUNCATION_TYPES = [
-        "Multimodal (LPIPS)",
-        "Multimodal (L2)",
-        "Global",
-    ]
-
-    def __init__(self):
-        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-        self._download_all_models()
-        self._download_all_cluster_centers()
-        self._download_all_cluster_center_images()
-
-        self.model_name = self.MODEL_NAMES[0]
-        self.model = self._load_model(self.model_name)
-        self.cluster_centers = self._load_cluster_centers(self.model_name)
-        self.cluster_center_images = self._load_cluster_center_images(self.model_name)
-
-        self.lpips = LPIPS()
-        self.cluster_center_lpips_feature_dict = self._compute_cluster_center_lpips_features()
-
-    def _load_model(self, model_name: str) -> nn.Module:
-        path = hf_hub_download("public-data/Self-Distilled-StyleGAN", f"models/{model_name}_pytorch.pkl")
-        with open(path, "rb") as f:
-            model = pickle.load(f)["G_ema"]
-        model.eval()
-        model.to(self.device)
-        return model
-
-    def _load_cluster_centers(self, model_name: str) -> torch.Tensor:
-        path = hf_hub_download("public-data/Self-Distilled-StyleGAN", f"cluster_centers/{model_name}.npy")
-        centers = np.load(path)
-        centers = torch.from_numpy(centers).float().to(self.device)
-        return centers
-
-    def _load_cluster_center_images(self, model_name: str) -> np.ndarray:
-        path = hf_hub_download("public-data/Self-Distilled-StyleGAN", f"cluster_center_images/{model_name}.npy")
-        return np.load(path)
-
-    def set_model(self, model_name: str) -> None:
-        if model_name == self.model_name:
-            return
-        self.model_name = model_name
-        self.model = self._load_model(model_name)
-        self.cluster_centers = self._load_cluster_centers(model_name)
-        self.cluster_center_images = self._load_cluster_center_images(model_name)
-
-    def _download_all_models(self):
-        for name in self.MODEL_NAMES:
-            self._load_model(name)
-
-    def _download_all_cluster_centers(self):
-        for name in self.MODEL_NAMES:
-            self._load_cluster_centers(name)
-
-    def _download_all_cluster_center_images(self):
-        for name in self.MODEL_NAMES:
-            self._load_cluster_center_images(name)
-
-    def generate_z(self, seed: int) -> torch.Tensor:
-        seed = int(np.clip(seed, 0, np.iinfo(np.uint32).max))
-        return torch.from_numpy(np.random.RandomState(seed).randn(1, self.model.z_dim)).float().to(self.device)
-
-    def compute_w(self, z: torch.Tensor) -> torch.Tensor:
-        label = torch.zeros((1, self.model.c_dim), device=self.device)
-        w = self.model.mapping(z, label)
-        return w
-
-    @staticmethod
-    def truncate_w(w_center: torch.Tensor, w: torch.Tensor, psi: float) -> torch.Tensor:
-        if psi == 1:
-            return w
-        return w_center.lerp(w, psi)
-
-    @torch.inference_mode()
-    def synthesize(self, w: torch.Tensor) -> torch.Tensor:
-        return self.model.synthesis(w)
-
-    def postprocess(self, tensor: torch.Tensor) -> np.ndarray:
-        tensor = (tensor.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255).to(torch.uint8)
-        return tensor.cpu().numpy()
-
-    def compute_lpips_features(self, image: np.ndarray) -> list[torch.Tensor]:
-        data = self.lpips.preprocess(image)
-        return self.lpips.compute_features(data)
-
-    def _compute_cluster_center_lpips_features(self) -> dict[str, list[list[torch.Tensor]]]:
-        res = dict()
-        for name in self.MODEL_NAMES:
-            images = self._load_cluster_center_images(name)
-            res[name] = [self.compute_lpips_features(image) for image in images]
-        return res
-
-    def compute_distance_to_cluster_centers(self, ws: torch.Tensor, distance_type: str) -> list[torch.Tensor]:
-        if distance_type == "l2":
-            return self._compute_l2_distance_to_cluster_centers(ws)
-        elif distance_type == "lpips":
-            return self._compute_lpips_distance_to_cluster_centers(ws)
-        else:
-            raise ValueError
-
-    def _compute_l2_distance_to_cluster_centers(self, ws: torch.Tensor) -> np.ndarray:
-        dist2 = ((self.cluster_centers - ws[0, 0]) ** 2).sum(dim=1)
-        return dist2.cpu().numpy()
-
-    def _compute_lpips_distance_to_cluster_centers(self, ws: torch.Tensor) -> np.ndarray:
-        x = self.synthesize(ws)
-        x = self.postprocess(x)[0]
-        feat0 = self.compute_lpips_features(x)
-        cluster_center_features = self.cluster_center_lpips_feature_dict[self.model_name]
-        distances = [self.lpips.compute_distance(feat0, feat1) for feat1 in cluster_center_features]
-        return np.asarray(distances)
-
-    def find_nearest_cluster_center(self, ws: torch.Tensor, distance_type: str) -> int:
-        distances = self.compute_distance_to_cluster_centers(ws, distance_type)
-        return int(np.argmin(distances))
-
-    def generate_image(self, seed: int, truncation_psi: float, truncation_type: str) -> np.ndarray:
-        z = self.generate_z(seed)
-        ws = self.compute_w(z)
-        if truncation_type == self.TRUNCATION_TYPES[2]:
-            w0 = self.model.mapping.w_avg
-        else:
-            if truncation_type == self.TRUNCATION_TYPES[0]:
-                distance_type = "lpips"
-            elif truncation_type == self.TRUNCATION_TYPES[1]:
-                distance_type = "l2"
-            else:
-                raise ValueError
-            cluster_index = self.find_nearest_cluster_center(ws, distance_type)
-            w0 = self.cluster_centers[cluster_index]
-        new_ws = self.truncate_w(w0, ws, truncation_psi)
-        out = self.synthesize(new_ws)
-        out = self.postprocess(out)
-        return out[0]
-
-    def set_model_and_generate_image(
-        self, model_name: str, seed: int, truncation_psi: float, truncation_type: str
-    ) -> np.ndarray:
-        self.set_model(model_name)
-        return self.generate_image(seed, truncation_psi, truncation_type)

requirements.txt

@@ -1,6 +1,5 @@
-lpips==0.1.4
-numpy==1.26.4
-Pillow==10.3.0
-scipy==1.13.1
-torch==2.0.1
-torchvision==0.15.2
+numpy==1.22.3
+Pillow==9.0.1
+scipy==1.8.0
+torch==1.11.0
+torchvision==0.12.0

style.css

@@ -1,8 +0,0 @@
-h1 {
-  text-align: center;
-  display: block;
-}
-div#result {
-  max-width: 600px;
-  max-height: 600px;
-}