TSE1966/Super-Resolution-Anime-Diffusion
0
1import os2import random3import zipfile4import findfile5import PIL.Image6import autocuda7from pyabsa.utils.pyabsa_utils import fprint8 9try:10 for z_file in findfile.find_cwd_files(and_key=['.zip'],11 exclude_key=['.ignore', 'git', 'SuperResolutionAnimeDiffusion'],12 recursive=10):13 fprint(f"Extracting {z_file}...")14 with zipfile.ZipFile(z_file, 'r') as zip_ref:15 zip_ref.extractall(os.path.dirname(z_file))16except Exception as e:17 os.system('unzip random_examples.zip')18 19from diffusers import (20 AutoencoderKL,21 UNet2DConditionModel,22 StableDiffusionPipeline,23 StableDiffusionImg2ImgPipeline,24 DPMSolverMultistepScheduler,25)26import gradio as gr27import torch28from PIL import Image29import utils30import datetime31import time32import psutil33from Waifu2x.magnify import ImageMagnifier34from RealESRGANv030.interface import realEsrgan35 36magnifier = ImageMagnifier()37 38start_time = time.time()39is_colab = utils.is_google_colab()40 41CUDA_VISIBLE_DEVICES = ""42device = autocuda.auto_cuda()43 44dtype = torch.float16 if device != "cpu" else torch.float3245 46 47 48class Model:49 def __init__(self, name, path="", prefix=""):50 self.name = name51 self.path = path52 self.prefix = prefix53 self.pipe_t2i = None54 self.pipe_i2i = None55 56 57models = [58 # Model("anything v3", "Linaqruf/anything-v3.0", "anything v3 style"),59 Model("anything v5", "stablediffusionapi/anything-v5", "anything v5 style"),60]61# Model("Spider-Verse", "nitrosocke/spider-verse-diffusion", "spiderverse style "),62# Model("Balloon Art", "Fictiverse/Stable_Diffusion_BalloonArt_Model", "BalloonArt "),63# Model("Elden Ring", "nitrosocke/elden-ring-diffusion", "elden ring style "),64# Model("Tron Legacy", "dallinmackay/Tron-Legacy-diffusion", "trnlgcy ")65# Model("Pokémon", "lambdalabs/sd-pokemon-diffusers", ""),66# Model("Pony Diffusion", "AstraliteHeart/pony-diffusion", ""),67# Model("Robo Diffusion", "nousr/robo-diffusion", ""),68 69scheduler = DPMSolverMultistepScheduler(70 beta_start=0.00085,71 beta_end=0.012,72 beta_schedule="scaled_linear",73 num_train_timesteps=1000,74 trained_betas=None,75 predict_epsilon=True,76 thresholding=False,77 algorithm_type="dpmsolver++",78 solver_type="midpoint",79 solver_order=2,80 # lower_order_final=True,81)82 83custom_model = None84if is_colab:85 models.insert(0, Model("Custom model"))86 custom_model = models[0]87 88last_mode = "txt2img"89current_model = models[1] if is_colab else models[0]90current_model_path = current_model.path91 92if is_colab:93 pipe = StableDiffusionPipeline.from_pretrained(94 current_model.path,95 torch_dtype=dtype,96 scheduler=scheduler,97 safety_checker=lambda images, clip_input: (images, False),98 )99 100else: # download all models101 print(f"{datetime.datetime.now()} Downloading vae...")102 vae = AutoencoderKL.from_pretrained(103 current_model.path, subfolder="vae", torch_dtype=dtype104 )105 for model in models:106 try:107 print(f"{datetime.datetime.now()} Downloading {model.name} model...")108 unet = UNet2DConditionModel.from_pretrained(109 model.path, subfolder="unet", torch_dtype=dtype110 )111 model.pipe_t2i = StableDiffusionPipeline.from_pretrained(112 model.path,113 unet=unet,114 vae=vae,115 torch_dtype=dtype,116 scheduler=scheduler,117 safety_checker=None,118 )119 model.pipe_i2i = StableDiffusionImg2ImgPipeline.from_pretrained(120 model.path,121 unet=unet,122 vae=vae,123 torch_dtype=dtype,124 scheduler=scheduler,125 safety_checker=None,126 )127 except Exception as e:128 print(129 f"{datetime.datetime.now()} Failed to load model "130 + model.name131 + ": "132 + str(e)133 )134 models.remove(model)135 pipe = models[0].pipe_t2i136 137# model.pipe_i2i = torch.compile(model.pipe_i2i)138# model.pipe_t2i = torch.compile(model.pipe_t2i)139if torch.cuda.is_available():140 pipe = pipe.to(device)141 142 143# device = "GPU 🔥" if torch.cuda.is_available() else "CPU 🥶"144 145 146def error_str(error, title="Error"):147 return (148 f"""#### {title}149 {error}"""150 if error151 else ""152 )153 154 155def custom_model_changed(path):156 models[0].path = path157 global current_model158 current_model = models[0]159 160 161def on_model_change(model_name):162 prefix = (163 'Enter prompt. "'164 + next((m.prefix for m in models if m.name == model_name), None)165 + '" is prefixed automatically'166 if model_name != models[0].name167 else "Don't forget to use the custom model prefix in the prompt!"168 )169 170 return (171 gr.update(visible=model_name == models[0].name),172 gr.update(placeholder=prefix),173 )174 175 176def inference(177 model_name,178 prompt,179 guidance,180 steps,181 width=512,182 height=512,183 seed=0,184 img=None,185 strength=0.5,186 neg_prompt="",187 scale="ESRGAN4x",188 scale_factor=2,189):190 fprint(psutil.virtual_memory()) # print memory usage191 192 fprint(f"Prompt: {prompt}")193 global current_model194 for model in models:195 if model.name == model_name:196 current_model = model197 model_path = current_model.path198 199 generator = torch.Generator(device).manual_seed(seed) if seed != 0 else None200 201 try:202 if img is not None:203 return (204 img_to_img(205 model_path,206 prompt,207 neg_prompt,208 img,209 strength,210 guidance,211 steps,212 width,213 height,214 generator,215 scale,216 scale_factor,217 ),218 None,219 )220 else:221 return (222 txt_to_img(223 model_path,224 prompt,225 neg_prompt,226 guidance,227 steps,228 width,229 height,230 generator,231 scale,232 scale_factor,233 ),234 None,235 )236 except Exception as e:237 return None, error_str(e)238 # if img is not None:239 # return img_to_img(model_path, prompt, neg_prompt, img, strength, guidance, steps, width, height,240 # generator, scale, scale_factor), None241 # else:242 # return txt_to_img(model_path, prompt, neg_prompt, guidance, steps, width, height, generator, scale, scale_factor), None243 244 245def txt_to_img(246 model_path,247 prompt,248 neg_prompt,249 guidance,250 steps,251 width,252 height,253 generator,254 scale,255 scale_factor,256):257 print(f"{datetime.datetime.now()} txt_to_img, model: {current_model.name}")258 259 global last_mode260 global pipe261 global current_model_path262 if model_path != current_model_path or last_mode != "txt2img":263 current_model_path = model_path264 265 if is_colab or current_model == custom_model:266 pipe = StableDiffusionPipeline.from_pretrained(267 current_model_path,268 torch_dtype=dtype,269 scheduler=scheduler,270 safety_checker=lambda images, clip_input: (images, False),271 )272 else:273 # pipe = pipe.to("cpu")274 pipe = current_model.pipe_t2i275 276 if torch.cuda.is_available():277 pipe = pipe.to(device)278 last_mode = "txt2img"279 280 prompt = current_model.prefix + prompt281 result = pipe(282 prompt,283 negative_prompt=neg_prompt,284 # num_images_per_prompt=n_images,285 num_inference_steps=int(steps),286 guidance_scale=guidance,287 width=width,288 height=height,289 generator=generator,290 )291 292 # result.images[0] = magnifier.magnify(result.images[0], scale_factor=scale_factor)293 # enhance resolution294 if scale_factor > 1:295 if scale == "ESRGAN4x":296 fp32 = True if device == "cpu" else False297 result.images[0] = realEsrgan(298 input_dir=result.images[0],299 suffix="",300 output_dir="imgs",301 fp32=fp32,302 outscale=scale_factor,303 )[0]304 else:305 result.images[0] = magnifier.magnify(306 result.images[0], scale_factor=scale_factor307 )308 # save image309 result.images[0].save(310 "imgs/result-{}.png".format(datetime.datetime.now().strftime("%Y%m%d-%H%M%S"))311 )312 return replace_nsfw_images(result)313 314 315def img_to_img(316 model_path,317 prompt,318 neg_prompt,319 img,320 strength,321 guidance,322 steps,323 width,324 height,325 generator,326 scale,327 scale_factor,328):329 fprint(f"{datetime.datetime.now()} img_to_img, model: {model_path}")330 331 global last_mode332 global pipe333 global current_model_path334 if model_path != current_model_path or last_mode != "img2img":335 current_model_path = model_path336 337 if is_colab or current_model == custom_model:338 pipe = StableDiffusionImg2ImgPipeline.from_pretrained(339 current_model_path,340 torch_dtype=dtype,341 scheduler=scheduler,342 safety_checker=lambda images, clip_input: (images, False),343 )344 else:345 # pipe = pipe.to("cpu")346 pipe = current_model.pipe_i2i347 348 if torch.cuda.is_available():349 pipe = pipe.to(device)350 last_mode = "img2img"351 352 prompt = current_model.prefix + prompt353 ratio = min(height / img.height, width / img.width)354 img = img.resize((int(img.width * ratio), int(img.height * ratio)), Image.LANCZOS)355 result = pipe(356 prompt,357 negative_prompt=neg_prompt,358 # num_images_per_prompt=n_images,359 image=img,360 num_inference_steps=int(steps),361 strength=strength,362 guidance_scale=guidance,363 # width=width,364 # height=height,365 generator=generator,366 )367 if scale_factor > 1:368 if scale == "ESRGAN4x":369 fp32 = True if device == "cpu" else False370 result.images[0] = realEsrgan(371 input_dir=result.images[0],372 suffix="",373 output_dir="imgs",374 fp32=fp32,375 outscale=scale_factor,376 )[0]377 else:378 result.images[0] = magnifier.magnify(379 result.images[0], scale_factor=scale_factor380 )381 # save image382 result.images[0].save(383 "imgs/result-{}.png".format(datetime.datetime.now().strftime("%Y%m%d-%H%M%S"))384 )385 return replace_nsfw_images(result)386 387 388def replace_nsfw_images(results):389 if is_colab:390 return results.images[0]391 if hasattr(results, "nsfw_content_detected") and results.nsfw_content_detected:392 for i in range(len(results.images)):393 if results.nsfw_content_detected[i]:394 results.images[i] = Image.open("nsfw.png")395 return results.images[0]396 397 398css = """.finetuned-diffusion-div div{display:inline-flex;align-items:center;gap:.8rem;font-size:1.75rem}.finetuned-diffusion-div div h1{font-weight:900;margin-bottom:7px}.finetuned-diffusion-div p{margin-bottom:10px;font-size:94%}a{text-decoration:underline}.tabs{margin-top:0;margin-bottom:0}#gallery{min-height:20rem}399"""400with gr.Blocks(css=css) as demo:401 if not os.path.exists("imgs"):402 os.mkdir("imgs")403 404 gr.Markdown("# Super Resolution Anime Diffusion")405 gr.Markdown(406 "## Author: [yangheng95](https://github.com/yangheng95) Github:[Github](https://github.com/yangheng95/stable-diffusion-webui)"407 )408 gr.Markdown(409 "### This demo is running on a CPU, so it will take at least 20 minutes. "410 "If you have a GPU, you can clone from [Github](https://github.com/yangheng95/SuperResolutionAnimeDiffusion) and run it locally."411 )412 gr.Markdown(413 "### FYI: to generate a 512*512 image and magnify 4x, it only takes 5~8 seconds on a RTX 2080 GPU"414 )415 gr.Markdown(416 "### You can duplicate this demo on HuggingFace Spaces, click [here](https://huggingface.co/spaces/yangheng/Super-Resolution-Anime-Diffusion?duplicate=true)"417 )418 419 with gr.Row():420 with gr.Column(scale=55):421 with gr.Group():422 gr.Markdown("Text to image")423 424 model_name = gr.Dropdown(425 label="Model",426 choices=[m.name for m in models],427 value=current_model.name,428 )429 430 with gr.Box(visible=False) as custom_model_group:431 custom_model_path = gr.Textbox(432 label="Custom model path",433 placeholder="Path to model, e.g. nitrosocke/Arcane-Diffusion",434 interactive=True,435 )436 gr.HTML(437 "<div><font size='2'>Custom models have to be downloaded first, so give it some time.</font></div>"438 )439 440 with gr.Row():441 prompt = gr.Textbox(442 label="Prompt",443 show_label=False,444 max_lines=2,445 placeholder="Enter prompt. Style applied automatically",446 ).style(container=False)447 with gr.Row():448 generate = gr.Button(value="Generate")449 450 with gr.Row():451 with gr.Group():452 neg_prompt = gr.Textbox(453 label="Negative prompt",454 value="bad result, worst, random, invalid, inaccurate, imperfect, blurry, deformed,"455 " disfigured, mutation, mutated, ugly, out of focus, bad anatomy, text, error,"456 " extra digit, fewer digits, worst quality, low quality, normal quality, noise, "457 "jpeg artifact, compression artifact, signature, watermark, username, logo, "458 "low resolution, worst resolution, bad resolution, normal resolution, bad detail,"459 " bad details, bad lighting, bad shadow, bad shading, bad background,"460 " worst background.",461 )462 463 image_out = gr.Image(height="auto", width="auto")464 error_output = gr.Markdown()465 466 with gr.Row():467 gr.Markdown(468 "# Random Image Generation Preview (512*768)x4 magnified"469 )470 for f_img in findfile.find_cwd_files(".png", recursive=2):471 with gr.Row():472 image = gr.Image(height=512, value=PIL.Image.open(f_img))473 # gallery = gr.Gallery(474 # label="Generated images", show_label=False, elem_id="gallery"475 # ).style(grid=[1], height="auto")476 477 with gr.Column(scale=45):478 with gr.Group():479 gr.Markdown("Image to Image")480 481 with gr.Row():482 with gr.Group():483 image = gr.Image(484 label="Image", height=256, tool="editor", type="pil"485 )486 strength = gr.Slider(487 label="Transformation strength",488 minimum=0,489 maximum=1,490 step=0.01,491 value=0.5,492 )493 494 with gr.Row():495 with gr.Group():496 # n_images = gr.Slider(label="Images", value=1, minimum=1, maximum=4, step=1)497 498 with gr.Row():499 guidance = gr.Slider(500 label="Guidance scale", value=7.5, maximum=15501 )502 steps = gr.Slider(503 label="Steps", value=15, minimum=2, maximum=75, step=1504 )505 506 with gr.Row():507 width = gr.Slider(508 label="Width",509 value=512,510 minimum=64,511 maximum=1024,512 step=8,513 )514 height = gr.Slider(515 label="Height",516 value=768,517 minimum=64,518 maximum=1024,519 step=8,520 )521 with gr.Row():522 scale = gr.Radio(523 label="Scale",524 choices=["Waifu2x", "ESRGAN4x"],525 value="Waifu2x",526 )527 with gr.Row():528 scale_factor = gr.Slider(529 1,530 8,531 label="Scale factor (to magnify image) (1, 2, 4, 8)",532 value=1,533 step=1,534 )535 536 seed = gr.Slider(537 0, 2147483647, label="Seed (0 = random)", value=0, step=1538 )539 540 if is_colab:541 model_name.change(542 on_model_change,543 inputs=model_name,544 outputs=[custom_model_group, prompt],545 queue=False,546 )547 custom_model_path.change(548 custom_model_changed, inputs=custom_model_path, outputs=None549 )550 # n_images.change(lambda n: gr.Gallery().style(grid=[2 if n > 1 else 1], height="auto"), inputs=n_images, outputs=gallery)551 552 gr.Markdown(553 "### based on [Anything V5]"554 )555 556 inputs = [557 model_name,558 prompt,559 guidance,560 steps,561 width,562 height,563 seed,564 image,565 strength,566 neg_prompt,567 scale,568 scale_factor,569 ]570 outputs = [image_out, error_output]571 prompt.submit(inference, inputs=inputs, outputs=outputs)572 generate.click(inference, inputs=inputs, outputs=outputs, api_name="generate")573 574 prompt_keys = [575 "girl",576 "lovely",577 "cute",578 "beautiful eyes",579 "cumulonimbus clouds",580 random.choice(["dress"]),581 random.choice(["white hair"]),582 random.choice(["blue eyes"]),583 random.choice(["flower meadow"]),584 random.choice(["Elif", "Angel"]),585 ]586 prompt.value = ",".join(prompt_keys)587 ex = gr.Examples(588 [589 [models[0].name, prompt.value, 7.5, 15],590 ],591 inputs=[model_name, prompt, guidance, steps, seed],592 outputs=outputs,593 fn=inference,594 cache_examples=False,595 )596 597print(f"Space built in {time.time() - start_time:.2f} seconds")598 599if not is_colab:600 demo.queue(concurrency_count=2)601demo.launch(debug=is_colab, enable_queue=True, share=is_colab)