Commit da464a3f authored by AUTOMATIC1111's avatar AUTOMATIC1111
Browse files

SDXL support

parent af081211

modules/launch_utils.py

+17 −0
Original line number Diff line number Diff line
@@ -224,6 +224,20 @@ def run_extensions_installers(settings_file): 
        run_extension_installer(os.path.join(extensions_dir, dirname_extension))





def mute_sdxl_imports():

    """create fake modules that SDXL wants to import but doesn't actually use for our purposes"""



    import importlib



    module = importlib.util.module_from_spec(importlib.machinery.ModuleSpec('taming.modules.losses.lpips', None))

    module.LPIPS = None

    sys.modules['taming.modules.losses.lpips'] = module



    module = importlib.util.module_from_spec(importlib.machinery.ModuleSpec('sgm.data', None))

    module.StableDataModuleFromConfig = None

    sys.modules['sgm.data'] = module





def prepare_environment():

    torch_index_url = os.environ.get('TORCH_INDEX_URL', "https://download.pytorch.org/whl/cu118")

    torch_command = os.environ.get('TORCH_COMMAND', f"pip install torch==2.0.1 torchvision==0.15.2 --extra-index-url {torch_index_url}")
@@ -319,11 +333,14 @@ def prepare_environment(): 
    if args.update_all_extensions:

        git_pull_recursive(extensions_dir)



    mute_sdxl_imports()



    if "--exit" in sys.argv:

        print("Exiting because of --exit argument")

        exit(0)







def configure_for_tests():

    if "--api" not in sys.argv:

        sys.argv.append("--api")

modules/lowvram.py

+37 −14
Original line number Diff line number Diff line
@@ -53,19 +53,46 @@ def setup_for_low_vram(sd_model, use_medvram): 
        send_me_to_gpu(first_stage_model, None)

        return first_stage_model_decode(z)



    # for SD1, cond_stage_model is CLIP and its NN is in the tranformer frield, but for SD2, it's open clip, and it's in model field

    if hasattr(sd_model.cond_stage_model, 'model'):

        sd_model.cond_stage_model.transformer = sd_model.cond_stage_model.model



    # remove several big modules: cond, first_stage, depth/embedder (if applicable), and unet from the model and then

    # send the model to GPU. Then put modules back. the modules will be in CPU.

    stored = sd_model.cond_stage_model.transformer, sd_model.first_stage_model, getattr(sd_model, 'depth_model', None), getattr(sd_model, 'embedder', None), sd_model.model

    sd_model.cond_stage_model.transformer, sd_model.first_stage_model, sd_model.depth_model, sd_model.embedder, sd_model.model = None, None, None, None, None

    to_remain_in_cpu = [

        (sd_model, 'first_stage_model'),

        (sd_model, 'depth_model'),

        (sd_model, 'embedder'),

        (sd_model, 'model'),

        (sd_model, 'embedder'),

    ]



    is_sdxl = hasattr(sd_model, 'conditioner')

    is_sd2 = not is_sdxl and hasattr(sd_model.cond_stage_model, 'model')



    if is_sdxl:

        to_remain_in_cpu.append((sd_model, 'conditioner'))

    elif is_sd2:

        to_remain_in_cpu.append((sd_model.cond_stage_model, 'model'))

    else:

        to_remain_in_cpu.append((sd_model.cond_stage_model, 'transformer'))



    # remove several big modules: cond, first_stage, depth/embedder (if applicable), and unet from the model

    stored = []

    for obj, field in to_remain_in_cpu:

        module = getattr(obj, field, None)

        stored.append(module)

        setattr(obj, field, None)



    # send the model to GPU.

    sd_model.to(devices.device)

    sd_model.cond_stage_model.transformer, sd_model.first_stage_model, sd_model.depth_model, sd_model.embedder, sd_model.model = stored



    # put modules back. the modules will be in CPU.

    for (obj, field), module in zip(to_remain_in_cpu, stored):

        setattr(obj, field, module)



    # register hooks for those the first three models

    if is_sdxl:

        sd_model.conditioner.register_forward_pre_hook(send_me_to_gpu)

    elif is_sd2:

        sd_model.cond_stage_model.model.register_forward_pre_hook(send_me_to_gpu)

    else:

        sd_model.cond_stage_model.transformer.register_forward_pre_hook(send_me_to_gpu)



    sd_model.first_stage_model.register_forward_pre_hook(send_me_to_gpu)

    sd_model.first_stage_model.encode = first_stage_model_encode_wrap

    sd_model.first_stage_model.decode = first_stage_model_decode_wrap
@@ -75,10 +102,6 @@ def setup_for_low_vram(sd_model, use_medvram): 
        sd_model.embedder.register_forward_pre_hook(send_me_to_gpu)

    parents[sd_model.cond_stage_model.transformer] = sd_model.cond_stage_model



    if hasattr(sd_model.cond_stage_model, 'model'):

        sd_model.cond_stage_model.model = sd_model.cond_stage_model.transformer

        del sd_model.cond_stage_model.transformer



    if use_medvram:

        sd_model.model.register_forward_pre_hook(send_me_to_gpu)

    else:

modules/paths.py

+8 −1
Original line number Diff line number Diff line
@@ -20,7 +20,7 @@ assert sd_path is not None, f"Couldn't find Stable Diffusion in any of: {possibl 


path_dirs = [

    (sd_path, 'ldm', 'Stable Diffusion', []),

    (os.path.join(sd_path, '../generative-models'), 'sgm', 'Stable Diffusion XL', []),

    (os.path.join(sd_path, '../generative-models'), 'sgm', 'Stable Diffusion XL', ["sgm"]),

    (os.path.join(sd_path, '../CodeFormer'), 'inference_codeformer.py', 'CodeFormer', []),

    (os.path.join(sd_path, '../BLIP'), 'models/blip.py', 'BLIP', []),

    (os.path.join(sd_path, '../k-diffusion'), 'k_diffusion/sampling.py', 'k_diffusion', ["atstart"]),
@@ -36,6 +36,13 @@ for d, must_exist, what, options in path_dirs: 
        d = os.path.abspath(d)

        if "atstart" in options:

            sys.path.insert(0, d)

        elif "sgm" in options:

            # Stable Diffusion XL repo has scripts dir with __init__.py in it which ruins every extension's scripts dir, so we

            # import sgm and remove it from sys.path so that when a script imports scripts.something, it doesbn't use sgm's scripts dir.



            sys.path.insert(0, d)

            import sgm

            sys.path.pop(0)

        else:

            sys.path.append(d)

        paths[what] = d

modules/processing.py

+5 −2
Original line number Diff line number Diff line
@@ -343,10 +343,13 @@ class StableDiffusionProcessing: 
        return cache[1]



    def setup_conds(self):

        prompts = prompt_parser.SdConditioning(self.prompts, width=self.width, height=self.height)

        negative_prompts = prompt_parser.SdConditioning(self.negative_prompts, width=self.width, height=self.height)



        sampler_config = sd_samplers.find_sampler_config(self.sampler_name)

        self.step_multiplier = 2 if sampler_config and sampler_config.options.get("second_order", False) else 1

        self.uc = self.get_conds_with_caching(prompt_parser.get_learned_conditioning, self.negative_prompts, self.steps * self.step_multiplier, [self.cached_uc], self.extra_network_data)

        self.c = self.get_conds_with_caching(prompt_parser.get_multicond_learned_conditioning, self.prompts, self.steps * self.step_multiplier, [self.cached_c], self.extra_network_data)

        self.uc = self.get_conds_with_caching(prompt_parser.get_learned_conditioning, negative_prompts, self.steps * self.step_multiplier, [self.cached_uc], self.extra_network_data)

        self.c = self.get_conds_with_caching(prompt_parser.get_multicond_learned_conditioning, prompts, self.steps * self.step_multiplier, [self.cached_c], self.extra_network_data)



    def parse_extra_network_prompts(self):

        self.prompts, self.extra_network_data = extra_networks.parse_prompts(self.prompts)

modules/prompt_parser.py

+20 −3
Original line number Diff line number Diff line 
from __future__ import annotations



import re

from collections import namedtuple

from typing import List
@@ -109,7 +111,19 @@ def get_learned_conditioning_prompt_schedules(prompts, steps): 
ScheduledPromptConditioning = namedtuple("ScheduledPromptConditioning", ["end_at_step", "cond"])





def get_learned_conditioning(model, prompts, steps):

class SdConditioning(list):

    """

    A list with prompts for stable diffusion's conditioner model.

    Can also specify width and height of created image - SDXL needs it.

    """

    def __init__(self, prompts, width=None, height=None):

        super().__init__()

        self.extend(prompts)

        self.width = width or getattr(prompts, 'width', None)

        self.height = height or getattr(prompts, 'height', None)





def get_learned_conditioning(model, prompts: SdConditioning | list[str], steps):

    """converts a list of prompts into a list of prompt schedules - each schedule is a list of ScheduledPromptConditioning, specifying the comdition (cond),

    and the sampling step at which this condition is to be replaced by the next one.
@@ -160,11 +174,13 @@ def get_learned_conditioning(model, prompts, steps): 
re_AND = re.compile(r"\bAND\b")

re_weight = re.compile(r"^(.*?)(?:\s*:\s*([-+]?(?:\d+\.?|\d*\.\d+)))?\s*$")



def get_multicond_prompt_list(prompts):



def get_multicond_prompt_list(prompts: SdConditioning | list[str]):

    res_indexes = []



    prompt_flat_list = []

    prompt_indexes = {}

    prompt_flat_list = SdConditioning(prompts)

    prompt_flat_list.clear()



    for prompt in prompts:

        subprompts = re_AND.split(prompt)
@@ -201,6 +217,7 @@ class MulticondLearnedConditioning: 
        self.shape: tuple = shape  # the shape field is needed to send this object to DDIM/PLMS

        self.batch: List[List[ComposableScheduledPromptConditioning]] = batch





def get_multicond_learned_conditioning(model, prompts, steps) -> MulticondLearnedConditioning:

    """same as get_learned_conditioning, but returns a list of ScheduledPromptConditioning along with the weight objects for each prompt.

    For each prompt, the list is obtained by splitting the prompt using the AND separator.

CRA Git | Maintained and supported by SUSTech CRA and CCSE