imaginAIry/imaginairy/weight_management/utils.py
Bryce f97f6a3b4b feature: use refiners library for generation
BREAKING CHANGE

  - stable diffusion 1.5 + inpainting working
  - self-attention guidance working. improves image generation quality
  - tile-mode working
  - inpainting self-attention guidance working

disable/broken features:
  - sd 1.4, 2.0, 2.1
  - most of the samplers
  - pix2pix edit
  - most of the controlnets
  - memory management
  - python 3.8 support

wip
2023-11-22 13:22:00 -08:00

49 lines
1.1 KiB
Python

import os.path
_base_dir = os.path.dirname(os.path.realpath(__file__))
WEIGHT_MAPS_PATH = os.path.join(_base_dir, "maps")
WEIGHT_INFO_PATH = os.path.join(_base_dir, "weight_info")
class MODEL_NAMES:
SD15 = "stable-diffusion-1-5"
class COMPONENT_NAMES:
VAE = "vae"
TEXT_ENCODER = "text"
UNET = "unet"
LORA = "lora"
class FORMAT_NAMES:
COMPVIS = "compvis"
DIFFUSERS = "diffusers"
REFINERS = "refiners"
def save_model_info(model_name, component_name, format_name, info_type, data):
import json
model_name = model_name.replace("_", "-")
component_name = component_name.replace("_", "-")
format_name = format_name.replace("_", "-")
filename = os.path.join(
WEIGHT_INFO_PATH,
f"{model_name}_{component_name}_{format_name}.{info_type}.json",
)
with open(filename, "w") as f:
f.write(json.dumps(data, indent=2))
def prefixes_only(keys):
new_keys = []
prev_key = None
for k in keys:
new_key = k.rsplit(".", 1)[0]
if new_key != prev_key:
new_keys.append(new_key)
prev_key = new_key
return new_keys