Merge branch 'main' of https://huggingface.co/spaces/akhaliq/BlendGAN into main

app.py

@@ -1,29 +1,20 @@
 import os 
-os.system("git clone https://github.com/onion-liu/BlendGAN.git")
-os.system("gdown https://drive.google.com/uc?id=1eF04jKMLAb9DvzI72m8Akn5ykWf3EafE")
-os.system("gdown https://drive.google.com/uc?id=14nevG94hNkkwaoK5eJLF1iv78cv5O8fN")
-from PIL import Image
-import torch
 import gradio as gr
-model2 = torch.hub.load(
-    "AK391/animegan2-pytorch:main",
-    "generator",
-    pretrained=True,
-    device="cuda",
-    progress=False
-)
-model1 = torch.hub.load("AK391/animegan2-pytorch:main", "generator", pretrained="face_paint_512_v1",  device="cuda")
-face2paint = torch.hub.load(
-    'AK391/animegan2-pytorch:main', 'face2paint', 
-    size=512, device="cuda",side_by_side=False
-)
-def inference(img, ver):
-    os.system("""python style_transfer_folder.py --size 1024 --ckpt ./pretrained_models/blendgan.pt --psp_encoder_ckpt ./pretrained_models/psp_encoder.pt --style_img_path /content/BlendGAN/style/ --input_img_path /content/BlendGAN/input/ --outdir results/style_transfer/""")
-    return out
+from PIL import Image
+
+os.system("wget https://www.dropbox.com/s/fgupbov77x4rrru/blendgan.pt")
+os.system("wget https://www.dropbox.com/s/v8q0dd3r4u20659/psp_encoder.pt")
+
+def inference(content, style, index):
+    content.save('content.png')
+    style.save('style.png')
+    os.system("""python style_transfer_folder.py --size 1024 --add_weight_index """+str(int(index))+""" --ckpt ./blendgan.pt --psp_encoder_ckpt ./psp_encoder.pt --style_img_path style.png --input_img_path content.png""")
+    return "out.jpg"
   
-title = "AnimeGANv2"
-description = "Gradio Demo for AnimeGanv2 Face Portrait v2. To use it, simply upload your image, or click one of the examples to load them. Read more at the links below. Please use a cropped portrait picture for best results similar to the examples below."
-article = "<p style='text-align: center'><a href='https://github.com/bryandlee/animegan2-pytorch' target='_blank'>Github Repo Pytorch</a> | <a href='https://github.com/Kazuhito00/AnimeGANv2-ONNX-Sample' target='_blank'>Github Repo ONNX</a></p><p style='text-align: center'>samples from repo: <img src='https://user-images.githubusercontent.com/26464535/129888683-98bb6283-7bb8-4d1a-a04a-e795f5858dcf.gif' alt='animation'/> <img src='https://user-images.githubusercontent.com/26464535/137619176-59620b59-4e20-4d98-9559-a424f86b7f24.jpg' alt='animation'/></p>"
-examples=[['groot.jpeg','version 2 (🔺 robustness,🔻 stylization)'],['bill.png','version 1 (🔺 stylization, 🔻 robustness)'],['tony.png','version 1 (🔺 stylization, 🔻 robustness)'],['elon.png','version 2 (🔺 robustness,🔻 stylization)'],['IU.png','version 1 (🔺 stylization, 🔻 robustness)'],['billie.png','version 2 (🔺 robustness,🔻 stylization)'],['will.png','version 2 (🔺 robustness,🔻 stylization)'],['beyonce.jpeg','version 1 (🔺 stylization, 🔻 robustness)'],['gongyoo.jpeg','version 1 (🔺 stylization, 🔻 robustness)']]
-gr.Interface(inference, [gr.inputs.Image(type="pil"),gr.inputs.Radio(['version 1 (🔺 stylization, 🔻 robustness)','version 2 (🔺 robustness,🔻 stylization)'], type="value", default='version 2 (🔺 robustness,🔻 stylization)', label='version')
-], gr.outputs.Image(type="pil"),title=title,description=description,article=article,enable_queue=True,examples=examples,allow_flagging=False).launch()
+title = "BlendGAN"
+description = "Gradio Demo for BlendGAN: Implicitly GAN Blending for Arbitrary Stylized Face Generation. To use it, simply upload your images, or click one of the examples to load them. Read more at the links below. Please use a cropped portrait picture for best results similar to the examples below."
+article = "<p style='text-align: center'><a href='https://arxiv.org/abs/2110.11728' target='_blank'>BlendGAN: Implicitly GAN Blending for Arbitrary Stylized Face Generation</a> | <a href='https://github.com/onion-liu/BlendGAN' target='_blank'>Github Repo</a></p><p style='text-align: center'>samples from repo: <img src='https://user-images.githubusercontent.com/6346064/142623312-3e6f09aa-ce88-465c-b956-a8b4db95b4da.gif' alt='animation'/> <img src='https://user-images.githubusercontent.com/6346064/142621044-086cde48-8604-467b-8c43-8768b6670ec2.gif' alt='animation'/></p>"
+
+examples=[['000000.png','100001.png',6]]
+gr.Interface(inference, [gr.inputs.Image(type="pil"),gr.inputs.Image(type="pil"),gr.inputs.Slider(minimum=1, maximum=30, step=1, default=6, label="Weight Index")
+], gr.outputs.Image(type="file"),title=title,description=description,article=article,enable_queue=True,examples=examples,allow_flagging=False).launch()

packages.txt

@@ -0,0 +1,4 @@
+wget
+ffmpeg 
+libsm6 
+libxext6

psp_encoder/psp_encoders.py

@@ -141,7 +141,7 @@ class PSPEncoder(Module):
         print('Loading psp encoders weights from irse50!')
         encoder_ckpt = torch.load(encoder_ckpt_path, map_location='cpu')
         self.encoder.load_state_dict(get_keys(encoder_ckpt, 'encoder'), strict=True)
-        self.latent_avg = encoder_ckpt['latent_avg'].cuda()
+        self.latent_avg = encoder_ckpt['latent_avg']
 
         self.face_pool = torch.nn.AdaptiveAvgPool2d((256, 256))
 

requirements.txt

@@ -1,4 +1,5 @@
 torch
 numpy
 opencv-python-headless
-gdown
+torchvision
+Pillow

style_transfer_folder.py

@@ -20,7 +20,7 @@ torch.cuda.manual_seed_all(seed)
 
 
 if __name__ == '__main__':
-    device = 'cuda'
+    device = 'cpu'
 
     parser = argparse.ArgumentParser()
 
@@ -38,9 +38,7 @@ if __name__ == '__main__':
 
     args = parser.parse_args()
 
-    outdir = args.outdir
-    if not os.path.exists(outdir):
-        os.makedirs(outdir, exist_ok=True)
+
 
     args.latent = 512
     args.n_mlp = 8
@@ -58,35 +56,29 @@ if __name__ == '__main__':
     psp_encoder = PSPEncoder(args.psp_encoder_ckpt, output_size=args.size).to(device)
     psp_encoder.eval()
 
-    input_img_paths = sorted(glob.glob(os.path.join(args.input_img_path, '*.*')))
-    style_img_paths = sorted(glob.glob(os.path.join(args.style_img_path, '*.*')))[:]
-
     num = 0
 
-    for input_img_path in input_img_paths:
-        print(num)
-        num += 1
-
-        name_in = os.path.splitext(os.path.basename(input_img_path))[0]
-        img_in = cv2.imread(input_img_path, 1)
-        img_in_ten = cv2ten(img_in, device)
-        img_in = cv2.resize(img_in, (args.size, args.size))
-
-        for style_img_path in style_img_paths:
-            name_style = os.path.splitext(os.path.basename(style_img_path))[0]
-            img_style = cv2.imread(style_img_path, 1)
-            img_style_ten = cv2ten(img_style, device)
-            img_style = cv2.resize(img_style, (args.size, args.size))
-
-            with torch.no_grad():
-                sample_style = g_ema.get_z_embed(img_style_ten)
-                sample_in = psp_encoder(img_in_ten)
-                img_out_ten, _ = g_ema([sample_in], z_embed=sample_style, add_weight_index=args.add_weight_index,
-                                       input_is_latent=True, return_latents=False, randomize_noise=False)
-                img_out = ten2cv(img_out_ten)
-            out = np.concatenate([img_in, img_style, img_out], axis=1)
-            # out = img_out
-            cv2.imwrite(f'{args.outdir}/{name_in}_v_{name_style}.jpg', out)
+   
+    print(num)
+    num += 1
+
+    img_in = cv2.imread(args.input_img_path)
+    img_in_ten = cv2ten(img_in, device)
+    img_in = cv2.resize(img_in, (args.size, args.size))
+
+  
+    img_style = cv2.imread(args.style_img_path)
+    img_style_ten = cv2ten(img_style, device)
+    img_style = cv2.resize(img_style, (args.size, args.size))
+
+    with torch.no_grad():
+        sample_style = g_ema.get_z_embed(img_style_ten)
+        sample_in = psp_encoder(img_in_ten)
+        img_out_ten, _ = g_ema([sample_in], z_embed=sample_style, add_weight_index=args.add_weight_index,
+                               input_is_latent=True, return_latents=False, randomize_noise=False)
+        img_out = ten2cv(img_out_ten)
+    #out = np.concatenate([img_in, img_style, img_out], axis=1)
+    cv2.imwrite('out.jpg', img_out)
 
     print('Done!')