Upload 19 files

gradio_tabs/animation.py

@@ -0,0 +1,295 @@
+import gradio as gr
+import os
+import torch
+import torchvision
+from PIL import Image
+import numpy as np
+import imageio
+
+extensions_dir = "./torch_extension/"
+os.environ["TORCH_EXTENSIONS_DIR"] = extensions_dir
+
+from networks.generator import Generator
+
+device = torch.device("cuda")
+ckpt_path = './models/lia-x.pt'
+gen = Generator(size=512, motion_dim=40, scale=2).to(device)
+gen.load_state_dict(torch.load(ckpt_path, weights_only=False))
+gen.eval()
+
+output_dir = "./res_gradio"
+os.makedirs(output_dir, exist_ok=True)
+
+# lables
+labels_k = [
+	'yaw1',
+	'yaw2',
+	'pitch',
+	'roll1',
+	'roll2',
+	'neck',
+
+	'pout',
+	'open->close',
+	'"O" Mouth',
+	'apple cheek',
+
+	'close->open',
+	'eyebrows',
+	'eyeballs1',
+	'eyeballs2',
+
+]
+
+labels_v = [
+	37, 39, 28, 15, 33, 31,
+	6, 25, 16, 19,
+	13, 24, 17, 26
+]
+
+
+def load_image(img, size):
+	# img = Image.open(filename).convert('RGB')
+	if not isinstance(img, np.ndarray):
+		img = Image.open(img).convert('RGB')
+		img = img.resize((size, size))
+		img = np.asarray(img)
+	img = np.transpose(img, (2, 0, 1))	# 3 x 256 x 256
+
+	return img / 255.0
+
+
+def img_preprocessing(img_path, size):
+	img = load_image(img_path, size)  # [0, 1]
+	img = torch.from_numpy(img).unsqueeze(0).float()  # [0, 1]
+	imgs_norm = (img - 0.5) * 2.0  # [-1, 1]
+
+	return imgs_norm
+
+
+def resize(img, size):
+	transform = torchvision.transforms.Compose([
+		torchvision.transforms.Resize(size, antialias=True),
+		torchvision.transforms.CenterCrop(size)
+	])
+
+	return transform(img)
+
+
+def vid_preprocessing(vid_path, size):
+	vid_dict = torchvision.io.read_video(vid_path, pts_unit='sec')
+	vid = vid_dict[0].permute(0, 3, 1, 2).unsqueeze(0)	# btchw
+	fps = vid_dict[2]['video_fps']
+	vid_norm = (vid / 255.0 - 0.5) * 2.0  # [-1, 1]
+
+	vid_norm = torch.cat([
+		resize(vid_norm[:, i, :, :, :], size).unsqueeze(1) for i in range(vid.size(1))
+	], dim=1)
+
+	return vid_norm, fps
+
+
+def img_denorm(img):
+	img = img.clamp(-1, 1).cpu()
+	img = (img - img.min()) / (img.max() - img.min())
+
+	return img
+
+
+def vid_denorm(vid):
+	vid = vid.clamp(-1, 1).cpu()
+	vid = (vid - vid.min()) / (vid.max() - vid.min())
+
+	return vid
+
+
+def img_postprocessing(image, output_path=output_dir + "/output_img.png"):
+
+	image = image.permute(0, 2, 3, 1)
+	edited_image = img_denorm(image)
+	img_output = (edited_image[0].numpy() * 255).astype(np.uint8)
+	imageio.imwrite(output_path, img_output, quality=6)
+
+	return output_path
+
+
+def vid_postprocessing(video, fps, output_path=output_dir + "/output_vid.mp4"):
+	# video: BCTHW
+
+	vid = video.permute(0, 2, 3, 4, 1)	# B T H W C
+	vid_np = (vid_denorm(vid[0]).numpy() * 255).astype('uint8')
+	imageio.mimwrite(output_path, vid_np, fps=fps, codec='libx264', quality=10)
+
+	return output_path
+
+
+@torch.no_grad()
+def edit_media(image, *selected_s):
+
+	image_tensor = img_preprocessing(image, 512)
+	image_tensor = image_tensor.to(device)
+
+	edited_image_tensor = gen.edit_img(image_tensor, labels_v, selected_s)
+
+	# de-norm
+	edited_image = img_postprocessing(edited_image_tensor)
+
+	return edited_image
+
+
+@torch.no_grad()
+def animate_media(image, video, *selected_s):
+
+	image_tensor = img_preprocessing(image, 512)
+	vid_target_tensor, fps = vid_preprocessing(video, 512)
+	image_tensor = image_tensor.to(device)
+	video_target_tensor = vid_target_tensor.to(device)
+
+	animated_video = gen.animate(image_tensor, video_target_tensor, labels_v, selected_s)
+
+	# postprocessing
+	animated_video = vid_postprocessing(animated_video, fps)
+
+	return animated_video
+
+
+def clear_media():
+	return None, None, *([0] * len(labels_k))
+
+
+image_output = gr.Image(label="Output Image", type='numpy', interactive=False, width=512)
+video_output = gr.Video(label="Output Video", width=512)
+
+
+def animation():
+	with gr.Tab("Animation & Image Editing"):
+
+		inputs_s = []
+
+		with gr.Row():
+			with gr.Column(scale=1):
+				with gr.Row():
+					with gr.Accordion(open=True, label="Source Image"):
+						image_input = gr.Image(type="filepath", width=512)	# , height=550)
+						gr.Examples(
+							examples=[
+								["./data/source/macron.png"],
+								["./data/source/einstein.png"],
+								["./data/source/taylor.png"],
+								["./data/source/portrait1.png"],
+								["./data/source/portrait2.png"],
+								["./data/source/portrait3.png"],
+							],
+							inputs=[image_input],
+							cache_examples=False,
+							visible=True,
+							)
+
+					with gr.Accordion(open=True, label="Driving Video"):
+						video_input = gr.Video(width=512)  # , height=550)
+						gr.Examples(
+							examples=[
+								["./data/driving/driving1.mp4"],
+								["./data/driving/driving2.mp4"],
+								["./data/driving/driving4.mp4"],
+								#["./data/driving/driving5.mp4"],
+								["./data/driving/driving6.mp4"],
+								#["./data/driving/driving7.mp4"],
+								["./data/driving/driving8.mov"],
+							],
+							inputs=[video_input],
+							cache_examples=False,
+							visible=True,
+							)
+
+				with gr.Row():
+					with gr.Column(scale=1):
+						with gr.Row():	# Buttons now within a single Row
+							edit_btn = gr.Button("Edit")
+							clear_btn = gr.Button("Clear")
+						with gr.Row():
+							animate_btn = gr.Button("Animate")
+
+
+
+			with gr.Column(scale=1):
+
+				with gr.Row():
+					with gr.Accordion(open=True, label="Edited Source Image"):
+						image_output.render()
+
+					with gr.Accordion(open=True, label="Animated Video"):
+						video_output.render()
+
+				with gr.Accordion("Control Panel", open=True):
+					with gr.Tab("Head"):
+						with gr.Row():
+							for k in labels_k[:3]:
+								slider = gr.Slider(minimum=-1.0, maximum=0.5, value=0, label=k)
+								inputs_s.append(slider)
+						with gr.Row():
+							for k in labels_k[3:6]:
+								slider = gr.Slider(minimum=-0.5, maximum=0.5, value=0, label=k)
+								inputs_s.append(slider)
+
+					with gr.Tab("Mouth"):
+						with gr.Row():
+							for k in labels_k[6:8]:
+								slider = gr.Slider(minimum=-0.4, maximum=0.4, value=0, label=k)
+								inputs_s.append(slider)
+						with gr.Row():
+							for k in labels_k[8:10]:
+								slider = gr.Slider(minimum=-0.4, maximum=0.4, value=0, label=k)
+								inputs_s.append(slider)
+
+					with gr.Tab("Eyes"):
+						with gr.Row():
+							for k in labels_k[10:12]:
+								slider = gr.Slider(minimum=-0.4, maximum=0.4, value=0, label=k)
+								inputs_s.append(slider)
+						with gr.Row():
+							for k in labels_k[12:14]:
+								slider = gr.Slider(minimum=-0.2, maximum=0.2, value=0, label=k)
+								inputs_s.append(slider)
+
+
+		edit_btn.click(
+			fn=edit_media,
+			inputs=[image_input] + inputs_s,
+			outputs=[image_output],
+			show_progress=True
+		)
+
+		animate_btn.click(
+			fn=animate_media,
+			inputs=[image_input, video_input] + inputs_s,  # [image_input, video_input] + inputs_s,
+			outputs=[video_output],
+		)
+
+		clear_btn.click(
+			fn=clear_media,
+			outputs=[image_output, video_output] + inputs_s
+		)
+
+		gr.Examples(
+			examples=[
+				['./data/source/macron.png', './data/driving/driving1.mp4', 0.14,0,-0.26,-0.29,-0.11,0,-0.13,-0.18,0,0,0,0,-0.02,0.07],
+				['./data/source/portrait1.png', './data/driving/driving2.mp4', -0.1, 0, 0, 0.17, 0.16, 0, 0.01, 0, 0.17,0.17, 0, 0, 0, 0],
+				['./data/source/macron.png', './data/driving/driving4.mp4', -0.24, -0.17, -0.15, 0, 0, 0, 0, -0.16,
+				 0.08, 0, 0, 0, 0, 0],
+				['./data/source/portrait2.png', './data/driving/driving3.mp4', 0.33, 0.38, -0.22, 0.25, -0.23, 0, -0.16,
+				 0, 0.06, 0, 0, 0, 0, 0],
+				['./data/source/portrait2.png', './data/driving/driving6.mp4', -0.27, -0.25, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+				 0, 0, 0],
+				['./data/source/portrait2.png','./data/driving/driving1.mp4',-0.03,0.21,-0.41,-0.29,-0.11,0,0,-0.23,0,0,0,0,-0.02,0.07],
+				['./data/source/portrait3.png','./data/driving/driving1.mp4',-0.03,0.21,-0.31,-0.12,-0.11,0,-0.05,-0.16,0,0,0,0,-0.02,0.07],
+				['./data/source/portrait1.png','./data/driving/driving1.mp4',-0.03,0.21,-0.31,-0.12,-0.11,0,-0.1,-0.12,0,0.11,0,0,-0.02,0.07],
+				['./data/source/einstein.png','./data/driving/driving2.mp4',-0.31,0,0,0.16,0.08,0,-0.07,0,0.13,0,0,0,0,0],
+				['./data/source/einstein.png', './data/driving/driving4.mp4',0,0,0,0,0,0,0,-0.14,0.1,0,0,0,0,0],
+				['./data/source/portrait1.png', './data/driving/driving4.mp4',0,0,0,0,0,0,0,-0.1,0.19,0,0,0,0,0],
+				['./data/source/macron.png', './data/driving/driving6.mp4',-0.37,-0.34,0,0,0,0,0,0,0,0,0,0,0,0],
+			],
+			inputs=[image_input, video_input] + inputs_s
+		)
+
+

gradio_tabs/vid_edit.py

@@ -0,0 +1,293 @@
+import gradio as gr
+import os
+import torch
+import torchvision
+from PIL import Image
+import numpy as np
+import imageio
+from einops import rearrange
+
+extensions_dir = "./torch_extension/"
+os.environ["TORCH_EXTENSIONS_DIR"] = extensions_dir
+
+from networks.generator import Generator
+
+device = torch.device("cuda")
+ckpt_path = './models/lia-x.pt'
+gen = Generator(size=512, motion_dim=40, scale=2).to(device)
+gen.load_state_dict(torch.load(ckpt_path, weights_only=False))
+gen.eval()
+
+output_dir = "./res_gradio"
+os.makedirs(output_dir, exist_ok=True)
+
+# lables
+labels_k = [
+	'yaw1',
+	'yaw2',
+	'pitch',
+	'roll1',
+	'roll2',
+	'neck',
+
+	'pout',
+	'open->close',
+	'"O" mouth',
+	'apple cheek',
+
+	'close->open',
+	'eyebrows',
+	'eyeballs1',
+	'eyeballs2',
+
+]
+
+labels_v = [
+	37, 39, 28, 15, 33, 31,
+	6, 25, 16, 19,
+	13, 24, 17, 26
+]
+
+
+def load_image(img, size):
+	# img = Image.open(filename).convert('RGB')
+	if not isinstance(img, np.ndarray):
+		img = Image.open(img).convert('RGB')
+		img = img.resize((size, size))
+		img = np.asarray(img)
+	img = np.transpose(img, (2, 0, 1))	# 3 x 256 x 256
+
+	return img / 255.0
+
+
+def img_preprocessing(img_path, size):
+	img = load_image(img_path, size)  # [0, 1]
+	img = torch.from_numpy(img).unsqueeze(0).float()  # [0, 1]
+	imgs_norm = (img - 0.5) * 2.0  # [-1, 1]
+
+	return imgs_norm
+
+
+def resize(img, size):
+	transform = torchvision.transforms.Compose([
+		torchvision.transforms.Resize(size, antialias=True),
+		torchvision.transforms.CenterCrop(size)
+	])
+
+	return transform(img)
+
+
+def vid_preprocessing(vid_path, size):
+	vid_dict = torchvision.io.read_video(vid_path, pts_unit='sec')
+	vid = vid_dict[0].permute(0, 3, 1, 2).unsqueeze(0)	# btchw
+	fps = vid_dict[2]['video_fps']
+	vid_norm = (vid / 255.0 - 0.5) * 2.0  # [-1, 1]
+
+	vid_norm = torch.cat([
+		resize(vid_norm[:, i, :, :, :], size).unsqueeze(1) for i in range(vid.size(1))
+	], dim=1)
+
+	return vid_norm, fps
+
+
+def img_denorm(img):
+	img = img.clamp(-1, 1).cpu()
+	img = (img - img.min()) / (img.max() - img.min())
+
+	return img
+
+
+def vid_denorm(vid):
+	vid = vid.clamp(-1, 1).cpu()
+	vid = (vid - vid.min()) / (vid.max() - vid.min())
+
+	return vid
+
+
+def img_postprocessing(image, output_path=output_dir + "/output_img.png"):
+	image = image.permute(0, 2, 3, 1)
+	edited_image = img_denorm(image)
+	img_output = (edited_image[0].numpy() * 255).astype(np.uint8)
+	imageio.imwrite(output_path, img_output, quality=6)
+
+	return output_path
+
+
+def vid_all_save(vid_d, vid_a, fps, output_path=output_dir + "/output_vid.mp4", output_all_path=output_dir + "/output_all_vid.mp4"):
+
+	vid_d = rearrange(vid_d, 'b t c h w -> b t h w c')
+	vid_a = rearrange(vid_a, 'b c t h w -> b t h w c')
+	vid_all = torch.cat([vid_d, vid_a], dim=3)
+
+	vid_a_np = (vid_denorm(vid_a[0]).numpy() * 255).astype('uint8')
+	vid_all_np = (vid_denorm(vid_all[0]).numpy() * 255).astype('uint8')
+
+	imageio.mimwrite(output_path, vid_a_np, fps=fps, codec='libx264', quality=8)
+	imageio.mimwrite(output_all_path, vid_all_np, fps=fps, codec='libx264', quality=8)
+
+	return output_path, output_all_path
+
+
+@torch.no_grad()
+def edit_img(video, *selected_s):
+
+	vid_target_tensor, fps = vid_preprocessing(video, 512)
+	video_target_tensor = vid_target_tensor.to(device)
+	image_tensor = video_target_tensor[:,0,:,:,:]
+
+	edited_image_tensor = gen.edit_img(image_tensor, labels_v, selected_s)
+
+	# de-norm
+	edited_image = img_postprocessing(edited_image_tensor)
+
+	return edited_image
+
+
+@torch.no_grad()
+def edit_vid(video, *selected_s):
+
+	video_target_tensor, fps = vid_preprocessing(video, 512)
+	video_target_tensor = video_target_tensor.to(device)
+
+	edited_video_tensor = gen.edit_vid(video_target_tensor, labels_v, selected_s)
+
+	# de-norm
+	animated_video, animated_all_video = vid_all_save(video_target_tensor, edited_video_tensor, fps)
+
+	return animated_video, animated_all_video
+
+
+def clear_media():
+	return None, None, None, *([0] * len(labels_k))
+
+
+image_output = gr.Image(label="Image", type='numpy', interactive=False, width=512)
+video_output = gr.Video(label="Video", width=512)
+video_all_output = gr.Video(label="Videos")
+
+
+def vid_edit():
+	with gr.Tab("Video Editing"):
+
+		inputs_c = []
+		inputs_s = []
+
+		with gr.Row():
+			with gr.Column(scale=1):
+				with gr.Row():
+					with gr.Accordion(open=True, label="Video"):
+						video_input = gr.Video(width=512)  # , height=550)
+						gr.Examples(
+							examples=[
+								["./data/driving/driving1.mp4"],
+								["./data/driving/driving2.mp4"],
+								["./data/driving/driving4.mp4"],
+								#["./data/driving/driving5.mp4"],
+								#["./data/driving/driving6.mp4"],
+								#["./data/driving/driving7.mp4"],
+								["./data/driving/driving3.mp4"],
+								["./data/driving/driving8.mov"],
+								["./data/driving/driving9.mov"],
+							],
+							inputs=[video_input],
+							cache_examples=False,
+							visible=True,
+						)
+
+				# with gr.Row():
+				# 	with gr.Column(scale=1):
+				# 		with gr.Row():	# Buttons now within a single Row
+				# 			edit_btn = gr.Button("Edit")
+				# 			clear_btn = gr.Button("Clear")
+				# 		with gr.Row():
+				# 			animate_btn = gr.Button("Generate")
+
+			with gr.Column(scale=2):
+
+				with gr.Row():
+					with gr.Accordion(open=True, label="Edited First Frame"):
+						image_output.render()
+
+					with gr.Accordion(open=True, label="Edited Video"):
+						video_output.render()
+				
+				with gr.Row():
+					with gr.Accordion(open=True, label="Original & Edited Videos"):
+						video_all_output.render()
+			
+			with gr.Column(scale=1):
+				with gr.Accordion("Control Panel", open=True):
+					with gr.Tab("Head"):
+						with gr.Row():
+							for k in labels_k[:3]:
+								slider = gr.Slider(minimum=-1.0, maximum=0.5, value=0, label=k)
+								inputs_s.append(slider)
+						with gr.Row():
+							for k in labels_k[3:6]:
+								slider = gr.Slider(minimum=-0.5, maximum=0.5, value=0, label=k)
+								inputs_s.append(slider)
+
+					with gr.Tab("Mouth"):
+						with gr.Row():
+							for k in labels_k[6:8]:
+								slider = gr.Slider(minimum=-0.4, maximum=0.4, value=0, label=k)
+								inputs_s.append(slider)
+						with gr.Row():
+							for k in labels_k[8:10]:
+								slider = gr.Slider(minimum=-0.4, maximum=0.4, value=0, label=k)
+								inputs_s.append(slider)
+
+					with gr.Tab("Eyes"):
+						with gr.Row():
+							for k in labels_k[10:12]:
+								slider = gr.Slider(minimum=-0.4, maximum=0.4, value=0, label=k)
+								inputs_s.append(slider)
+						with gr.Row():
+							for k in labels_k[12:14]:
+								slider = gr.Slider(minimum=-0.2, maximum=0.2, value=0, label=k)
+								inputs_s.append(slider)
+
+				with gr.Row():
+					with gr.Column(scale=1):
+						with gr.Row():	# Buttons now within a single Row
+							edit_btn = gr.Button("Edit")
+							clear_btn = gr.Button("Clear")
+						with gr.Row():
+							animate_btn = gr.Button("Generate")
+
+		edit_btn.click(
+			fn=edit_img,
+			inputs=[video_input] + inputs_s,
+			outputs=[image_output],
+			show_progress=True
+		)
+
+		animate_btn.click(
+			fn=edit_vid,
+			inputs=[video_input] + inputs_s,  # [image_input, video_input] + inputs_s,
+			outputs=[video_output, video_all_output],
+		)
+
+		clear_btn.click(
+			fn=clear_media,
+			outputs=[image_output, video_output, video_all_output] + inputs_s
+		)
+
+		gr.Examples(
+			examples=[
+				['./data/driving/driving1.mp4', 0.5, 0.5, 0, 0, 0, 0, 0,
+				 0, 0, 0, 0, 0, 0, 0],
+				['./data/driving/driving2.mp4', 0.5, 0.5, 0, 0, 0, 0, 0, 0, 0,
+				 0, 0, 0, 0, 0],
+				['./data/driving/driving1.mp4', 0, 0, 0, 0, 0, 0, 0,
+				 0, 0, 0, 0, -0.3, 0, 0],
+				['./data/driving/driving3.mp4', -0.6, 0, 0, 0, 0, 0, 0,
+				 0, 0, 0, 0, 0, 0, 0],
+				['./data/driving/driving9.mov', 0, 0, 0, 0, 0, 0, 0,
+				 0, 0, 0, 0, 0, -0.1, 0.07],
+			],
+			inputs=[video_input] + inputs_s
+		)
+
+
+

networks/decoder.py

@@ -0,0 +1,287 @@
+import math
+import torch
+from torch import nn
+import torch.nn.functional as F
+from .ops import (ConstantInput, ConvLayer, StyledConv, ToFlow, ToRGB, Direction)
+
+
+class FlowResBlock(nn.Module):
+	def __init__(self, in_channel, out_channel, style_dim):
+		super().__init__()
+
+		self.norm = nn.GroupNorm(32, out_channel)
+
+		self.conv1 = StyledConv(in_channel, out_channel, 3, style_dim, False)
+		self.conv2 = StyledConv(out_channel, out_channel, 3, style_dim, False)
+
+		self.gamma = nn.Parameter(1e-5 * torch.ones([1, out_channel, 1, 1]))
+
+	def forward(self, x, style):
+		h = x
+		h = self.conv1(h, style)
+		skip = h
+
+		h = self.norm(h)
+		h = self.conv2(h, style)
+		h = self.gamma * h
+
+		return h + skip
+
+
+class ResBlock(nn.Module):
+	def __init__(self, in_channel, out_channel):
+		super().__init__()
+
+		self.conv1 = ConvLayer(in_channel, out_channel, 3, upsample=False)
+		self.conv2 = ConvLayer(out_channel, out_channel, 3, upsample=False)
+
+		if in_channel != out_channel:
+			self.skip = ConvLayer(in_channel, out_channel, 1, upsample=False, activate=False, bias=False)
+		else:
+			self.skip = torch.nn.Identity()
+
+	def forward(self, x):
+
+		h = x
+		h = self.conv1(h)
+		h = self.conv2(h)
+		skip = self.skip(x)
+
+		return (h + skip) / math.sqrt(2)
+
+
+class Decoder(nn.Module):
+	def __init__(self, style_dim, motion_dim, scale=1):
+		super().__init__()
+		
+		channels = [512*scale, 256 * scale, 128 * scale, 64 * scale]
+
+		self.direction = Direction(style_dim, motion_dim)
+
+		self.input = ConstantInput(channels[0], size=4)  # 4
+
+		# block1, 4
+		self.conv1 = StyledConv(channels[0], channels[0], 3, style_dim, False)
+		
+		# for 512
+		self.conv_512_1 = StyledConv(channels[0], channels[0], 3, style_dim, True)
+		self.conv_512_2 = nn.ModuleList([
+			FlowResBlock(channels[0], channels[0], style_dim),
+			FlowResBlock(channels[0], channels[0], style_dim),
+			FlowResBlock(channels[0], channels[0], style_dim),
+			FlowResBlock(channels[0], channels[0], style_dim),
+		])
+		self.conv_512_2_rgb = nn.ModuleList([
+			ResBlock(channels[0], channels[0]),
+			ResBlock(channels[0], channels[0]),
+			ResBlock(channels[0], channels[0]),
+			ResBlock(channels[0], channels[0]),
+		])
+		self.rgb_512 = ToRGB(channels[0])
+		self.flow_512 = ToFlow(channels[0], style_dim)	# 16	
+
+		# block2, 8
+		self.conv2_1 = StyledConv(channels[0], channels[0], 3, style_dim, True)
+		self.conv2_2 = nn.ModuleList([
+			FlowResBlock(channels[0], channels[0], style_dim),
+			FlowResBlock(channels[0], channels[0], style_dim),
+			FlowResBlock(channels[0], channels[0], style_dim),
+			FlowResBlock(channels[0], channels[0], style_dim),
+		])
+		self.conv2_2_up = ConvLayer(channels[0], channels[0], 3, upsample=True)
+		self.conv2_2_rgb = nn.ModuleList([
+			ResBlock(channels[0], channels[0]),
+			ResBlock(channels[0], channels[0]),
+			ResBlock(channels[0], channels[0]),
+			ResBlock(channels[0], channels[0]),
+		])
+		self.rgb2 = ToRGB(channels[0])
+		self.flow2 = ToFlow(channels[0], style_dim)  # 16
+
+		# block3, 16
+		self.conv3_1 = StyledConv(channels[0], channels[0], 3, style_dim, True)
+		self.conv3_2 = nn.ModuleList([
+			FlowResBlock(channels[0], channels[0], style_dim),
+			FlowResBlock(channels[0], channels[0], style_dim),
+			FlowResBlock(channels[0], channels[0], style_dim),
+			FlowResBlock(channels[0], channels[0], style_dim),
+		])
+		self.conv3_2_up = ConvLayer(channels[0], channels[0], 3, upsample=True)
+		self.conv3_2_rgb = nn.ModuleList([
+			ResBlock(channels[0], channels[0]),
+			ResBlock(channels[0], channels[0]),
+			ResBlock(channels[0], channels[0]),
+			ResBlock(channels[0], channels[0]),
+		])
+		self.rgb3 = ToRGB(channels[0])
+		self.flow3 = ToFlow(channels[0], style_dim)  # 32
+
+		# block4, 32
+		self.conv4_1 = StyledConv(channels[0], channels[0], 3, style_dim, True)
+		self.conv4_2 = nn.ModuleList([
+			FlowResBlock(channels[0], channels[0], style_dim),
+			FlowResBlock(channels[0], channels[0], style_dim),
+			FlowResBlock(channels[0], channels[0], style_dim),
+			FlowResBlock(channels[0], channels[0], style_dim),
+		])
+		self.conv4_2_up = ConvLayer(channels[0], channels[0], 3, upsample=True)
+		self.conv4_2_rgb = nn.ModuleList([
+			ResBlock(channels[0], channels[0]),
+			ResBlock(channels[0], channels[0]),
+			ResBlock(channels[0], channels[0]),
+			ResBlock(channels[0], channels[0]),
+		]) 
+		self.rgb4 = ToRGB(channels[0])
+		self.flow4 = ToFlow(channels[0], style_dim)  # 64
+
+		# block5, 64
+		self.conv5_1 = StyledConv(channels[0], channels[1], 3, style_dim, True)
+		self.conv5_2 = nn.ModuleList([
+			FlowResBlock(channels[1], channels[1], style_dim),
+			FlowResBlock(channels[1], channels[1], style_dim),
+			FlowResBlock(channels[1], channels[1], style_dim),
+			FlowResBlock(channels[1], channels[1], style_dim),
+		])
+		self.conv5_2_up = ConvLayer(channels[0], channels[1], 3, upsample=True)
+		self.conv5_2_rgb = nn.ModuleList([
+			ResBlock(channels[1], channels[1]),
+			ResBlock(channels[1], channels[1]),
+			ResBlock(channels[1], channels[1]),
+			ResBlock(channels[1], channels[1]),
+		])
+		self.rgb5 = ToRGB(channels[1])
+		self.flow5 = ToFlow(channels[1], style_dim)  # 128
+
+		# block6, 128
+		self.conv6_1 = StyledConv(channels[1], channels[2], 3, style_dim, True)
+		self.conv6_2 = nn.ModuleList([
+			FlowResBlock(channels[2], channels[2], style_dim),
+			FlowResBlock(channels[2], channels[2], style_dim),
+			FlowResBlock(channels[2], channels[2], style_dim),
+			FlowResBlock(channels[2], channels[2], style_dim),
+		])
+		self.conv6_2_up = ConvLayer(channels[1], channels[2], 3, upsample=True)
+		self.conv6_2_rgb = nn.ModuleList([
+			ResBlock(channels[2], channels[2]),
+			ResBlock(channels[2], channels[2]),
+			ResBlock(channels[2], channels[2]),
+			ResBlock(channels[2], channels[2]),
+		])
+		self.rgb6 = ToRGB(channels[2])
+		self.flow6 = ToFlow(channels[2], style_dim)  # 128
+
+		# block7, 256
+		self.conv7_1 = StyledConv(channels[2], channels[3], 3, style_dim, True)
+		self.conv7_2 = nn.ModuleList([
+			FlowResBlock(channels[3], channels[3], style_dim),
+			FlowResBlock(channels[3], channels[3], style_dim),
+			FlowResBlock(channels[3], channels[3], style_dim),
+			FlowResBlock(channels[3], channels[3], style_dim),
+		])
+		self.conv7_2_up = ConvLayer(channels[2], channels[3], 3, upsample=True)
+		self.conv7_2_rgb = nn.ModuleList([
+			ResBlock(channels[3], channels[3]),
+			ResBlock(channels[3], channels[3]),
+			ResBlock(channels[3], channels[3]),
+			ResBlock(channels[3], channels[3]),
+		])
+		self.rgb7 = ToRGB(channels[3])
+		self.flow7 = ToFlow(channels[3], style_dim)  # 128
+
+	def navigation(self, z_s2r, alpha):
+
+		if alpha is not None:
+			# generating moving directions
+			if len(alpha) > 1:
+				z_r2t = self.direction(alpha[0])  # target
+				z_r2s = self.direction(alpha[1])  # source
+				z_start = self.direction(alpha[2])	# start
+				z_s2t = z_s2r + (z_r2t - z_start) + z_r2s
+			else:
+				z_r2t = self.direction(alpha[0])
+				z_s2t = z_s2r + z_r2t  # wa + directions
+		else:
+			z_s2t = z_s2r
+
+		return z_s2t
+
+	def apply_flow(self, h, mask, flow, feat):
+
+		feat_warp = F.grid_sample(feat, flow) * mask
+		h = feat_warp + (1 - mask) * h
+
+		return feat_warp, h
+
+	def forward(self, z_s2r, alpha, feats):
+		# z_s2r: bs x style_dim
+		# alpha: bs x style_dim
+
+		z_s2t = self.navigation(z_s2r, alpha)
+
+		h = self.input(z_s2t)
+		h = self.conv1(h, z_s2t)
+		
+		#for 512
+		h = self.conv_512_1(h, z_s2t)
+		for conv in self.conv_512_2:
+			h = conv(h, z_s2t)
+		h_warp_512, h, h_flow_512 = self.flow_512(h, z_s2t, feats[0])
+		for conv in self.conv_512_2_rgb:
+			h_warp_512 = conv(h_warp_512)
+		rgb_512 = self.rgb_512(h_warp_512)
+
+		h = self.conv2_1(h, z_s2t)
+		for conv in self.conv2_2:
+			h = conv(h, z_s2t)
+		h_warp2, h, h_flow2 = self.flow2(h, z_s2t, feats[1], h_flow_512)
+		h_warp2 = h_warp2 + self.conv2_2_up(h_warp_512)
+		for conv in self.conv2_2_rgb:
+			h_warp2 = conv(h_warp2)
+		rgb2 = self.rgb2(h_warp2, rgb_512)
+
+		h = self.conv3_1(h, z_s2t)
+		for conv in self.conv3_2:
+			h = conv(h, z_s2t)
+		h_warp3, h, h_flow3 = self.flow3(h, z_s2t, feats[2], h_flow2)
+		h_warp3 = h_warp3 + self.conv3_2_up(h_warp2)
+		for conv in self.conv3_2_rgb:
+			h_warp3 = conv(h_warp3)
+		rgb3 = self.rgb3(h_warp3, rgb2)
+
+		h = self.conv4_1(h, z_s2t)
+		for conv in self.conv4_2:
+			h = conv(h, z_s2t)
+		h_warp4, h, h_flow4 = self.flow4(h, z_s2t, feats[3], h_flow3)
+		h_warp4 = h_warp4 + self.conv4_2_up(h_warp3)
+		for conv in self.conv4_2_rgb:
+			h_warp4 = conv(h_warp4)
+		rgb4 = self.rgb4(h_warp4, rgb3)
+
+		h = self.conv5_1(h, z_s2t)
+		for conv in self.conv5_2:
+			h = conv(h, z_s2t)
+		h_warp5, h, h_flow5 = self.flow5(h, z_s2t, feats[4], h_flow4)
+		h_warp5 = h_warp5 + self.conv5_2_up(h_warp4)
+		for conv in self.conv5_2_rgb:
+			h_warp5 = conv(h_warp5)
+		rgb5 = self.rgb5(h_warp5, rgb4)
+
+		h = self.conv6_1(h, z_s2t)
+		for conv in self.conv6_2:
+			h = conv(h, z_s2t)
+		h_warp6, h, h_flow6 = self.flow6(h, z_s2t, feats[5], h_flow5)
+		h_warp6 = h_warp6 + self.conv6_2_up(h_warp5)
+		for conv in self.conv6_2_rgb:
+			h_warp6 = conv(h_warp6)
+		rgb6 = self.rgb6(h_warp6, rgb5)
+
+		h = self.conv7_1(h, z_s2t)
+		for conv in self.conv7_2:
+			h = conv(h, z_s2t)
+		h_warp7, h, h_flow7 = self.flow7(h, z_s2t, feats[6], h_flow6)
+		h_warp7 = h_warp7 + self.conv7_2_up(h_warp6)
+		for conv in self.conv7_2_rgb:
+			h_warp7 = conv(h_warp7)
+		out = self.rgb7(h_warp7, rgb6)
+
+		return out

networks/encoder.py

@@ -0,0 +1,141 @@
+import math
+import torch
+from torch import nn
+from torch.nn import functional as F
+from .ops import (EqualConv2d, EqualLinear, ConvLayer)
+
+
+class ResBlock(nn.Module):
+	def __init__(self, in_channel, out_channel):
+		super().__init__()
+
+		self.conv1 = ConvLayer(in_channel, out_channel, 3)
+		self.conv2 = ConvLayer(out_channel, out_channel, 3, downsample=True)
+
+		self.skip = ConvLayer(in_channel, out_channel, 1, downsample=True, activate=False, bias=False)
+
+
+	def forward(self, x):
+		
+		h = x
+		
+		h = self.conv1(h)
+		h = self.conv2(h)
+
+		skip = self.skip(x)
+		h = (h + skip) / math.sqrt(2)
+
+		return h
+
+
+class Encoder2R(nn.Module):
+	def __init__(self, latent_dim=512, scale=1):
+		super(Encoder2R, self).__init__()
+		
+		channels = [64*scale, 128*scale, 256*scale, 512*scale]
+
+		# version1
+		self.block1 = ConvLayer(3, channels[0], 1) # 256, 3 -> 64
+		self.block2 = nn.Sequential(
+			ResBlock(channels[0], channels[1])
+		) # 64 -> 128
+		self.block3 = nn.Sequential(
+			ResBlock(channels[1], channels[2])
+		) # 128 -> 256
+		self.block4 = nn.Sequential(
+			ResBlock(channels[2], channels[3])
+		) # 256 -> 512
+		self.block5 = nn.Sequential(
+			ResBlock(channels[3], channels[3])
+		) # 512 -> 512
+		self.block6 = nn.Sequential(
+			ResBlock(channels[3], channels[3])
+		) # 512 -> 512
+		self.block7 = nn.Sequential(
+			ResBlock(channels[3], channels[3])
+		) # 512 -> 512
+		
+		self.block_512 = ResBlock(channels[3], channels[3])
+		self.block8 = EqualConv2d(channels[3], latent_dim, 4, padding=0, bias=False)
+
+	def forward(self, x):
+
+		res = []
+		h = x
+		h = self.block1(h) # 256
+		res.append(h)
+		h = self.block2(h) # 128
+		res.append(h)
+		h = self.block3(h) # 64
+		res.append(h)
+		h = self.block4(h) # 32
+		res.append(h)
+		h = self.block5(h) # 16
+		res.append(h)
+		h = self.block6(h) # 8
+		res.append(h)
+		h = self.block7(h) # 4
+		res.append(h)
+		h = self.block_512(h)
+		h = self.block8(h) # 1
+
+		return h.squeeze(-1).squeeze(-1), res[::-1]
+
+
+class Encoder(nn.Module):
+	def __init__(self, dim=512, dim_motion=20, scale=1):
+		super(Encoder, self).__init__()
+
+		# 2R netmork
+		self.enc_2r = Encoder2R(dim, scale)
+
+		# R2T
+		self.enc_r2t = nn.Sequential(
+			EqualLinear(dim, dim_motion)
+		)
+	
+	def enc_motion(self, x):
+
+		z_t2r, _ = self.enc_2r(x)
+		alpha_r2t = self.enc_r2t(z_t2r)
+
+		return alpha_r2t
+
+
+	def enc_transfer_img(self, z_s2r, d_l, s_l):
+
+		alpha_r2s = self.enc_r2t(z_s2r)
+		alpha_r2s[:, d_l] = alpha_r2s[:, d_l] + torch.FloatTensor(s_l).unsqueeze(0).to('cuda')
+		alpha = [alpha_r2s]
+
+		return alpha
+
+	def enc_transfer_vid(self, alpha_r2s, input_target, alpha_start):
+
+		z_t2r, _ = self.enc_2r(input_target)
+		alpha_r2t = self.enc_r2t(z_t2r)
+		alpha = [alpha_r2t, alpha_r2s, alpha_start]
+
+		return alpha
+
+
+	def forward(self, input_source, input_target, alpha_start=None):
+
+		if input_target is not None:
+
+			z_s2r, feats = self.enc_2r(input_source)
+			z_t2r, _ = self.enc_2r(input_target)
+
+			alpha_r2t = self.enc_r2t(z_t2r)
+
+			if alpha_start is not None:
+				alpha_r2s = self.enc_r2t(z_s2r)
+				alpha = [alpha_r2t, alpha_r2s, alpha_start]
+			else:
+				alpha = [alpha_r2t]
+
+			return z_s2r, alpha, feats
+		else:
+			z_s2r, feats = self.enc_2r(input_source)
+
+			return z_s2r, None, feats

networks/generator.py

@@ -0,0 +1,121 @@
+import torch
+from torch import nn
+from networks.encoder import Encoder
+from networks.decoder import Decoder
+import numpy as np
+from tqdm import tqdm
+
+
+class Generator(nn.Module):
+	def __init__(self, size, style_dim=512, motion_dim=40, scale=1):
+		super(Generator, self).__init__()
+
+		style_dim = style_dim * scale
+
+		# encoder
+		self.enc = Encoder(style_dim, motion_dim, scale)
+		self.dec = Decoder(style_dim, motion_dim, scale)
+	
+	def get_alpha(self, x):
+		return self.enc.enc_motion(x)
+
+	def edit_img(self, img_source, d_l, v_l):
+
+		z_s2r, feat_rgb = self.enc.enc_2r(img_source)
+		alpha_r2s = self.enc.enc_r2t(z_s2r)
+		alpha_r2s[:, d_l] = alpha_r2s[:, d_l] + torch.FloatTensor(v_l).unsqueeze(0).to('cuda')
+		img_recon = self.dec(z_s2r, [alpha_r2s], feat_rgb)
+
+		return img_recon
+
+	def animate(self, img_source, vid_target, d_l, v_l):
+
+		alpha_start = self.get_alpha(vid_target[:, 0, :, :, :])
+
+		vid_target_recon = []
+		z_s2r, feat_rgb = self.enc.enc_2r(img_source)
+		alpha_r2s = self.enc.enc_r2t(z_s2r)
+		alpha_r2s[:, d_l] = alpha_r2s[:, d_l] + torch.FloatTensor(v_l).unsqueeze(0).to('cuda')
+
+		for i in tqdm(range(vid_target.size(1))):
+			img_target = vid_target[:, i, :, :, :]
+			alpha = self.enc.enc_transfer_vid(alpha_r2s, img_target, alpha_start)
+			img_recon = self.dec(z_s2r, alpha, feat_rgb)
+			vid_target_recon.append(img_recon.unsqueeze(2))
+		vid_target_recon = torch.cat(vid_target_recon, dim=2) # BCTHW
+
+		return vid_target_recon
+
+	def edit_vid(self, vid_target, d_l, v_l):
+
+		img_source = vid_target[:, 0, :, :, :]
+		alpha_start = self.get_alpha(vid_target[:, 0, :, :, :])
+
+		vid_target_recon = []
+		z_s2r, feat_rgb = self.enc.enc_2r(img_source)
+		alpha_r2s = self.enc.enc_r2t(z_s2r)
+		alpha_r2s[:, d_l] = alpha_r2s[:, d_l] + torch.FloatTensor(v_l).unsqueeze(0).to('cuda')
+
+		for i in tqdm(range(vid_target.size(1))):
+			img_target = vid_target[:, i, :, :, :]
+			alpha = self.enc.enc_transfer_vid(alpha_r2s, img_target, alpha_start)
+			img_recon = self.dec(z_s2r, alpha, feat_rgb)
+			vid_target_recon.append(img_recon.unsqueeze(2))
+		vid_target_recon = torch.cat(vid_target_recon, dim=2) # BCTHW
+
+		return vid_target_recon
+
+
+	def interpolate_img(self, img_source, d_l, v_l):
+
+		vid_target_recon = []
+
+		step = 16
+		v_start = np.array([0.] * len(v_l))
+		v_end = np.array(v_l)
+		stride = (v_end - v_start) / step
+
+		z_s2r, feat_rgb = self.enc.enc_2r(img_source)
+
+		v_tmp = v_start
+		for i in range(step):
+			v_tmp = v_tmp + stride
+			alpha = self.enc.enc_transfer_img(z_s2r, d_l, v_tmp)
+			img_recon = self.dec(z_s2r, alpha, feat_rgb)
+			vid_target_recon.append(img_recon.unsqueeze(2))
+
+		for i in range(step):
+			v_tmp = v_tmp - stride
+			alpha = self.enc.enc_transfer_img(z_s2r, d_l, v_tmp)
+			img_recon = self.dec(z_s2r, alpha, feat_rgb)
+			vid_target_recon.append(img_recon.unsqueeze(2))
+
+		if (v_l[6]!=0) or (v_l[7]!=0) or (v_l[8]!=0) or (v_l[9]!=0):
+			for i in range(step):
+				v_tmp = v_tmp + stride
+				alpha = self.enc.enc_transfer_img(z_s2r, d_l, v_tmp)
+				img_recon = self.dec(z_s2r, alpha, feat_rgb)
+				vid_target_recon.append(img_recon.unsqueeze(2))
+
+			for i in range(step):
+				v_tmp = v_tmp - stride
+				alpha = self.enc.enc_transfer_img(z_s2r, d_l, v_tmp)
+				img_recon = self.dec(z_s2r, alpha, feat_rgb)
+				vid_target_recon.append(img_recon.unsqueeze(2))
+		else:
+			for i in range(step):
+				v_tmp = v_tmp - stride
+				alpha = self.enc.enc_transfer_img(z_s2r, d_l, v_tmp)
+				img_recon = self.dec(z_s2r, alpha, feat_rgb)
+				vid_target_recon.append(img_recon.unsqueeze(2))
+
+			for i in range(step):
+				v_tmp = v_tmp + stride
+				alpha = self.enc.enc_transfer_img(z_s2r, d_l, v_tmp)
+				img_recon = self.dec(z_s2r, alpha, feat_rgb)
+				vid_target_recon.append(img_recon.unsqueeze(2))
+
+		vid_target_recon = torch.cat(vid_target_recon, dim=2)  # BCTHW
+
+		return vid_target_recon
+

networks/op/__init__.py

@@ -0,0 +1,2 @@
+from .fused_act import FusedLeakyReLU, fused_leaky_relu
+from .upfirdn2d import upfirdn2d

networks/op/conv2d_gradfix.py

@@ -0,0 +1,227 @@
+import contextlib
+import warnings
+
+import torch
+from torch import autograd
+from torch.nn import functional as F
+
+enabled = True
+weight_gradients_disabled = False
+
+
+@contextlib.contextmanager
+def no_weight_gradients():
+    global weight_gradients_disabled
+
+    old = weight_gradients_disabled
+    weight_gradients_disabled = True
+    yield
+    weight_gradients_disabled = old
+
+
+def conv2d(input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1):
+    if could_use_op(input):
+        return conv2d_gradfix(
+            transpose=False,
+            weight_shape=weight.shape,
+            stride=stride,
+            padding=padding,
+            output_padding=0,
+            dilation=dilation,
+            groups=groups,
+        ).apply(input, weight, bias)
+
+    return F.conv2d(
+        input=input,
+        weight=weight,
+        bias=bias,
+        stride=stride,
+        padding=padding,
+        dilation=dilation,
+        groups=groups,
+    )
+
+
+def conv_transpose2d(
+    input,
+    weight,
+    bias=None,
+    stride=1,
+    padding=0,
+    output_padding=0,
+    groups=1,
+    dilation=1,
+):
+    if could_use_op(input):
+        return conv2d_gradfix(
+            transpose=True,
+            weight_shape=weight.shape,
+            stride=stride,
+            padding=padding,
+            output_padding=output_padding,
+            groups=groups,
+            dilation=dilation,
+        ).apply(input, weight, bias)
+
+    return F.conv_transpose2d(
+        input=input,
+        weight=weight,
+        bias=bias,
+        stride=stride,
+        padding=padding,
+        output_padding=output_padding,
+        dilation=dilation,
+        groups=groups,
+    )
+
+
+def could_use_op(input):
+    if (not enabled) or (not torch.backends.cudnn.enabled):
+        return False
+
+    if input.device.type != "cuda":
+        return False
+
+    if any(torch.__version__.startswith(x) for x in ["1.7.", "1.8."]):
+        return True
+
+    warnings.warn(
+        f"conv2d_gradfix not supported on PyTorch {torch.__version__}. Falling back to torch.nn.functional.conv2d()."
+    )
+
+    return False
+
+
+def ensure_tuple(xs, ndim):
+    xs = tuple(xs) if isinstance(xs, (tuple, list)) else (xs,) * ndim
+
+    return xs
+
+
+conv2d_gradfix_cache = dict()
+
+
+def conv2d_gradfix(
+    transpose, weight_shape, stride, padding, output_padding, dilation, groups
+):
+    ndim = 2
+    weight_shape = tuple(weight_shape)
+    stride = ensure_tuple(stride, ndim)
+    padding = ensure_tuple(padding, ndim)
+    output_padding = ensure_tuple(output_padding, ndim)
+    dilation = ensure_tuple(dilation, ndim)
+
+    key = (transpose, weight_shape, stride, padding, output_padding, dilation, groups)
+    if key in conv2d_gradfix_cache:
+        return conv2d_gradfix_cache[key]
+
+    common_kwargs = dict(
+        stride=stride, padding=padding, dilation=dilation, groups=groups
+    )
+
+    def calc_output_padding(input_shape, output_shape):
+        if transpose:
+            return [0, 0]
+
+        return [
+            input_shape[i + 2]
+            - (output_shape[i + 2] - 1) * stride[i]
+            - (1 - 2 * padding[i])
+            - dilation[i] * (weight_shape[i + 2] - 1)
+            for i in range(ndim)
+        ]
+
+    class Conv2d(autograd.Function):
+        @staticmethod
+        def forward(ctx, input, weight, bias):
+            if not transpose:
+                out = F.conv2d(input=input, weight=weight, bias=bias, **common_kwargs)
+
+            else:
+                out = F.conv_transpose2d(
+                    input=input,
+                    weight=weight,
+                    bias=bias,
+                    output_padding=output_padding,
+                    **common_kwargs,
+                )
+
+            ctx.save_for_backward(input, weight)
+
+            return out
+
+        @staticmethod
+        def backward(ctx, grad_output):
+            input, weight = ctx.saved_tensors
+            grad_input, grad_weight, grad_bias = None, None, None
+
+            if ctx.needs_input_grad[0]:
+                p = calc_output_padding(
+                    input_shape=input.shape, output_shape=grad_output.shape
+                )
+                grad_input = conv2d_gradfix(
+                    transpose=(not transpose),
+                    weight_shape=weight_shape,
+                    output_padding=p,
+                    **common_kwargs,
+                ).apply(grad_output, weight, None)
+
+            if ctx.needs_input_grad[1] and not weight_gradients_disabled:
+                grad_weight = Conv2dGradWeight.apply(grad_output, input)
+
+            if ctx.needs_input_grad[2]:
+                grad_bias = grad_output.sum((0, 2, 3))
+
+            return grad_input, grad_weight, grad_bias
+
+    class Conv2dGradWeight(autograd.Function):
+        @staticmethod
+        def forward(ctx, grad_output, input):
+            op = torch._C._jit_get_operation(
+                "aten::cudnn_convolution_backward_weight"
+                if not transpose
+                else "aten::cudnn_convolution_transpose_backward_weight"
+            )
+            flags = [
+                torch.backends.cudnn.benchmark,
+                torch.backends.cudnn.deterministic,
+                torch.backends.cudnn.allow_tf32,
+            ]
+            grad_weight = op(
+                weight_shape,
+                grad_output,
+                input,
+                padding,
+                stride,
+                dilation,
+                groups,
+                *flags,
+            )
+            ctx.save_for_backward(grad_output, input)
+
+            return grad_weight
+
+        @staticmethod
+        def backward(ctx, grad_grad_weight):
+            grad_output, input = ctx.saved_tensors
+            grad_grad_output, grad_grad_input = None, None
+
+            if ctx.needs_input_grad[0]:
+                grad_grad_output = Conv2d.apply(input, grad_grad_weight, None)
+
+            if ctx.needs_input_grad[1]:
+                p = calc_output_padding(
+                    input_shape=input.shape, output_shape=grad_output.shape
+                )
+                grad_grad_input = conv2d_gradfix(
+                    transpose=(not transpose),
+                    weight_shape=weight_shape,
+                    output_padding=p,
+                    **common_kwargs,
+                ).apply(grad_output, grad_grad_weight, None)
+
+            return grad_grad_output, grad_grad_input
+
+    conv2d_gradfix_cache[key] = Conv2d
+
+    return Conv2d

networks/op/fused_act.py

@@ -0,0 +1,127 @@
+import os
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+from torch.autograd import Function
+from torch.utils.cpp_extension import load
+
+
+module_path = os.path.dirname(__file__)
+fused = load(
+    "fused",
+    sources=[
+        os.path.join(module_path, "fused_bias_act.cpp"),
+        os.path.join(module_path, "fused_bias_act_kernel.cu"),
+    ],
+)
+
+
+class FusedLeakyReLUFunctionBackward(Function):
+    @staticmethod
+    def forward(ctx, grad_output, out, bias, negative_slope, scale):
+        ctx.save_for_backward(out)
+        ctx.negative_slope = negative_slope
+        ctx.scale = scale
+
+        empty = grad_output.new_empty(0)
+
+        grad_input = fused.fused_bias_act(
+            grad_output.contiguous(), empty, out, 3, 1, negative_slope, scale
+        )
+
+        dim = [0]
+
+        if grad_input.ndim > 2:
+            dim += list(range(2, grad_input.ndim))
+
+        if bias:
+            grad_bias = grad_input.sum(dim).detach()
+
+        else:
+            grad_bias = empty
+
+        return grad_input, grad_bias
+
+    @staticmethod
+    def backward(ctx, gradgrad_input, gradgrad_bias):
+        out, = ctx.saved_tensors
+        gradgrad_out = fused.fused_bias_act(
+            gradgrad_input.contiguous(),
+            gradgrad_bias,
+            out,
+            3,
+            1,
+            ctx.negative_slope,
+            ctx.scale,
+        )
+
+        return gradgrad_out, None, None, None, None
+
+
+class FusedLeakyReLUFunction(Function):
+    @staticmethod
+    def forward(ctx, input, bias, negative_slope, scale):
+        empty = input.new_empty(0)
+
+        ctx.bias = bias is not None
+
+        if bias is None:
+            bias = empty
+
+        out = fused.fused_bias_act(input, bias, empty, 3, 0, negative_slope, scale)
+        ctx.save_for_backward(out)
+        ctx.negative_slope = negative_slope
+        ctx.scale = scale
+
+        return out
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        out, = ctx.saved_tensors
+
+        grad_input, grad_bias = FusedLeakyReLUFunctionBackward.apply(
+            grad_output, out, ctx.bias, ctx.negative_slope, ctx.scale
+        )
+
+        if not ctx.bias:
+            grad_bias = None
+
+        return grad_input, grad_bias, None, None
+
+
+class FusedLeakyReLU(nn.Module):
+    def __init__(self, channel, bias=True, negative_slope=0.2, scale=2 ** 0.5):
+        super().__init__()
+
+        if bias:
+            self.bias = nn.Parameter(torch.zeros(channel))
+
+        else:
+            self.bias = None
+
+        self.negative_slope = negative_slope
+        self.scale = scale
+
+    def forward(self, input):
+        return fused_leaky_relu(input, self.bias, self.negative_slope, self.scale)
+
+
+def fused_leaky_relu(input, bias=None, negative_slope=0.2, scale=2 ** 0.5):
+    if input.device.type == "cpu":
+        if bias is not None:
+            rest_dim = [1] * (input.ndim - bias.ndim - 1)
+            return (
+                F.leaky_relu(
+                    input + bias.view(1, bias.shape[0], *rest_dim), negative_slope=0.2
+                )
+                * scale
+            )
+
+        else:
+            return F.leaky_relu(input, negative_slope=0.2) * scale
+
+    else:
+        return FusedLeakyReLUFunction.apply(
+            input.contiguous(), bias, negative_slope, scale
+        )

networks/op/fused_bias_act.cpp

@@ -0,0 +1,32 @@
+
+#include <ATen/ATen.h>
+#include <torch/extension.h>
+
+torch::Tensor fused_bias_act_op(const torch::Tensor &input,
+                                const torch::Tensor &bias,
+                                const torch::Tensor &refer, int act, int grad,
+                                float alpha, float scale);
+
+#define CHECK_CUDA(x)                                                          \
+  TORCH_CHECK(x.type().is_cuda(), #x " must be a CUDA tensor")
+#define CHECK_CONTIGUOUS(x)                                                    \
+  TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
+#define CHECK_INPUT(x)                                                         \
+  CHECK_CUDA(x);                                                               \
+  CHECK_CONTIGUOUS(x)
+
+torch::Tensor fused_bias_act(const torch::Tensor &input,
+                             const torch::Tensor &bias,
+                             const torch::Tensor &refer, int act, int grad,
+                             float alpha, float scale) {
+  CHECK_INPUT(input);
+  CHECK_INPUT(bias);
+
+  at::DeviceGuard guard(input.device());
+
+  return fused_bias_act_op(input, bias, refer, act, grad, alpha, scale);
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.def("fused_bias_act", &fused_bias_act, "fused bias act (CUDA)");
+}

networks/op/fused_bias_act_kernel.cu

@@ -0,0 +1,105 @@
+// Copyright (c) 2019, NVIDIA Corporation. All rights reserved.
+//
+// This work is made available under the Nvidia Source Code License-NC.
+// To view a copy of this license, visit
+// https://nvlabs.github.io/stylegan2/license.html
+
+#include <torch/types.h>
+
+#include <ATen/ATen.h>
+#include <ATen/AccumulateType.h>
+#include <ATen/cuda/CUDAApplyUtils.cuh>
+#include <ATen/cuda/CUDAContext.h>
+
+
+#include <cuda.h>
+#include <cuda_runtime.h>
+
+template <typename scalar_t>
+static __global__ void
+fused_bias_act_kernel(scalar_t *out, const scalar_t *p_x, const scalar_t *p_b,
+                      const scalar_t *p_ref, int act, int grad, scalar_t alpha,
+                      scalar_t scale, int loop_x, int size_x, int step_b,
+                      int size_b, int use_bias, int use_ref) {
+  int xi = blockIdx.x * loop_x * blockDim.x + threadIdx.x;
+
+  scalar_t zero = 0.0;
+
+  for (int loop_idx = 0; loop_idx < loop_x && xi < size_x;
+       loop_idx++, xi += blockDim.x) {
+    scalar_t x = p_x[xi];
+
+    if (use_bias) {
+      x += p_b[(xi / step_b) % size_b];
+    }
+
+    scalar_t ref = use_ref ? p_ref[xi] : zero;
+
+    scalar_t y;
+
+    switch (act * 10 + grad) {
+    default:
+    case 10:
+      y = x;
+      break;
+    case 11:
+      y = x;
+      break;
+    case 12:
+      y = 0.0;
+      break;
+
+    case 30:
+      y = (x > 0.0) ? x : x * alpha;
+      break;
+    case 31:
+      y = (ref > 0.0) ? x : x * alpha;
+      break;
+    case 32:
+      y = 0.0;
+      break;
+    }
+
+    out[xi] = y * scale;
+  }
+}
+
+torch::Tensor fused_bias_act_op(const torch::Tensor &input,
+                                const torch::Tensor &bias,
+                                const torch::Tensor &refer, int act, int grad,
+                                float alpha, float scale) {
+  int curDevice = -1;
+  cudaGetDevice(&curDevice);
+  cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+
+  auto x = input.contiguous();
+  auto b = bias.contiguous();
+  auto ref = refer.contiguous();
+
+  int use_bias = b.numel() ? 1 : 0;
+  int use_ref = ref.numel() ? 1 : 0;
+
+  int size_x = x.numel();
+  int size_b = b.numel();
+  int step_b = 1;
+
+  for (int i = 1 + 1; i < x.dim(); i++) {
+    step_b *= x.size(i);
+  }
+
+  int loop_x = 4;
+  int block_size = 4 * 32;
+  int grid_size = (size_x - 1) / (loop_x * block_size) + 1;
+
+  auto y = torch::empty_like(x);
+
+  AT_DISPATCH_FLOATING_TYPES_AND_HALF(
+      x.scalar_type(), "fused_bias_act_kernel", [&] {
+        fused_bias_act_kernel<scalar_t><<<grid_size, block_size, 0, stream>>>(
+            y.data_ptr<scalar_t>(), x.data_ptr<scalar_t>(),
+            b.data_ptr<scalar_t>(), ref.data_ptr<scalar_t>(), act, grad, alpha,
+            scale, loop_x, size_x, step_b, size_b, use_bias, use_ref);
+      });
+
+  return y;
+}

networks/op/upfirdn2d.cpp

@@ -0,0 +1,31 @@
+#include <ATen/ATen.h>
+#include <torch/extension.h>
+
+torch::Tensor upfirdn2d_op(const torch::Tensor &input,
+                           const torch::Tensor &kernel, int up_x, int up_y,
+                           int down_x, int down_y, int pad_x0, int pad_x1,
+                           int pad_y0, int pad_y1);
+
+#define CHECK_CUDA(x)                                                          \
+  TORCH_CHECK(x.type().is_cuda(), #x " must be a CUDA tensor")
+#define CHECK_CONTIGUOUS(x)                                                    \
+  TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
+#define CHECK_INPUT(x)                                                         \
+  CHECK_CUDA(x);                                                               \
+  CHECK_CONTIGUOUS(x)
+
+torch::Tensor upfirdn2d(const torch::Tensor &input, const torch::Tensor &kernel,
+                        int up_x, int up_y, int down_x, int down_y, int pad_x0,
+                        int pad_x1, int pad_y0, int pad_y1) {
+  CHECK_INPUT(input);
+  CHECK_INPUT(kernel);
+
+  at::DeviceGuard guard(input.device());
+
+  return upfirdn2d_op(input, kernel, up_x, up_y, down_x, down_y, pad_x0, pad_x1,
+                      pad_y0, pad_y1);
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.def("upfirdn2d", &upfirdn2d, "upfirdn2d (CUDA)");
+}

networks/op/upfirdn2d.py

@@ -0,0 +1,209 @@
+from collections import abc
+import os
+
+import torch
+from torch.nn import functional as F
+from torch.autograd import Function
+from torch.utils.cpp_extension import load
+
+
+module_path = os.path.dirname(__file__)
+upfirdn2d_op = load(
+    "upfirdn2d",
+    sources=[
+        os.path.join(module_path, "upfirdn2d.cpp"),
+        os.path.join(module_path, "upfirdn2d_kernel.cu"),
+    ],
+)
+
+
+class UpFirDn2dBackward(Function):
+    @staticmethod
+    def forward(
+        ctx, grad_output, kernel, grad_kernel, up, down, pad, g_pad, in_size, out_size
+    ):
+
+        up_x, up_y = up
+        down_x, down_y = down
+        g_pad_x0, g_pad_x1, g_pad_y0, g_pad_y1 = g_pad
+
+        grad_output = grad_output.reshape(-1, out_size[0], out_size[1], 1)
+
+        grad_input = upfirdn2d_op.upfirdn2d(
+            grad_output,
+            grad_kernel,
+            down_x,
+            down_y,
+            up_x,
+            up_y,
+            g_pad_x0,
+            g_pad_x1,
+            g_pad_y0,
+            g_pad_y1,
+        )
+        grad_input = grad_input.view(in_size[0], in_size[1], in_size[2], in_size[3])
+
+        ctx.save_for_backward(kernel)
+
+        pad_x0, pad_x1, pad_y0, pad_y1 = pad
+
+        ctx.up_x = up_x
+        ctx.up_y = up_y
+        ctx.down_x = down_x
+        ctx.down_y = down_y
+        ctx.pad_x0 = pad_x0
+        ctx.pad_x1 = pad_x1
+        ctx.pad_y0 = pad_y0
+        ctx.pad_y1 = pad_y1
+        ctx.in_size = in_size
+        ctx.out_size = out_size
+
+        return grad_input
+
+    @staticmethod
+    def backward(ctx, gradgrad_input):
+        kernel, = ctx.saved_tensors
+
+        gradgrad_input = gradgrad_input.reshape(-1, ctx.in_size[2], ctx.in_size[3], 1)
+
+        gradgrad_out = upfirdn2d_op.upfirdn2d(
+            gradgrad_input,
+            kernel,
+            ctx.up_x,
+            ctx.up_y,
+            ctx.down_x,
+            ctx.down_y,
+            ctx.pad_x0,
+            ctx.pad_x1,
+            ctx.pad_y0,
+            ctx.pad_y1,
+        )
+        # gradgrad_out = gradgrad_out.view(ctx.in_size[0], ctx.out_size[0], ctx.out_size[1], ctx.in_size[3])
+        gradgrad_out = gradgrad_out.view(
+            ctx.in_size[0], ctx.in_size[1], ctx.out_size[0], ctx.out_size[1]
+        )
+
+        return gradgrad_out, None, None, None, None, None, None, None, None
+
+
+class UpFirDn2d(Function):
+    @staticmethod
+    def forward(ctx, input, kernel, up, down, pad):
+        up_x, up_y = up
+        down_x, down_y = down
+        pad_x0, pad_x1, pad_y0, pad_y1 = pad
+
+        kernel_h, kernel_w = kernel.shape
+        batch, channel, in_h, in_w = input.shape
+        ctx.in_size = input.shape
+
+        input = input.reshape(-1, in_h, in_w, 1)
+
+        ctx.save_for_backward(kernel, torch.flip(kernel, [0, 1]))
+
+        out_h = (in_h * up_y + pad_y0 + pad_y1 - kernel_h + down_y) // down_y
+        out_w = (in_w * up_x + pad_x0 + pad_x1 - kernel_w + down_x) // down_x
+        ctx.out_size = (out_h, out_w)
+
+        ctx.up = (up_x, up_y)
+        ctx.down = (down_x, down_y)
+        ctx.pad = (pad_x0, pad_x1, pad_y0, pad_y1)
+
+        g_pad_x0 = kernel_w - pad_x0 - 1
+        g_pad_y0 = kernel_h - pad_y0 - 1
+        g_pad_x1 = in_w * up_x - out_w * down_x + pad_x0 - up_x + 1
+        g_pad_y1 = in_h * up_y - out_h * down_y + pad_y0 - up_y + 1
+
+        ctx.g_pad = (g_pad_x0, g_pad_x1, g_pad_y0, g_pad_y1)
+
+        out = upfirdn2d_op.upfirdn2d(
+            input, kernel, up_x, up_y, down_x, down_y, pad_x0, pad_x1, pad_y0, pad_y1
+        )
+        # out = out.view(major, out_h, out_w, minor)
+        out = out.view(-1, channel, out_h, out_w)
+
+        return out
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        kernel, grad_kernel = ctx.saved_tensors
+
+        grad_input = None
+
+        if ctx.needs_input_grad[0]:
+            grad_input = UpFirDn2dBackward.apply(
+                grad_output,
+                kernel,
+                grad_kernel,
+                ctx.up,
+                ctx.down,
+                ctx.pad,
+                ctx.g_pad,
+                ctx.in_size,
+                ctx.out_size,
+            )
+
+        return grad_input, None, None, None, None
+
+
+def upfirdn2d(input, kernel, up=1, down=1, pad=(0, 0)):
+    if not isinstance(up, abc.Iterable):
+        up = (up, up)
+
+    if not isinstance(down, abc.Iterable):
+        down = (down, down)
+
+    if len(pad) == 2:
+        pad = (pad[0], pad[1], pad[0], pad[1])
+
+    if input.device.type == "cpu":
+        out = upfirdn2d_native(input, kernel, *up, *down, *pad)
+
+    else:
+        out = UpFirDn2d.apply(input, kernel, up, down, pad)
+
+    return out
+
+
+def upfirdn2d_native(
+    input, kernel, up_x, up_y, down_x, down_y, pad_x0, pad_x1, pad_y0, pad_y1
+):
+    _, channel, in_h, in_w = input.shape
+    input = input.reshape(-1, in_h, in_w, 1)
+
+    _, in_h, in_w, minor = input.shape
+    kernel_h, kernel_w = kernel.shape
+
+    out = input.view(-1, in_h, 1, in_w, 1, minor)
+    out = F.pad(out, [0, 0, 0, up_x - 1, 0, 0, 0, up_y - 1])
+    out = out.view(-1, in_h * up_y, in_w * up_x, minor)
+
+    out = F.pad(
+        out, [0, 0, max(pad_x0, 0), max(pad_x1, 0), max(pad_y0, 0), max(pad_y1, 0)]
+    )
+    out = out[
+        :,
+        max(-pad_y0, 0) : out.shape[1] - max(-pad_y1, 0),
+        max(-pad_x0, 0) : out.shape[2] - max(-pad_x1, 0),
+        :,
+    ]
+
+    out = out.permute(0, 3, 1, 2)
+    out = out.reshape(
+        [-1, 1, in_h * up_y + pad_y0 + pad_y1, in_w * up_x + pad_x0 + pad_x1]
+    )
+    w = torch.flip(kernel, [0, 1]).view(1, 1, kernel_h, kernel_w)
+    out = F.conv2d(out, w)
+    out = out.reshape(
+        -1,
+        minor,
+        in_h * up_y + pad_y0 + pad_y1 - kernel_h + 1,
+        in_w * up_x + pad_x0 + pad_x1 - kernel_w + 1,
+    )
+    out = out.permute(0, 2, 3, 1)
+    out = out[:, ::down_y, ::down_x, :]
+
+    out_h = (in_h * up_y + pad_y0 + pad_y1 - kernel_h + down_y) // down_y
+    out_w = (in_w * up_x + pad_x0 + pad_x1 - kernel_w + down_x) // down_x
+
+    return out.view(-1, channel, out_h, out_w)

networks/op/upfirdn2d_kernel.cu

@@ -0,0 +1,369 @@
+// Copyright (c) 2019, NVIDIA Corporation. All rights reserved.
+//
+// This work is made available under the Nvidia Source Code License-NC.
+// To view a copy of this license, visit
+// https://nvlabs.github.io/stylegan2/license.html
+
+#include <torch/types.h>
+
+#include <ATen/ATen.h>
+#include <ATen/AccumulateType.h>
+#include <ATen/cuda/CUDAApplyUtils.cuh>
+#include <ATen/cuda/CUDAContext.h>
+
+#include <cuda.h>
+#include <cuda_runtime.h>
+
+static __host__ __device__ __forceinline__ int floor_div(int a, int b) {
+  int c = a / b;
+
+  if (c * b > a) {
+    c--;
+  }
+
+  return c;
+}
+
+struct UpFirDn2DKernelParams {
+  int up_x;
+  int up_y;
+  int down_x;
+  int down_y;
+  int pad_x0;
+  int pad_x1;
+  int pad_y0;
+  int pad_y1;
+
+  int major_dim;
+  int in_h;
+  int in_w;
+  int minor_dim;
+  int kernel_h;
+  int kernel_w;
+  int out_h;
+  int out_w;
+  int loop_major;
+  int loop_x;
+};
+
+template <typename scalar_t>
+__global__ void upfirdn2d_kernel_large(scalar_t *out, const scalar_t *input,
+                                       const scalar_t *kernel,
+                                       const UpFirDn2DKernelParams p) {
+  int minor_idx = blockIdx.x * blockDim.x + threadIdx.x;
+  int out_y = minor_idx / p.minor_dim;
+  minor_idx -= out_y * p.minor_dim;
+  int out_x_base = blockIdx.y * p.loop_x * blockDim.y + threadIdx.y;
+  int major_idx_base = blockIdx.z * p.loop_major;
+
+  if (out_x_base >= p.out_w || out_y >= p.out_h ||
+      major_idx_base >= p.major_dim) {
+    return;
+  }
+
+  int mid_y = out_y * p.down_y + p.up_y - 1 - p.pad_y0;
+  int in_y = min(max(floor_div(mid_y, p.up_y), 0), p.in_h);
+  int h = min(max(floor_div(mid_y + p.kernel_h, p.up_y), 0), p.in_h) - in_y;
+  int kernel_y = mid_y + p.kernel_h - (in_y + 1) * p.up_y;
+
+  for (int loop_major = 0, major_idx = major_idx_base;
+       loop_major < p.loop_major && major_idx < p.major_dim;
+       loop_major++, major_idx++) {
+    for (int loop_x = 0, out_x = out_x_base;
+         loop_x < p.loop_x && out_x < p.out_w; loop_x++, out_x += blockDim.y) {
+      int mid_x = out_x * p.down_x + p.up_x - 1 - p.pad_x0;
+      int in_x = min(max(floor_div(mid_x, p.up_x), 0), p.in_w);
+      int w = min(max(floor_div(mid_x + p.kernel_w, p.up_x), 0), p.in_w) - in_x;
+      int kernel_x = mid_x + p.kernel_w - (in_x + 1) * p.up_x;
+
+      const scalar_t *x_p =
+          &input[((major_idx * p.in_h + in_y) * p.in_w + in_x) * p.minor_dim +
+                 minor_idx];
+      const scalar_t *k_p = &kernel[kernel_y * p.kernel_w + kernel_x];
+      int x_px = p.minor_dim;
+      int k_px = -p.up_x;
+      int x_py = p.in_w * p.minor_dim;
+      int k_py = -p.up_y * p.kernel_w;
+
+      scalar_t v = 0.0f;
+
+      for (int y = 0; y < h; y++) {
+        for (int x = 0; x < w; x++) {
+          v += static_cast<scalar_t>(*x_p) * static_cast<scalar_t>(*k_p);
+          x_p += x_px;
+          k_p += k_px;
+        }
+
+        x_p += x_py - w * x_px;
+        k_p += k_py - w * k_px;
+      }
+
+      out[((major_idx * p.out_h + out_y) * p.out_w + out_x) * p.minor_dim +
+          minor_idx] = v;
+    }
+  }
+}
+
+template <typename scalar_t, int up_x, int up_y, int down_x, int down_y,
+          int kernel_h, int kernel_w, int tile_out_h, int tile_out_w>
+__global__ void upfirdn2d_kernel(scalar_t *out, const scalar_t *input,
+                                 const scalar_t *kernel,
+                                 const UpFirDn2DKernelParams p) {
+  const int tile_in_h = ((tile_out_h - 1) * down_y + kernel_h - 1) / up_y + 1;
+  const int tile_in_w = ((tile_out_w - 1) * down_x + kernel_w - 1) / up_x + 1;
+
+  __shared__ volatile float sk[kernel_h][kernel_w];
+  __shared__ volatile float sx[tile_in_h][tile_in_w];
+
+  int minor_idx = blockIdx.x;
+  int tile_out_y = minor_idx / p.minor_dim;
+  minor_idx -= tile_out_y * p.minor_dim;
+  tile_out_y *= tile_out_h;
+  int tile_out_x_base = blockIdx.y * p.loop_x * tile_out_w;
+  int major_idx_base = blockIdx.z * p.loop_major;
+
+  if (tile_out_x_base >= p.out_w | tile_out_y >= p.out_h |
+      major_idx_base >= p.major_dim) {
+    return;
+  }
+
+  for (int tap_idx = threadIdx.x; tap_idx < kernel_h * kernel_w;
+       tap_idx += blockDim.x) {
+    int ky = tap_idx / kernel_w;
+    int kx = tap_idx - ky * kernel_w;
+    scalar_t v = 0.0;
+
+    if (kx < p.kernel_w & ky < p.kernel_h) {
+      v = kernel[(p.kernel_h - 1 - ky) * p.kernel_w + (p.kernel_w - 1 - kx)];
+    }
+
+    sk[ky][kx] = v;
+  }
+
+  for (int loop_major = 0, major_idx = major_idx_base;
+       loop_major < p.loop_major & major_idx < p.major_dim;
+       loop_major++, major_idx++) {
+    for (int loop_x = 0, tile_out_x = tile_out_x_base;
+         loop_x < p.loop_x & tile_out_x < p.out_w;
+         loop_x++, tile_out_x += tile_out_w) {
+      int tile_mid_x = tile_out_x * down_x + up_x - 1 - p.pad_x0;
+      int tile_mid_y = tile_out_y * down_y + up_y - 1 - p.pad_y0;
+      int tile_in_x = floor_div(tile_mid_x, up_x);
+      int tile_in_y = floor_div(tile_mid_y, up_y);
+
+      __syncthreads();
+
+      for (int in_idx = threadIdx.x; in_idx < tile_in_h * tile_in_w;
+           in_idx += blockDim.x) {
+        int rel_in_y = in_idx / tile_in_w;
+        int rel_in_x = in_idx - rel_in_y * tile_in_w;
+        int in_x = rel_in_x + tile_in_x;
+        int in_y = rel_in_y + tile_in_y;
+
+        scalar_t v = 0.0;
+
+        if (in_x >= 0 & in_y >= 0 & in_x < p.in_w & in_y < p.in_h) {
+          v = input[((major_idx * p.in_h + in_y) * p.in_w + in_x) *
+                        p.minor_dim +
+                    minor_idx];
+        }
+
+        sx[rel_in_y][rel_in_x] = v;
+      }
+
+      __syncthreads();
+      for (int out_idx = threadIdx.x; out_idx < tile_out_h * tile_out_w;
+           out_idx += blockDim.x) {
+        int rel_out_y = out_idx / tile_out_w;
+        int rel_out_x = out_idx - rel_out_y * tile_out_w;
+        int out_x = rel_out_x + tile_out_x;
+        int out_y = rel_out_y + tile_out_y;
+
+        int mid_x = tile_mid_x + rel_out_x * down_x;
+        int mid_y = tile_mid_y + rel_out_y * down_y;
+        int in_x = floor_div(mid_x, up_x);
+        int in_y = floor_div(mid_y, up_y);
+        int rel_in_x = in_x - tile_in_x;
+        int rel_in_y = in_y - tile_in_y;
+        int kernel_x = (in_x + 1) * up_x - mid_x - 1;
+        int kernel_y = (in_y + 1) * up_y - mid_y - 1;
+
+        scalar_t v = 0.0;
+
+#pragma unroll
+        for (int y = 0; y < kernel_h / up_y; y++)
+#pragma unroll
+          for (int x = 0; x < kernel_w / up_x; x++)
+            v += sx[rel_in_y + y][rel_in_x + x] *
+                 sk[kernel_y + y * up_y][kernel_x + x * up_x];
+
+        if (out_x < p.out_w & out_y < p.out_h) {
+          out[((major_idx * p.out_h + out_y) * p.out_w + out_x) * p.minor_dim +
+              minor_idx] = v;
+        }
+      }
+    }
+  }
+}
+
+torch::Tensor upfirdn2d_op(const torch::Tensor &input,
+                           const torch::Tensor &kernel, int up_x, int up_y,
+                           int down_x, int down_y, int pad_x0, int pad_x1,
+                           int pad_y0, int pad_y1) {
+  int curDevice = -1;
+  cudaGetDevice(&curDevice);
+  cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+
+  UpFirDn2DKernelParams p;
+
+  auto x = input.contiguous();
+  auto k = kernel.contiguous();
+
+  p.major_dim = x.size(0);
+  p.in_h = x.size(1);
+  p.in_w = x.size(2);
+  p.minor_dim = x.size(3);
+  p.kernel_h = k.size(0);
+  p.kernel_w = k.size(1);
+  p.up_x = up_x;
+  p.up_y = up_y;
+  p.down_x = down_x;
+  p.down_y = down_y;
+  p.pad_x0 = pad_x0;
+  p.pad_x1 = pad_x1;
+  p.pad_y0 = pad_y0;
+  p.pad_y1 = pad_y1;
+
+  p.out_h = (p.in_h * p.up_y + p.pad_y0 + p.pad_y1 - p.kernel_h + p.down_y) /
+            p.down_y;
+  p.out_w = (p.in_w * p.up_x + p.pad_x0 + p.pad_x1 - p.kernel_w + p.down_x) /
+            p.down_x;
+
+  auto out =
+      at::empty({p.major_dim, p.out_h, p.out_w, p.minor_dim}, x.options());
+
+  int mode = -1;
+
+  int tile_out_h = -1;
+  int tile_out_w = -1;
+
+  if (p.up_x == 1 && p.up_y == 1 && p.down_x == 1 && p.down_y == 1 &&
+      p.kernel_h <= 4 && p.kernel_w <= 4) {
+    mode = 1;
+    tile_out_h = 16;
+    tile_out_w = 64;
+  }
+
+  if (p.up_x == 1 && p.up_y == 1 && p.down_x == 1 && p.down_y == 1 &&
+      p.kernel_h <= 3 && p.kernel_w <= 3) {
+    mode = 2;
+    tile_out_h = 16;
+    tile_out_w = 64;
+  }
+
+  if (p.up_x == 2 && p.up_y == 2 && p.down_x == 1 && p.down_y == 1 &&
+      p.kernel_h <= 4 && p.kernel_w <= 4) {
+    mode = 3;
+    tile_out_h = 16;
+    tile_out_w = 64;
+  }
+
+  if (p.up_x == 2 && p.up_y == 2 && p.down_x == 1 && p.down_y == 1 &&
+      p.kernel_h <= 2 && p.kernel_w <= 2) {
+    mode = 4;
+    tile_out_h = 16;
+    tile_out_w = 64;
+  }
+
+  if (p.up_x == 1 && p.up_y == 1 && p.down_x == 2 && p.down_y == 2 &&
+      p.kernel_h <= 4 && p.kernel_w <= 4) {
+    mode = 5;
+    tile_out_h = 8;
+    tile_out_w = 32;
+  }
+
+  if (p.up_x == 1 && p.up_y == 1 && p.down_x == 2 && p.down_y == 2 &&
+      p.kernel_h <= 2 && p.kernel_w <= 2) {
+    mode = 6;
+    tile_out_h = 8;
+    tile_out_w = 32;
+  }
+
+  dim3 block_size;
+  dim3 grid_size;
+
+  if (tile_out_h > 0 && tile_out_w > 0) {
+    p.loop_major = (p.major_dim - 1) / 16384 + 1;
+    p.loop_x = 1;
+    block_size = dim3(32 * 8, 1, 1);
+    grid_size = dim3(((p.out_h - 1) / tile_out_h + 1) * p.minor_dim,
+                     (p.out_w - 1) / (p.loop_x * tile_out_w) + 1,
+                     (p.major_dim - 1) / p.loop_major + 1);
+  } else {
+    p.loop_major = (p.major_dim - 1) / 16384 + 1;
+    p.loop_x = 4;
+    block_size = dim3(4, 32, 1);
+    grid_size = dim3((p.out_h * p.minor_dim - 1) / block_size.x + 1,
+                     (p.out_w - 1) / (p.loop_x * block_size.y) + 1,
+                     (p.major_dim - 1) / p.loop_major + 1);
+  }
+
+  AT_DISPATCH_FLOATING_TYPES_AND_HALF(x.scalar_type(), "upfirdn2d_cuda", [&] {
+    switch (mode) {
+    case 1:
+      upfirdn2d_kernel<scalar_t, 1, 1, 1, 1, 4, 4, 16, 64>
+          <<<grid_size, block_size, 0, stream>>>(out.data_ptr<scalar_t>(),
+                                                 x.data_ptr<scalar_t>(),
+                                                 k.data_ptr<scalar_t>(), p);
+
+      break;
+
+    case 2:
+      upfirdn2d_kernel<scalar_t, 1, 1, 1, 1, 3, 3, 16, 64>
+          <<<grid_size, block_size, 0, stream>>>(out.data_ptr<scalar_t>(),
+                                                 x.data_ptr<scalar_t>(),
+                                                 k.data_ptr<scalar_t>(), p);
+
+      break;
+
+    case 3:
+      upfirdn2d_kernel<scalar_t, 2, 2, 1, 1, 4, 4, 16, 64>
+          <<<grid_size, block_size, 0, stream>>>(out.data_ptr<scalar_t>(),
+                                                 x.data_ptr<scalar_t>(),
+                                                 k.data_ptr<scalar_t>(), p);
+
+      break;
+
+    case 4:
+      upfirdn2d_kernel<scalar_t, 2, 2, 1, 1, 2, 2, 16, 64>
+          <<<grid_size, block_size, 0, stream>>>(out.data_ptr<scalar_t>(),
+                                                 x.data_ptr<scalar_t>(),
+                                                 k.data_ptr<scalar_t>(), p);
+
+      break;
+
+    case 5:
+      upfirdn2d_kernel<scalar_t, 1, 1, 2, 2, 4, 4, 8, 32>
+          <<<grid_size, block_size, 0, stream>>>(out.data_ptr<scalar_t>(),
+                                                 x.data_ptr<scalar_t>(),
+                                                 k.data_ptr<scalar_t>(), p);
+
+      break;
+
+    case 6:
+      upfirdn2d_kernel<scalar_t, 1, 1, 2, 2, 4, 4, 8, 32>
+          <<<grid_size, block_size, 0, stream>>>(out.data_ptr<scalar_t>(),
+                                                 x.data_ptr<scalar_t>(),
+                                                 k.data_ptr<scalar_t>(), p);
+
+      break;
+
+    default:
+      upfirdn2d_kernel_large<scalar_t><<<grid_size, block_size, 0, stream>>>(
+          out.data_ptr<scalar_t>(), x.data_ptr<scalar_t>(),
+          k.data_ptr<scalar_t>(), p);
+    }
+  });
+
+  return out;
+}

networks/op/upfirdn2d_new.py

@@ -0,0 +1,230 @@
+import os
+
+import torch
+import torch.nn.functional as F
+from torch.autograd import Function
+from torch.utils.cpp_extension import load
+#from util import is_custom_kernel_supported as is_custom_kernel_supported
+
+def is_custom_kernel_supported():
+    version_str = str(torch.version.cuda).split(".")
+    major = version_str[0]
+    minor = version_str[1]
+    return int(major) >= 10 and int(minor) >= 1
+
+
+if is_custom_kernel_supported():
+    print("Loading custom kernel...")
+    module_path = os.path.dirname(__file__)
+    upfirdn2d_op = load(
+        'upfirdn2d',
+        sources=[
+            os.path.join(module_path, 'upfirdn2d.cpp'),
+            os.path.join(module_path, 'upfirdn2d_kernel.cu'),
+        ],
+        verbose=True
+    )
+
+use_custom_kernel = is_custom_kernel_supported()
+
+
+class UpFirDn2dBackward(Function):
+    @staticmethod
+    def forward(
+        ctx, grad_output, kernel, grad_kernel, up, down, pad, g_pad, in_size, out_size
+    ):
+
+        up_x, up_y = up
+        down_x, down_y = down
+        g_pad_x0, g_pad_x1, g_pad_y0, g_pad_y1 = g_pad
+
+        grad_output = grad_output.reshape(-1, out_size[0], out_size[1], 1)
+
+        grad_input = upfirdn2d_op.upfirdn2d(
+            grad_output,
+            grad_kernel,
+            down_x,
+            down_y,
+            up_x,
+            up_y,
+            g_pad_x0,
+            g_pad_x1,
+            g_pad_y0,
+            g_pad_y1,
+        )
+        grad_input = grad_input.view(in_size[0], in_size[1], in_size[2], in_size[3])
+
+        ctx.save_for_backward(kernel)
+
+        pad_x0, pad_x1, pad_y0, pad_y1 = pad
+
+        ctx.up_x = up_x
+        ctx.up_y = up_y
+        ctx.down_x = down_x
+        ctx.down_y = down_y
+        ctx.pad_x0 = pad_x0
+        ctx.pad_x1 = pad_x1
+        ctx.pad_y0 = pad_y0
+        ctx.pad_y1 = pad_y1
+        ctx.in_size = in_size
+        ctx.out_size = out_size
+
+        return grad_input
+
+    @staticmethod
+    def backward(ctx, gradgrad_input):
+        kernel, = ctx.saved_tensors
+
+        gradgrad_input = gradgrad_input.reshape(-1, ctx.in_size[2], ctx.in_size[3], 1)
+
+        gradgrad_out = upfirdn2d_op.upfirdn2d(
+            gradgrad_input,
+            kernel,
+            ctx.up_x,
+            ctx.up_y,
+            ctx.down_x,
+            ctx.down_y,
+            ctx.pad_x0,
+            ctx.pad_x1,
+            ctx.pad_y0,
+            ctx.pad_y1,
+        )
+        # gradgrad_out = gradgrad_out.view(ctx.in_size[0], ctx.out_size[0], ctx.out_size[1], ctx.in_size[3])
+        gradgrad_out = gradgrad_out.view(
+            ctx.in_size[0], ctx.in_size[1], ctx.out_size[0], ctx.out_size[1]
+        )
+
+        return gradgrad_out, None, None, None, None, None, None, None, None
+
+
+class UpFirDn2d(Function):
+    @staticmethod
+    def forward(ctx, input, kernel, up, down, pad):
+        up_x, up_y = up
+        down_x, down_y = down
+        pad_x0, pad_x1, pad_y0, pad_y1 = pad
+
+        kernel_h, kernel_w = kernel.shape
+        batch, channel, in_h, in_w = input.shape
+        ctx.in_size = input.shape
+
+        input = input.reshape(-1, in_h, in_w, 1)
+
+        ctx.save_for_backward(kernel, torch.flip(kernel, [0, 1]))
+
+        out_h = (in_h * up_y + pad_y0 + pad_y1 - kernel_h) // down_y + 1
+        out_w = (in_w * up_x + pad_x0 + pad_x1 - kernel_w) // down_x + 1
+        ctx.out_size = (out_h, out_w)
+
+        ctx.up = (up_x, up_y)
+        ctx.down = (down_x, down_y)
+        ctx.pad = (pad_x0, pad_x1, pad_y0, pad_y1)
+
+        g_pad_x0 = kernel_w - pad_x0 - 1
+        g_pad_y0 = kernel_h - pad_y0 - 1
+        g_pad_x1 = in_w * up_x - out_w * down_x + pad_x0 - up_x + 1
+        g_pad_y1 = in_h * up_y - out_h * down_y + pad_y0 - up_y + 1
+
+        ctx.g_pad = (g_pad_x0, g_pad_x1, g_pad_y0, g_pad_y1)
+
+        out = upfirdn2d_op.upfirdn2d(
+            input, kernel, up_x, up_y, down_x, down_y, pad_x0, pad_x1, pad_y0, pad_y1
+        )
+        # out = out.view(major, out_h, out_w, minor)
+        out = out.view(-1, channel, out_h, out_w)
+
+        return out
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        kernel, grad_kernel = ctx.saved_tensors
+
+        grad_input = UpFirDn2dBackward.apply(
+            grad_output,
+            kernel,
+            grad_kernel,
+            ctx.up,
+            ctx.down,
+            ctx.pad,
+            ctx.g_pad,
+            ctx.in_size,
+            ctx.out_size,
+        )
+
+        return grad_input, None, None, None, None
+
+
+def upfirdn2d(input, kernel, up=1, down=1, pad=(0, 0)):
+    global use_custom_kernel
+    if use_custom_kernel:
+        out = UpFirDn2d.apply(
+            input, kernel, (up, up), (down, down), (pad[0], pad[1], pad[0], pad[1])
+        )
+    else:
+        out = upfirdn2d_native(input, kernel, up, up, down, down, pad[0], pad[1], pad[0], pad[1])
+
+    return out
+
+
+def upfirdn2d_native(
+    input, kernel, up_x, up_y, down_x, down_y, pad_x0, pad_x1, pad_y0, pad_y1
+):
+    bs, ch, in_h, in_w = input.shape
+    minor = 1
+    kernel_h, kernel_w = kernel.shape
+
+    #assert kernel_h == 1 and kernel_w == 1
+
+    #print("original shape ", input.shape, up_x, down_x, pad_x0, pad_x1)
+
+    out = input.view(-1, in_h, 1, in_w, 1, minor)
+    if up_x > 1 or up_y > 1:
+        out = F.pad(out, [0, 0, 0, up_x - 1, 0, 0, 0, up_y - 1])
+
+    #print("after padding ", out.shape)
+    out = out.view(-1, in_h * up_y, in_w * up_x, minor)
+
+    #print("after reshaping ", out.shape)
+
+    if pad_x0 > 0 or pad_x1 > 0 or pad_y0 > 0 or pad_y1 > 0:
+        out = F.pad(out, [0, 0, max(pad_x0, 0), max(pad_x1, 0), max(pad_y0, 0), max(pad_y1, 0)])
+
+    #print("after second padding ", out.shape)
+    out = out[
+        :,
+        max(-pad_y0, 0) : out.shape[1] - max(-pad_y1, 0),
+        max(-pad_x0, 0) : out.shape[2] - max(-pad_x1, 0),
+        :,
+    ]
+
+    #print("after trimming ", out.shape)
+
+    out = out.permute(0, 3, 1, 2)
+    out = out.reshape(
+        [-1, 1, in_h * up_y + pad_y0 + pad_y1, in_w * up_x + pad_x0 + pad_x1]
+    )
+
+    #print("after reshaping", out.shape)
+    w = torch.flip(kernel, [0, 1]).view(1, 1, kernel_h, kernel_w)
+    out = F.conv2d(out, w)
+
+    #print("after conv ", out.shape)
+    out = out.reshape(
+        -1,
+        minor,
+        in_h * up_y + pad_y0 + pad_y1 - kernel_h + 1,
+        in_w * up_x + pad_x0 + pad_x1 - kernel_w + 1,
+    )
+
+    out = out.permute(0, 2, 3, 1)
+
+    #print("after permuting ", out.shape)
+
+    out = out[:, ::down_y, ::down_x, :]
+
+    out = out.view(bs, ch, out.size(1), out.size(2))
+
+    #print("final shape ", out.shape)
+
+    return out
+

networks/ops.py

@@ -0,0 +1,490 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+from .op import (FusedLeakyReLU, fused_leaky_relu, upfirdn2d)
+import numpy as np
+
+
+def make_kernel(k):
+	k = torch.tensor(k, dtype=torch.float32)
+
+	if k.ndim == 1:
+		k = k[None, :] * k[:, None]
+
+	k /= k.sum()
+
+	return k
+
+
+class Blur(nn.Module):
+	def __init__(self, kernel, pad, upsample_factor=1):
+		super().__init__()
+
+		kernel = make_kernel(kernel)
+
+		if upsample_factor > 1:
+			kernel = kernel * (upsample_factor ** 2)
+
+		self.register_buffer('kernel', kernel)
+
+		self.pad = pad
+
+	def forward(self, input):
+		return upfirdn2d(input, self.kernel, pad=self.pad)
+
+
+class ScaledLeakyReLU(nn.Module):
+	def __init__(self, negative_slope=0.2):
+		super().__init__()
+
+		self.negative_slope = negative_slope
+
+	def forward(self, input):
+		return F.leaky_relu(input, negative_slope=self.negative_slope)
+
+
+class EqualConv2d(nn.Module):
+	def __init__(self, in_channel, out_channel, kernel_size, stride=1, padding=0, bias=True):
+		super().__init__()
+
+		self.weight = nn.Parameter(torch.randn(out_channel, in_channel, kernel_size, kernel_size))
+		self.scale = 1 / math.sqrt(in_channel * kernel_size ** 2)
+
+		self.stride = stride
+		self.padding = padding
+
+		if bias:
+			self.bias = nn.Parameter(torch.zeros(out_channel))
+		else:
+			self.bias = None
+
+	def forward(self, input):
+
+		return F.conv2d(input, self.weight * self.scale, bias=self.bias, stride=self.stride, padding=self.padding)
+
+	def __repr__(self):
+		return (
+			f'{self.__class__.__name__}({self.weight.shape[1]}, {self.weight.shape[0]},'
+			f' {self.weight.shape[2]}, stride={self.stride}, padding={self.padding})'
+		)
+
+
+class EqualLinear(nn.Module):
+	def __init__(self, in_dim, out_dim, bias=True, bias_init=0, lr_mul=1, activation=None):
+		super().__init__()
+
+		self.weight = nn.Parameter(torch.randn(out_dim, in_dim).div_(lr_mul))
+		
+		bias_init = np.broadcast_to(np.asarray(bias_init, dtype=np.float32), [out_dim])
+		if bias:
+			self.bias = nn.Parameter(torch.from_numpy(bias_init / lr_mul))
+			#self.bias = nn.Parameter(torch.zeros(out_dim).fill_(bias_init))
+		else:
+			self.bias = None
+
+		self.activation = activation
+
+		self.scale = (1 / math.sqrt(in_dim)) * lr_mul
+		self.lr_mul = lr_mul
+
+	def forward(self, input):
+
+		if self.activation:
+			out = F.linear(input, self.weight * self.scale)
+			out = fused_leaky_relu(out, self.bias * self.lr_mul)
+		else:
+			out = F.linear(input, self.weight * self.scale, bias=self.bias * self.lr_mul)
+
+		return out
+
+	def __repr__(self):
+		return (f'{self.__class__.__name__}({self.weight.shape[1]}, {self.weight.shape[0]})')
+
+
+class ConvLayer(nn.Sequential):
+	def __init__(
+			self,
+			in_channel,
+			out_channel,
+			kernel_size,
+			downsample=False,
+			upsample=False,
+			blur_kernel=[1, 3, 3, 1],
+			bias=True,
+			activate=True,
+	):
+		layers = []
+
+		if downsample:
+			factor = 2
+			p = (len(blur_kernel) - factor) + (kernel_size - 1)
+			pad0 = (p + 1) // 2
+			pad1 = p // 2
+
+			layers.append(Blur(blur_kernel, pad=(pad0, pad1)))
+
+			stride = 2
+			self.padding = 0
+		
+		elif upsample:
+			layers.append(Upsample(blur_kernel))
+			
+			stride = 1
+			self.padding = kernel_size // 2	
+		else:
+			stride = 1
+			self.padding = kernel_size // 2
+
+		layers.append(EqualConv2d(in_channel, out_channel, kernel_size, padding=self.padding, stride=stride,
+								  bias=bias and not activate))
+
+		if activate:
+			if bias:
+				layers.append(FusedLeakyReLU(out_channel))
+			else:
+				layers.append(ScaledLeakyReLU(0.2))
+
+		super().__init__(*layers)
+
+
+class ResBlock(nn.Module):
+	def __init__(self, in_channel, out_channel, blur_kernel=[1, 3, 3, 1]):
+		super().__init__()
+
+		self.conv1 = ConvLayer(in_channel, in_channel, 3)
+		self.conv2 = ConvLayer(in_channel, out_channel, 3)
+		self.skip = nn.Identity()
+
+	def forward(self, input):
+		out = self.conv1(input)
+		out = self.conv2(out)
+
+		skip = self.skip(input)
+		out = (out + skip) / math.sqrt(2)
+
+		return out
+
+
+class ResDownBlock(nn.Module):
+	def __init__(self, in_channel, out_channel, blur_kernel=[1, 3, 3, 1]):
+		super().__init__()
+
+		self.conv1 = ConvLayer(in_channel, in_channel, 3)
+		self.conv2 = ConvLayer(in_channel, out_channel, 3, downsample=True)
+
+		self.skip = ConvLayer(in_channel, out_channel, 1, downsample=True, activate=False, bias=False)	
+
+	def forward(self, input):
+		out = self.conv1(input)
+		out = self.conv2(out)
+
+		skip = self.skip(input)
+		out = (out + skip) / math.sqrt(2)
+
+		return out
+
+
+class ResUpBlock(nn.Module):
+	def __init__(self, in_channel, out_channel, blur_kernel=[1, 3, 3, 1]):
+		super().__init__()
+
+		self.conv1 = ConvLayer(in_channel, out_channel, 3, upsample=True)
+		self.conv2 = ConvLayer(out_channel, out_channel, 3, upsample=False)
+		
+		if in_channel != out_channel:
+			self.skip = ConvLayer(in_channel, out_channel, 1, upsample=True, activate=False, bias=False)
+		else:
+			self.skip = torch.nn.Identity()	
+
+	def forward(self, x):
+		out = self.conv1(x)
+		out = self.conv2(out)
+
+		skip = self.skip(x)
+		out = (out + skip) / math.sqrt(2)
+
+		return out
+
+
+class Upsample(nn.Module):
+	def __init__(self, kernel, factor=2):
+		super().__init__()
+
+		self.factor = factor
+		kernel = make_kernel(kernel) * (factor ** 2)
+		self.register_buffer('kernel', kernel)
+
+		p = kernel.shape[0] - factor
+
+		pad0 = (p + 1) // 2 + factor - 1
+		pad1 = p // 2
+
+		self.pad = (pad0, pad1)
+
+	def forward(self, input):
+		return upfirdn2d(input, self.kernel, up=self.factor, down=1, pad=self.pad)
+
+
+class Downsample(nn.Module):
+	def __init__(self, kernel, factor=2):
+		super().__init__()
+
+		self.factor = factor
+		kernel = make_kernel(kernel)
+		self.register_buffer('kernel', kernel)
+
+		p = kernel.shape[0] - factor
+
+		pad0 = (p + 1) // 2
+		pad1 = p // 2
+
+		self.pad = (pad0, pad1)
+
+	def forward(self, input):
+		return upfirdn2d(input, self.kernel, up=1, down=self.factor, pad=self.pad)
+
+
+class ModulatedConv2d(nn.Module):
+	def __init__(self, in_channel, out_channel, kernel_size, style_dim, demodulate=True, upsample=False,
+				 downsample=False, blur_kernel=[1, 3, 3, 1], ):
+		super().__init__()
+
+		self.eps = 1e-8
+		self.kernel_size = kernel_size
+		self.in_channel = in_channel
+		self.out_channel = out_channel
+		self.upsample = upsample
+		self.downsample = downsample
+
+		if upsample:
+			factor = 2
+			p = (len(blur_kernel) - factor) - (kernel_size - 1)
+			pad0 = (p + 1) // 2 + factor - 1
+			pad1 = p // 2 + 1
+
+			self.blur = Blur(blur_kernel, pad=(pad0, pad1), upsample_factor=factor)
+
+		if downsample:
+			factor = 2
+			p = (len(blur_kernel) - factor) + (kernel_size - 1)
+			pad0 = (p + 1) // 2
+			pad1 = p // 2
+
+			self.blur = Blur(blur_kernel, pad=(pad0, pad1))
+
+		fan_in = in_channel * kernel_size ** 2
+		self.scale = 1 / math.sqrt(fan_in)
+		self.padding = kernel_size // 2
+
+		self.weight = nn.Parameter(torch.randn(1, out_channel, in_channel, kernel_size, kernel_size))
+
+		self.modulation = EqualLinear(style_dim, in_channel, bias_init=1)
+		self.demodulate = demodulate
+	
+	def __repr__(self):
+		return (
+			f'{self.__class__.__name__}({self.in_channel}, {self.out_channel}, {self.kernel_size}, '
+			f'upsample={self.upsample}, downsample={self.downsample})'
+		)
+
+	def forward(self, input, style):
+		batch, in_channel, height, width = input.shape
+
+		style = self.modulation(style).view(batch, 1, in_channel, 1, 1)
+		weight = self.scale * self.weight * style
+
+		if self.demodulate:
+			demod = torch.rsqrt(weight.pow(2).sum([2, 3, 4]) + 1e-8)
+			weight = weight * demod.view(batch, self.out_channel, 1, 1, 1)
+
+		weight = weight.view(batch * self.out_channel, in_channel, self.kernel_size, self.kernel_size)
+
+		if self.upsample:
+			input = input.view(1, batch * in_channel, height, width)
+			weight = weight.view(batch, self.out_channel, in_channel, self.kernel_size, self.kernel_size)
+			weight = weight.transpose(1, 2).reshape(batch * in_channel, self.out_channel, self.kernel_size,
+													self.kernel_size)
+			out = F.conv_transpose2d(input, weight, padding=0, stride=2, groups=batch)
+			_, _, height, width = out.shape
+			out = out.view(batch, self.out_channel, height, width)
+			out = self.blur(out)
+		elif self.downsample:
+			input = self.blur(input)
+			_, _, height, width = input.shape
+			input = input.view(1, batch * in_channel, height, width)
+			out = F.conv2d(input, weight, padding=0, stride=2, groups=batch)
+			_, _, height, width = out.shape
+			out = out.view(batch, self.out_channel, height, width)
+		else:
+			input = input.view(1, batch * in_channel, height, width)
+			out = F.conv2d(input, weight, padding=self.padding, groups=batch)
+			_, _, height, width = out.shape
+			out = out.view(batch, self.out_channel, height, width)
+
+		return out
+
+
+class ConstantInput(nn.Module):
+	def __init__(self, channel, size=4):
+		super().__init__()
+
+		self.input = nn.Parameter(torch.randn(1, channel, size, size))
+
+	def forward(self, input):
+		batch = input.shape[0]
+		out = self.input.repeat(batch, 1, 1, 1)
+
+		return out
+
+class StyledConv(nn.Module):
+	def __init__(self, in_channel, out_channel, kernel_size, style_dim, upsample=False, demodulate=True):
+		super().__init__()
+
+		self.conv = ModulatedConv2d(
+			in_channel,
+			out_channel,
+			kernel_size,
+			style_dim,
+			upsample=upsample,
+			blur_kernel=[1,3,3,1],
+			demodulate=demodulate,
+		)
+
+		self.activate = FusedLeakyReLU(out_channel)
+
+	def forward(self, input, style):
+		out = self.conv(input, style)
+		out = self.activate(out)
+
+		return out
+
+class ToRGB(nn.Module):
+	def __init__(self, in_channel, upsample=True, blur_kernel=[1, 3, 3, 1]):
+		super().__init__()
+		
+		self.upsample = upsample
+
+		if upsample:
+			self.up = Upsample(blur_kernel)
+
+		self.conv = ConvLayer(in_channel, 3, 1)
+		self.bias = nn.Parameter(torch.zeros(1, 3, 1, 1))
+
+	def forward(self, input, skip=None):
+		out = self.conv(input)
+		out = out + self.bias
+
+		if skip is not None:
+			skip = self.up(skip)
+			out = out + skip
+
+		return out
+
+class ToFlow(nn.Module):
+	def __init__(self, in_channel, style_dim, upsample=True, blur_kernel=[1, 3, 3, 1]):
+		super().__init__()
+		
+		self.upsample = upsample
+		if upsample:
+			self.up = Upsample(blur_kernel)
+
+		self.conv = ModulatedConv2d(in_channel, 3, 1, style_dim, demodulate=False)
+		self.bias = nn.Parameter(torch.zeros(1, 3, 1, 1))
+
+	def forward(self, h, style, feat, skip=None):
+
+		out = self.conv(h, style)
+		out = out + self.bias
+
+		if skip is not None:
+			if self.upsample:
+				skip = self.up(skip)
+			out = out + skip
+		
+		xs = torch.linspace(-1, 1, out.size(2)).to(h.device)
+		xs = torch.meshgrid(xs, xs, indexing='xy')
+		xs = torch.stack(xs, 2)
+		xs = xs.unsqueeze(0).repeat(out.size(0), 1, 1, 1)
+
+		sampler = torch.tanh(out[:, 0:2, :, :])
+		mask = torch.sigmoid(out[:, 2:3, :, :])
+		flow = sampler.permute(0, 2, 3, 1) + xs	
+		
+		feat_warp = F.grid_sample(feat, flow, align_corners=True) * mask
+		h = feat_warp + (1 - mask) * h
+
+		#return h, out
+		return feat_warp, h, out
+
+
+class Direction(nn.Module):
+	def __init__(self, style_dim, motion_dim):
+		super(Direction, self).__init__()
+
+		self.weight = nn.Parameter(torch.randn(style_dim, motion_dim))
+
+	def forward(self, input):
+		# input: (bs*t) x 512
+
+		weight = self.weight + 1e-8
+		Q, R = torch.linalg.qr(weight)	# get eignvector, orthogonal [n1, n2, n3, n4]
+
+		input_diag = torch.diag_embed(input)  # alpha, diagonal matrix
+		out = torch.matmul(input_diag, Q.T)
+		out = torch.sum(out, dim=1)
+
+		return out
+
+
+
+
+			
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+

utils/data_processing.py

@@ -0,0 +1,141 @@
+import os
+import torch
+import torchvision
+from PIL import Image
+import numpy as np
+import imageio
+from einops import rearrange, repeat
+
+
+def load_image(img, size):
+    # img = Image.open(filename).convert('RGB')
+    if not isinstance(img, np.ndarray):
+        img = Image.open(img).convert('RGB')
+        img = img.resize((size, size))
+        img = np.asarray(img)
+    img = np.transpose(img, (2, 0, 1))  # 3 x 256 x 256
+
+    return img / 255.0
+
+
+def img_preprocessing(img_path, size):
+    img = load_image(img_path, size)  # [0, 1]
+    img = torch.from_numpy(img).unsqueeze(0).float()  # [0, 1]
+    imgs_norm = (img - 0.5) * 2.0  # [-1, 1]
+
+    return imgs_norm
+
+
+def resize(img, size):
+    transform = torchvision.transforms.Compose([
+        torchvision.transforms.Resize(size, antialias=True),
+        torchvision.transforms.CenterCrop(size)
+    ])
+
+    return transform(img)
+
+
+def vid_preprocessing(vid_path, size):
+    vid_dict = torchvision.io.read_video(vid_path, pts_unit='sec')
+    vid = vid_dict[0].permute(0, 3, 1, 2).unsqueeze(0)  # btchw
+    fps = vid_dict[2]['video_fps']
+    vid_norm = (vid / 255.0 - 0.5) * 2.0  # [-1, 1]
+
+    vid_norm = torch.cat([
+        resize(vid_norm[:, i, :, :, :], size).unsqueeze(1) for i in range(vid.size(1))
+    ], dim=1)
+
+    return vid_norm, fps
+
+
+def img_denorm(img):
+    img = img.clamp(-1, 1).cpu()
+    img = (img - img.min()) / (img.max() - img.min())
+
+    return img
+
+
+def vid_denorm(vid):
+    vid = vid.clamp(-1, 1).cpu()
+    vid = (vid - vid.min()) / (vid.max() - vid.min())
+
+    return vid
+
+
+def save_img_edit(save_dir, img, img_e):
+    # img: BCHW
+    # img_e: BCHW
+
+    output_img_path = os.path.join(save_dir, "img_edit.png")
+    output_img_all_path = os.path.join(save_dir, "img_all.png")
+
+    img = rearrange(img, 'b c h w -> b h w c')
+    img_e = rearrange(img_e, 'b c h w -> b h w c')
+    img_all = torch.cat([img, img_e], dim=2)
+
+    img_e_np = (img_denorm(img_e[0]).numpy() * 255).astype('uint8')
+    img_all_np = (img_denorm(img_all[0]).numpy() * 255).astype('uint8')
+
+    imageio.imwrite(output_img_path, img_e_np, quality=8)
+    imageio.imwrite(output_img_all_path, img_all_np, quality=8)
+
+    return
+
+
+def save_vid_edit(save_dir, vid_d, vid_a, fps):
+    # img_s: BCHW
+    # vid_d: BTCHW
+    # vid_a: BCTHW
+
+    output_vid_a_path = os.path.join(save_dir, "vid_animation.mp4")
+    output_vid_all_path = os.path.join(save_dir, "vid_all.mp4")
+
+    vid_d = rearrange(vid_d, 'b t c h w -> b t h w c')
+    vid_a = rearrange(vid_a, 'b c t h w -> b t h w c')
+    vid_all = torch.cat([vid_d, vid_a], dim=3)
+
+    vid_a_np = (vid_denorm(vid_a[0]).numpy() * 255).astype('uint8')
+    vid_all_np = (vid_denorm(vid_all[0]).numpy() * 255).astype('uint8')
+
+    imageio.mimwrite(output_vid_a_path, vid_a_np, fps=fps, codec='libx264', quality=8)
+    imageio.mimwrite(output_vid_all_path, vid_all_np, fps=fps, codec='libx264', quality=8)
+
+    return
+
+
+def save_animation(save_dir, img_s, vid_d, vid_a, fps):
+    # img_s: BCHW
+    # vid_d: BTCHW
+    # vid_a: BCTHW
+
+    output_vid_a_path = os.path.join(save_dir, "vid_animation.mp4")
+    output_img_e_path = os.path.join(save_dir, "img_edit.png")
+    output_vid_all_path = os.path.join(save_dir, "vid_all.mp4")
+
+    vid_d = rearrange(vid_d, 'b t c h w -> b t h w c')
+    vid_a = rearrange(vid_a, 'b c t h w -> b t h w c')
+    img_s = repeat(rearrange(img_s, 'b c h w -> b h w c'), 'b h w c -> b t h w c', t=vid_d.size(1))
+    vid_all = torch.cat([img_s, vid_d, vid_a], dim=3)
+
+    vid_a_np = (vid_denorm(vid_a[0]).numpy() * 255).astype('uint8')
+    img_e_np = vid_a_np[0]
+    vid_all_np = (vid_denorm(vid_all[0]).numpy() * 255).astype('uint8')
+
+    imageio.mimwrite(output_vid_a_path, vid_a_np, fps=fps, codec='libx264', quality=8)
+    imageio.mimwrite(output_vid_all_path, vid_all_np, fps=fps, codec='libx264', quality=8)
+    imageio.imwrite(output_img_e_path, img_e_np, quality=8)
+
+    return
+
+
+def save_linear_manipulation(save_dir, vid, fps):
+    # vid: BCTHW
+
+    output_vid_path = os.path.join(save_dir, "vid_interpolation.mp4")
+
+    vid = rearrange(vid, 'b c t h w -> b t h w c')
+    vid_np = (vid_denorm(vid[0]).numpy() * 255).astype('uint8')
+
+    imageio.mimwrite(output_vid_path, vid_np, fps=fps, codec='libx264', quality=8)
+
+    return