EDIT: add super-resolution model and inpainting sd

handler.py

@@ -1,12 +1,11 @@
 from typing import  Dict, List, Any
 import torch
-from diffusers import DPMSolverMultistepScheduler, StableDiffusionInpaintPipeline
+from diffusers import DPMSolverMultistepScheduler, StableDiffusionInpaintPipeline, EulerAncestralDiscreteScheduler
 from PIL import Image
 import base64
 from io import BytesIO
-from datetime import datetime
-import time
-
+import numpy as np
+from RealESRGAN import RealESRGAN
 
 # set device
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
@@ -14,20 +13,19 @@ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 if device.type != 'cuda':
     raise ValueError("need to run on GPU")
 
-begin_runtime = datetime.now()
-    
-def print_current_time(tag: str):
-    global begin_runtime
-    print(str(tag).upper() + ": " + str(datetime.now() - begin_runtime))
-
 class EndpointHandler():
     def __init__(self, path=""):
         # load StableDiffusionInpaintPipeline pipeline
         self.pipe = StableDiffusionInpaintPipeline.from_pretrained(path, torch_dtype=torch.float16)
         # use DPMSolverMultistepScheduler
-        self.pipe.scheduler = DPMSolverMultistepScheduler.from_config(self.pipe.scheduler.config)
+        self.pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
+        # pipe.enable_sequential_cpu_offload()
         # move to device
-        self.pipe = self.pipe.to(device)
+        self.pipe.to(device)
+        self.pipe.enable_xformers_memory_efficient_attention()
+
+        self.upscaler = RealESRGAN(device, scale=4)
+        self.upscaler.load_weights('weights/RealESRGAN_x4.pth', download=True)
 
 
     def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
@@ -45,14 +43,12 @@ class EndpointHandler():
 
         # hyperparamters
         num_inference_steps = data.pop("num_inference_steps", 50)
-        # guidance_scale = data.pop("guidance_scale", 7.5)
+        guidance_scale = data.pop("guidance_scale", 7.5)
         negative_prompt = data.pop("negative_prompt", None)
         height = data.pop("height", None)
         width = data.pop("width", None)
 
         # process image
-        print_current_time("Start decoding")
-        
         if encoded_image is not None and encoded_mask_image is not None:
             image = self.decode_base64_image(encoded_image)
             mask_image = self.decode_base64_image(encoded_mask_image)
@@ -60,29 +56,31 @@ class EndpointHandler():
             image = None
             mask_image = None
 
-        print_current_time("Finish decoding")
-            
         # run inference pipeline
         out = self.pipe(inputs,
                         image=image,
                         mask_image=mask_image,
                         num_inference_steps=num_inference_steps,
-                        # guidance_scale=guidance_scale,
+                        guidance_scale=guidance_scale,
                         num_images_per_prompt=num_images,
                         negative_prompt=negative_prompt,
                         height=height,
                         width=width
-        )
+        ).images
+
+        for i in range(len(out)):
+          gen_img = Image.composite(out[i], image.resize(out[i].size), mask_image.resize(out[i].size))
+          gen_img = self.upscaler.predict(gen_img)
+          gen_img = Image.composite(gen_img, image.resize(gen_img.size), mask_image.resize(gen_img.size))
+          out[i] = gen_img
 
         # return first generate PIL image
         json_imgs = {}
-        for i in range(len(out.images)):
+        for i in range(len(out)):
           buffered = BytesIO()
-          out.images[i].save(buffered, format="PNG")
+          out[i].save(buffered, format="PNG")
           img_str = base64.b64encode(buffered.getvalue())
           json_imgs[f"{i}"] = img_str.decode()
-
-        print_current_time("Complete Stable diffusion")
         return json_imgs
 
     # helper to decode input image

model_index.json

@@ -1,38 +0,0 @@
-{
-  "_class_name": "StableDiffusionInpaintPipeline",
-  "_diffusers_version": "0.27.2",
-  "_name_or_path": "stabilityai/stable-diffusion-2-1",
-  "feature_extractor": [
-    null,
-    null
-  ],
-  "image_encoder": [
-    null,
-    null
-  ],
-  "requires_safety_checker": false,
-  "safety_checker": [
-    null,
-    null
-  ],
-  "scheduler": [
-    "diffusers",
-    "DDIMScheduler"
-  ],
-  "text_encoder": [
-    "transformers",
-    "CLIPTextModel"
-  ],
-  "tokenizer": [
-    "transformers",
-    "CLIPTokenizer"
-  ],
-  "unet": [
-    "diffusers",
-    "UNet2DConditionModel"
-  ],
-  "vae": [
-    "diffusers",
-    "AutoencoderKL"
-  ]
-}

requirements.txt

@@ -0,0 +1 @@
+git+https://github.com/sberbank-ai/Real-ESRGAN.git

scheduler/scheduler_config.json

@@ -1,20 +0,0 @@
-{
-  "_class_name": "DDIMScheduler",
-  "_diffusers_version": "0.27.2",
-  "beta_end": 0.012,
-  "beta_schedule": "scaled_linear",
-  "beta_start": 0.00085,
-  "clip_sample": false,
-  "clip_sample_range": 1.0,
-  "dynamic_thresholding_ratio": 0.995,
-  "num_train_timesteps": 1000,
-  "prediction_type": "v_prediction",
-  "rescale_betas_zero_snr": false,
-  "sample_max_value": 1.0,
-  "set_alpha_to_one": false,
-  "skip_prk_steps": true,
-  "steps_offset": 1,
-  "thresholding": false,
-  "timestep_spacing": "leading",
-  "trained_betas": null
-}

text_encoder/config.json

@@ -1,25 +0,0 @@
-{
-  "_name_or_path": "/root/.cache/huggingface/hub/models--stabilityai--stable-diffusion-2-1/snapshots/f7f33030acc57428be85fbec092c37a78231d75a/text_encoder",
-  "architectures": [
-    "CLIPTextModel"
-  ],
-  "attention_dropout": 0.0,
-  "bos_token_id": 0,
-  "dropout": 0.0,
-  "eos_token_id": 2,
-  "hidden_act": "gelu",
-  "hidden_size": 1024,
-  "initializer_factor": 1.0,
-  "initializer_range": 0.02,
-  "intermediate_size": 4096,
-  "layer_norm_eps": 1e-05,
-  "max_position_embeddings": 77,
-  "model_type": "clip_text_model",
-  "num_attention_heads": 16,
-  "num_hidden_layers": 23,
-  "pad_token_id": 1,
-  "projection_dim": 512,
-  "torch_dtype": "float16",
-  "transformers_version": "4.38.2",
-  "vocab_size": 49408
-}

text_encoder/model.safetensors

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:bc1827c465450322616f06dea41596eac7d493f4e95904dcb51f0fc745c4e13f
-size 680820392

tokenizer/special_tokens_map.json

@@ -1,24 +0,0 @@
-{
-  "bos_token": {
-    "content": "<|startoftext|>",
-    "lstrip": false,
-    "normalized": true,
-    "rstrip": false,
-    "single_word": false
-  },
-  "eos_token": {
-    "content": "<|endoftext|>",
-    "lstrip": false,
-    "normalized": true,
-    "rstrip": false,
-    "single_word": false
-  },
-  "pad_token": "!",
-  "unk_token": {
-    "content": "<|endoftext|>",
-    "lstrip": false,
-    "normalized": true,
-    "rstrip": false,
-    "single_word": false
-  }
-}

tokenizer/tokenizer_config.json

@@ -1,38 +0,0 @@
-{
-  "add_prefix_space": false,
-  "added_tokens_decoder": {
-    "0": {
-      "content": "!",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "49406": {
-      "content": "<|startoftext|>",
-      "lstrip": false,
-      "normalized": true,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "49407": {
-      "content": "<|endoftext|>",
-      "lstrip": false,
-      "normalized": true,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    }
-  },
-  "bos_token": "<|startoftext|>",
-  "clean_up_tokenization_spaces": true,
-  "do_lower_case": true,
-  "eos_token": "<|endoftext|>",
-  "errors": "replace",
-  "model_max_length": 77,
-  "pad_token": "!",
-  "tokenizer_class": "CLIPTokenizer",
-  "unk_token": "<|endoftext|>"
-}

unet/config.json

@@ -1,73 +0,0 @@
-{
-  "_class_name": "UNet2DConditionModel",
-  "_diffusers_version": "0.27.2",
-  "_name_or_path": "/root/.cache/huggingface/hub/models--stabilityai--stable-diffusion-2-1/snapshots/f7f33030acc57428be85fbec092c37a78231d75a/unet",
-  "act_fn": "silu",
-  "addition_embed_type": null,
-  "addition_embed_type_num_heads": 64,
-  "addition_time_embed_dim": null,
-  "attention_head_dim": [
-    5,
-    10,
-    20,
-    20
-  ],
-  "attention_type": "default",
-  "block_out_channels": [
-    320,
-    640,
-    1280,
-    1280
-  ],
-  "center_input_sample": false,
-  "class_embed_type": null,
-  "class_embeddings_concat": false,
-  "conv_in_kernel": 3,
-  "conv_out_kernel": 3,
-  "cross_attention_dim": 1024,
-  "cross_attention_norm": null,
-  "down_block_types": [
-    "CrossAttnDownBlock2D",
-    "CrossAttnDownBlock2D",
-    "CrossAttnDownBlock2D",
-    "DownBlock2D"
-  ],
-  "downsample_padding": 1,
-  "dropout": 0.0,
-  "dual_cross_attention": false,
-  "encoder_hid_dim": null,
-  "encoder_hid_dim_type": null,
-  "flip_sin_to_cos": true,
-  "freq_shift": 0,
-  "in_channels": 4,
-  "layers_per_block": 2,
-  "mid_block_only_cross_attention": null,
-  "mid_block_scale_factor": 1,
-  "mid_block_type": "UNetMidBlock2DCrossAttn",
-  "norm_eps": 1e-05,
-  "norm_num_groups": 32,
-  "num_attention_heads": null,
-  "num_class_embeds": null,
-  "only_cross_attention": false,
-  "out_channels": 4,
-  "projection_class_embeddings_input_dim": null,
-  "resnet_out_scale_factor": 1.0,
-  "resnet_skip_time_act": false,
-  "resnet_time_scale_shift": "default",
-  "reverse_transformer_layers_per_block": null,
-  "sample_size": 96,
-  "time_cond_proj_dim": null,
-  "time_embedding_act_fn": null,
-  "time_embedding_dim": null,
-  "time_embedding_type": "positional",
-  "timestep_post_act": null,
-  "transformer_layers_per_block": 1,
-  "up_block_types": [
-    "UpBlock2D",
-    "CrossAttnUpBlock2D",
-    "CrossAttnUpBlock2D",
-    "CrossAttnUpBlock2D"
-  ],
-  "upcast_attention": true,
-  "use_linear_projection": true
-}

unet/diffusion_pytorch_model.safetensors

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:8a3a4d7978884c5e4ef00b62641b1b544b257be2f6715d984188610ad6475ad2
-size 1731904736

vae/config.json

@@ -1,34 +0,0 @@
-{
-  "_class_name": "AutoencoderKL",
-  "_diffusers_version": "0.27.2",
-  "_name_or_path": "/root/.cache/huggingface/hub/models--stabilityai--stable-diffusion-2-1/snapshots/f7f33030acc57428be85fbec092c37a78231d75a/vae",
-  "act_fn": "silu",
-  "block_out_channels": [
-    128,
-    256,
-    512,
-    512
-  ],
-  "down_block_types": [
-    "DownEncoderBlock2D",
-    "DownEncoderBlock2D",
-    "DownEncoderBlock2D",
-    "DownEncoderBlock2D"
-  ],
-  "force_upcast": true,
-  "in_channels": 3,
-  "latent_channels": 4,
-  "latents_mean": null,
-  "latents_std": null,
-  "layers_per_block": 2,
-  "norm_num_groups": 32,
-  "out_channels": 3,
-  "sample_size": 768,
-  "scaling_factor": 0.18215,
-  "up_block_types": [
-    "UpDecoderBlock2D",
-    "UpDecoderBlock2D",
-    "UpDecoderBlock2D",
-    "UpDecoderBlock2D"
-  ]
-}

vae/diffusion_pytorch_model.safetensors

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:3e4c08995484ee61270175e9e7a072b66a6e4eeb5f0c266667fe1f45b90daf9a
-size 167335342