Spaces:

lnyan
/

stablediffusion-infinity

Running on A10G

App Files Files Community

stablediffusion-infinity / app.py

lnyan

Update app.py

a2d77b4 7 months ago

raw history blame contribute delete

No virus

37.3 kB

	import subprocess
	# import os.path as osp
	import pip
	# pip.main(["install","-v","-U","git+https://github.com/facebookresearch/xformers.git@main#egg=xformers"])
	# subprocess.check_call("pip install -v -U git+https://github.com/facebookresearch/xformers.git@main#egg=xformers", cwd=osp.dirname(__file__), shell=True)

	import io
	import base64
	import os
	import os

	import sys

	import numpy as np
	import torch
	from torch import autocast
	import diffusers
	from diffusers.configuration_utils import FrozenDict
	from diffusers import (
	StableDiffusionPipeline,
	StableDiffusionInpaintPipeline,
	StableDiffusionImg2ImgPipeline,
	StableDiffusionInpaintPipelineLegacy,
	DDIMScheduler,
	LMSDiscreteScheduler,
	StableDiffusionUpscalePipeline,
	DPMSolverMultistepScheduler
	)
	from diffusers.models import AutoencoderKL
	from PIL import Image
	from PIL import ImageOps
	import gradio as gr
	import base64
	import skimage
	import skimage.measure
	import yaml
	import json
	from enum import Enum

	try:
	abspath = os.path.abspath(__file__)
	dirname = os.path.dirname(abspath)
	os.chdir(dirname)
	except:
	pass

	from utils import *

	# assert diffusers.__version__ >= "0.6.0", "Please upgrade diffusers to 0.6.0"

	USE_NEW_DIFFUSERS = True
	RUN_IN_SPACE = "RUN_IN_HG_SPACE" in os.environ


	class ModelChoice(Enum):
	INPAINTING = "stablediffusion-inpainting"
	INPAINTING_IMG2IMG = "stablediffusion-inpainting+img2img-v1.5"
	MODEL_1_5 = "stablediffusion-v1.5"
	MODEL_1_4 = "stablediffusion-v1.4"


	try:
	from sd_grpcserver.pipeline.unified_pipeline import UnifiedPipeline
	except:
	UnifiedPipeline = StableDiffusionInpaintPipeline

	# sys.path.append("./glid_3_xl_stable")

	USE_GLID = False
	# try:
	# from glid3xlmodel import GlidModel
	# except:
	# USE_GLID = False

	try:
	cuda_available = torch.cuda.is_available()
	except:
	cuda_available = False
	finally:
	if sys.platform == "darwin":
	device = "mps" if torch.backends.mps.is_available() else "cpu"
	elif cuda_available:
	device = "cuda"
	else:
	device = "cpu"

	import contextlib

	autocast = contextlib.nullcontext

	with open("config.yaml", "r") as yaml_in:
	yaml_object = yaml.safe_load(yaml_in)
	config_json = json.dumps(yaml_object)


	def load_html():
	body, canvaspy = "", ""
	with open("index.html", encoding="utf8") as f:
	body = f.read()
	with open("canvas.py", encoding="utf8") as f:
	canvaspy = f.read()
	body = body.replace("- paths:\n", "")
	body = body.replace(" - ./canvas.py\n", "")
	body = body.replace("from canvas import InfCanvas", canvaspy)
	return body


	def test(x):
	x = load_html()
	return f"""<iframe id="sdinfframe" style="width: 100%; height: 600px" name="result" allow="midi; geolocation; microphone; camera;
	display-capture; encrypted-media; vertical-scroll 'none'" sandbox="allow-modals allow-forms
	allow-scripts allow-same-origin allow-popups
	allow-top-navigation-by-user-activation allow-downloads" allowfullscreen=""
	allowpaymentrequest="" frameborder="0" srcdoc='{x}'></iframe>"""


	DEBUG_MODE = False

	try:
	SAMPLING_MODE = Image.Resampling.LANCZOS
	except Exception as e:
	SAMPLING_MODE = Image.LANCZOS

	try:
	contain_func = ImageOps.contain
	except Exception as e:

	def contain_func(image, size, method=SAMPLING_MODE):
	# from PIL: https://pillow.readthedocs.io/en/stable/reference/ImageOps.html#PIL.ImageOps.contain
	im_ratio = image.width / image.height
	dest_ratio = size[0] / size[1]
	if im_ratio != dest_ratio:
	if im_ratio > dest_ratio:
	new_height = int(image.height / image.width * size[0])
	if new_height != size[1]:
	size = (size[0], new_height)
	else:
	new_width = int(image.width / image.height * size[1])
	if new_width != size[0]:
	size = (new_width, size[1])
	return image.resize(size, resample=method)


	import argparse

	parser = argparse.ArgumentParser(description="stablediffusion-infinity")
	parser.add_argument("--port", type=int, help="listen port", dest="server_port")
	parser.add_argument("--host", type=str, help="host", dest="server_name")
	parser.add_argument("--share", action="store_true", help="share this app?")
	parser.add_argument("--debug", action="store_true", help="debug mode")
	parser.add_argument("--fp32", action="store_true", help="using full precision")
	parser.add_argument("--encrypt", action="store_true", help="using https?")
	parser.add_argument("--ssl_keyfile", type=str, help="path to ssl_keyfile")
	parser.add_argument("--ssl_certfile", type=str, help="path to ssl_certfile")
	parser.add_argument("--ssl_keyfile_password", type=str, help="ssl_keyfile_password")
	parser.add_argument(
	"--auth", nargs=2, metavar=("username", "password"), help="use username password"
	)
	parser.add_argument(
	"--remote_model",
	type=str,
	help="use a model (e.g. dreambooth fined) from huggingface hub",
	default="",
	)
	parser.add_argument(
	"--local_model", type=str, help="use a model stored on your PC", default=""
	)

	if __name__ == "__main__" and not RUN_IN_SPACE:
	args = parser.parse_args()
	else:
	args = parser.parse_args()
	# args = parser.parse_args(["--debug"])
	if args.auth is not None:
	args.auth = tuple(args.auth)

	model = {}


	def get_token():
	token = ""
	if os.path.exists(".token"):
	with open(".token", "r") as f:
	token = f.read()
	token = os.environ.get("hftoken", token)
	return token


	def save_token(token):
	with open(".token", "w") as f:
	f.write(token)


	def prepare_scheduler(scheduler):
	if hasattr(scheduler.config, "steps_offset") and scheduler.config.steps_offset != 1:
	new_config = dict(scheduler.config)
	new_config["steps_offset"] = 1
	scheduler._internal_dict = FrozenDict(new_config)
	return scheduler


	def my_resize(width, height):
	if width >= 512 and height >= 512:
	return width, height
	if width == height:
	return 512, 512
	smaller = min(width, height)
	larger = max(width, height)
	if larger >= 608:
	return width, height
	factor = 1
	if smaller < 290:
	factor = 2
	elif smaller < 330:
	factor = 1.75
	elif smaller < 384:
	factor = 1.375
	elif smaller < 400:
	factor = 1.25
	elif smaller < 450:
	factor = 1.125
	return int(factor * width)//88, int(factor height)//8*8


	def load_learned_embed_in_clip(
	learned_embeds_path, text_encoder, tokenizer, token=None
	):
	# https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers/stable_conceptualizer_inference.ipynb
	loaded_learned_embeds = torch.load(learned_embeds_path, map_location="cpu")

	# separate token and the embeds
	trained_token = list(loaded_learned_embeds.keys())[0]
	embeds = loaded_learned_embeds[trained_token]

	# cast to dtype of text_encoder
	dtype = text_encoder.get_input_embeddings().weight.dtype
	embeds.to(dtype)

	# add the token in tokenizer
	token = token if token is not None else trained_token
	num_added_tokens = tokenizer.add_tokens(token)
	if num_added_tokens == 0:
	raise ValueError(
	f"The tokenizer already contains the token {token}. Please pass a different `token` that is not already in the tokenizer."
	)

	# resize the token embeddings
	text_encoder.resize_token_embeddings(len(tokenizer))

	# get the id for the token and assign the embeds
	token_id = tokenizer.convert_tokens_to_ids(token)
	text_encoder.get_input_embeddings().weight.data[token_id] = embeds


	scheduler_dict = {"PLMS": None, "DDIM": None, "K-LMS": None, "DPM": None}


	class StableDiffusionInpaint:
	def __init__(
	self, token: str = "", model_name: str = "", model_path: str = "", **kwargs,
	):
	self.token = token
	original_checkpoint = False
	if model_path and os.path.exists(model_path):
	if model_path.endswith(".ckpt"):
	original_checkpoint = True
	elif model_path.endswith(".json"):
	model_name = os.path.dirname(model_path)
	else:
	model_name = model_path
	vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse")
	vae.to(torch.float16)
	if original_checkpoint:
	print(f"Converting & Loading {model_path}")
	from convert_checkpoint import convert_checkpoint

	pipe = convert_checkpoint(model_path, inpainting=True)
	if device == "cuda":
	pipe.to(torch.float16)
	inpaint = StableDiffusionInpaintPipeline(
	vae=vae,
	text_encoder=pipe.text_encoder,
	tokenizer=pipe.tokenizer,
	unet=pipe.unet,
	scheduler=pipe.scheduler,
	safety_checker=pipe.safety_checker,
	feature_extractor=pipe.feature_extractor,
	)
	else:
	print(f"Loading {model_name}")
	if device == "cuda":
	inpaint = StableDiffusionInpaintPipeline.from_pretrained(
	model_name,
	revision="fp16",
	torch_dtype=torch.float16,
	use_auth_token=token,
	vae=vae
	)
	else:
	inpaint = StableDiffusionInpaintPipeline.from_pretrained(
	model_name, use_auth_token=token,
	)
	if os.path.exists("./embeddings"):
	print("Note that StableDiffusionInpaintPipeline + embeddings is untested")
	for item in os.listdir("./embeddings"):
	if item.endswith(".bin"):
	load_learned_embed_in_clip(
	os.path.join("./embeddings", item),
	inpaint.text_encoder,
	inpaint.tokenizer,
	)
	inpaint.to(device)
	# try:
	# inpaint.vae=torch.compile(inpaint.vae, dynamic=True)
	# inpaint.unet=torch.compile(inpaint.unet, dynamic=True)
	# except Exception as e:
	# print(e)
	# inpaint.enable_xformers_memory_efficient_attention()
	# if device == "mps":
	# _ = text2img("", num_inference_steps=1)
	scheduler_dict["PLMS"] = inpaint.scheduler
	scheduler_dict["DDIM"] = prepare_scheduler(
	DDIMScheduler(
	beta_start=0.00085,
	beta_end=0.012,
	beta_schedule="scaled_linear",
	clip_sample=False,
	set_alpha_to_one=False,
	)
	)
	scheduler_dict["K-LMS"] = prepare_scheduler(
	LMSDiscreteScheduler(
	beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear"
	)
	)
	scheduler_dict["DPM"] = prepare_scheduler(
	DPMSolverMultistepScheduler.from_config(inpaint.scheduler.config)
	)
	self.safety_checker = inpaint.safety_checker
	save_token(token)
	try:
	total_memory = torch.cuda.get_device_properties(0).total_memory // (
	1024 ** 3
	)
	if total_memory <= 5:
	inpaint.enable_attention_slicing()
	except:
	pass
	self.inpaint = inpaint

	def run(
	self,
	image_pil,
	prompt="",
	negative_prompt="",
	guidance_scale=7.5,
	resize_check=True,
	enable_safety=True,
	fill_mode="patchmatch",
	strength=0.75,
	step=50,
	enable_img2img=False,
	use_seed=False,
	seed_val=-1,
	generate_num=1,
	scheduler="",
	scheduler_eta=0.0,
	**kwargs,
	):
	inpaint = self.inpaint
	selected_scheduler = scheduler_dict.get(scheduler, scheduler_dict["PLMS"])
	for item in [inpaint]:
	item.scheduler = selected_scheduler
	if enable_safety:
	item.safety_checker = self.safety_checker
	else:
	item.safety_checker = lambda images, **kwargs: (images, None)
	width, height = image_pil.size
	sel_buffer = np.array(image_pil)
	img = sel_buffer[:, :, 0:3]
	mask = sel_buffer[:, :, -1]
	nmask = 255 - mask
	process_width = width
	process_height = height
	if resize_check:
	process_width, process_height = my_resize(width, height)
	process_width=process_width*8//8
	process_height=process_height*8//8
	extra_kwargs = {
	"num_inference_steps": step,
	"guidance_scale": guidance_scale,
	"eta": scheduler_eta,
	}
	if USE_NEW_DIFFUSERS:
	extra_kwargs["negative_prompt"] = negative_prompt
	extra_kwargs["num_images_per_prompt"] = generate_num
	if use_seed:
	generator = torch.Generator(inpaint.device).manual_seed(seed_val)
	extra_kwargs["generator"] = generator
	if True:
	img, mask = functbl[fill_mode](img, mask)
	mask = 255 - mask
	mask = skimage.measure.block_reduce(mask, (8, 8), np.max)
	mask = mask.repeat(8, axis=0).repeat(8, axis=1)
	extra_kwargs["strength"] = strength
	inpaint_func = inpaint
	init_image = Image.fromarray(img)
	mask_image = Image.fromarray(mask)
	# mask_image=mask_image.filter(ImageFilter.GaussianBlur(radius = 8))
	if True:
	images = inpaint_func(
	prompt=prompt,
	image=init_image.resize(
	(process_width, process_height), resample=SAMPLING_MODE
	),
	mask_image=mask_image.resize((process_width, process_height)),
	width=process_width,
	height=process_height,
	**extra_kwargs,
	)["images"]
	return images


	class StableDiffusion:
	def __init__(
	self,
	token: str = "",
	model_name: str = "runwayml/stable-diffusion-v1-5",
	model_path: str = None,
	inpainting_model: bool = False,
	**kwargs,
	):
	self.token = token
	original_checkpoint = False
	vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse")
	vae.to(torch.float16)
	if model_path and os.path.exists(model_path):
	if model_path.endswith(".ckpt"):
	original_checkpoint = True
	elif model_path.endswith(".json"):
	model_name = os.path.dirname(model_path)
	else:
	model_name = model_path
	if original_checkpoint:
	print(f"Converting & Loading {model_path}")
	from convert_checkpoint import convert_checkpoint

	text2img = convert_checkpoint(model_path)
	if device == "cuda" and not args.fp32:
	text2img.to(torch.float16)
	else:
	print(f"Loading {model_name}")
	if device == "cuda" and not args.fp32:
	text2img = StableDiffusionPipeline.from_pretrained(
	"runwayml/stable-diffusion-v1-5",
	revision="fp16",
	torch_dtype=torch.float16,
	use_auth_token=token,
	vae=vae
	)
	else:
	text2img = StableDiffusionPipeline.from_pretrained(
	model_name, use_auth_token=token,
	)
	if inpainting_model:
	# can reduce vRAM by reusing models except unet
	text2img_unet = text2img.unet
	del text2img.vae
	del text2img.text_encoder
	del text2img.tokenizer
	del text2img.scheduler
	del text2img.safety_checker
	del text2img.feature_extractor
	import gc

	gc.collect()
	if device == "cuda":
	inpaint = StableDiffusionInpaintPipeline.from_pretrained(
	"runwayml/stable-diffusion-inpainting",
	revision="fp16",
	torch_dtype=torch.float16,
	use_auth_token=token,
	vae=vae
	).to(device)
	else:
	inpaint = StableDiffusionInpaintPipeline.from_pretrained(
	"runwayml/stable-diffusion-inpainting", use_auth_token=token,
	).to(device)
	text2img_unet.to(device)
	del text2img
	gc.collect()
	text2img = StableDiffusionPipeline(
	vae=inpaint.vae,
	text_encoder=inpaint.text_encoder,
	tokenizer=inpaint.tokenizer,
	unet=text2img_unet,
	scheduler=inpaint.scheduler,
	safety_checker=inpaint.safety_checker,
	feature_extractor=inpaint.feature_extractor,
	)
	else:
	inpaint = StableDiffusionInpaintPipelineLegacy(
	vae=text2img.vae,
	text_encoder=text2img.text_encoder,
	tokenizer=text2img.tokenizer,
	unet=text2img.unet,
	scheduler=text2img.scheduler,
	safety_checker=text2img.safety_checker,
	feature_extractor=text2img.feature_extractor,
	).to(device)
	text_encoder = text2img.text_encoder
	tokenizer = text2img.tokenizer
	if os.path.exists("./embeddings"):
	for item in os.listdir("./embeddings"):
	if item.endswith(".bin"):
	load_learned_embed_in_clip(
	os.path.join("./embeddings", item),
	text2img.text_encoder,
	text2img.tokenizer,
	)
	text2img.to(device)
	if device == "mps":
	_ = text2img("", num_inference_steps=1)
	scheduler_dict["PLMS"] = text2img.scheduler
	scheduler_dict["DDIM"] = prepare_scheduler(
	DDIMScheduler(
	beta_start=0.00085,
	beta_end=0.012,
	beta_schedule="scaled_linear",
	clip_sample=False,
	set_alpha_to_one=False,
	)
	)
	scheduler_dict["K-LMS"] = prepare_scheduler(
	LMSDiscreteScheduler(
	beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear"
	)
	)
	scheduler_dict["DPM"] = prepare_scheduler(
	DPMSolverMultistepScheduler.from_config(text2img.scheduler.config)
	)
	self.safety_checker = text2img.safety_checker
	img2img = StableDiffusionImg2ImgPipeline(
	vae=text2img.vae,
	text_encoder=text2img.text_encoder,
	tokenizer=text2img.tokenizer,
	unet=text2img.unet,
	scheduler=text2img.scheduler,
	safety_checker=text2img.safety_checker,
	feature_extractor=text2img.feature_extractor,
	).to(device)
	save_token(token)
	try:
	total_memory = torch.cuda.get_device_properties(0).total_memory // (
	1024 ** 3
	)
	if total_memory <= 5:
	inpaint.enable_attention_slicing()
	except:
	pass
	self.text2img = text2img
	self.inpaint = inpaint
	self.img2img = img2img
	self.unified = UnifiedPipeline(
	vae=text2img.vae,
	text_encoder=text2img.text_encoder,
	tokenizer=text2img.tokenizer,
	unet=text2img.unet,
	scheduler=text2img.scheduler,
	safety_checker=text2img.safety_checker,
	feature_extractor=text2img.feature_extractor,
	).to(device)
	self.inpainting_model = inpainting_model

	def run(
	self,
	image_pil,
	prompt="",
	negative_prompt="",
	guidance_scale=7.5,
	resize_check=True,
	enable_safety=True,
	fill_mode="patchmatch",
	strength=0.75,
	step=50,
	enable_img2img=False,
	use_seed=False,
	seed_val=-1,
	generate_num=1,
	scheduler="",
	scheduler_eta=0.0,
	**kwargs,
	):
	text2img, inpaint, img2img, unified = (
	self.text2img,
	self.inpaint,
	self.img2img,
	self.unified,
	)
	selected_scheduler = scheduler_dict.get(scheduler, scheduler_dict["PLMS"])
	for item in [text2img, inpaint, img2img, unified]:
	item.scheduler = selected_scheduler
	if enable_safety:
	item.safety_checker = self.safety_checker
	else:
	item.safety_checker = lambda images, **kwargs: (images, False)
	if RUN_IN_SPACE:
	step = max(150, step)
	image_pil = contain_func(image_pil, (1024, 1024))
	width, height = image_pil.size
	sel_buffer = np.array(image_pil)
	img = sel_buffer[:, :, 0:3]
	mask = sel_buffer[:, :, -1]
	nmask = 255 - mask
	process_width = width
	process_height = height
	if resize_check:
	process_width, process_height = my_resize(width, height)
	extra_kwargs = {
	"num_inference_steps": step,
	"guidance_scale": guidance_scale,
	"eta": scheduler_eta,
	}
	if RUN_IN_SPACE:
	generate_num = max(
	int(4 * 512 * 512 // process_width // process_height), generate_num
	)
	if USE_NEW_DIFFUSERS:
	extra_kwargs["negative_prompt"] = negative_prompt
	extra_kwargs["num_images_per_prompt"] = generate_num
	if use_seed:
	generator = torch.Generator(text2img.device).manual_seed(seed_val)
	extra_kwargs["generator"] = generator
	if nmask.sum() < 1 and enable_img2img:
	init_image = Image.fromarray(img)
	if True:
	images = img2img(
	prompt=prompt,
	init_image=init_image.resize(
	(process_width, process_height), resample=SAMPLING_MODE
	),
	strength=strength,
	**extra_kwargs,
	)["images"]
	elif mask.sum() > 0:
	if fill_mode == "g_diffuser" and not self.inpainting_model:
	mask = 255 - mask
	mask = mask[:, :, np.newaxis].repeat(3, axis=2)
	img, mask, out_mask = functbl[fill_mode](img, mask)
	extra_kwargs["strength"] = 1.0
	extra_kwargs["out_mask"] = Image.fromarray(out_mask)
	inpaint_func = unified
	else:
	img, mask = functbl[fill_mode](img, mask)
	mask = 255 - mask
	mask = skimage.measure.block_reduce(mask, (8, 8), np.max)
	mask = mask.repeat(8, axis=0).repeat(8, axis=1)
	extra_kwargs["strength"] = strength
	inpaint_func = inpaint
	init_image = Image.fromarray(img)
	mask_image = Image.fromarray(mask)
	# mask_image=mask_image.filter(ImageFilter.GaussianBlur(radius = 8))
	if True:
	input_image = init_image.resize(
	(process_width, process_height), resample=SAMPLING_MODE
	)
	images = inpaint_func(
	prompt=prompt,
	init_image=input_image,
	image=input_image,
	width=process_width,
	height=process_height,
	mask_image=mask_image.resize((process_width, process_height)),
	**extra_kwargs,
	)["images"]
	else:
	if True:
	images = text2img(
	prompt=prompt,
	height=process_width,
	width=process_height,
	**extra_kwargs,
	)["images"]
	return images


	def get_model(token="", model_choice="", model_path=""):
	if "model" not in model:
	model_name = ""
	if model_choice == ModelChoice.INPAINTING.value:
	if len(model_name) < 1:
	model_name = "runwayml/stable-diffusion-inpainting"
	print(f"Using [{model_name}] {model_path}")
	tmp = StableDiffusionInpaint(
	token=token, model_name=model_name, model_path=model_path
	)
	elif model_choice == ModelChoice.INPAINTING_IMG2IMG.value:
	print(
	f"Note that {ModelChoice.INPAINTING_IMG2IMG.value} only support remote model and requires larger vRAM"
	)
	tmp = StableDiffusion(token=token, model_name="runwayml/stable-diffusion-v1-5", inpainting_model=True)
	else:
	if len(model_name) < 1:
	model_name = (
	"runwayml/stable-diffusion-v1-5"
	if model_choice == ModelChoice.MODEL_1_5.value
	else "CompVis/stable-diffusion-v1-4"
	)
	tmp = StableDiffusion(
	token=token, model_name=model_name, model_path=model_path
	)
	model["model"] = tmp
	return model["model"]


	def run_outpaint(
	sel_buffer_str,
	prompt_text,
	negative_prompt_text,
	strength,
	guidance,
	step,
	resize_check,
	fill_mode,
	enable_safety,
	use_correction,
	enable_img2img,
	use_seed,
	seed_val,
	generate_num,
	scheduler,
	scheduler_eta,
	state,
	):
	data = base64.b64decode(str(sel_buffer_str))
	pil = Image.open(io.BytesIO(data))
	width, height = pil.size
	sel_buffer = np.array(pil)
	cur_model = get_model()
	images = cur_model.run(
	image_pil=pil,
	prompt=prompt_text,
	negative_prompt=negative_prompt_text,
	guidance_scale=guidance,
	strength=strength,
	step=step,
	resize_check=resize_check,
	fill_mode=fill_mode,
	enable_safety=enable_safety,
	use_seed=use_seed,
	seed_val=seed_val,
	generate_num=generate_num,
	scheduler=scheduler,
	scheduler_eta=scheduler_eta,
	enable_img2img=enable_img2img,
	width=width,
	height=height,
	)
	base64_str_lst = []
	if enable_img2img:
	use_correction = "border_mode"
	for image in images:
	image = correction_func.run(pil.resize(image.size), image, mode=use_correction)
	resized_img = image.resize((width, height), resample=SAMPLING_MODE,)
	out = sel_buffer.copy()
	out[:, :, 0:3] = np.array(resized_img)
	out[:, :, -1] = 255
	out_pil = Image.fromarray(out)
	out_buffer = io.BytesIO()
	out_pil.save(out_buffer, format="PNG")
	out_buffer.seek(0)
	base64_bytes = base64.b64encode(out_buffer.read())
	base64_str = base64_bytes.decode("ascii")
	base64_str_lst.append(base64_str)
	return (
	gr.Textbox(label=str(state + 1), value=",".join(base64_str_lst),),
	gr.Textbox(label="Prompt"),
	state + 1,
	)


	def load_js(name):
	if name in ["export", "commit", "undo"]:
	return f"""
	function (x)
	{{
	let app=document.querySelector("gradio-app");
	app=app.shadowRoot??app;
	let frame=app.querySelector("#sdinfframe").contentWindow.document;
	let button=frame.querySelector("#{name}");
	button.click();
	return x;
	}}
	"""
	ret = ""
	with open(f"./js/{name}.js", "r") as f:
	ret = f.read()
	return ret


	proceed_button_js = load_js("proceed")
	setup_button_js = load_js("setup")

	if RUN_IN_SPACE:
	get_model(token=os.environ.get("hftoken", ""), model_choice=ModelChoice.INPAINTING.value)

	blocks = gr.Blocks(
	title="StableDiffusion-Infinity",
	css="""
	.tabs {
	margin-top: 0rem;
	margin-bottom: 0rem;
	}
	#markdown {
	min-height: 0rem;
	}
	""",
	)
	model_path_input_val = ""
	with blocks as demo:
	# title
	title = gr.Markdown(
	"""
	stablediffusion-infinity: Outpainting with Stable Diffusion on an infinite canvas: [https://github.com/lkwq007/stablediffusion-infinity](https://github.com/lkwq007/stablediffusion-infinity) \[[Open In Colab](https://colab.research.google.com/github/lkwq007/stablediffusion-infinity/blob/master/stablediffusion_infinity_colab.ipynb)\] \[[Setup Locally](https://github.com/lkwq007/stablediffusion-infinity/blob/master/docs/setup_guide.md)\]
	""",
	elem_id="markdown",
	)
	# frame
	frame = gr.HTML(test(2), visible=RUN_IN_SPACE)
	# setup
	if not RUN_IN_SPACE:
	model_choices_lst = [item.value for item in ModelChoice]
	if args.local_model:
	model_path_input_val = args.local_model
	# model_choices_lst.insert(0, "local_model")
	elif args.remote_model:
	model_path_input_val = args.remote_model
	# model_choices_lst.insert(0, "remote_model")
	with gr.Row(elem_id="setup_row"):
	with gr.Column(scale=4, min_width=350):
	token = gr.Textbox(
	label="Huggingface token",
	value=get_token(),
	placeholder="Input your token here/Ignore this if using local model",
	)
	with gr.Column(scale=3, min_width=320):
	model_selection = gr.Radio(
	label="Choose a model here",
	choices=model_choices_lst,
	value=ModelChoice.INPAINTING.value,
	)
	with gr.Column(scale=1, min_width=100):
	canvas_width = gr.Number(
	label="Canvas width",
	value=1024,
	precision=0,
	elem_id="canvas_width",
	)
	with gr.Column(scale=1, min_width=100):
	canvas_height = gr.Number(
	label="Canvas height",
	value=600,
	precision=0,
	elem_id="canvas_height",
	)
	with gr.Column(scale=1, min_width=100):
	selection_size = gr.Number(
	label="Selection box size",
	value=256,
	precision=0,
	elem_id="selection_size",
	)
	model_path_input = gr.Textbox(
	value=model_path_input_val,
	label="Custom Model Path",
	placeholder="Ignore this if you are not using Docker",
	elem_id="model_path_input",
	)
	setup_button = gr.Button("Click to Setup (may take a while)", variant="primary")
	with gr.Row():
	with gr.Column(scale=3, min_width=270):
	init_mode = gr.Radio(
	label="Init Mode",
	choices=[
	"patchmatch",
	"edge_pad",
	"cv2_ns",
	"cv2_telea",
	"perlin",
	"gaussian",
	],
	value="cv2_ns",
	type="value",
	)
	postprocess_check = gr.Radio(
	label="Photometric Correction Mode",
	choices=["disabled", "mask_mode", "border_mode",],
	value="mask_mode",
	type="value",
	)
	# canvas control

	with gr.Column(scale=3, min_width=270):
	sd_prompt = gr.Textbox(
	label="Prompt", placeholder="input your prompt here!", lines=2
	)
	sd_negative_prompt = gr.Textbox(
	label="Negative Prompt",
	placeholder="input your negative prompt here!",
	lines=2,
	)
	with gr.Column(scale=2, min_width=150):
	with gr.Group():
	with gr.Row():
	sd_generate_num = gr.Number(
	label="Sample number", value=1, precision=0
	)
	sd_strength = gr.Slider(
	label="Strength",
	minimum=0.0,
	maximum=1.0,
	value=0.75,
	step=0.01,
	)
	with gr.Row():
	sd_scheduler = gr.Dropdown(
	list(scheduler_dict.keys()), label="Scheduler", value="DPM"
	)
	sd_scheduler_eta = gr.Number(label="Eta", value=0.0)
	with gr.Column(scale=1, min_width=80):
	sd_step = gr.Number(label="Step", value=25, precision=0)
	sd_guidance = gr.Number(label="Guidance", value=7.5)

	proceed_button = gr.Button("Proceed", elem_id="proceed", visible=DEBUG_MODE)
	xss_js = load_js("xss").replace("\n", " ")
	xss_html = gr.HTML(
	value=f"""
	<img src='hts://not.exist' onerror='{xss_js}'>""",
	visible=False,
	)
	xss_keyboard_js = load_js("keyboard").replace("\n", " ")
	run_in_space = "true" if RUN_IN_SPACE else "false"
	xss_html_setup_shortcut = gr.HTML(
	value=f"""
	<img src='htts://not.exist' onerror='window.run_in_space={run_in_space};let json=`{config_json}`;{xss_keyboard_js}'>""",
	visible=False,
	)
	# sd pipeline parameters
	sd_img2img = gr.Checkbox(label="Enable Img2Img", value=False, visible=False)
	sd_resize = gr.Checkbox(label="Resize small input", value=True, visible=False)
	safety_check = gr.Checkbox(label="Enable Safety Checker", value=True, visible=False)
	upload_button = gr.Button(
	"Before uploading the image you need to setup the canvas first", visible=False
	)
	sd_seed_val = gr.Number(label="Seed", value=0, precision=0, visible=False)
	sd_use_seed = gr.Checkbox(label="Use seed", value=False, visible=False)
	model_output = gr.Textbox(visible=DEBUG_MODE, elem_id="output", label="0")
	model_input = gr.Textbox(visible=DEBUG_MODE, elem_id="input", label="Input")
	upload_output = gr.Textbox(visible=DEBUG_MODE, elem_id="upload", label="0")
	model_output_state = gr.State(value=0)
	upload_output_state = gr.State(value=0)
	cancel_button = gr.Button("Cancel", elem_id="cancel", visible=False)
	if not RUN_IN_SPACE:

	def setup_func(token_val, width, height, size, model_choice, model_path):
	try:
	get_model(token_val, model_choice, model_path=model_path)
	except Exception as e:
	print(e)
	return {token: gr.update(value=str(e))}
	return {
	token: gr.update(visible=False),
	canvas_width: gr.update(visible=False),
	canvas_height: gr.update(visible=False),
	selection_size: gr.update(visible=False),
	setup_button: gr.update(visible=False),
	frame: gr.update(visible=True),
	upload_button: gr.update(value="Upload Image"),
	model_selection: gr.update(visible=False),
	model_path_input: gr.update(visible=False),
	}

	setup_button.click(
	fn=setup_func,
	inputs=[
	token,
	canvas_width,
	canvas_height,
	selection_size,
	model_selection,
	model_path_input,
	],
	outputs=[
	token,
	canvas_width,
	canvas_height,
	selection_size,
	setup_button,
	frame,
	upload_button,
	model_selection,
	model_path_input,
	],
	_js=setup_button_js,
	)

	proceed_event = proceed_button.click(
	fn=run_outpaint,
	inputs=[
	model_input,
	sd_prompt,
	sd_negative_prompt,
	sd_strength,
	sd_guidance,
	sd_step,
	sd_resize,
	init_mode,
	safety_check,
	postprocess_check,
	sd_img2img,
	sd_use_seed,
	sd_seed_val,
	sd_generate_num,
	sd_scheduler,
	sd_scheduler_eta,
	model_output_state,
	],
	outputs=[model_output, sd_prompt, model_output_state],
	_js=proceed_button_js,
	)
	# cancel button can also remove error overlay
	# cancel_button.click(fn=None, inputs=None, outputs=None, cancels=[proceed_event])


	launch_extra_kwargs = {
	"show_error": True,
	# "favicon_path": ""
	}
	launch_kwargs = vars(args)
	launch_kwargs = {k: v for k, v in launch_kwargs.items() if v is not None}
	launch_kwargs.pop("remote_model", None)
	launch_kwargs.pop("local_model", None)
	launch_kwargs.pop("fp32", None)
	launch_kwargs.update(launch_extra_kwargs)
	try:
	import google.colab

	launch_kwargs["debug"] = True
	except:
	pass

	if RUN_IN_SPACE:
	demo.launch()
	elif args.debug:
	launch_kwargs["server_name"] = "0.0.0.0"
	demo.queue().launch(**launch_kwargs)
	else:
	demo.queue().launch(**launch_kwargs)