leditsplusplus

Running on A10G

App Files Files Community

leditsplusplus / app.py

linoyts HF staff

Update app.py (#16)

535a876 verified 3 months ago

raw history blame contribute delete

No virus

40.5 kB

	import gradio as gr
	import torch
	import numpy as np
	import requests
	import random
	from io import BytesIO
	from utils import *
	from constants import *
	from pipeline_semantic_stable_diffusion_img2img_solver import SemanticStableDiffusionImg2ImgPipeline_DPMSolver
	from torch import autocast, inference_mode
	from diffusers import StableDiffusionPipeline, AutoencoderKL
	from diffusers.schedulers import DDIMScheduler
	from scheduling_dpmsolver_multistep_inject import DPMSolverMultistepSchedulerInject
	from transformers import AutoProcessor, BlipForConditionalGeneration
	from share_btn import community_icon_html, loading_icon_html, share_js

	# load pipelines
	# sd_model_id = "runwayml/stable-diffusion-v1-5"
	sd_model_id = "stabilityai/stable-diffusion-2-1-base"
	device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
	vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse", torch_dtype=torch.float16)
	pipe = SemanticStableDiffusionImg2ImgPipeline_DPMSolver.from_pretrained(sd_model_id,vae=vae,torch_dtype=torch.float16, safety_checker=None, requires_safety_checker=False).to(device)
	pipe.scheduler = DPMSolverMultistepSchedulerInject.from_pretrained(sd_model_id, subfolder="scheduler"
	, algorithm_type="sde-dpmsolver++", solver_order=2)

	blip_processor = AutoProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
	blip_model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base",torch_dtype=torch.float16).to(device)

	## IMAGE CPATIONING ##
	def caption_image(input_image):
	inputs = blip_processor(images=input_image, return_tensors="pt").to(device, torch.float16)
	pixel_values = inputs.pixel_values

	generated_ids = blip_model.generate(pixel_values=pixel_values, max_length=50)
	generated_caption = blip_processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
	return generated_caption, generated_caption

	def sample(zs, wts, attention_store, text_cross_attention_maps, prompt_tar="", cfg_scale_tar=15, skip=36, eta=1):
	latents = wts[-1].expand(1, -1, -1, -1)
	img, attention_store, text_cross_attention_maps = pipe(
	prompt=prompt_tar,
	init_latents=latents,
	guidance_scale=cfg_scale_tar,
	# num_images_per_prompt=1,
	# num_inference_steps=steps,
	# use_ddpm=True,
	# wts=wts.value,
	attention_store = attention_store, text_cross_attention_maps=text_cross_attention_maps,
	zs=zs,
	)
	return img.images[0], attention_store, text_cross_attention_maps


	def reconstruct(
	tar_prompt,
	image_caption,
	tar_cfg_scale,
	skip,
	wts,
	zs,
	attention_store,
	text_cross_attention_maps,
	do_reconstruction,
	reconstruction,
	reconstruct_button,
	):
	if reconstruct_button == "Hide Reconstruction":
	return (
	reconstruction,
	reconstruction,
	gr.update(visible=False),
	do_reconstruction,
	"Show Reconstruction",
	)

	else:
	if do_reconstruction:
	if (
	image_caption.lower() == tar_prompt.lower()
	): # if image caption was not changed, run actual reconstruction
	tar_prompt = ""
	latents = wts[-1].expand(1, -1, -1, -1)
	reconstruction, attention_store, text_cross_attention_maps = sample(
	zs, wts, attention_store=attention_store, text_cross_attention_maps=text_cross_attention_maps,prompt_tar=tar_prompt, skip=skip, cfg_scale_tar=tar_cfg_scale
	)
	do_reconstruction = False
	return (
	reconstruction,
	reconstruction,
	gr.update(visible=True),
	do_reconstruction,
	"Hide Reconstruction",
	)


	def load_and_invert(
	input_image,
	do_inversion,
	seed,
	randomize_seed,
	wts,
	zs,
	src_prompt="",
	# tar_prompt="",
	steps=30,
	src_cfg_scale=3.5,
	skip=15,
	tar_cfg_scale=15,
	progress=gr.Progress(track_tqdm=True),
	):
	# x0 = load_512(input_image, device=device).to(torch.float16)

	if do_inversion or randomize_seed:
	seed = randomize_seed_fn(seed, randomize_seed)
	seed_everything(seed)
	# invert and retrieve noise maps and latent
	zs_tensor, wts_tensor = pipe.invert(
	image_path=input_image,
	source_prompt=src_prompt,
	source_guidance_scale=src_cfg_scale,
	num_inversion_steps=steps,
	skip=skip,
	eta=1.0,
	)
	wts = wts_tensor
	zs = zs_tensor
	do_inversion = False

	return wts, zs, do_inversion, gr.update(visible=False)

	## SEGA ##

	def edit(input_image,
	wts, zs, attention_store, text_cross_attention_maps,
	tar_prompt,
	image_caption,
	steps,
	skip,
	tar_cfg_scale,
	edit_concept_1,edit_concept_2,edit_concept_3,
	guidnace_scale_1,guidnace_scale_2,guidnace_scale_3,
	warmup_1, warmup_2, warmup_3,
	neg_guidance_1, neg_guidance_2, neg_guidance_3,
	threshold_1, threshold_2, threshold_3,
	do_reconstruction,
	reconstruction,
	# for inversion in case it needs to be re computed (and avoid delay):
	do_inversion,
	seed,
	randomize_seed,
	src_prompt,
	src_cfg_scale,
	mask_type,
	progress=gr.Progress(track_tqdm=True)):
	show_share_button = gr.update(visible=True)
	if(mask_type == "No mask"):
	use_cross_attn_mask = False
	use_intersect_mask = False
	elif(mask_type=="Cross Attention Mask"):
	use_cross_attn_mask = True
	use_intersect_mask = False
	elif(mask_type=="Intersect Mask"):
	use_cross_attn_mask = False
	use_intersect_mask = True

	if randomize_seed:
	seed = randomize_seed_fn(seed, randomize_seed)
	seed_everything(seed)

	if do_inversion or randomize_seed:
	zs_tensor, wts_tensor = pipe.invert(
	image_path = input_image,
	source_prompt =src_prompt,
	source_guidance_scale= src_cfg_scale,
	num_inversion_steps = steps,
	skip = skip,
	eta = 1.0,
	)
	wts = wts_tensor
	zs = zs_tensor
	do_inversion = False

	if image_caption.lower() == tar_prompt.lower(): # if image caption was not changed, run pure sega
	tar_prompt = ""

	if edit_concept_1 != "" or edit_concept_2 != "" or edit_concept_3 != "":
	editing_args = dict(
	editing_prompt = [edit_concept_1,edit_concept_2,edit_concept_3],
	reverse_editing_direction = [ neg_guidance_1, neg_guidance_2, neg_guidance_3,],
	edit_warmup_steps=[warmup_1, warmup_2, warmup_3,],
	edit_guidance_scale=[guidnace_scale_1,guidnace_scale_2,guidnace_scale_3],
	edit_threshold=[threshold_1, threshold_2, threshold_3],
	edit_momentum_scale=0,
	edit_mom_beta=0,
	eta=1,
	use_cross_attn_mask=use_cross_attn_mask,
	use_intersect_mask=use_intersect_mask
	)

	latnets = wts[-1].expand(1, -1, -1, -1)
	sega_out, attention_store, text_cross_attention_maps = pipe(prompt=tar_prompt,
	init_latents=latnets,
	guidance_scale = tar_cfg_scale,
	# num_images_per_prompt=1,
	# num_inference_steps=steps,
	# use_ddpm=True,
	# wts=wts.value,
	zs=zs, attention_store=attention_store, text_cross_attention_maps=text_cross_attention_maps, **editing_args)

	return sega_out.images[0], gr.update(visible=True), do_reconstruction, reconstruction, wts, zs, attention_store, text_cross_attention_maps, do_inversion, show_share_button


	else: # if sega concepts were not added, performs regular ddpm sampling

	if do_reconstruction: # if ddpm sampling wasn't computed
	pure_ddpm_img, attention_store, text_cross_attention_maps = sample(zs, wts, attention_store=attention_store, text_cross_attention_maps=text_cross_attention_maps, prompt_tar=tar_prompt, skip=skip, cfg_scale_tar=tar_cfg_scale)
	reconstruction = pure_ddpm_img
	do_reconstruction = False
	return pure_ddpm_img, gr.update(visible=False), do_reconstruction, reconstruction, wts, zs, attention_store, text_cross_attention_maps, do_inversion, show_share_button

	return reconstruction, gr.update(visible=False), do_reconstruction, reconstruction, wts, zs, attention_store, text_cross_attention_maps, do_inversion, show_share_button


	def randomize_seed_fn(seed, is_random):
	if is_random:
	seed = random.randint(0, np.iinfo(np.int32).max)
	return seed

	def seed_everything(seed):
	torch.manual_seed(seed)
	torch.cuda.manual_seed(seed)
	random.seed(seed)
	np.random.seed(seed)

	def crop_image(image):
	h, w, c = image.shape
	if h < w:
	offset = (w - h) // 2
	image = image[:, offset:offset + h]
	elif w < h:
	offset = (h - w) // 2
	image = image[offset:offset + w]
	image = np.array(Image.fromarray(image).resize((512, 512)))
	return image


	def get_example():
	case = [
	[
	'examples/car_input.png',
	# '',
	'cherry blossom', 'green cabriolet','yellow car',

	'examples/car_output.png',


	13,11,7,
	2,2,2,
	False, False, True,
	50,
	25,
	7.5,
	0.65, 0.8, 0.8,
	890000000

	],
	[
	'examples/girl_with_pearl_earring_input.png',
	# '',
	'glasses', '','',

	'examples/girl_with_pearl_earring_output.png',


	4,7,0,
	3,2,2,
	False,False,False,
	50,
	25,
	5,
	0.97, 0.95,0.95,
	1900000000

	],

	[
	'examples/flower_field_input.jpg',
	# '',
	'pink tulips', 'red flowers',
	'van gogh painting',
	'examples/flower_field_output.png',


	20,7,10,
	1,1,1,
	False,True,False,
	50,
	25,
	7,
	0.9, 0.9,0.8,
	1900000000


	],

	]
	return case


	def swap_visibilities(input_image,
	edit_concept_1,
	edit_concept_2,
	edit_concept_3,
	sega_edited_image,
	guidnace_scale_1,
	guidnace_scale_2,
	guidnace_scale_3,
	warmup_1,
	warmup_2,
	warmup_3,
	neg_guidance_1,
	neg_guidance_2,
	neg_guidance_3,
	steps,
	skip,
	tar_cfg_scale,
	threshold_1,
	threshold_2,
	threshold_3,
	sega_concepts_counter

	):
	sega_concepts_counter=0
	concept1_update = update_display_concept("Remove" if neg_guidance_1 else "Add", edit_concept_1, neg_guidance_1, sega_concepts_counter)
	if(edit_concept_2 != ""):
	concept2_update = update_display_concept("Remove" if neg_guidance_2 else "Add", edit_concept_2, neg_guidance_2, sega_concepts_counter+1)
	else:
	concept2_update = gr.update(visible=False), gr.update(visible=False),gr.update(visible=False), gr.update(value=neg_guidance_2),gr.update(visible=True),gr.update(visible=False),sega_concepts_counter+1

	return (gr.update(visible=True), concept1_update[:-1], concept2_update)



	########
	# demo #
	########


	intro = """
	<div style="display: flex;align-items: center;justify-content: center">
	<img src="https://huggingface.co/spaces/editing-images/leditsplusplus/resolve/main/magician.png" width="100" style="display: inline-block">
	<h1 style="margin-left: 12px;text-align: center;margin-bottom: 7px;display: inline-block">LEDITS++</h1>
	<h3 style="display: inline-block;margin-left: 10px;margin-top: 6px;font-weight: 500">Limitless Image Editing using Text-to-Image Models</h3>
	</div>

	<p style="font-size: 0.95rem;margin: 0rem;line-height: 1.2em;margin-top:1em;display: inline-block">
	<a href="https://leditsplusplus-project.static.hf.space" target="_blank">project page</a> \| <a href="https://arxiv.org/abs/2311.16711" target="_blank">paper</a>
	\|
	<a href="https://huggingface.co/spaces/leditsplusplus/demo?duplicate=true" target="_blank" style="
	display: inline-block;
	">
	<img style="margin-top: -1em;margin-bottom: 0em;position: absolute;" src="https://bit.ly/3CWLGkA" alt="Duplicate Space"></a>
	</p>
	"""


	with gr.Blocks(css="style.css") as demo:
	def update_counter(sega_concepts_counter, concept1, concept2, concept3):
	if sega_concepts_counter == "":
	sega_concepts_counter = sum(1 for concept in (concept1, concept2, concept3) if concept != '')
	return sega_concepts_counter
	def remove_concept(sega_concepts_counter, row_triggered):
	sega_concepts_counter -= 1
	rows_visibility = [gr.update(visible=False) for _ in range(4)]

	if(row_triggered-1 > sega_concepts_counter):
	rows_visibility[sega_concepts_counter] = gr.update(visible=True)
	else:
	rows_visibility[row_triggered-1] = gr.update(visible=True)

	row1_visibility, row2_visibility, row3_visibility, row4_visibility = rows_visibility

	guidance_scale_label = "Concept Guidance Scale"
	# enable_interactive = gr.update(interactive=True)
	return (gr.update(visible=False),
	gr.update(visible=False, value="",),
	gr.update(interactive=True, value=""),
	gr.update(visible=False,label = guidance_scale_label),
	gr.update(interactive=True, value =False),
	gr.update(value=DEFAULT_WARMUP_STEPS),
	gr.update(value=DEFAULT_THRESHOLD),
	gr.update(visible=True),
	gr.update(interactive=True, value="custom"),
	row1_visibility,
	row2_visibility,
	row3_visibility,
	row4_visibility,
	sega_concepts_counter
	)



	def update_display_concept(button_label, edit_concept, neg_guidance, sega_concepts_counter):
	sega_concepts_counter += 1
	guidance_scale_label = "Concept Guidance Scale"
	if(button_label=='Remove'):
	neg_guidance = True
	guidance_scale_label = "Negative Guidance Scale"

	return (gr.update(visible=True), #boxn
	gr.update(visible=True, value=edit_concept), #concept_n
	gr.update(visible=True,label = guidance_scale_label), #guidance_scale_n
	gr.update(value=neg_guidance),#neg_guidance_n
	gr.update(visible=False), #row_n
	gr.update(visible=True), #row_n+1
	sega_concepts_counter
	)


	def display_editing_options(run_button, clear_button, sega_tab):
	return gr.update(visible=True), gr.update(visible=True), gr.update(visible=True)

	def update_interactive_mode(add_button_label):
	if add_button_label == "Clear":
	return gr.update(interactive=False), gr.update(interactive=False)
	else:
	return gr.update(interactive=True), gr.update(interactive=True)

	def update_dropdown_parms(dropdown):
	if dropdown == 'custom':
	return DEFAULT_SEGA_CONCEPT_GUIDANCE_SCALE,DEFAULT_WARMUP_STEPS, DEFAULT_THRESHOLD
	elif dropdown =='style':
	return STYLE_SEGA_CONCEPT_GUIDANCE_SCALE,STYLE_WARMUP_STEPS, STYLE_THRESHOLD
	elif dropdown =='object':
	return OBJECT_SEGA_CONCEPT_GUIDANCE_SCALE,OBJECT_WARMUP_STEPS, OBJECT_THRESHOLD
	elif dropdown =='faces':
	return FACE_SEGA_CONCEPT_GUIDANCE_SCALE,FACE_WARMUP_STEPS, FACE_THRESHOLD


	def reset_do_inversion():
	return True

	def reset_do_reconstruction():
	do_reconstruction = True
	return do_reconstruction

	def reset_image_caption():
	return ""

	def update_inversion_progress_visibility(input_image, do_inversion):
	if do_inversion and not input_image is None:
	return gr.update(visible=True)
	else:
	return gr.update(visible=False)

	def update_edit_progress_visibility(input_image, do_inversion):
	# if do_inversion and not input_image is None:
	# return inversion_progress.update(visible=True)
	# else:
	return gr.update(visible=True)


	gr.HTML(intro)
	wts = gr.State()
	zs = gr.State()
	attention_store=gr.State()
	text_cross_attention_maps = gr.State()
	reconstruction = gr.State()
	do_inversion = gr.State(value=True)
	do_reconstruction = gr.State(value=True)
	sega_concepts_counter = gr.State(0)
	image_caption = gr.State(value="")

	with gr.Row():
	input_image = gr.Image(label="Input Image", interactive=True, elem_id="input_image")
	ddpm_edited_image = gr.Image(label=f"Pure DDPM Inversion Image", interactive=False, visible=False)
	sega_edited_image = gr.Image(label=f"LEDITS Edited Image", interactive=False, elem_id="output_image")

	with gr.Group(visible=False, elem_id="share-btn-wrapper") as share_btn_container:
	with gr.Group(elem_id="share-btn-container"):
	community_icon = gr.HTML(community_icon_html, visible=True)
	loading_icon = gr.HTML(loading_icon_html, visible=False)
	share_button = gr.Button("Share to community", elem_id="share-btn", visible=True)

	with gr.Row():
	with gr.Group(visible=False, elem_id="box1") as box1:
	with gr.Row():
	concept_1 = gr.Button(scale=3, value="")
	remove_concept1 = gr.Button("x", scale=1, min_width=10)
	with gr.Row():
	guidnace_scale_1 = gr.Slider(label='Concept Guidance Scale', minimum=1, maximum=30,
	info="How strongly the concept should modify the image",
	value=DEFAULT_SEGA_CONCEPT_GUIDANCE_SCALE,
	step=0.5, interactive=True)
	with gr.Group(visible=False, elem_id="box2") as box2:
	with gr.Row():
	concept_2 = gr.Button(scale=3, value="")
	remove_concept2 = gr.Button("x", scale=1, min_width=10)
	with gr.Row():
	guidnace_scale_2 = gr.Slider(label='Concept Guidance Scale', minimum=1, maximum=30,
	info="How strongly the concept should modify the image",
	value=DEFAULT_SEGA_CONCEPT_GUIDANCE_SCALE,
	step=0.5, interactive=True)
	with gr.Group(visible=False, elem_id="box3") as box3:
	with gr.Row():
	concept_3 = gr.Button(scale=3, value="")
	remove_concept3 = gr.Button("x", scale=1, min_width=10)
	with gr.Row():
	guidnace_scale_3 = gr.Slider(label='Concept Guidance Scale', minimum=1, maximum=30,
	info="How strongly the concept should modify the image",
	value=DEFAULT_SEGA_CONCEPT_GUIDANCE_SCALE,
	step=0.5, interactive=True)


	with gr.Row():
	inversion_progress = gr.Textbox(visible=False, label="Inversion progress")

	with gr.Group():
	intro_segs = gr.Markdown("Add/Remove Concepts from your Image <span style=\"font-size: 12px; color: rgb(156, 163, 175)\">with Semantic Guidance</span>")
	# 1st SEGA concept
	with gr.Row() as row1:
	with gr.Column(scale=3, min_width=100):
	with gr.Row():
	# with gr.Column(scale=3, min_width=100):
	edit_concept_1 = gr.Textbox(
	label="Concept",
	show_label=True,
	max_lines=1, value="",
	placeholder="E.g.: Sunglasses",
	)
	# with gr.Column(scale=2, min_width=100):# better mobile ui
	dropdown1 = gr.Dropdown(label = "Edit Type", value ='custom' , choices=['custom','style', 'object', 'faces'])


	with gr.Column(scale=1, min_width=100, visible=False):
	neg_guidance_1 = gr.Checkbox(
	label='Remove Concept?')

	with gr.Column(scale=1, min_width=100):
	with gr.Row(): # better mobile ui
	with gr.Column():
	add_1 = gr.Button('Add')
	remove_1 = gr.Button('Remove')


	# 2nd SEGA concept
	with gr.Row(visible=False) as row2:
	with gr.Column(scale=3, min_width=100):
	with gr.Row(): #better mobile UI
	# with gr.Column(scale=3, min_width=100):
	edit_concept_2 = gr.Textbox(
	label="Concept",
	show_label=True,
	max_lines=1,
	placeholder="E.g.: Realistic",
	)
	# with gr.Column(scale=2, min_width=100):# better mobile ui
	dropdown2 = gr.Dropdown(label = "Edit Type", value ='custom' , choices=['custom','style', 'object', 'faces'])

	with gr.Column(scale=1, min_width=100, visible=False):
	neg_guidance_2 = gr.Checkbox(
	label='Remove Concept?')

	with gr.Column(scale=1, min_width=100):
	with gr.Row(): # better mobile ui
	with gr.Column():
	add_2 = gr.Button('Add')
	remove_2 = gr.Button('Remove')

	# 3rd SEGA concept
	with gr.Row(visible=False) as row3:
	with gr.Column(scale=3, min_width=100):
	with gr.Row(): #better mobile UI
	# with gr.Column(scale=3, min_width=100):
	edit_concept_3 = gr.Textbox(
	label="Concept",
	show_label=True,
	max_lines=1,
	placeholder="E.g.: orange",
	)
	# with gr.Column(scale=2, min_width=100):
	dropdown3 = gr.Dropdown(label = "Edit Type", value ='custom' , choices=['custom','style', 'object', 'faces'])

	with gr.Column(scale=1, min_width=100, visible=False):
	neg_guidance_3 = gr.Checkbox(
	label='Remove Concept?',visible=True)

	with gr.Column(scale=1, min_width=100):
	with gr.Row(): # better mobile ui
	with gr.Column():
	add_3 = gr.Button('Add')
	remove_3 = gr.Button('Remove')

	with gr.Row(visible=False) as row4:
	gr.Markdown("### Max of 3 concepts reached. Remove a concept to add more")

	#with gr.Row(visible=False).style(mobile_collapse=False, equal_height=True):
	# add_concept_button = gr.Button("+1 concept")




	# caption_button = gr.Button("Caption Image", scale=1)


	with gr.Row():
	run_button = gr.Button("Edit your image!", visible=True)


	with gr.Accordion("Advanced Options", open=False):
	with gr.Row():
	tar_prompt = gr.Textbox(
	label="Describe your edited image (optional)",
	elem_id="target_prompt",
	# show_label=False,
	max_lines=1, value="", scale=3,
	placeholder="Target prompt, DDPM Inversion", info = "DPM Solver++ Inversion Prompt. Can help with global changes, modify to what you would like to see"
	)
	with gr.Tabs() as tabs:

	with gr.TabItem('General options', id=2):
	with gr.Row():
	with gr.Column(min_width=100):
	clear_button = gr.Button("Clear", visible=True)
	src_prompt = gr.Textbox(lines=1, label="Source Prompt", interactive=True, placeholder="")
	steps = gr.Number(value=50, precision=0, label="Num Diffusion Steps", interactive=True)
	src_cfg_scale = gr.Number(value=3.5, label=f"Source Guidance Scale", interactive=True)
	mask_type = gr.Radio(choices=["No mask", "Cross Attention Mask", "Intersect Mask"], value="Intersect Mask", label="Mask type")

	with gr.Column(min_width=100):
	reconstruct_button = gr.Button("Show Reconstruction", visible=False)
	skip = gr.Slider(minimum=0, maximum=95, value=25, step=1, label="Skip Steps", interactive=True, info = "Percentage of skipped denoising steps. Bigger values increase fidelity to input image")
	tar_cfg_scale = gr.Slider(minimum=1, maximum=30,value=7.5, label=f"Guidance Scale", interactive=True)
	seed = gr.Slider(minimum=0, maximum=np.iinfo(np.int32).max, label="Seed", interactive=True, randomize=True)
	randomize_seed = gr.Checkbox(label='Randomize seed', value=False)

	with gr.TabItem('SEGA options', id=3) as sega_advanced_tab:
	# 1st SEGA concept
	gr.Markdown("1st concept")
	with gr.Row():
	warmup_1 = gr.Slider(label='Warmup', minimum=0, maximum=50,
	value=DEFAULT_WARMUP_STEPS,
	step=1, interactive=True, info="At which step to start applying semantic guidance. Bigger values reduce edit concept's effect")
	threshold_1 = gr.Slider(label='Threshold', minimum=0, maximum=0.99,
	value=DEFAULT_THRESHOLD, step=0.01, interactive=True,
	info = "Lower the threshold for more effect (e.g. ~0.9 for style transfer)")

	# 2nd SEGA concept
	gr.Markdown("2nd concept")
	with gr.Row() as row2_advanced:
	warmup_2 = gr.Slider(label='Warmup', minimum=0, maximum=50,
	value=DEFAULT_WARMUP_STEPS,
	step=1, interactive=True, info="At which step to start applying semantic guidance. Bigger values reduce edit concept's effect")
	threshold_2 = gr.Slider(label='Threshold', minimum=0, maximum=0.99,
	value=DEFAULT_THRESHOLD,
	step=0.01, interactive=True,
	info = "Lower the threshold for more effect (e.g. ~0.9 for style transfer)")
	# 3rd SEGA concept
	gr.Markdown("3rd concept")
	with gr.Row() as row3_advanced:
	warmup_3 = gr.Slider(label='Warmup', minimum=0, maximum=50,
	value=DEFAULT_WARMUP_STEPS, step=1,
	interactive=True, info="At which step to start applying semantic guidance. Bigger values reduce edit concept's effect")
	threshold_3 = gr.Slider(label='Threshold', minimum=0, maximum=0.99,
	value=DEFAULT_THRESHOLD, step=0.01,
	interactive=True,
	info = "Lower the threshold for more effect (e.g. ~0.9 for style transfer)")

	# caption_button.click(
	# fn = caption_image,
	# inputs = [input_image],
	# outputs = [tar_prompt]
	# )
	#neg_guidance_1.change(fn = update_label, inputs=[neg_guidance_1], outputs=[add_1])
	#neg_guidance_2.change(fn = update_label, inputs=[neg_guidance_2], outputs=[add_2])
	#neg_guidance_3.change(fn = update_label, inputs=[neg_guidance_3], outputs=[add_3])
	add_1.click(fn=update_counter,
	inputs=[sega_concepts_counter,edit_concept_1,edit_concept_2,edit_concept_3],
	outputs=sega_concepts_counter,queue=False).then(fn = update_display_concept, inputs=[add_1, edit_concept_1, neg_guidance_1, sega_concepts_counter], outputs=[box1, concept_1, guidnace_scale_1,neg_guidance_1,row1, row2, sega_concepts_counter],queue=False)
	add_2.click(fn=update_counter,inputs=[sega_concepts_counter,edit_concept_1,edit_concept_2,edit_concept_3], outputs=sega_concepts_counter,queue=False).then(fn = update_display_concept, inputs=[add_2, edit_concept_2, neg_guidance_2, sega_concepts_counter], outputs=[box2, concept_2, guidnace_scale_2,neg_guidance_2,row2, row3, sega_concepts_counter],queue=False)
	add_3.click(fn=update_counter,inputs=[sega_concepts_counter,edit_concept_1,edit_concept_2,edit_concept_3], outputs=sega_concepts_counter,queue=False).then(fn = update_display_concept, inputs=[add_3, edit_concept_3, neg_guidance_3, sega_concepts_counter], outputs=[box3, concept_3, guidnace_scale_3,neg_guidance_3,row3, row4, sega_concepts_counter],queue=False)

	remove_1.click(fn = update_display_concept, inputs=[remove_1, edit_concept_1, neg_guidance_1, sega_concepts_counter], outputs=[box1, concept_1, guidnace_scale_1,neg_guidance_1,row1, row2, sega_concepts_counter],queue=False)
	remove_2.click(fn = update_display_concept, inputs=[remove_2, edit_concept_2, neg_guidance_2 ,sega_concepts_counter], outputs=[box2, concept_2, guidnace_scale_2,neg_guidance_2,row2, row3,sega_concepts_counter],queue=False)
	remove_3.click(fn = update_display_concept, inputs=[remove_3, edit_concept_3, neg_guidance_3, sega_concepts_counter], outputs=[box3, concept_3, guidnace_scale_3,neg_guidance_3, row3, row4, sega_concepts_counter],queue=False)

	remove_concept1.click(
	fn=update_counter,inputs=[sega_concepts_counter,edit_concept_1,edit_concept_2,edit_concept_3], outputs=sega_concepts_counter,queue=False).then(
	fn = remove_concept, inputs=[sega_concepts_counter,gr.State(1)], outputs= [box1, concept_1, edit_concept_1, guidnace_scale_1,neg_guidance_1,warmup_1, threshold_1, add_1, dropdown1, row1, row2, row3, row4, sega_concepts_counter],queue=False)
	remove_concept2.click(
	fn=update_counter,inputs=[sega_concepts_counter,edit_concept_1,edit_concept_2,edit_concept_3], outputs=sega_concepts_counter,queue=False).then(
	fn = remove_concept, inputs=[sega_concepts_counter,gr.State(2)], outputs=[box2, concept_2, edit_concept_2, guidnace_scale_2,neg_guidance_2, warmup_2, threshold_2, add_2 , dropdown2, row1, row2, row3, row4, sega_concepts_counter],queue=False)
	remove_concept3.click(
	fn=update_counter,inputs=[sega_concepts_counter,edit_concept_1,edit_concept_2,edit_concept_3], outputs=sega_concepts_counter,queue=False).then(
	fn = remove_concept,inputs=[sega_concepts_counter,gr.State(3)], outputs=[box3, concept_3, edit_concept_3, guidnace_scale_3,neg_guidance_3,warmup_3, threshold_3, add_3, dropdown3, row1, row2, row3, row4, sega_concepts_counter],queue=False)

	#add_concept_button.click(fn = update_display_concept, inputs=sega_concepts_counter,
	# outputs= [row2, row2_advanced, row3, row3_advanced, add_concept_button, sega_concepts_counter], queue = False)

	run_button.click(
	fn=edit,
	inputs=[input_image,
	wts, zs, attention_store,
	text_cross_attention_maps,
	tar_prompt,
	image_caption,
	steps,
	skip,
	tar_cfg_scale,
	edit_concept_1,edit_concept_2,edit_concept_3,
	guidnace_scale_1,guidnace_scale_2,guidnace_scale_3,
	warmup_1, warmup_2, warmup_3,
	neg_guidance_1, neg_guidance_2, neg_guidance_3,
	threshold_1, threshold_2, threshold_3, do_reconstruction, reconstruction,
	do_inversion,
	seed,
	randomize_seed,
	src_prompt,
	src_cfg_scale,
	mask_type
	],
	outputs=[sega_edited_image, reconstruct_button, do_reconstruction, reconstruction, wts, zs,attention_store, text_cross_attention_maps, do_inversion, share_btn_container]

	)
	# .success(fn=update_gallery_display, inputs= [prev_output_image, sega_edited_image], outputs = [gallery, gallery, prev_output_image])


	input_image.change(
	fn = reset_do_inversion,
	outputs = [do_inversion],
	queue=False,
	concurrency_limit=None
	).then(
	fn = randomize_seed_fn,
	inputs = [seed, randomize_seed],
	outputs = [seed],
	queue=False,
	concurrency_limit=None
	)

	# Automatically start inverting upon input_image change
	input_image.upload(
	fn = crop_image,
	inputs = [input_image],
	outputs = [input_image],
	queue=False,
	concurrency_limit=None,
	).then(
	fn = reset_do_inversion,
	outputs = [do_inversion],
	queue=False,
	concurrency_limit=None
	).then(
	fn = randomize_seed_fn,
	inputs = [seed, randomize_seed],
	outputs = [seed],
	queue=False,
	concurrency_limit=None
	).then(fn = caption_image,
	inputs = [input_image],
	outputs = [tar_prompt, image_caption],
	queue=False,
	concurrency_limit=None
	)

	# Repeat inversion (and reconstruction) when these params are changed:
	src_prompt.change(
	fn = reset_do_inversion,
	outputs = [do_inversion],
	queue = False
	).then(
	fn = reset_do_reconstruction,
	outputs = [do_reconstruction],
	queue = False
	)

	steps.change(
	fn = reset_do_inversion,
	outputs = [do_inversion],
	queue = False
	).then(
	fn = reset_do_reconstruction,
	outputs = [do_reconstruction],
	queue = False
	)

	src_cfg_scale.change(
	fn = reset_do_inversion,
	outputs = [do_inversion],
	queue = False
	).then(
	fn = reset_do_reconstruction,
	outputs = [do_reconstruction],
	queue = False
	)

	# Repeat only reconstruction these params are changed:
	tar_prompt.change(
	fn = reset_do_reconstruction,
	outputs = [do_reconstruction],
	queue = False
	)

	tar_cfg_scale.change(
	fn = reset_do_reconstruction,
	outputs = [do_reconstruction],
	queue = False
	)

	skip.change(
	fn = reset_do_inversion,
	outputs = [do_inversion],
	queue = False
	).then(
	fn = reset_do_reconstruction,
	outputs = [do_reconstruction],
	queue = False
	)

	seed.change(
	fn=reset_do_inversion,
	outputs=[do_inversion],
	queue=False
	).then(
	fn=reset_do_reconstruction,
	outputs=[do_reconstruction],
	queue=False
	)

	dropdown1.change(fn=update_dropdown_parms, inputs = [dropdown1], outputs = [guidnace_scale_1,warmup_1, threshold_1], queue=False)
	dropdown2.change(fn=update_dropdown_parms, inputs = [dropdown2], outputs = [guidnace_scale_2,warmup_2, threshold_2], queue=False)
	dropdown3.change(fn=update_dropdown_parms, inputs = [dropdown3], outputs = [guidnace_scale_3,warmup_3, threshold_3], queue=False)

	clear_components = [input_image,ddpm_edited_image,ddpm_edited_image,sega_edited_image, do_inversion,
	src_prompt, steps, src_cfg_scale, seed,
	tar_prompt, skip, tar_cfg_scale, reconstruct_button,reconstruct_button,
	edit_concept_1, guidnace_scale_1,guidnace_scale_1,warmup_1, threshold_1, neg_guidance_1,dropdown1, concept_1, concept_1, row1,
	edit_concept_2, guidnace_scale_2,guidnace_scale_2,warmup_2, threshold_2, neg_guidance_2,dropdown2, concept_2, concept_2, row2,
	edit_concept_3, guidnace_scale_3,guidnace_scale_3,warmup_3, threshold_3, neg_guidance_3,dropdown3, concept_3,concept_3, row3,
	row4,sega_concepts_counter, box1, box2, box3 ]

	clear_components_output_vals = [None, None,gr.update(visible=False), None, True,
	"", DEFAULT_DIFFUSION_STEPS, DEFAULT_SOURCE_GUIDANCE_SCALE, DEFAULT_SEED,
	"", DEFAULT_SKIP_STEPS, DEFAULT_TARGET_GUIDANCE_SCALE, gr.update(value="Show Reconstruction"),gr.update(visible=False),
	"", DEFAULT_SEGA_CONCEPT_GUIDANCE_SCALE,gr.update(visible=False), DEFAULT_WARMUP_STEPS, DEFAULT_THRESHOLD, DEFAULT_NEGATIVE_GUIDANCE, "custom","", gr.update(visible=False), gr.update(visible=True),
	"", DEFAULT_SEGA_CONCEPT_GUIDANCE_SCALE,gr.update(visible=False), DEFAULT_WARMUP_STEPS, DEFAULT_THRESHOLD, DEFAULT_NEGATIVE_GUIDANCE, "custom","", gr.update(visible=False), gr.update(visible=False),
	"", DEFAULT_SEGA_CONCEPT_GUIDANCE_SCALE,gr.update(visible=False), DEFAULT_WARMUP_STEPS, DEFAULT_THRESHOLD, DEFAULT_NEGATIVE_GUIDANCE, "custom","",gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(value=0),
	gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)]


	clear_button.click(lambda: clear_components_output_vals, outputs = clear_components)

	reconstruct_button.click(lambda: ddpm_edited_image.update(visible=True), outputs=[ddpm_edited_image]).then(fn = reconstruct,
	inputs = [tar_prompt,
	image_caption,
	tar_cfg_scale,
	skip,
	wts, zs,
	do_reconstruction,
	reconstruction,
	reconstruct_button],
	outputs = [ddpm_edited_image,reconstruction, ddpm_edited_image, do_reconstruction, reconstruct_button])

	randomize_seed.change(
	fn = randomize_seed_fn,
	inputs = [seed, randomize_seed],
	outputs = [seed],
	queue = False)

	share_button.click(None, [], [], js=share_js)

	gr.Examples(
	label='Examples',
	fn=swap_visibilities,
	run_on_click=True,
	examples=get_example(),
	inputs=[input_image,
	edit_concept_1,
	edit_concept_2,
	edit_concept_3,
	sega_edited_image,
	guidnace_scale_1,
	guidnace_scale_2,
	guidnace_scale_3,
	warmup_1,
	warmup_2,
	warmup_3,
	neg_guidance_1,
	neg_guidance_2,
	neg_guidance_3,
	steps,
	skip,
	tar_cfg_scale,
	threshold_1,
	threshold_2,
	threshold_3,
	seed,
	sega_concepts_counter
	],
	outputs=[share_btn_container, box1, concept_1, guidnace_scale_1,neg_guidance_1, row1, row2,box2, concept_2, guidnace_scale_2,neg_guidance_2,row2, row3,sega_concepts_counter],
	cache_examples=True
	)

	demo.queue(default_concurrency_limit=1)
	demo.launch()