add support for diffusers checkpoint loading

.gitignore

@@ -4,4 +4,6 @@
 
 outputs/
 gradio_tmp/
-__pycache__/
+__pycache__/
+
+checkpoints/

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2023 Picsart AI Research (PAIR)
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

app.py

@@ -1,40 +1,44 @@
 import os
+import sys
+from pathlib import Path
 from collections import OrderedDict
 
 import gradio as gr
 import shutil
 import uuid
 import torch
-from pathlib import Path
-from lib.utils.iimage import IImage
 from PIL import Image
 
-from lib import models
-from lib.methods import rasg, sd, sr
-from lib.utils import poisson_blend, image_from_url_text
+demo_path = Path(__file__).resolve().parent
+root_path = demo_path
+sys.path.append(str(root_path))
+from src import models
+from src.methods import rasg, sd, sr
+from src.utils import IImage, poisson_blend, image_from_url_text
 
 
-TMP_DIR = 'gradio_tmp'
-if Path(TMP_DIR).exists():
-    shutil.rmtree(TMP_DIR)
-Path(TMP_DIR).mkdir(exist_ok=True, parents=True)
+TMP_DIR = root_path / 'gradio_tmp'
+if TMP_DIR.exists():
+    shutil.rmtree(str(TMP_DIR))
+TMP_DIR.mkdir(exist_ok=True, parents=True)
 
-os.environ['GRADIO_TEMP_DIR'] = TMP_DIR
+os.environ['GRADIO_TEMP_DIR'] = str(TMP_DIR)
 
 on_huggingspace = os.environ.get("SPACE_AUTHOR_NAME") == "PAIR"
 
 negative_prompt_str = "text, bad anatomy, bad proportions, blurry, cropped, deformed, disfigured, duplicate, error, extra limbs, gross proportions, jpeg artifacts, long neck, low quality, lowres, malformed, morbid, mutated, mutilated, out of frame, ugly, worst quality"
 positive_prompt_str = "Full HD, 4K, high quality, high resolution"
 
+examples_path = root_path / '__assets__/demo/examples'
 example_inputs = [
-    ['assets/examples/images_1024/a40.jpg', 'assets/examples/images_2048/a40.jpg', 'medieval castle'],
-    ['assets/examples/images_1024/a4.jpg', 'assets/examples/images_2048/a4.jpg', 'parrot'],
-    ['assets/examples/images_1024/a65.jpg', 'assets/examples/images_2048/a65.jpg', 'hoodie'],
-    ['assets/examples/images_1024/a54.jpg', 'assets/examples/images_2048/a54.jpg', 'salad'],
-    ['assets/examples/images_1024/a51.jpg', 'assets/examples/images_2048/a51.jpg', 'space helmet'],
-    ['assets/examples/images_1024/a46.jpg', 'assets/examples/images_2048/a46.jpg', 'stack of books'],
-    ['assets/examples/images_1024/a19.jpg', 'assets/examples/images_2048/a19.jpg', 'antique greek vase'],
-    ['assets/examples/images_1024/a2.jpg', 'assets/examples/images_2048/a2.jpg', 'sunglasses'],
+    [f'{examples_path}/images_1024/a40.jpg', f'{examples_path}/images_2048/a40.jpg', 'medieval castle'],
+    [f'{examples_path}/images_1024/a4.jpg', f'{examples_path}/images_2048/a4.jpg', 'parrot'],
+    [f'{examples_path}/images_1024/a65.jpg', f'{examples_path}/images_2048/a65.jpg', 'hoodie'],
+    [f'{examples_path}/images_1024/a54.jpg', f'{examples_path}/images_2048/a54.jpg', 'salad'],
+    [f'{examples_path}/images_1024/a51.jpg', f'{examples_path}/images_2048/a51.jpg', 'space helmet'],
+    [f'{examples_path}/images_1024/a46.jpg', f'{examples_path}/images_2048/a46.jpg', 'stack of books'],
+    [f'{examples_path}/images_1024/a19.jpg', f'{examples_path}/images_2048/a19.jpg', 'antique greek vase'],
+    [f'{examples_path}/images_1024/a2.jpg', f'{examples_path}/images_2048/a2.jpg', 'sunglasses'],
 ]
 
 thumbnails = [
@@ -60,27 +64,35 @@ example_previews = [
 ]
 
 # Load models
+models.pre_download_inpainting_models()
 inpainting_models = OrderedDict([
-    ("Dreamshaper Inpainting V8", models.ds_inp.load_model()),
-    ("Stable-Inpainting 2.0", models.sd2_inp.load_model()),
-    ("Stable-Inpainting 1.5", models.sd15_inp.load_model())
+    ("Dreamshaper Inpainting V8", 'ds8_inp'),
+    ("Stable-Inpainting 2.0", 'sd2_inp'),
+    ("Stable-Inpainting 1.5", 'sd15_inp')
 ])
 sr_model = models.sd2_sr.load_model(device='cuda:1')
 sam_predictor = models.sam.load_model(device='cuda:0')
 
-inp_model = inpainting_models[list(inpainting_models.keys())[0]]
-def set_model_from_name(inp_model_name):
+inp_model_name = list(inpainting_models.keys())[0]
+inp_model = models.load_inpainting_model(
+    inpainting_models[inp_model_name], device='cuda:0', cache=False)
+
+
+def set_model_from_name(new_inp_model_name):
     global inp_model
-    print (f"Activating Inpaintng Model: {inp_model_name}")
-    inp_model = inpainting_models[inp_model_name]
+    global inp_model_name
+    if new_inp_model_name != inp_model_name:
+        print (f"Activating Inpaintng Model: {new_inp_model_name}")
+        inp_model = models.load_inpainting_model(
+            inpainting_models[new_inp_model_name], device='cuda:0', cache=False)
+        inp_model_name = new_inp_model_name
 
 
 def save_user_session(hr_image, hr_mask, lr_results, prompt, session_id=None):
     if session_id == '':
         session_id = str(uuid.uuid4())
     
-    tmp_dir = Path(TMP_DIR)
-    session_dir = tmp_dir / session_id
+    session_dir = TMP_DIR / session_id
     session_dir.mkdir(exist_ok=True, parents=True)
     
     hr_image.save(session_dir / 'hr_image.png')
@@ -103,8 +115,7 @@ def recover_user_session(session_id):
     if session_id == '':
         return None, None, [], ''
     
-    tmp_dir = Path(TMP_DIR)
-    session_dir = tmp_dir / session_id
+    session_dir = TMP_DIR / session_id
     lr_results_dir = session_dir / 'lr_results'
 
     hr_image = Image.open(session_dir / 'hr_image.png')
@@ -121,64 +132,22 @@ def recover_user_session(session_id):
     return hr_image, hr_mask, gallery, prompt
 
 
-def rasg_run(
-    use_painta, prompt, imageMask, hr_image, seed, eta,
-    negative_prompt, positive_prompt, ddim_steps,
-    guidance_scale=7.5,
-    batch_size=1, session_id=''
+def inpainting_run(model_name, use_rasg, use_painta, prompt, imageMask,
+    hr_image, seed, eta, negative_prompt, positive_prompt, ddim_steps,
+    guidance_scale=7.5, batch_size=1, session_id=''
 ):
     torch.cuda.empty_cache()
+    set_model_from_name(model_name)
 
-    seed = int(seed)
-    batch_size = max(1, min(int(batch_size), 4))
-
-    image = IImage(hr_image).resize(512)
-    mask = IImage(imageMask['mask']).rgb().resize(512)
-
-    method = ['rasg']
+    method = ['default']
     if use_painta: method.append('painta')
+    if use_rasg: method.append('rasg')
     method = '-'.join(method)
 
-    inpainted_images = []
-    blended_images = []
-    for i in range(batch_size):
-        seed = seed + i * 1000
-
-        inpainted_image = rasg.run(
-            ddim=inp_model,
-            method=method,
-            prompt=prompt,
-            image=image,
-            mask=mask,
-            seed=seed,
-            eta=eta,
-            negative_prompt=negative_prompt,
-            positive_prompt=positive_prompt,
-            num_steps=ddim_steps,
-            guidance_scale=guidance_scale
-        ).crop(image.size)
-        
-        blended_image = poisson_blend(
-            orig_img=image.data[0],
-            fake_img=inpainted_image.data[0],
-            mask=mask.data[0],
-            dilation=12
-        )
-        blended_images.append(blended_image)
-        inpainted_images.append(inpainted_image.pil())
-
-    session_id = save_user_session(
-        hr_image, imageMask['mask'], inpainted_images, prompt, session_id=session_id)
-
-    return blended_images, session_id
-
-
-def sd_run(use_painta, prompt, imageMask, hr_image, seed, eta,
-    negative_prompt, positive_prompt, ddim_steps,
-    guidance_scale=7.5,
-    batch_size=1, session_id=''
-):
-    torch.cuda.empty_cache()
+    if use_rasg:
+        inpainting_f = rasg.run
+    else:
+        inpainting_f = sd.run
 
     seed = int(seed)
     batch_size = max(1, min(int(batch_size), 4))
@@ -195,7 +164,7 @@ def sd_run(use_painta, prompt, imageMask, hr_image, seed, eta,
     for i in range(batch_size):
         seed = seed + i * 1000
 
-        inpainted_image = sd.run(
+        inpainted_image = inpainting_f(
             ddim=inp_model,
             method=method,
             prompt=prompt,
@@ -226,13 +195,12 @@ def sd_run(use_painta, prompt, imageMask, hr_image, seed, eta,
 
 def upscale_run(
     ddim_steps, seed, use_sam_mask, session_id, img_index,
-    negative_prompt='',
-    positive_prompt=', high resolution professional photo'
+    negative_prompt='', positive_prompt='high resolution professional photo'
 ):
     hr_image, hr_mask, gallery, prompt = recover_user_session(session_id)
 
     if len(gallery) == 0:
-        return Image.open('./assets/sr_info.png')
+        return Image.open(root_path / '__assets__/sr_info.png')
 
     torch.cuda.empty_cache()
 
@@ -249,7 +217,7 @@ def upscale_run(
         inpainted_image,
         hr_image,
         hr_mask,
-        prompt=prompt + positive_prompt,
+        prompt=f'{prompt}, {positive_prompt}',
         noise_level=20,
         blend_trick=True,
         blend_output=True,
@@ -261,14 +229,7 @@ def upscale_run(
     return output_image
 
 
-def switch_run(use_rasg, model_name, *args):
-    set_model_from_name(model_name)
-    if use_rasg:
-        return rasg_run(*args)
-    return sd_run(*args)
-
-
-with gr.Blocks(css='style.css') as demo:
+with gr.Blocks(css=demo_path / 'style.css') as demo:
     gr.HTML(
         """
         <div style="text-align: center; max-width: 1200px; margin: 20px auto;">
@@ -300,7 +261,7 @@ with gr.Blocks(css='style.css') as demo:
         <img style="margin-top: 0em; margin-bottom: 0em" src="https://bit.ly/3gLdBN6" alt="Duplicate Space"></a>
         </p>""")
 
-    with open('script.js', 'r') as f:
+    with open(demo_path / 'script.js', 'r') as f:
         js_str = f.read()
 
     demo.load(_js=js_str)
@@ -380,10 +341,10 @@ with gr.Blocks(css='style.css') as demo:
     html_info = gr.HTML(elem_id=f'html_info', elem_classes="infotext")
 
     inpaint_btn.click(
-        fn=switch_run, 
+        fn=inpainting_run, 
         inputs=[
-            use_rasg,
             model_picker,
+            use_rasg,
             use_painta,
             prompt,
             imageMask,
@@ -415,4 +376,4 @@ with gr.Blocks(css='style.css') as demo:
     )
 
 demo.queue(max_size=20)
-demo.launch(share=True, allowed_paths=[TMP_DIR])
+demo.launch(share=True, allowed_paths=[str(TMP_DIR)])

assets/.gitignore

@@ -1 +0,0 @@
-models/

config/ddpm/v1.yaml

@@ -0,0 +1,14 @@
+linear_start: 0.00085
+linear_end: 0.0120
+num_timesteps_cond: 1
+log_every_t: 200
+timesteps: 1000
+first_stage_key: "jpg"
+cond_stage_key: "txt"
+image_size: 64
+channels: 4
+cond_stage_trainable: false
+conditioning_key: crossattn
+monitor: val/loss_simple_ema
+scale_factor: 0.18215
+use_ema: False # we set this to false because this is an inference only config

config/ddpm/v2-upsample.yaml

@@ -0,0 +1,24 @@
+parameterization: "v"
+low_scale_key: "lr"
+linear_start: 0.0001
+linear_end: 0.02
+num_timesteps_cond: 1
+log_every_t: 200
+timesteps: 1000
+first_stage_key: "jpg"
+cond_stage_key: "txt"
+image_size: 128
+channels: 4
+cond_stage_trainable: false
+conditioning_key: "hybrid-adm"
+monitor: val/loss_simple_ema
+scale_factor: 0.08333
+use_ema: False
+
+low_scale_config:
+  target: ldm.modules.diffusionmodules.upscaling.ImageConcatWithNoiseAugmentation
+  params:
+    noise_schedule_config: # image space
+      linear_start: 0.0001
+      linear_end: 0.02
+    max_noise_level: 350

config/encoders/clip.yaml

@@ -0,0 +1 @@
+__class__: smplfusion.models.encoders.clip_embedder.FrozenCLIPEmbedder

config/encoders/openclip.yaml

@@ -0,0 +1,4 @@
+__class__: smplfusion.models.encoders.open_clip_embedder.FrozenOpenCLIPEmbedder
+__init__:
+  freeze: True
+  layer: "penultimate"

config/unet/inpainting/v1.yaml

@@ -0,0 +1,15 @@
+__class__: smplfusion.models.unet.UNetModel
+__init__:
+  image_size: 32 # unused
+  in_channels: 9  # 4 data + 4 downscaled image + 1 mask
+  out_channels: 4
+  model_channels: 320
+  attention_resolutions: [ 4, 2, 1 ]
+  num_res_blocks: 2
+  channel_mult: [ 1, 2, 4, 4 ]
+  num_heads: 8
+  use_spatial_transformer: True
+  transformer_depth: 1
+  context_dim: 768
+  use_checkpoint: False
+  legacy: False

config/unet/inpainting/v2.yaml

@@ -0,0 +1,16 @@
+__class__: smplfusion.models.unet.UNetModel
+__init__:
+  use_checkpoint: False
+  image_size: 32 # unused
+  in_channels: 9
+  out_channels: 4
+  model_channels: 320
+  attention_resolutions: [ 4, 2, 1 ]
+  num_res_blocks: 2
+  channel_mult: [ 1, 2, 4, 4 ]
+  num_head_channels: 64 # need to fix for flash-attn
+  use_spatial_transformer: True
+  use_linear_in_transformer: True
+  transformer_depth: 1
+  context_dim: 1024
+  legacy: False

config/unet/upsample/v2.yaml

@@ -0,0 +1,19 @@
+__class__: smplfusion.models.unet.UNetModel
+__init__:
+  use_checkpoint: False
+  num_classes: 1000  # timesteps for noise conditioning (here constant, just need one)
+  image_size: 128
+  in_channels: 7
+  out_channels: 4
+  model_channels: 256
+  attention_resolutions: [ 2,4,8]
+  num_res_blocks: 2
+  channel_mult: [ 1, 2, 2, 4]
+  disable_self_attentions: [True, True, True, False]
+  disable_middle_self_attn: False
+  num_heads: 8
+  use_spatial_transformer: True
+  transformer_depth: 1
+  context_dim: 1024
+  legacy: False
+  use_linear_in_transformer: True

config/vae-upsample.yaml

@@ -0,0 +1,16 @@
+__class__: smplfusion.models.vae.AutoencoderKL
+__init__:
+  embed_dim: 4
+  ddconfig:
+    double_z: True
+    z_channels: 4
+    resolution: 256
+    in_channels: 3
+    out_ch: 3
+    ch: 128
+    ch_mult: [ 1,2,4 ]
+    num_res_blocks: 2
+    attn_resolutions: [ ]
+    dropout: 0.0
+  lossconfig:
+    target: torch.nn.Identity

config/vae.yaml

@@ -0,0 +1,17 @@
+__class__: smplfusion.models.vae.AutoencoderKL
+__init__:
+  embed_dim: 4
+  monitor: val/rec_loss
+  ddconfig:
+    double_z: true
+    z_channels: 4
+    resolution: 256
+    in_channels: 3
+    out_ch: 3
+    ch: 128
+    ch_mult: [1,2,4,4]
+    num_res_blocks: 2
+    attn_resolutions: []
+    dropout: 0.0
+  lossconfig:
+    target: torch.nn.Identity

lib/models/__init__.py

@@ -1 +0,0 @@
-from . import sd2_inp, ds_inp, sd15_inp, sd2_sr, sam

lib/models/common.py

@@ -1,49 +0,0 @@
-import importlib
-import requests
-from pathlib import Path
-from os.path import dirname
-
-from omegaconf import OmegaConf
-from tqdm import tqdm
-
-
-PROJECT_DIR = dirname(dirname(dirname(__file__)))
-CONFIG_FOLDER =  f'{PROJECT_DIR}/assets/config'
-MODEL_FOLDER =  f'{PROJECT_DIR}/assets/models'
-
-
-def download_file(url, save_path, chunk_size=1024):
-    try:
-        save_path = Path(save_path)
-        if save_path.exists():
-            print(f'{save_path.name} exists')
-            return
-        save_path.parent.mkdir(exist_ok=True, parents=True)
-        resp = requests.get(url, stream=True)
-        total = int(resp.headers.get('content-length', 0))
-        with open(save_path, 'wb') as file, tqdm(
-            desc=save_path.name,
-            total=total,
-            unit='iB',
-            unit_scale=True,
-            unit_divisor=1024,
-        ) as bar:
-            for data in resp.iter_content(chunk_size=chunk_size):
-                size = file.write(data)
-                bar.update(size)
-        print(f'{save_path.name} download finished')
-    except Exception as e:
-        raise Exception(f"Download failed: {e}")
-
-
-def get_obj_from_str(string):
-    module, cls = string.rsplit(".", 1)
-    try:
-        return getattr(importlib.import_module(module, package=None), cls)
-    except:
-        return getattr(importlib.import_module('lib.' + module, package=None), cls)
-
-
-def load_obj(path):
-    objyaml = OmegaConf.load(path)
-    return get_obj_from_str(objyaml['__class__'])(**objyaml.get("__init__", {}))

lib/models/ds_inp.py

@@ -1,51 +0,0 @@
-import importlib
-from omegaconf import OmegaConf
-import torch
-import safetensors
-import safetensors.torch
-
-from lib.smplfusion import DDIM, share, scheduler
-from .common import *
-
-
-MODEL_PATH = f'{MODEL_FOLDER}/dreamshaper/dreamshaper_8Inpainting.safetensors'
-DOWNLOAD_URL = 'https://civitai.com/api/download/models/131004'
-
-# pre-download
-download_file(DOWNLOAD_URL, MODEL_PATH)
-
-
-def load_model(dtype=torch.float16):
-    print ("Loading model: Dreamshaper Inpainting V8")
-
-    download_file(DOWNLOAD_URL, MODEL_PATH)
-    
-    state_dict = safetensors.torch.load_file(MODEL_PATH)
-
-    config = OmegaConf.load(f'{CONFIG_FOLDER}/ddpm/v1.yaml')
-    unet = load_obj(f'{CONFIG_FOLDER}/unet/inpainting/v1.yaml').eval().cuda()
-    vae = load_obj(f'{CONFIG_FOLDER}/vae.yaml').eval().cuda()
-    encoder = load_obj(f'{CONFIG_FOLDER}/encoders/clip.yaml').eval().cuda()
-
-    extract = lambda state_dict, model: {x[len(model)+1:]:y for x,y in state_dict.items() if model in x}
-    unet_state = extract(state_dict, 'model.diffusion_model')
-    encoder_state = extract(state_dict, 'cond_stage_model')
-    vae_state = extract(state_dict, 'first_stage_model')
-
-    unet.load_state_dict(unet_state)
-    encoder.load_state_dict(encoder_state)
-    vae.load_state_dict(vae_state)
-
-    if dtype == torch.float16:
-        unet.convert_to_fp16()
-    vae.to(dtype)
-    encoder.to(dtype)
-
-    unet = unet.requires_grad_(False)
-    encoder = encoder.requires_grad_(False)
-    vae = vae.requires_grad_(False)
-    
-    ddim = DDIM(config, vae, encoder, unet)
-    share.schedule = scheduler.linear(config.timesteps, config.linear_start, config.linear_end)
-
-    return ddim