Delete app.py

app.py

@@ -1,237 +0,0 @@
-import os
-
-import random
-import torch
-import gradio as gr
-
-from e4e.models.psp import pSp
-from util import *
-from huggingface_hub import hf_hub_download
-
-import tempfile
-from argparse import Namespace
-import shutil
-
-import dlib
-import numpy as np
-import torchvision.transforms as transforms
-from torchvision import utils
-
-from model.sg2_model import Generator
-from generate_videos import project_code_by_edit_name
-import urllib.request
-import clip
-
-# Fetch image for analysis
-img_url = "http://claireye.com.tw/img/230212a.jpg"
-urllib.request.urlretrieve(img_url, "pose.jpg")
-model_dir = "models"
-os.makedirs(model_dir, exist_ok=True)
-
-model_repos = {
-    "e4e": ("aijack/e4e", "e4e.pt"),
-               "dlib": ("aijack/jojogan", "face_landmarks.dat"),
-               "base": ("aijack/stylegan2", "stylegan2.pt"),
-               "sketch": ("aijack/sketch", "sketch.pt"),
-                "jojo": ("aijack/jojo", "jojo.pt"),
-                "art": ("aijack/art", "art.pt"),
-                "arcane": ("aijack/arcane", "arcane.pt")
-
-}
-
-interface_gan_map = {"None": None, "Masculine": ("gender", 1.0), "Feminine": ("gender", -1.0),
-                     "Smiling": ("smile", 1.0),
-                     "Frowning": ("smile", -1.0), "Young": ("age", -1.0), "Old": ("age", 1.0),
-                     "Long Hair": ("hair_length", -1.0), "Short Hair": ("hair_length", 1.0)}
-
-
-def get_models():
-    os.makedirs(model_dir, exist_ok=True)
-
-    model_paths = {}
-
-    for model_name, repo_details in model_repos.items():
-        download_path = hf_hub_download(repo_id=repo_details[0], filename=repo_details[1])
-        model_paths[model_name] = download_path
-
-    return model_paths
-
-
-model_paths = get_models()
-
-
-class ImageEditor(object):
-    def __init__(self):
-        self.device = "cuda" if torch.cuda.is_available() else "cpu"
-
-        latent_size = 512
-        n_mlp = 8
-        channel_mult = 2
-        model_size = 1024
-
-        self.generators = {}
-
-        self.model_list = [name for name in model_paths.keys() if name not in ["e4e", "dlib"]]
-
-        for model in self.model_list:
-            g_ema = Generator(
-                model_size, latent_size, n_mlp, channel_multiplier=channel_mult
-            ).to(self.device)
-
-            checkpoint = torch.load(model_paths[model], map_location=self.device)
-
-            g_ema.load_state_dict(checkpoint["g_ema"], strict=False)
-
-            self.generators[model] = g_ema
-
-        self.experiment_args = {"model_path": model_paths["e4e"]}
-        self.experiment_args["transform"] = transforms.Compose(
-            [
-                transforms.Resize((256, 256)),
-                transforms.ToTensor(),
-                transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]),
-            ]
-        )
-        self.resize_dims = (256, 256)
-
-        model_path = self.experiment_args["model_path"]
-
-        ckpt = torch.load(model_path, map_location="cuda:0" if torch.cuda.is_available() else "cpu")
-        opts = ckpt["opts"]
-
-        opts["checkpoint_path"] = model_path
-        opts = Namespace(**opts)
-
-        self.e4e_net = pSp(opts, self.device)
-        self.e4e_net.eval()
-
-        self.shape_predictor = dlib.shape_predictor(
-            model_paths["dlib"]
-        )
-
-
-        self.clip_model, _ = clip.load("ViT-B/32", device=self.device)
-
-        print("setup complete")
-
-    def get_style_list(self):
-        style_list = []
-
-        for key in self.generators:
-            style_list.append(key)
-
-        return style_list
-
-    def invert_image(self, input_image):
-        input_image = self.run_alignment(str(input_image))
-
-        input_image = input_image.resize(self.resize_dims)
-
-        img_transforms = self.experiment_args["transform"]
-        transformed_image = img_transforms(input_image)
-
-        with torch.no_grad():
-            images, latents = self.run_on_batch(transformed_image.unsqueeze(0))
-            result_image, latent = images[0], latents[0]
-
-        inverted_latent = latent.unsqueeze(0).unsqueeze(1)
-
-        return inverted_latent
-
-    def get_generators_for_styles(self, output_styles, loop_styles=False):
-
-        if "base" in output_styles:  # always start with base if chosen
-            output_styles.insert(0, output_styles.pop(output_styles.index("base")))
-        if loop_styles:
-            output_styles.append(output_styles[0])
-
-        return [self.generators[style] for style in output_styles]
-
-
-
-    def get_target_latent(self, source_latent, alter, generators):
-        np_source_latent = source_latent.squeeze(0).cpu().detach().numpy()
-        if alter == "None":
-            return random.choice([source_latent.squeeze(0),] * max((len(generators) - 1), 1))
-        edit = interface_gan_map[alter]
-        projected_code_np = project_code_by_edit_name(np_source_latent, edit[0], edit[1])
-        return torch.from_numpy(projected_code_np).float().to(self.device)
-
-    def edit_image(self, input, output_styles, edit_choices):
-        return self.predict(input, output_styles, edit_choices=edit_choices)
-
-    def predict(
-            self,
-            input,  # Input image path
-            output_styles,  # Style checkbox options.
-            loop_styles=False,  # Loop back to the initial style
-            edit_choices=None,  # Optional dictionary with edit choice arguments
-    ):
-
-        if edit_choices is None:
-            edit_choices = {"edit_type": "None"}
-
-        # @title Align image
-        out_dir = tempfile.mkdtemp()
-
-        inverted_latent = self.invert_image(input)
-        generators = self.get_generators_for_styles(output_styles, loop_styles)
-        output_paths = []
-
-        with torch.no_grad():
-            for g_ema in generators:
-                latent_for_gen = self.get_target_latent(inverted_latent, edit_choices, generators)
-
-                img, _ = g_ema([latent_for_gen], input_is_latent=True, truncation=1, randomize_noise=False)
-
-                output_path = os.path.join(out_dir, f"out_{len(output_paths)}.jpg")
-                utils.save_image(img, output_path, nrow=1, normalize=True, range=(-1, 1))
-
-                output_paths.append(output_path)
-
-        return output_paths
-
-
-    def run_alignment(self, image_path):
-        aligned_image = align_face(filepath=image_path, predictor=self.shape_predictor)
-        print("Aligned image has shape: {}".format(aligned_image.size))
-        return aligned_image
-
-    def run_on_batch(self, inputs):
-        images, latents = self.e4e_net(
-            inputs.to(self.device).float(), randomize_noise=False, return_latents=True
-        )
-        return images, latents
-
-
-editor = ImageEditor()
-
-blocks = gr.Blocks(theme="darkdefault")
-
-with blocks:
-    gr.Markdown("<h1><center>Holiday Filters </center></h1>")
-    gr.Markdown(
-        "<div>Upload an image of your face, pick your desired output styles, pick any modifiers, and apply StyleGAN-based editing.</div>"
-    )
-    with gr.Row():
-        with gr.Column():
-            input_img = gr.Image(type="filepath", label="Input image")
-        with gr.Column():
-            style_choice = gr.CheckboxGroup(choices=editor.get_style_list(), value=editor.get_style_list(), type="value",                     label="Styles")
-            alter = gr.Dropdown(
-                choices=["None", "Masculine", "Feminine", "Smiling", "Frowning", "Young", "Old", "Short Hair",
-                         "Long Hair"], value="None", label="Additional Modifiers")
-            img_button = gr.Button("Edit Image")
-    with gr.Row():
-            img_output = gr.Gallery(label="Output Images")
-            img_output.style(grid=(3, 3, 4, 4, 6, 6))
-
-    img_button.click(fn=editor.edit_image, inputs=[input_img, style_choice, alter], outputs=img_output)
-    ex = gr.Examples(examples=[['pose.jpg', editor.get_style_list(), "Smiling"],['pose.jpg', editor.get_style_list(), "Long Hair"]], fn=editor.edit_image, inputs=[input_img, style_choice, alter],
-                     outputs=[img_output], cache_examples=True,
-                     run_on_click=True)
-    ex.dataset.headers = [""]
-    article = "<p style='text-align: center'><a href='http://claireye.com.tw'>Claireye</a> | 2023</p>"
-    gr.Markdown(article)
-
-blocks.launch(enable_queue=True)