Spaces:

garibida
/

ReNoise-Inversion

Runtime error

App Files Files Community

ReNoise-Inversion / gradio_app.py

garibida

Add space

657e152 2 months ago

raw history blame contribute delete

No virus

16.1 kB

	from __future__ import annotations

	import gradio as gr
	import spaces
	from PIL import Image
	import torch

	from src.eunms import Model_Type, Scheduler_Type, Gradient_Averaging_Type, Epsilon_Update_Type
	from src.enums_utils import model_type_to_size, get_pipes
	from src.config import RunConfig
	from main import run as run_model


	DESCRIPTION = '''# ReNoise: Real Image Inversion Through Iterative Noising
	This is a demo for our ''ReNoise: Real Image Inversion Through Iterative Noising'' [paper](https://garibida.github.io/ReNoise-Inversion/). Code is available [here](https://github.com/garibida/ReNoise-Inversion)
	Our ReNoise inversion technique can be applied to various diffusion models, including recent few-step ones such as SDXL-Turbo.
	This demo preform real image editing using our ReNoise inversion. The input image is resize to size of 512x512, the optimal size of SDXL Turbo.
	'''

	device = 'cuda' if torch.cuda.is_available() else 'cpu'
	model_type = Model_Type.SDXL_Turbo
	scheduler_type = Scheduler_Type.EULER
	image_size = model_type_to_size(Model_Type.SDXL_Turbo)
	pipe_inversion, pipe_inference = get_pipes(model_type, scheduler_type, device=device)

	cache_size = 25
	prev_configs = [None for i in range(cache_size)]
	prev_inv_latents = [None for i in range(cache_size)]
	prev_images = [None for i in range(cache_size)]
	prev_noises = [None for i in range(cache_size)]

	@spaces.GPU
	def main_pipeline(
	input_image: str,
	src_prompt: str,
	tgt_prompt: str,
	edit_cfg: float,
	number_of_renoising_iterations: int,
	inersion_strength: float,
	avg_gradients: bool,
	first_step_range_start: int,
	first_step_range_end: int,
	rest_step_range_start: int,
	rest_step_range_end: int,
	lambda_ac: float,
	lambda_kl: float,
	noise_correction: bool):

	global prev_configs, prev_inv_latents, prev_images, prev_noises

	update_epsilon_type = Epsilon_Update_Type.OPTIMIZE if noise_correction else Epsilon_Update_Type.NONE
	avg_gradients_type = Gradient_Averaging_Type.ON_END if avg_gradients else Gradient_Averaging_Type.NONE

	first_step_range = (first_step_range_start, first_step_range_end)
	rest_step_range = (rest_step_range_start, rest_step_range_end)

	config = RunConfig(model_type = model_type,
	num_inference_steps = 4,
	num_inversion_steps = 4,
	guidance_scale = 0.0,
	max_num_aprox_steps_first_step = first_step_range_end+1,
	num_aprox_steps = number_of_renoising_iterations,
	inversion_max_step = inersion_strength,
	gradient_averaging_type = avg_gradients_type,
	gradient_averaging_first_step_range = first_step_range,
	gradient_averaging_step_range = rest_step_range,
	scheduler_type = scheduler_type,
	num_reg_steps = 4,
	num_ac_rolls = 5,
	lambda_ac = lambda_ac,
	lambda_kl = lambda_kl,
	update_epsilon_type = update_epsilon_type,
	do_reconstruction = True)
	config.prompt = src_prompt

	inv_latent = None
	noise_list = None
	for i in range(cache_size):
	if prev_configs[i] is not None and prev_configs[i] == config and prev_images[i] == input_image:
	print(f"Using cache for config #{i}")
	inv_latent = prev_inv_latents[i]
	noise_list = prev_noises[i]
	prev_configs.pop(i)
	prev_inv_latents.pop(i)
	prev_images.pop(i)
	prev_noises.pop(i)
	break

	original_image = Image.open(input_image).convert("RGB").resize(image_size)

	res_image, inv_latent, noise, all_latents = run_model(original_image,
	config,
	latents=inv_latent,
	pipe_inversion=pipe_inversion,
	pipe_inference=pipe_inference,
	edit_prompt=tgt_prompt,
	noise=noise_list,
	edit_cfg=edit_cfg)

	prev_configs.append(config)
	prev_inv_latents.append(inv_latent)
	prev_images.append(input_image)
	prev_noises.append(noise)

	if len(prev_configs) > cache_size:
	print("Popping cache")
	prev_configs.pop(0)
	prev_inv_latents.pop(0)
	prev_images.pop(0)
	prev_noises.pop(0)

	return res_image


	with gr.Blocks(css='style.css') as demo:
	gr.Markdown(DESCRIPTION)

	gr.HTML(
	'''<a href="https://huggingface.co/spaces/garibida/ReNoise-Inversion?duplicate=true">
	<img src="https://bit.ly/3gLdBN6" alt="Duplicate Space"></a>Duplicate the Space to run privately without waiting in queue''')

	with gr.Row():
	with gr.Column():
	input_image = gr.Image(
	label="Input image",
	type="filepath",
	height=image_size[0],
	width=image_size[1]
	)
	src_prompt = gr.Text(
	label='Source Prompt',
	max_lines=1,
	placeholder='A kitten is sitting in a basket on a branch',
	)
	tgt_prompt = gr.Text(
	label='Target Prompt',
	max_lines=1,
	placeholder='A plush toy kitten is sitting in a basket on a branch',
	)
	with gr.Accordion("Advanced Options", open=False):
	edit_cfg = gr.Slider(
	label='Denoise Classifier-Free Guidence Scale',
	minimum=1.0,
	maximum=3.5,
	value=1.0,
	step=0.1
	)
	number_of_renoising_iterations = gr.Slider(
	label='Number of ReNoise Iterations',
	minimum=0,
	maximum=20,
	value=9,
	step=1
	)
	inersion_strength = gr.Slider(
	label='Inversion Strength',
	info="Indicates how much to invert the reference image. The number of denoising steps depends on the amount of noise initially added. When strength is 1, the image will be inverted to complete noise and the denoising process will run for the full number of steps (4). When strength is 0.5, the image will be inverted to half noise and the denoising process will run for 2 steps.",
	minimum=0.0,
	maximum=1.0,
	value=1.0,
	step=0.25
	)
	avg_gradients = gr.Checkbox(
	label="Preform Estimation Averaging",
	info="IMPROVES RECONSTRUCTION. Averagin the estination over multiple ReNoise iterations can improve the quality of the reconstruction. The Next 4 sliders control the range of steps to average over. The first two sliders control the range of steps to average over for the first inversion step (t < 250). The last two sliders control the range of steps to average over for the rest of the inversion step (t > 250).",
	value=True
	)
	first_step_range_start = gr.Slider(
	label='First Estimation in Average (t < 250)',
	minimum=0,
	maximum=21,
	value=0,
	step=1
	)
	first_step_range_end = gr.Slider(
	label='Last Estimation in Average (t < 250)',
	minimum=0,
	maximum=21,
	value=5,
	step=1
	)
	rest_step_range_start = gr.Slider(
	label='First Estimation in Average (t > 250)',
	minimum=0,
	maximum=21,
	value=8,
	step=1
	)
	rest_step_range_end = gr.Slider(
	label='Last Estimation in Average (t > 250)',
	minimum=0,
	maximum=21,
	value=10,
	step=1
	)
	num_reg_steps = 4
	num_ac_rolls = 5
	lambda_ac = gr.Slider(
	label='Labmda AC',
	info="IMPROVES EDITABILITY. The weight of the pair loss in the noise prediction regulariztion. This loss encourages the inversion to predict more editable noise. A higher value allows more significant changes to the image (higher editability), but may result in less faithful reconstructions.",
	minimum=0.0,
	maximum=50.0,
	value=20.0,
	step=1.0
	)
	lambda_kl = gr.Slider(
	label='Labmda Patch KL',
	info="IMPROVES EDITABILITY. This weight controls the strength of the patch-level KL divergence term in the noise prediction regularization. While it encourages editable noise like the Labmda AC, it often has a less detrimental effect on reconstruction fidelity.",
	minimum=0.0,
	maximum=0.4,
	value=0.065,
	step=0.005
	)
	noise_correction = gr.Checkbox(
	label="Preform Noise Correction",
	info="IMPROVES RECONSTRUCTION. Performs noise correction to improve the reconstruction of the image.",
	value=True
	)

	run_button = gr.Button('Edit')
	with gr.Column():
	# result = gr.Gallery(label='Result')
	result = gr.Image(
	label="Result",
	type="pil",
	height=image_size[0],
	width=image_size[1]
	)

	examples = [
	[
	"example_images/kitten.jpg", #input_image
	"A kitten is sitting in a basket on a branch", #src_prompt
	"a lego kitten is sitting in a basket on a branch", #tgt_prompt
	1.0, #edit_cfg
	9, #number_of_renoising_iterations
	1.0, #inersion_strength
	True, #avg_gradients
	0, #first_step_range_start
	5, #first_step_range_end
	8, #rest_step_range_start
	10, #rest_step_range_end
	20.0, #lambda_ac
	0.055, #lambda_kl
	False #noise_correction
	],
	[
	"example_images/kitten.jpg", #input_image
	"A kitten is sitting in a basket on a branch", #src_prompt
	"a brokkoli is sitting in a basket on a branch", #tgt_prompt
	1.0, #edit_cfg
	9, #number_of_renoising_iterations
	1.0, #inersion_strength
	True, #avg_gradients
	0, #first_step_range_start
	5, #first_step_range_end
	8, #rest_step_range_start
	10, #rest_step_range_end
	20.0, #lambda_ac
	0.055, #lambda_kl
	False #noise_correction
	],
	[
	"example_images/kitten.jpg", #input_image
	"A kitten is sitting in a basket on a branch", #src_prompt
	"a dog is sitting in a basket on a branch", #tgt_prompt
	1.0, #edit_cfg
	9, #number_of_renoising_iterations
	1.0, #inersion_strength
	True, #avg_gradients
	0, #first_step_range_start
	5, #first_step_range_end
	8, #rest_step_range_start
	10, #rest_step_range_end
	20.0, #lambda_ac
	0.055, #lambda_kl
	False #noise_correction
	],
	[
	"example_images/monkey.jpeg", #input_image
	"a monkey sitting on a tree branch in the forest", #src_prompt
	"a beaver sitting on a tree branch in the forest", #tgt_prompt
	1.0, #edit_cfg
	9, #number_of_renoising_iterations
	1.0, #inersion_strength
	True, #avg_gradients
	0, #first_step_range_start
	5, #first_step_range_end
	8, #rest_step_range_start
	10, #rest_step_range_end
	20.0, #lambda_ac
	0.055, #lambda_kl
	True #noise_correction
	],
	[
	"example_images/monkey.jpeg", #input_image
	"a monkey sitting on a tree branch in the forest", #src_prompt
	"a raccoon sitting on a tree branch in the forest", #tgt_prompt
	1.0, #edit_cfg
	9, #number_of_renoising_iterations
	1.0, #inersion_strength
	True, #avg_gradients
	0, #first_step_range_start
	5, #first_step_range_end
	8, #rest_step_range_start
	10, #rest_step_range_end
	20.0, #lambda_ac
	0.055, #lambda_kl
	True #noise_correction
	],
	[
	"example_images/lion.jpeg", #input_image
	"a lion is sitting in the grass at sunset", #src_prompt
	"a tiger is sitting in the grass at sunset", #tgt_prompt
	1.0, #edit_cfg
	9, #number_of_renoising_iterations
	1.0, #inersion_strength
	True, #avg_gradients
	0, #first_step_range_start
	5, #first_step_range_end
	8, #rest_step_range_start
	10, #rest_step_range_end
	20.0, #lambda_ac
	0.055, #lambda_kl
	True #noise_correction
	]
	]

	gr.Examples(examples=examples,
	inputs=[
	input_image,
	src_prompt,
	tgt_prompt,
	edit_cfg,
	number_of_renoising_iterations,
	inersion_strength,
	avg_gradients,
	first_step_range_start,
	first_step_range_end,
	rest_step_range_start,
	rest_step_range_end,
	lambda_ac,
	lambda_kl,
	noise_correction
	],
	outputs=[
	result
	],
	fn=main_pipeline,
	cache_examples=True)


	inputs = [
	input_image,
	src_prompt,
	tgt_prompt,
	edit_cfg,
	number_of_renoising_iterations,
	inersion_strength,
	avg_gradients,
	first_step_range_start,
	first_step_range_end,
	rest_step_range_start,
	rest_step_range_end,
	lambda_ac,
	lambda_kl,
	noise_correction
	]
	outputs = [
	result
	]
	run_button.click(fn=main_pipeline, inputs=inputs, outputs=outputs)

	demo.queue(max_size=50).launch(share=False)