DragGan

Running

DragGan / visualizer_drag_gradio.py

download from hub

58845cd 12 months ago

32.4 kB

	import os
	import os.path as osp
	from argparse import ArgumentParser
	from functools import partial
	from huggingface_hub import snapshot_download
	from pathlib import Path

	import gradio as gr
	import numpy as np
	import torch
	from PIL import Image

	import dnnlib
	from gradio_utils import (ImageMask, draw_mask_on_image, draw_points_on_image,
	get_latest_points_pair, get_valid_mask,
	on_change_single_global_state)
	from viz.renderer import Renderer, add_watermark_np


	# download models from hub
	model_dir = Path('./checkpoints')
	snapshot_download('radames/DragGan', repo_type='model', local_dir=model_dir)

	parser = ArgumentParser()
	parser.add_argument('--share', action='store_true')
	parser.add_argument('--cache-dir', type=str, default='./checkpoints')
	args = parser.parse_args()

	cache_dir = args.cache_dir

	device = 'cuda'


	def reverse_point_pairs(points):
	new_points = []
	for p in points:
	new_points.append([p[1], p[0]])
	return new_points


	def clear_state(global_state, target=None):
	"""Clear target history state from global_state
	If target is not defined, points and mask will be both removed.
	1. set global_state['points'] as empty dict
	2. set global_state['mask'] as full-one mask.
	"""
	if target is None:
	target = ['point', 'mask']
	if not isinstance(target, list):
	target = [target]
	if 'point' in target:
	global_state['points'] = dict()
	print('Clear Points State!')
	if 'mask' in target:
	image_raw = global_state["images"]["image_raw"]
	global_state['mask'] = np.ones((image_raw.size[1], image_raw.size[0]),
	dtype=np.uint8)
	print('Clear mask State!')

	return global_state


	def init_images(global_state):
	"""This function is called only ones with Gradio App is started.
	0. pre-process global_state, unpack value from global_state of need
	1. Re-init renderer
	2. run `renderer._render_drag_impl` with `is_drag=False` to generate
	new image
	3. Assign images to global state and re-generate mask
	"""

	if isinstance(global_state, gr.State):
	state = global_state.value
	else:
	state = global_state

	state['renderer'].init_network(
	state['generator_params'], # res
	valid_checkpoints_dict[state['pretrained_weight']], # pkl
	state['params']['seed'], # w0_seed,
	None, # w_load
	state['params']['latent_space'] == 'w+', # w_plus
	'const',
	state['params']['trunc_psi'], # trunc_psi,
	state['params']['trunc_cutoff'], # trunc_cutoff,
	None, # input_transform
	state['params']['lr'] # lr,
	)

	state['renderer']._render_drag_impl(state['generator_params'],
	is_drag=False,
	to_pil=True)

	init_image = state['generator_params'].image
	state['images']['image_orig'] = init_image
	state['images']['image_raw'] = init_image
	state['images']['image_show'] = Image.fromarray(
	add_watermark_np(np.array(init_image)))
	state['mask'] = np.ones((init_image.size[1], init_image.size[0]),
	dtype=np.uint8)
	return global_state


	def update_image_draw(image, points, mask, show_mask, global_state=None):

	image_draw = draw_points_on_image(image, points)
	if show_mask and mask is not None and not (mask == 0).all() and not (
	mask == 1).all():
	image_draw = draw_mask_on_image(image_draw, mask)

	image_draw = Image.fromarray(add_watermark_np(np.array(image_draw)))
	if global_state is not None:
	global_state['images']['image_show'] = image_draw
	return image_draw


	def preprocess_mask_info(global_state, image):
	"""Function to handle mask information.
	1. last_mask is None: Do not need to change mask, return mask
	2. last_mask is not None:
	2.1 global_state is remove_mask:
	2.2 global_state is add_mask:
	"""
	if isinstance(image, dict):
	last_mask = get_valid_mask(image['mask'])
	else:
	last_mask = None
	mask = global_state['mask']

	# mask in global state is a placeholder with all 1.
	if (mask == 1).all():
	mask = last_mask

	# last_mask = global_state['last_mask']
	editing_mode = global_state['editing_state']

	if last_mask is None:
	return global_state

	if editing_mode == 'remove_mask':
	updated_mask = np.clip(mask - last_mask, 0, 1)
	print(f'Last editing_state is {editing_mode}, do remove.')
	elif editing_mode == 'add_mask':
	updated_mask = np.clip(mask + last_mask, 0, 1)
	print(f'Last editing_state is {editing_mode}, do add.')
	else:
	updated_mask = mask
	print(f'Last editing_state is {editing_mode}, '
	'do nothing to mask.')

	global_state['mask'] = updated_mask
	# global_state['last_mask'] = None # clear buffer
	return global_state


	valid_checkpoints_dict = {
	f.split('/')[-1].split('.')[0]: osp.join(cache_dir, f)
	for f in os.listdir(cache_dir)
	if (f.endswith('pkl') and osp.exists(osp.join(cache_dir, f)))
	}
	print(f'File under cache_dir ({cache_dir}):')
	print(os.listdir(cache_dir))
	print('Valid checkpoint file:')
	print(valid_checkpoints_dict)

	init_pkl = 'stylegan_human_v2_512'

	with gr.Blocks() as app:

	# renderer = Renderer()
	global_state = gr.State({
	"images": {
	# image_orig: the original image, change with seed/model is changed
	# image_raw: image with mask and points, change durning optimization
	# image_show: image showed on screen
	},
	"temporal_params": {
	# stop
	},
	'mask':
	None, # mask for visualization, 1 for editing and 0 for unchange
	'last_mask': None, # last edited mask
	'show_mask': True, # add button
	"generator_params": dnnlib.EasyDict(),
	"params": {
	"seed": 0,
	"motion_lambda": 20,
	"r1_in_pixels": 3,
	"r2_in_pixels": 12,
	"magnitude_direction_in_pixels": 1.0,
	"latent_space": "w+",
	"trunc_psi": 0.7,
	"trunc_cutoff": None,
	"lr": 0.001,
	},
	"device": device,
	"draw_interval": 1,
	"renderer": Renderer(disable_timing=True),
	"points": {},
	"curr_point": None,
	"curr_type_point": "start",
	'editing_state': 'add_points',
	'pretrained_weight': init_pkl
	})

	# init image
	global_state = init_images(global_state)

	with gr.Row():

	with gr.Row():

	# Left --> tools
	with gr.Column(scale=3):

	# Pickle
	with gr.Row():

	with gr.Column(scale=1, min_width=10):
	gr.Markdown(value='Pickle', show_label=False)

	with gr.Column(scale=4, min_width=10):
	form_pretrained_dropdown = gr.Dropdown(
	choices=list(valid_checkpoints_dict.keys()),
	label="Pretrained Model",
	value=init_pkl,
	)

	# Latent
	with gr.Row():
	with gr.Column(scale=1, min_width=10):
	gr.Markdown(value='Latent', show_label=False)

	with gr.Column(scale=4, min_width=10):
	form_seed_number = gr.Number(
	value=global_state.value['params']['seed'],
	interactive=True,
	label="Seed",
	)
	form_lr_number = gr.Number(
	value=global_state.value["params"]["lr"],
	interactive=True,
	label="Step Size")

	with gr.Row():
	with gr.Column(scale=2, min_width=10):
	form_reset_image = gr.Button("Reset Image")
	with gr.Column(scale=3, min_width=10):
	form_latent_space = gr.Radio(
	['w', 'w+'],
	value=global_state.value['params']
	['latent_space'],
	interactive=True,
	label='Latent space to optimize',
	show_label=False,
	)

	# Drag
	with gr.Row():
	with gr.Column(scale=1, min_width=10):
	gr.Markdown(value='Drag', show_label=False)
	with gr.Column(scale=4, min_width=10):
	with gr.Row():
	with gr.Column(scale=1, min_width=10):
	enable_add_points = gr.Button('Add Points')
	with gr.Column(scale=1, min_width=10):
	undo_points = gr.Button('Reset Points')
	with gr.Row():
	with gr.Column(scale=1, min_width=10):
	form_start_btn = gr.Button("Start")
	with gr.Column(scale=1, min_width=10):
	form_stop_btn = gr.Button("Stop")

	form_steps_number = gr.Number(value=0,
	label="Steps",
	interactive=False)

	# Mask
	with gr.Row():
	with gr.Column(scale=1, min_width=10):
	gr.Markdown(value='Mask', show_label=False)
	with gr.Column(scale=4, min_width=10):
	enable_add_mask = gr.Button('Edit Flexible Area')
	with gr.Row():
	with gr.Column(scale=1, min_width=10):
	form_reset_mask_btn = gr.Button("Reset mask")
	with gr.Column(scale=1, min_width=10):
	show_mask = gr.Checkbox(
	label='Show Mask',
	value=global_state.value['show_mask'],
	show_label=False)

	with gr.Row():
	form_lambda_number = gr.Number(
	value=global_state.value["params"]
	["motion_lambda"],
	interactive=True,
	label="Lambda",
	)

	form_draw_interval_number = gr.Number(
	value=global_state.value["draw_interval"],
	label="Draw Interval (steps)",
	interactive=True,
	visible=False)

	# Right --> Image
	with gr.Column(scale=8):
	form_image = ImageMask(
	value=global_state.value['images']['image_show'],
	brush_radius=20).style(
	width=768,
	height=768) # NOTE: hard image size code here.
	gr.Markdown("""
	## Quick Start

	1. Select desired `Pretrained Model` and adjust `Seed` to generate an
	initial image.
	2. Click on image to add control points.
	3. Click `Start` and enjoy it!

	## Advance Usage

	1. Change `Step Size` to adjust learning rate in drag optimization.
	2. Select `w` or `w+` to change latent space to optimize:
	* Optimize on `w` space may cause greater influence to the image.
	* Optimize on `w+` space may work slower than `w`, but usually achieve
	better results.
	* Note that changing the latent space will reset the image, points and
	mask (this has the same effect as `Reset Image` button).
	3. Click `Edit Flexible Area` to create a mask and constrain the
	unmasked region to remain unchanged.
	""")
	gr.HTML("""
	<style>
	.container {
	position: absolute;
	height: 50px;
	text-align: center;
	line-height: 50px;
	width: 100%;
	}
	</style>
	<div class="container">
	Gradio demo supported by
	<img src="https://avatars.githubusercontent.com/u/10245193?s=200&v=4" height="20" width="20" style="display:inline;">
	<a href="https://github.com/open-mmlab/mmagic">OpenMMLab MMagic</a>
	</div>
	""")

	# Network & latents tab listeners
	def on_change_pretrained_dropdown(pretrained_value, global_state):
	"""Function to handle model change.
	1. Set pretrained value to global_state
	2. Re-init images and clear all states
	"""

	global_state['pretrained_weight'] = pretrained_value
	init_images(global_state)
	clear_state(global_state)

	return global_state, global_state["images"]['image_show']

	form_pretrained_dropdown.change(
	on_change_pretrained_dropdown,
	inputs=[form_pretrained_dropdown, global_state],
	outputs=[global_state, form_image],
	)

	def on_click_reset_image(global_state):
	"""Reset image to the original one and clear all states
	1. Re-init images
	2. Clear all states
	"""

	init_images(global_state)
	clear_state(global_state)

	return global_state, global_state['images']['image_show']

	form_reset_image.click(
	on_click_reset_image,
	inputs=[global_state],
	outputs=[global_state, form_image],
	)

	# Update parameters
	def on_change_update_image_seed(seed, global_state):
	"""Function to handle generation seed change.
	1. Set seed to global_state
	2. Re-init images and clear all states
	"""

	global_state["params"]["seed"] = int(seed)
	init_images(global_state)
	clear_state(global_state)

	return global_state, global_state['images']['image_show']

	form_seed_number.change(
	on_change_update_image_seed,
	inputs=[form_seed_number, global_state],
	outputs=[global_state, form_image],
	)

	def on_click_latent_space(latent_space, global_state):
	"""Function to reset latent space to optimize.
	NOTE: this function we reset the image and all controls
	1. Set latent-space to global_state
	2. Re-init images and clear all state
	"""

	global_state['params']['latent_space'] = latent_space
	init_images(global_state)
	clear_state(global_state)

	return global_state, global_state['images']['image_show']

	form_latent_space.change(on_click_latent_space,
	inputs=[form_latent_space, global_state],
	outputs=[global_state, form_image])

	# ==== Params
	form_lambda_number.change(
	partial(on_change_single_global_state, ["params", "motion_lambda"]),
	inputs=[form_lambda_number, global_state],
	outputs=[global_state],
	)

	def on_change_lr(lr, global_state):
	if lr == 0:
	print('lr is 0, do nothing.')
	return global_state
	else:
	global_state["params"]["lr"] = lr
	renderer = global_state['renderer']
	renderer.update_lr(lr)
	print('New optimizer: ')
	print(renderer.w_optim)
	return global_state

	form_lr_number.change(
	on_change_lr,
	inputs=[form_lr_number, global_state],
	outputs=[global_state],
	)

	def on_click_start(global_state, image):
	p_in_pixels = []
	t_in_pixels = []
	valid_points = []

	# handle of start drag in mask editing mode
	global_state = preprocess_mask_info(global_state, image)

	# Prepare the points for the inference
	if len(global_state["points"]) == 0:
	# yield on_click_start_wo_points(global_state, image)
	image_raw = global_state['images']['image_raw']
	update_image_draw(
	image_raw,
	global_state['points'],
	global_state['mask'],
	global_state['show_mask'],
	global_state,
	)

	yield (
	global_state,
	0,
	global_state['images']['image_show'],
	# gr.File.update(visible=False),
	gr.Button.update(interactive=True),
	gr.Button.update(interactive=True),
	gr.Button.update(interactive=True),
	gr.Button.update(interactive=True),
	gr.Button.update(interactive=True),
	# latent space
	gr.Radio.update(interactive=True),
	gr.Button.update(interactive=True),
	# NOTE: disable stop button
	gr.Button.update(interactive=False),

	# update other comps
	gr.Dropdown.update(interactive=True),
	gr.Number.update(interactive=True),
	gr.Number.update(interactive=True),
	gr.Button.update(interactive=True),
	gr.Button.update(interactive=True),
	gr.Checkbox.update(interactive=True),
	# gr.Number.update(interactive=True),
	gr.Number.update(interactive=True),
	)
	else:

	# Transform the points into torch tensors
	for key_point, point in global_state["points"].items():
	try:
	p_start = point.get("start_temp", point["start"])
	p_end = point["target"]

	if p_start is None or p_end is None:
	continue

	except KeyError:
	continue

	p_in_pixels.append(p_start)
	t_in_pixels.append(p_end)
	valid_points.append(key_point)

	mask = torch.tensor(global_state['mask']).float()
	drag_mask = 1 - mask

	renderer: Renderer = global_state["renderer"]
	global_state['temporal_params']['stop'] = False
	global_state['editing_state'] = 'running'

	# reverse points order
	p_to_opt = reverse_point_pairs(p_in_pixels)
	t_to_opt = reverse_point_pairs(t_in_pixels)
	print('Running with:')
	print(f' Source: {p_in_pixels}')
	print(f' Target: {t_in_pixels}')
	step_idx = 0
	while True:
	if global_state["temporal_params"]["stop"]:
	break

	# do drage here!
	renderer._render_drag_impl(
	global_state['generator_params'],
	p_to_opt, # point
	t_to_opt, # target
	drag_mask, # mask,
	global_state['params']['motion_lambda'], # lambda_mask
	reg=0,
	feature_idx=5, # NOTE: do not support change for now
	r1=global_state['params']['r1_in_pixels'], # r1
	r2=global_state['params']['r2_in_pixels'], # r2
	# random_seed = 0,
	# noise_mode = 'const',
	trunc_psi=global_state['params']['trunc_psi'],
	# force_fp32 = False,
	# layer_name = None,
	# sel_channels = 3,
	# base_channel = 0,
	# img_scale_db = 0,
	# img_normalize = False,
	# untransform = False,
	is_drag=True,
	to_pil=True)

	if step_idx % global_state['draw_interval'] == 0:
	print('Current Source:')
	for key_point, p_i, t_i in zip(valid_points, p_to_opt,
	t_to_opt):
	global_state["points"][key_point]["start_temp"] = [
	p_i[1],
	p_i[0],
	]
	global_state["points"][key_point]["target"] = [
	t_i[1],
	t_i[0],
	]
	start_temp = global_state["points"][key_point][
	"start_temp"]
	print(f' {start_temp}')

	image_result = global_state['generator_params']['image']
	image_draw = update_image_draw(
	image_result,
	global_state['points'],
	global_state['mask'],
	global_state['show_mask'],
	global_state,
	)
	global_state['images']['image_raw'] = image_result

	yield (
	global_state,
	step_idx,
	global_state['images']['image_show'],
	# gr.File.update(visible=False),
	gr.Button.update(interactive=False),
	gr.Button.update(interactive=False),
	gr.Button.update(interactive=False),
	gr.Button.update(interactive=False),
	gr.Button.update(interactive=False),
	# latent space
	gr.Radio.update(interactive=False),
	gr.Button.update(interactive=False),
	# enable stop button in loop
	gr.Button.update(interactive=True),

	# update other comps
	gr.Dropdown.update(interactive=False),
	gr.Number.update(interactive=False),
	gr.Number.update(interactive=False),
	gr.Button.update(interactive=False),
	gr.Button.update(interactive=False),
	gr.Checkbox.update(interactive=False),
	# gr.Number.update(interactive=False),
	gr.Number.update(interactive=False),
	)

	# increate step
	step_idx += 1

	image_result = global_state['generator_params']['image']
	global_state['images']['image_raw'] = image_result
	image_draw = update_image_draw(image_result,
	global_state['points'],
	global_state['mask'],
	global_state['show_mask'],
	global_state)

	# fp = NamedTemporaryFile(suffix=".png", delete=False)
	# image_result.save(fp, "PNG")

	global_state['editing_state'] = 'add_points'

	yield (
	global_state,
	0, # reset step to 0 after stop.
	global_state['images']['image_show'],
	# gr.File.update(visible=True, value=fp.name),
	gr.Button.update(interactive=True),
	gr.Button.update(interactive=True),
	gr.Button.update(interactive=True),
	gr.Button.update(interactive=True),
	gr.Button.update(interactive=True),
	# latent space
	gr.Radio.update(interactive=True),
	gr.Button.update(interactive=True),
	# NOTE: disable stop button with loop finish
	gr.Button.update(interactive=False),

	# update other comps
	gr.Dropdown.update(interactive=True),
	gr.Number.update(interactive=True),
	gr.Number.update(interactive=True),
	gr.Checkbox.update(interactive=True),
	gr.Number.update(interactive=True),
	)

	form_start_btn.click(
	on_click_start,
	inputs=[global_state, form_image],
	outputs=[
	global_state,
	form_steps_number,
	form_image,
	# form_download_result_file,
	# >>> buttons
	form_reset_image,
	enable_add_points,
	enable_add_mask,
	undo_points,
	form_reset_mask_btn,
	form_latent_space,
	form_start_btn,
	form_stop_btn,
	# <<< buttonm
	# >>> inputs comps
	form_pretrained_dropdown,
	form_seed_number,
	form_lr_number,
	show_mask,
	form_lambda_number,
	],
	)

	def on_click_stop(global_state):
	"""Function to handle stop button is clicked.
	1. send a stop signal by set global_state["temporal_params"]["stop"] as True
	2. Disable Stop button
	"""
	global_state["temporal_params"]["stop"] = True

	return global_state, gr.Button.update(interactive=False)

	form_stop_btn.click(on_click_stop,
	inputs=[global_state],
	outputs=[global_state, form_stop_btn])

	form_draw_interval_number.change(
	partial(
	on_change_single_global_state,
	"draw_interval",
	map_transform=lambda x: int(x),
	),
	inputs=[form_draw_interval_number, global_state],
	outputs=[global_state],
	)

	def on_click_remove_point(global_state):
	choice = global_state["curr_point"]
	del global_state["points"][choice]

	choices = list(global_state["points"].keys())

	if len(choices) > 0:
	global_state["curr_point"] = choices[0]

	return (
	gr.Dropdown.update(choices=choices, value=choices[0]),
	global_state,
	)

	# Mask
	def on_click_reset_mask(global_state):
	global_state['mask'] = np.ones(
	(
	global_state["images"]["image_raw"].size[1],
	global_state["images"]["image_raw"].size[0],
	),
	dtype=np.uint8,
	)
	image_draw = update_image_draw(global_state['images']['image_raw'],
	global_state['points'],
	global_state['mask'],
	global_state['show_mask'], global_state)
	return global_state, image_draw

	form_reset_mask_btn.click(
	on_click_reset_mask,
	inputs=[global_state],
	outputs=[global_state, form_image],
	)

	# Image
	def on_click_enable_draw(global_state, image):
	"""Function to start add mask mode.
	1. Preprocess mask info from last state
	2. Change editing state to add_mask
	3. Set curr image with points and mask
	"""
	global_state = preprocess_mask_info(global_state, image)
	global_state['editing_state'] = 'add_mask'
	image_raw = global_state['images']['image_raw']
	image_draw = update_image_draw(image_raw, global_state['points'],
	global_state['mask'], True,
	global_state)
	return (global_state,
	gr.Image.update(value=image_draw, interactive=True))

	def on_click_remove_draw(global_state, image):
	"""Function to start remove mask mode.
	1. Preprocess mask info from last state
	2. Change editing state to remove_mask
	3. Set curr image with points and mask
	"""
	global_state = preprocess_mask_info(global_state, image)
	global_state['edinting_state'] = 'remove_mask'
	image_raw = global_state['images']['image_raw']
	image_draw = update_image_draw(image_raw, global_state['points'],
	global_state['mask'], True,
	global_state)
	return (global_state,
	gr.Image.update(value=image_draw, interactive=True))

	enable_add_mask.click(on_click_enable_draw,
	inputs=[global_state, form_image],
	outputs=[
	global_state,
	form_image,
	])

	def on_click_add_point(global_state, image: dict):
	"""Function switch from add mask mode to add points mode.
	1. Updaste mask buffer if need
	2. Change global_state['editing_state'] to 'add_points'
	3. Set current image with mask
	"""

	global_state = preprocess_mask_info(global_state, image)
	global_state['editing_state'] = 'add_points'
	mask = global_state['mask']
	image_raw = global_state['images']['image_raw']
	image_draw = update_image_draw(image_raw, global_state['points'], mask,
	global_state['show_mask'], global_state)

	return (global_state,
	gr.Image.update(value=image_draw, interactive=False))

	enable_add_points.click(on_click_add_point,
	inputs=[global_state, form_image],
	outputs=[global_state, form_image])

	def on_click_image(global_state, evt: gr.SelectData):
	"""This function only support click for point selection
	"""
	xy = evt.index
	if global_state['editing_state'] != 'add_points':
	print(f'In {global_state["editing_state"]} state. '
	'Do not add points.')

	return global_state, global_state['images']['image_show']

	points = global_state["points"]

	point_idx = get_latest_points_pair(points)
	if point_idx is None:
	points[0] = {'start': xy, 'target': None}
	print(f'Click Image - Start - {xy}')
	elif points[point_idx].get('target', None) is None:
	points[point_idx]['target'] = xy
	print(f'Click Image - Target - {xy}')
	else:
	points[point_idx + 1] = {'start': xy, 'target': None}
	print(f'Click Image - Start - {xy}')

	image_raw = global_state['images']['image_raw']
	image_draw = update_image_draw(
	image_raw,
	global_state['points'],
	global_state['mask'],
	global_state['show_mask'],
	global_state,
	)

	return global_state, image_draw

	form_image.select(
	on_click_image,
	inputs=[global_state],
	outputs=[global_state, form_image],
	)

	def on_click_clear_points(global_state):
	"""Function to handle clear all control points
	1. clear global_state['points'] (clear_state)
	2. re-init network
	2. re-draw image
	"""
	clear_state(global_state, target='point')

	renderer: Renderer = global_state["renderer"]
	renderer.feat_refs = None

	image_raw = global_state['images']['image_raw']
	image_draw = update_image_draw(image_raw, {}, global_state['mask'],
	global_state['show_mask'], global_state)
	return global_state, image_draw

	undo_points.click(on_click_clear_points,
	inputs=[global_state],
	outputs=[global_state, form_image])

	def on_click_show_mask(global_state, show_mask):
	"""Function to control whether show mask on image."""
	global_state['show_mask'] = show_mask

	image_raw = global_state['images']['image_raw']
	image_draw = update_image_draw(
	image_raw,
	global_state['points'],
	global_state['mask'],
	global_state['show_mask'],
	global_state,
	)
	return global_state, image_draw

	show_mask.change(
	on_click_show_mask,
	inputs=[global_state, show_mask],
	outputs=[global_state, form_image],
	)

	gr.close_all()
	app.queue(concurrency_count=5, max_size=20)
	app.launch(share=args.share)