lnyan/stablediffusion-infinity
807
1import subprocess2# import os.path as osp3import pip4# pip.main(["install","-v","-U","git+https://github.com/facebookresearch/xformers.git@main#egg=xformers"])5# subprocess.check_call("pip install -v -U git+https://github.com/facebookresearch/xformers.git@main#egg=xformers", cwd=osp.dirname(__file__), shell=True)6 7import io8import base649import os10import os11 12import sys13 14import numpy as np15import torch16from torch import autocast17import diffusers18from diffusers.configuration_utils import FrozenDict19from diffusers import (20 StableDiffusionPipeline,21 StableDiffusionInpaintPipeline,22 StableDiffusionImg2ImgPipeline,23 StableDiffusionInpaintPipelineLegacy,24 DDIMScheduler,25 LMSDiscreteScheduler,26 StableDiffusionUpscalePipeline,27 DPMSolverMultistepScheduler28)29from diffusers.models import AutoencoderKL30from PIL import Image31from PIL import ImageOps32import gradio as gr33import base6434import skimage35import skimage.measure36import yaml37import json38from enum import Enum39 40try:41 abspath = os.path.abspath(__file__)42 dirname = os.path.dirname(abspath)43 os.chdir(dirname)44except:45 pass46 47from utils import *48 49# assert diffusers.__version__ >= "0.6.0", "Please upgrade diffusers to 0.6.0"50 51USE_NEW_DIFFUSERS = True52RUN_IN_SPACE = "RUN_IN_HG_SPACE" in os.environ53 54 55class ModelChoice(Enum):56 INPAINTING = "stablediffusion-inpainting"57 INPAINTING_IMG2IMG = "stablediffusion-inpainting+img2img-v1.5"58 MODEL_1_5 = "stablediffusion-v1.5"59 MODEL_1_4 = "stablediffusion-v1.4"60 61 62try:63 from sd_grpcserver.pipeline.unified_pipeline import UnifiedPipeline64except:65 UnifiedPipeline = StableDiffusionInpaintPipeline66 67# sys.path.append("./glid_3_xl_stable")68 69USE_GLID = False70# try:71# from glid3xlmodel import GlidModel72# except:73# USE_GLID = False74 75try:76 cuda_available = torch.cuda.is_available()77except:78 cuda_available = False79finally:80 if sys.platform == "darwin":81 device = "mps" if torch.backends.mps.is_available() else "cpu"82 elif cuda_available:83 device = "cuda"84 else:85 device = "cpu"86 87import contextlib88 89autocast = contextlib.nullcontext90 91with open("config.yaml", "r") as yaml_in:92 yaml_object = yaml.safe_load(yaml_in)93 config_json = json.dumps(yaml_object)94 95 96def load_html():97 body, canvaspy = "", ""98 with open("index.html", encoding="utf8") as f:99 body = f.read()100 with open("canvas.py", encoding="utf8") as f:101 canvaspy = f.read()102 body = body.replace("- paths:\n", "")103 body = body.replace(" - ./canvas.py\n", "")104 body = body.replace("from canvas import InfCanvas", canvaspy)105 return body106 107 108def test(x):109 x = load_html()110 return f"""<iframe id="sdinfframe" style="width: 100%; height: 600px" name="result" allow="midi; geolocation; microphone; camera; 111 display-capture; encrypted-media; vertical-scroll 'none'" sandbox="allow-modals allow-forms 112 allow-scripts allow-same-origin allow-popups 113 allow-top-navigation-by-user-activation allow-downloads" allowfullscreen="" 114 allowpaymentrequest="" frameborder="0" srcdoc='{x}'></iframe>"""115 116 117DEBUG_MODE = False118 119try:120 SAMPLING_MODE = Image.Resampling.LANCZOS121except Exception as e:122 SAMPLING_MODE = Image.LANCZOS123 124try:125 contain_func = ImageOps.contain126except Exception as e:127 128 def contain_func(image, size, method=SAMPLING_MODE):129 # from PIL: https://pillow.readthedocs.io/en/stable/reference/ImageOps.html#PIL.ImageOps.contain130 im_ratio = image.width / image.height131 dest_ratio = size[0] / size[1]132 if im_ratio != dest_ratio:133 if im_ratio > dest_ratio:134 new_height = int(image.height / image.width * size[0])135 if new_height != size[1]:136 size = (size[0], new_height)137 else:138 new_width = int(image.width / image.height * size[1])139 if new_width != size[0]:140 size = (new_width, size[1])141 return image.resize(size, resample=method)142 143 144import argparse145 146parser = argparse.ArgumentParser(description="stablediffusion-infinity")147parser.add_argument("--port", type=int, help="listen port", dest="server_port")148parser.add_argument("--host", type=str, help="host", dest="server_name")149parser.add_argument("--share", action="store_true", help="share this app?")150parser.add_argument("--debug", action="store_true", help="debug mode")151parser.add_argument("--fp32", action="store_true", help="using full precision")152parser.add_argument("--encrypt", action="store_true", help="using https?")153parser.add_argument("--ssl_keyfile", type=str, help="path to ssl_keyfile")154parser.add_argument("--ssl_certfile", type=str, help="path to ssl_certfile")155parser.add_argument("--ssl_keyfile_password", type=str, help="ssl_keyfile_password")156parser.add_argument(157 "--auth", nargs=2, metavar=("username", "password"), help="use username password"158)159parser.add_argument(160 "--remote_model",161 type=str,162 help="use a model (e.g. dreambooth fined) from huggingface hub",163 default="",164)165parser.add_argument(166 "--local_model", type=str, help="use a model stored on your PC", default=""167)168 169if __name__ == "__main__" and not RUN_IN_SPACE:170 args = parser.parse_args()171else:172 args = parser.parse_args()173# args = parser.parse_args(["--debug"])174if args.auth is not None:175 args.auth = tuple(args.auth)176 177model = {}178 179 180def get_token():181 token = ""182 if os.path.exists(".token"):183 with open(".token", "r") as f:184 token = f.read()185 token = os.environ.get("hftoken", token)186 return token187 188 189def save_token(token):190 with open(".token", "w") as f:191 f.write(token)192 193 194def prepare_scheduler(scheduler):195 if hasattr(scheduler.config, "steps_offset") and scheduler.config.steps_offset != 1:196 new_config = dict(scheduler.config)197 new_config["steps_offset"] = 1198 scheduler._internal_dict = FrozenDict(new_config)199 return scheduler200 201 202def my_resize(width, height):203 if width >= 512 and height >= 512:204 return width, height205 if width == height:206 return 512, 512207 smaller = min(width, height)208 larger = max(width, height)209 if larger >= 608:210 return width, height211 factor = 1212 if smaller < 290:213 factor = 2214 elif smaller < 330:215 factor = 1.75216 elif smaller < 384:217 factor = 1.375218 elif smaller < 400:219 factor = 1.25220 elif smaller < 450:221 factor = 1.125222 return int(factor * width)//8*8, int(factor * height)//8*8223 224 225def load_learned_embed_in_clip(226 learned_embeds_path, text_encoder, tokenizer, token=None227):228 # https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers/stable_conceptualizer_inference.ipynb229 loaded_learned_embeds = torch.load(learned_embeds_path, map_location="cpu")230 231 # separate token and the embeds232 trained_token = list(loaded_learned_embeds.keys())[0]233 embeds = loaded_learned_embeds[trained_token]234 235 # cast to dtype of text_encoder236 dtype = text_encoder.get_input_embeddings().weight.dtype237 embeds.to(dtype)238 239 # add the token in tokenizer240 token = token if token is not None else trained_token241 num_added_tokens = tokenizer.add_tokens(token)242 if num_added_tokens == 0:243 raise ValueError(244 f"The tokenizer already contains the token {token}. Please pass a different `token` that is not already in the tokenizer."245 )246 247 # resize the token embeddings248 text_encoder.resize_token_embeddings(len(tokenizer))249 250 # get the id for the token and assign the embeds251 token_id = tokenizer.convert_tokens_to_ids(token)252 text_encoder.get_input_embeddings().weight.data[token_id] = embeds253 254 255scheduler_dict = {"PLMS": None, "DDIM": None, "K-LMS": None, "DPM": None}256 257 258class StableDiffusionInpaint:259 def __init__(260 self, token: str = "", model_name: str = "", model_path: str = "", **kwargs,261 ):262 self.token = token263 original_checkpoint = False264 if model_path and os.path.exists(model_path):265 if model_path.endswith(".ckpt"):266 original_checkpoint = True267 elif model_path.endswith(".json"):268 model_name = os.path.dirname(model_path)269 else:270 model_name = model_path271 vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse")272 vae.to(torch.float16)273 if original_checkpoint:274 print(f"Converting & Loading {model_path}")275 from convert_checkpoint import convert_checkpoint276 277 pipe = convert_checkpoint(model_path, inpainting=True)278 if device == "cuda":279 pipe.to(torch.float16)280 inpaint = StableDiffusionInpaintPipeline(281 vae=vae,282 text_encoder=pipe.text_encoder,283 tokenizer=pipe.tokenizer,284 unet=pipe.unet,285 scheduler=pipe.scheduler,286 safety_checker=pipe.safety_checker,287 feature_extractor=pipe.feature_extractor,288 )289 else:290 print(f"Loading {model_name}")291 if device == "cuda":292 inpaint = StableDiffusionInpaintPipeline.from_pretrained(293 model_name,294 revision="fp16",295 torch_dtype=torch.float16,296 use_auth_token=token,297 vae=vae298 )299 else:300 inpaint = StableDiffusionInpaintPipeline.from_pretrained(301 model_name, use_auth_token=token,302 )303 if os.path.exists("./embeddings"):304 print("Note that StableDiffusionInpaintPipeline + embeddings is untested")305 for item in os.listdir("./embeddings"):306 if item.endswith(".bin"):307 load_learned_embed_in_clip(308 os.path.join("./embeddings", item),309 inpaint.text_encoder,310 inpaint.tokenizer,311 )312 inpaint.to(device)313 # try:314 # inpaint.vae=torch.compile(inpaint.vae, dynamic=True)315 # inpaint.unet=torch.compile(inpaint.unet, dynamic=True)316 # except Exception as e:317 # print(e)318 # inpaint.enable_xformers_memory_efficient_attention()319 # if device == "mps":320 # _ = text2img("", num_inference_steps=1)321 scheduler_dict["PLMS"] = inpaint.scheduler322 scheduler_dict["DDIM"] = prepare_scheduler(323 DDIMScheduler(324 beta_start=0.00085,325 beta_end=0.012,326 beta_schedule="scaled_linear",327 clip_sample=False,328 set_alpha_to_one=False,329 )330 )331 scheduler_dict["K-LMS"] = prepare_scheduler(332 LMSDiscreteScheduler(333 beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear"334 )335 )336 scheduler_dict["DPM"] = prepare_scheduler(337 DPMSolverMultistepScheduler.from_config(inpaint.scheduler.config)338 )339 self.safety_checker = inpaint.safety_checker340 save_token(token)341 try:342 total_memory = torch.cuda.get_device_properties(0).total_memory // (343 1024 ** 3344 )345 if total_memory <= 5:346 inpaint.enable_attention_slicing()347 except:348 pass349 self.inpaint = inpaint350 351 def run(352 self,353 image_pil,354 prompt="",355 negative_prompt="",356 guidance_scale=7.5,357 resize_check=True,358 enable_safety=True,359 fill_mode="patchmatch",360 strength=0.75,361 step=50,362 enable_img2img=False,363 use_seed=False,364 seed_val=-1,365 generate_num=1,366 scheduler="",367 scheduler_eta=0.0,368 **kwargs,369 ):370 inpaint = self.inpaint371 selected_scheduler = scheduler_dict.get(scheduler, scheduler_dict["PLMS"])372 for item in [inpaint]:373 item.scheduler = selected_scheduler374 if enable_safety:375 item.safety_checker = self.safety_checker376 else:377 item.safety_checker = lambda images, **kwargs: (images, None)378 width, height = image_pil.size379 sel_buffer = np.array(image_pil)380 img = sel_buffer[:, :, 0:3]381 mask = sel_buffer[:, :, -1]382 nmask = 255 - mask383 process_width = width384 process_height = height385 if resize_check:386 process_width, process_height = my_resize(width, height)387 process_width=process_width*8//8388 process_height=process_height*8//8389 extra_kwargs = {390 "num_inference_steps": step,391 "guidance_scale": guidance_scale,392 "eta": scheduler_eta,393 }394 if USE_NEW_DIFFUSERS:395 extra_kwargs["negative_prompt"] = negative_prompt396 extra_kwargs["num_images_per_prompt"] = generate_num397 if use_seed:398 generator = torch.Generator(inpaint.device).manual_seed(seed_val)399 extra_kwargs["generator"] = generator400 if True:401 img, mask = functbl[fill_mode](img, mask)402 mask = 255 - mask403 mask = skimage.measure.block_reduce(mask, (8, 8), np.max)404 mask = mask.repeat(8, axis=0).repeat(8, axis=1)405 extra_kwargs["strength"] = strength406 inpaint_func = inpaint407 init_image = Image.fromarray(img)408 mask_image = Image.fromarray(mask)409 # mask_image=mask_image.filter(ImageFilter.GaussianBlur(radius = 8))410 if True:411 images = inpaint_func(412 prompt=prompt,413 image=init_image.resize(414 (process_width, process_height), resample=SAMPLING_MODE415 ),416 mask_image=mask_image.resize((process_width, process_height)),417 width=process_width,418 height=process_height,419 **extra_kwargs,420 )["images"]421 return images422 423 424class StableDiffusion:425 def __init__(426 self,427 token: str = "",428 model_name: str = "runwayml/stable-diffusion-v1-5",429 model_path: str = None,430 inpainting_model: bool = False,431 **kwargs,432 ):433 self.token = token434 original_checkpoint = False435 vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse")436 vae.to(torch.float16)437 if model_path and os.path.exists(model_path):438 if model_path.endswith(".ckpt"):439 original_checkpoint = True440 elif model_path.endswith(".json"):441 model_name = os.path.dirname(model_path)442 else:443 model_name = model_path444 if original_checkpoint:445 print(f"Converting & Loading {model_path}")446 from convert_checkpoint import convert_checkpoint447 448 text2img = convert_checkpoint(model_path)449 if device == "cuda" and not args.fp32:450 text2img.to(torch.float16)451 else:452 print(f"Loading {model_name}")453 if device == "cuda" and not args.fp32:454 text2img = StableDiffusionPipeline.from_pretrained(455 "runwayml/stable-diffusion-v1-5",456 revision="fp16",457 torch_dtype=torch.float16,458 use_auth_token=token,459 vae=vae460 )461 else:462 text2img = StableDiffusionPipeline.from_pretrained(463 model_name, use_auth_token=token,464 )465 if inpainting_model:466 # can reduce vRAM by reusing models except unet467 text2img_unet = text2img.unet468 del text2img.vae469 del text2img.text_encoder470 del text2img.tokenizer471 del text2img.scheduler472 del text2img.safety_checker473 del text2img.feature_extractor474 import gc475 476 gc.collect()477 if device == "cuda":478 inpaint = StableDiffusionInpaintPipeline.from_pretrained(479 "runwayml/stable-diffusion-inpainting",480 revision="fp16",481 torch_dtype=torch.float16,482 use_auth_token=token,483 vae=vae484 ).to(device)485 else:486 inpaint = StableDiffusionInpaintPipeline.from_pretrained(487 "runwayml/stable-diffusion-inpainting", use_auth_token=token,488 ).to(device)489 text2img_unet.to(device)490 del text2img491 gc.collect()492 text2img = StableDiffusionPipeline(493 vae=inpaint.vae,494 text_encoder=inpaint.text_encoder,495 tokenizer=inpaint.tokenizer,496 unet=text2img_unet,497 scheduler=inpaint.scheduler,498 safety_checker=inpaint.safety_checker,499 feature_extractor=inpaint.feature_extractor,500 )501 else:502 inpaint = StableDiffusionInpaintPipelineLegacy(503 vae=text2img.vae,504 text_encoder=text2img.text_encoder,505 tokenizer=text2img.tokenizer,506 unet=text2img.unet,507 scheduler=text2img.scheduler,508 safety_checker=text2img.safety_checker,509 feature_extractor=text2img.feature_extractor,510 ).to(device)511 text_encoder = text2img.text_encoder512 tokenizer = text2img.tokenizer513 if os.path.exists("./embeddings"):514 for item in os.listdir("./embeddings"):515 if item.endswith(".bin"):516 load_learned_embed_in_clip(517 os.path.join("./embeddings", item),518 text2img.text_encoder,519 text2img.tokenizer,520 )521 text2img.to(device)522 if device == "mps":523 _ = text2img("", num_inference_steps=1)524 scheduler_dict["PLMS"] = text2img.scheduler525 scheduler_dict["DDIM"] = prepare_scheduler(526 DDIMScheduler(527 beta_start=0.00085,528 beta_end=0.012,529 beta_schedule="scaled_linear",530 clip_sample=False,531 set_alpha_to_one=False,532 )533 )534 scheduler_dict["K-LMS"] = prepare_scheduler(535 LMSDiscreteScheduler(536 beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear"537 )538 )539 scheduler_dict["DPM"] = prepare_scheduler(540 DPMSolverMultistepScheduler.from_config(text2img.scheduler.config)541 )542 self.safety_checker = text2img.safety_checker543 img2img = StableDiffusionImg2ImgPipeline(544 vae=text2img.vae,545 text_encoder=text2img.text_encoder,546 tokenizer=text2img.tokenizer,547 unet=text2img.unet,548 scheduler=text2img.scheduler,549 safety_checker=text2img.safety_checker,550 feature_extractor=text2img.feature_extractor,551 ).to(device)552 save_token(token)553 try:554 total_memory = torch.cuda.get_device_properties(0).total_memory // (555 1024 ** 3556 )557 if total_memory <= 5:558 inpaint.enable_attention_slicing()559 except:560 pass561 self.text2img = text2img562 self.inpaint = inpaint563 self.img2img = img2img564 self.unified = UnifiedPipeline(565 vae=text2img.vae,566 text_encoder=text2img.text_encoder,567 tokenizer=text2img.tokenizer,568 unet=text2img.unet,569 scheduler=text2img.scheduler,570 safety_checker=text2img.safety_checker,571 feature_extractor=text2img.feature_extractor,572 ).to(device)573 self.inpainting_model = inpainting_model574 575 def run(576 self,577 image_pil,578 prompt="",579 negative_prompt="",580 guidance_scale=7.5,581 resize_check=True,582 enable_safety=True,583 fill_mode="patchmatch",584 strength=0.75,585 step=50,586 enable_img2img=False,587 use_seed=False,588 seed_val=-1,589 generate_num=1,590 scheduler="",591 scheduler_eta=0.0,592 **kwargs,593 ):594 text2img, inpaint, img2img, unified = (595 self.text2img,596 self.inpaint,597 self.img2img,598 self.unified,599 )600 selected_scheduler = scheduler_dict.get(scheduler, scheduler_dict["PLMS"])601 for item in [text2img, inpaint, img2img, unified]:602 item.scheduler = selected_scheduler603 if enable_safety:604 item.safety_checker = self.safety_checker605 else:606 item.safety_checker = lambda images, **kwargs: (images, False)607 if RUN_IN_SPACE:608 step = max(150, step)609 image_pil = contain_func(image_pil, (1024, 1024))610 width, height = image_pil.size611 sel_buffer = np.array(image_pil)612 img = sel_buffer[:, :, 0:3]613 mask = sel_buffer[:, :, -1]614 nmask = 255 - mask615 process_width = width616 process_height = height617 if resize_check:618 process_width, process_height = my_resize(width, height)619 extra_kwargs = {620 "num_inference_steps": step,621 "guidance_scale": guidance_scale,622 "eta": scheduler_eta,623 }624 if RUN_IN_SPACE:625 generate_num = max(626 int(4 * 512 * 512 // process_width // process_height), generate_num627 )628 if USE_NEW_DIFFUSERS:629 extra_kwargs["negative_prompt"] = negative_prompt630 extra_kwargs["num_images_per_prompt"] = generate_num631 if use_seed:632 generator = torch.Generator(text2img.device).manual_seed(seed_val)633 extra_kwargs["generator"] = generator634 if nmask.sum() < 1 and enable_img2img:635 init_image = Image.fromarray(img)636 if True:637 images = img2img(638 prompt=prompt,639 init_image=init_image.resize(640 (process_width, process_height), resample=SAMPLING_MODE641 ),642 strength=strength,643 **extra_kwargs,644 )["images"]645 elif mask.sum() > 0:646 if fill_mode == "g_diffuser" and not self.inpainting_model:647 mask = 255 - mask648 mask = mask[:, :, np.newaxis].repeat(3, axis=2)649 img, mask, out_mask = functbl[fill_mode](img, mask)650 extra_kwargs["strength"] = 1.0651 extra_kwargs["out_mask"] = Image.fromarray(out_mask)652 inpaint_func = unified653 else:654 img, mask = functbl[fill_mode](img, mask)655 mask = 255 - mask656 mask = skimage.measure.block_reduce(mask, (8, 8), np.max)657 mask = mask.repeat(8, axis=0).repeat(8, axis=1)658 extra_kwargs["strength"] = strength659 inpaint_func = inpaint660 init_image = Image.fromarray(img)661 mask_image = Image.fromarray(mask)662 # mask_image=mask_image.filter(ImageFilter.GaussianBlur(radius = 8))663 if True:664 input_image = init_image.resize(665 (process_width, process_height), resample=SAMPLING_MODE666 )667 images = inpaint_func(668 prompt=prompt,669 init_image=input_image,670 image=input_image,671 width=process_width,672 height=process_height,673 mask_image=mask_image.resize((process_width, process_height)),674 **extra_kwargs,675 )["images"]676 else:677 if True:678 images = text2img(679 prompt=prompt,680 height=process_width,681 width=process_height,682 **extra_kwargs,683 )["images"]684 return images685 686 687def get_model(token="", model_choice="", model_path=""):688 if "model" not in model:689 model_name = ""690 if model_choice == ModelChoice.INPAINTING.value:691 if len(model_name) < 1:692 model_name = "runwayml/stable-diffusion-inpainting"693 print(f"Using [{model_name}] {model_path}")694 tmp = StableDiffusionInpaint(695 token=token, model_name=model_name, model_path=model_path696 )697 elif model_choice == ModelChoice.INPAINTING_IMG2IMG.value:698 print(699 f"Note that {ModelChoice.INPAINTING_IMG2IMG.value} only support remote model and requires larger vRAM"700 )701 tmp = StableDiffusion(token=token, model_name="runwayml/stable-diffusion-v1-5", inpainting_model=True)702 else:703 if len(model_name) < 1:704 model_name = (705 "runwayml/stable-diffusion-v1-5"706 if model_choice == ModelChoice.MODEL_1_5.value707 else "CompVis/stable-diffusion-v1-4"708 )709 tmp = StableDiffusion(710 token=token, model_name=model_name, model_path=model_path711 )712 model["model"] = tmp713 return model["model"]714 715 716def run_outpaint(717 sel_buffer_str,718 prompt_text,719 negative_prompt_text,720 strength,721 guidance,722 step,723 resize_check,724 fill_mode,725 enable_safety,726 use_correction,727 enable_img2img,728 use_seed,729 seed_val,730 generate_num,731 scheduler,732 scheduler_eta,733 state,734):735 data = base64.b64decode(str(sel_buffer_str))736 pil = Image.open(io.BytesIO(data))737 width, height = pil.size738 sel_buffer = np.array(pil)739 cur_model = get_model()740 images = cur_model.run(741 image_pil=pil,742 prompt=prompt_text,743 negative_prompt=negative_prompt_text,744 guidance_scale=guidance,745 strength=strength,746 step=step,747 resize_check=resize_check,748 fill_mode=fill_mode,749 enable_safety=enable_safety,750 use_seed=use_seed,751 seed_val=seed_val,752 generate_num=generate_num,753 scheduler=scheduler,754 scheduler_eta=scheduler_eta,755 enable_img2img=enable_img2img,756 width=width,757 height=height,758 )759 base64_str_lst = []760 if enable_img2img:761 use_correction = "border_mode"762 for image in images:763 image = correction_func.run(pil.resize(image.size), image, mode=use_correction)764 resized_img = image.resize((width, height), resample=SAMPLING_MODE,)765 out = sel_buffer.copy()766 out[:, :, 0:3] = np.array(resized_img)767 out[:, :, -1] = 255768 out_pil = Image.fromarray(out)769 out_buffer = io.BytesIO()770 out_pil.save(out_buffer, format="PNG")771 out_buffer.seek(0)772 base64_bytes = base64.b64encode(out_buffer.read())773 base64_str = base64_bytes.decode("ascii")774 base64_str_lst.append(base64_str)775 return (776 gr.Textbox(label=str(state + 1), value=",".join(base64_str_lst),),777 gr.Textbox(label="Prompt"),778 state + 1,779 )780 781 782def load_js(name):783 if name in ["export", "commit", "undo"]:784 return f"""785function (x)786{{ 787 let app=document.querySelector("gradio-app");788 app=app.shadowRoot??app;789 let frame=app.querySelector("#sdinfframe").contentWindow.document;790 let button=frame.querySelector("#{name}");791 button.click();792 return x;793}}794"""795 ret = ""796 with open(f"./js/{name}.js", "r") as f:797 ret = f.read()798 return ret799 800 801proceed_button_js = load_js("proceed")802setup_button_js = load_js("setup")803 804if RUN_IN_SPACE:805 get_model(token=os.environ.get("hftoken", ""), model_choice=ModelChoice.INPAINTING.value)806 807blocks = gr.Blocks(808 title="StableDiffusion-Infinity",809 css="""810.tabs {811margin-top: 0rem;812margin-bottom: 0rem;813}814#markdown {815min-height: 0rem;816}817""",818)819model_path_input_val = ""820with blocks as demo:821 # title822 title = gr.Markdown(823 """824 **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)\] 825 """,826 elem_id="markdown",827 )828 # frame829 frame = gr.HTML(test(2), visible=RUN_IN_SPACE)830 # setup831 if not RUN_IN_SPACE:832 model_choices_lst = [item.value for item in ModelChoice]833 if args.local_model:834 model_path_input_val = args.local_model835 # model_choices_lst.insert(0, "local_model")836 elif args.remote_model:837 model_path_input_val = args.remote_model838 # model_choices_lst.insert(0, "remote_model")839 with gr.Row(elem_id="setup_row"):840 with gr.Column(scale=4, min_width=350):841 token = gr.Textbox(842 label="Huggingface token",843 value=get_token(),844 placeholder="Input your token here/Ignore this if using local model",845 )846 with gr.Column(scale=3, min_width=320):847 model_selection = gr.Radio(848 label="Choose a model here",849 choices=model_choices_lst,850 value=ModelChoice.INPAINTING.value,851 )852 with gr.Column(scale=1, min_width=100):853 canvas_width = gr.Number(854 label="Canvas width",855 value=1024,856 precision=0,857 elem_id="canvas_width",858 )859 with gr.Column(scale=1, min_width=100):860 canvas_height = gr.Number(861 label="Canvas height",862 value=600,863 precision=0,864 elem_id="canvas_height",865 )866 with gr.Column(scale=1, min_width=100):867 selection_size = gr.Number(868 label="Selection box size",869 value=256,870 precision=0,871 elem_id="selection_size",872 )873 model_path_input = gr.Textbox(874 value=model_path_input_val,875 label="Custom Model Path",876 placeholder="Ignore this if you are not using Docker",877 elem_id="model_path_input",878 )879 setup_button = gr.Button("Click to Setup (may take a while)", variant="primary")880 with gr.Row():881 with gr.Column(scale=3, min_width=270):882 init_mode = gr.Radio(883 label="Init Mode",884 choices=[885 "patchmatch",886 "edge_pad",887 "cv2_ns",888 "cv2_telea",889 "perlin",890 "gaussian",891 ],892 value="cv2_ns",893 type="value",894 )895 postprocess_check = gr.Radio(896 label="Photometric Correction Mode",897 choices=["disabled", "mask_mode", "border_mode",],898 value="mask_mode",899 type="value",900 )901 # canvas control902 903 with gr.Column(scale=3, min_width=270):904 sd_prompt = gr.Textbox(905 label="Prompt", placeholder="input your prompt here!", lines=2906 )907 sd_negative_prompt = gr.Textbox(908 label="Negative Prompt",909 placeholder="input your negative prompt here!",910 lines=2,911 )912 with gr.Column(scale=2, min_width=150):913 with gr.Group():914 with gr.Row():915 sd_generate_num = gr.Number(916 label="Sample number", value=1, precision=0917 )918 sd_strength = gr.Slider(919 label="Strength",920 minimum=0.0,921 maximum=1.0,922 value=0.75,923 step=0.01,924 )925 with gr.Row():926 sd_scheduler = gr.Dropdown(927 list(scheduler_dict.keys()), label="Scheduler", value="DPM"928 )929 sd_scheduler_eta = gr.Number(label="Eta", value=0.0)930 with gr.Column(scale=1, min_width=80):931 sd_step = gr.Number(label="Step", value=25, precision=0)932 sd_guidance = gr.Number(label="Guidance", value=7.5)933 934 proceed_button = gr.Button("Proceed", elem_id="proceed", visible=DEBUG_MODE)935 xss_js = load_js("xss").replace("\n", " ")936 xss_html = gr.HTML(937 value=f"""938 <img src='hts://not.exist' onerror='{xss_js}'>""",939 visible=False,940 )941 xss_keyboard_js = load_js("keyboard").replace("\n", " ")942 run_in_space = "true" if RUN_IN_SPACE else "false"943 xss_html_setup_shortcut = gr.HTML(944 value=f"""945 <img src='htts://not.exist' onerror='window.run_in_space={run_in_space};let json=`{config_json}`;{xss_keyboard_js}'>""",946 visible=False,947 )948 # sd pipeline parameters949 sd_img2img = gr.Checkbox(label="Enable Img2Img", value=False, visible=False)950 sd_resize = gr.Checkbox(label="Resize small input", value=True, visible=False)951 safety_check = gr.Checkbox(label="Enable Safety Checker", value=True, visible=False)952 upload_button = gr.Button(953 "Before uploading the image you need to setup the canvas first", visible=False954 )955 sd_seed_val = gr.Number(label="Seed", value=0, precision=0, visible=False)956 sd_use_seed = gr.Checkbox(label="Use seed", value=False, visible=False)957 model_output = gr.Textbox(visible=DEBUG_MODE, elem_id="output", label="0")958 model_input = gr.Textbox(visible=DEBUG_MODE, elem_id="input", label="Input")959 upload_output = gr.Textbox(visible=DEBUG_MODE, elem_id="upload", label="0")960 model_output_state = gr.State(value=0)961 upload_output_state = gr.State(value=0)962 cancel_button = gr.Button("Cancel", elem_id="cancel", visible=False)963 if not RUN_IN_SPACE:964 965 def setup_func(token_val, width, height, size, model_choice, model_path):966 try:967 get_model(token_val, model_choice, model_path=model_path)968 except Exception as e:969 print(e)970 return {token: gr.update(value=str(e))}971 return {972 token: gr.update(visible=False),973 canvas_width: gr.update(visible=False),974 canvas_height: gr.update(visible=False),975 selection_size: gr.update(visible=False),976 setup_button: gr.update(visible=False),977 frame: gr.update(visible=True),978 upload_button: gr.update(value="Upload Image"),979 model_selection: gr.update(visible=False),980 model_path_input: gr.update(visible=False),981 }982 983 setup_button.click(984 fn=setup_func,985 inputs=[986 token,987 canvas_width,988 canvas_height,989 selection_size,990 model_selection,991 model_path_input,992 ],993 outputs=[994 token,995 canvas_width,996 canvas_height,997 selection_size,998 setup_button,999 frame,1000 upload_button,1001 model_selection,1002 model_path_input,1003 ],1004 _js=setup_button_js,1005 )1006 1007 proceed_event = proceed_button.click(1008 fn=run_outpaint,1009 inputs=[1010 model_input,1011 sd_prompt,1012 sd_negative_prompt,1013 sd_strength,1014 sd_guidance,1015 sd_step,1016 sd_resize,1017 init_mode,1018 safety_check,1019 postprocess_check,1020 sd_img2img,1021 sd_use_seed,1022 sd_seed_val,1023 sd_generate_num,1024 sd_scheduler,1025 sd_scheduler_eta,1026 model_output_state,1027 ],1028 outputs=[model_output, sd_prompt, model_output_state],1029 _js=proceed_button_js,1030 )1031 # cancel button can also remove error overlay1032 # cancel_button.click(fn=None, inputs=None, outputs=None, cancels=[proceed_event])1033 1034 1035launch_extra_kwargs = {1036 "show_error": True,1037 # "favicon_path": ""1038}1039launch_kwargs = vars(args)1040launch_kwargs = {k: v for k, v in launch_kwargs.items() if v is not None}1041launch_kwargs.pop("remote_model", None)1042launch_kwargs.pop("local_model", None)1043launch_kwargs.pop("fp32", None)1044launch_kwargs.update(launch_extra_kwargs)1045try:1046 import google.colab1047 1048 launch_kwargs["debug"] = True1049except:1050 pass1051 1052if RUN_IN_SPACE:1053 demo.launch()1054elif args.debug:1055 launch_kwargs["server_name"] = "0.0.0.0"1056 demo.queue().launch(**launch_kwargs)1057else:1058 demo.queue().launch(**launch_kwargs)1059 1060