fluxdev/stable-diffusion-webui-forge

1import torch2import contextlib3 4from ldm_patched.modules import model_management5from ldm_patched.modules import model_detection6 7from ldm_patched.modules.sd import VAE, CLIP, load_model_weights8import ldm_patched.modules.model_patcher9import ldm_patched.modules.utils10import ldm_patched.modules.clip_vision11 12from omegaconf import OmegaConf13from modules.sd_models_config import find_checkpoint_config14from modules.shared import cmd_opts15from modules import sd_hijack16from modules.sd_models_xl import extend_sdxl17from ldm.util import instantiate_from_config18from modules_forge import forge_clip19from modules_forge.unet_patcher import UnetPatcher20from ldm_patched.modules.model_base import model_sampling, ModelType21 22import open_clip23from transformers import CLIPTextModel, CLIPTokenizer24 25 26class FakeObject:27    def __init__(self, *args, **kwargs):28        super().__init__()29        self.visual = None30        return31 32    def eval(self, *args, **kwargs):33        return self34 35    def parameters(self, *args, **kwargs):36        return []37 38 39class ForgeSD:40    def __init__(self, unet, clip, vae, clipvision):41        self.unet = unet42        self.clip = clip43        self.vae = vae44        self.clipvision = clipvision45 46    def shallow_copy(self):47        return ForgeSD(48            self.unet,49            self.clip,50            self.vae,51            self.clipvision52        )53 54 55@contextlib.contextmanager56def no_clip():57    backup_openclip = open_clip.create_model_and_transforms58    backup_CLIPTextModel = CLIPTextModel.from_pretrained59    backup_CLIPTokenizer = CLIPTokenizer.from_pretrained60 61    try:62        open_clip.create_model_and_transforms = lambda *args, **kwargs: (FakeObject(), None, None)63        CLIPTextModel.from_pretrained = lambda *args, **kwargs: FakeObject()64        CLIPTokenizer.from_pretrained = lambda *args, **kwargs: FakeObject()65        yield66 67    finally:68        open_clip.create_model_and_transforms = backup_openclip69        CLIPTextModel.from_pretrained = backup_CLIPTextModel70        CLIPTokenizer.from_pretrained = backup_CLIPTokenizer71    return72 73 74def load_checkpoint_guess_config(sd, output_vae=True, output_clip=True, output_clipvision=False, embedding_directory=None, output_model=True):75    sd_keys = sd.keys()76    clip = None77    clipvision = None78    vae = None79    model = None80    model_patcher = None81    clip_target = None82 83    parameters = ldm_patched.modules.utils.calculate_parameters(sd, "model.diffusion_model.")84    unet_dtype = model_management.unet_dtype(model_params=parameters)85    load_device = model_management.get_torch_device()86    manual_cast_dtype = model_management.unet_manual_cast(unet_dtype, load_device)87 88    class WeightsLoader(torch.nn.Module):89        pass90 91    model_config = model_detection.model_config_from_unet(sd, "model.diffusion_model.", unet_dtype)92    model_config.set_manual_cast(manual_cast_dtype)93 94    if model_config is None:95        raise RuntimeError("ERROR: Could not detect model type")96 97    if model_config.clip_vision_prefix is not None:98        if output_clipvision:99            clipvision = ldm_patched.modules.clip_vision.load_clipvision_from_sd(sd, model_config.clip_vision_prefix, True)100 101    if output_model:102        inital_load_device = model_management.unet_inital_load_device(parameters, unet_dtype)103        offload_device = model_management.unet_offload_device()104        model = model_config.get_model(sd, "model.diffusion_model.", device=inital_load_device)105        model.load_model_weights(sd, "model.diffusion_model.")106 107    if output_vae:108        vae_sd = ldm_patched.modules.utils.state_dict_prefix_replace(sd, {"first_stage_model.": ""}, filter_keys=True)109        vae_sd = model_config.process_vae_state_dict(vae_sd)110        vae = VAE(sd=vae_sd)111 112    if output_clip:113        w = WeightsLoader()114        clip_target = model_config.clip_target()115        if clip_target is not None:116            clip = CLIP(clip_target, embedding_directory=embedding_directory)117            w.cond_stage_model = clip.cond_stage_model118            sd = model_config.process_clip_state_dict(sd)119            load_model_weights(w, sd)120 121    left_over = sd.keys()122    if len(left_over) > 0:123        print("left over keys:", left_over)124 125    if output_model:126        model_patcher = UnetPatcher(model, load_device=load_device, offload_device=model_management.unet_offload_device(), current_device=inital_load_device)127        if inital_load_device != torch.device("cpu"):128            print("loaded straight to GPU")129            model_management.load_model_gpu(model_patcher)130 131    return ForgeSD(model_patcher, clip, vae, clipvision)132 133 134@torch.no_grad()135def load_model_for_a1111(timer, checkpoint_info=None, state_dict=None):136    a1111_config_filename = find_checkpoint_config(state_dict, checkpoint_info)137    a1111_config = OmegaConf.load(a1111_config_filename)138    timer.record("forge solving config")139 140    if hasattr(a1111_config.model.params, 'network_config'):141        a1111_config.model.params.network_config.target = 'modules_forge.forge_loader.FakeObject'142 143    if hasattr(a1111_config.model.params, 'unet_config'):144        a1111_config.model.params.unet_config.target = 'modules_forge.forge_loader.FakeObject'145 146    if hasattr(a1111_config.model.params, 'first_stage_config'):147        a1111_config.model.params.first_stage_config.target = 'modules_forge.forge_loader.FakeObject'148 149    with no_clip():150        sd_model = instantiate_from_config(a1111_config.model)151 152    timer.record("forge instantiate config")153 154    forge_objects = load_checkpoint_guess_config(155        state_dict,156        output_vae=True,157        output_clip=True,158        output_clipvision=True,159        embedding_directory=cmd_opts.embeddings_dir,160        output_model=True161    )162    sd_model.forge_objects = forge_objects163    sd_model.forge_objects_original = forge_objects.shallow_copy()164    sd_model.forge_objects_after_applying_lora = forge_objects.shallow_copy()165    timer.record("forge load real models")166 167    sd_model.first_stage_model = forge_objects.vae.first_stage_model168    sd_model.model.diffusion_model = forge_objects.unet.model.diffusion_model169 170    conditioner = getattr(sd_model, 'conditioner', None)171    if conditioner:172        text_cond_models = []173 174        for i in range(len(conditioner.embedders)):175            embedder = conditioner.embedders[i]176            typename = type(embedder).__name__177            if typename == 'FrozenCLIPEmbedder':  # SDXL Clip L178                embedder.tokenizer = forge_objects.clip.tokenizer.clip_l.tokenizer179                embedder.transformer = forge_objects.clip.cond_stage_model.clip_l.transformer180                model_embeddings = embedder.transformer.text_model.embeddings181                model_embeddings.token_embedding = sd_hijack.EmbeddingsWithFixes(182                    model_embeddings.token_embedding, sd_hijack.model_hijack)183                embedder = forge_clip.CLIP_SD_XL_L(embedder, sd_hijack.model_hijack)184                conditioner.embedders[i] = embedder185                text_cond_models.append(embedder)186            elif typename == 'FrozenOpenCLIPEmbedder2':  # SDXL Clip G187                embedder.tokenizer = forge_objects.clip.tokenizer.clip_g.tokenizer188                embedder.transformer = forge_objects.clip.cond_stage_model.clip_g.transformer189                embedder.text_projection = forge_objects.clip.cond_stage_model.clip_g.text_projection190                model_embeddings = embedder.transformer.text_model.embeddings191                model_embeddings.token_embedding = sd_hijack.EmbeddingsWithFixes(192                    model_embeddings.token_embedding, sd_hijack.model_hijack, textual_inversion_key='clip_g')193                embedder = forge_clip.CLIP_SD_XL_G(embedder, sd_hijack.model_hijack)194                conditioner.embedders[i] = embedder195                text_cond_models.append(embedder)196 197        if len(text_cond_models) == 1:198            sd_model.cond_stage_model = text_cond_models[0]199        else:200            sd_model.cond_stage_model = conditioner201    elif type(sd_model.cond_stage_model).__name__ == 'FrozenCLIPEmbedder':  # SD15 Clip202        sd_model.cond_stage_model.tokenizer = forge_objects.clip.tokenizer.clip_l.tokenizer203        sd_model.cond_stage_model.transformer = forge_objects.clip.cond_stage_model.clip_l.transformer204        model_embeddings = sd_model.cond_stage_model.transformer.text_model.embeddings205        model_embeddings.token_embedding = sd_hijack.EmbeddingsWithFixes(206            model_embeddings.token_embedding, sd_hijack.model_hijack)207        sd_model.cond_stage_model = forge_clip.CLIP_SD_15_L(sd_model.cond_stage_model, sd_hijack.model_hijack)208    elif type(sd_model.cond_stage_model).__name__ == 'FrozenOpenCLIPEmbedder':  # SD21 Clip209        sd_model.cond_stage_model.tokenizer = forge_objects.clip.tokenizer.clip_h.tokenizer210        sd_model.cond_stage_model.transformer = forge_objects.clip.cond_stage_model.clip_h.transformer211        model_embeddings = sd_model.cond_stage_model.transformer.text_model.embeddings212        model_embeddings.token_embedding = sd_hijack.EmbeddingsWithFixes(213            model_embeddings.token_embedding, sd_hijack.model_hijack)214        sd_model.cond_stage_model = forge_clip.CLIP_SD_21_H(sd_model.cond_stage_model, sd_hijack.model_hijack)215    else:216        raise NotImplementedError('Bad Clip Class Name:' + type(sd_model.cond_stage_model).__name__)217 218    timer.record("forge set components")219 220    sd_model_hash = checkpoint_info.calculate_shorthash()221    timer.record("calculate hash")222 223    if getattr(sd_model, 'parameterization', None) == 'v':224        sd_model.forge_objects.unet.model.model_sampling = model_sampling(sd_model.forge_objects.unet.model.model_config, ModelType.V_PREDICTION)225 226    sd_model.is_sdxl = conditioner is not None227    sd_model.is_sd2 = not sd_model.is_sdxl and hasattr(sd_model.cond_stage_model, 'model')228    sd_model.is_sd1 = not sd_model.is_sdxl and not sd_model.is_sd2229    sd_model.is_ssd = sd_model.is_sdxl and 'model.diffusion_model.middle_block.1.transformer_blocks.0.attn1.to_q.weight' not in sd_model.state_dict().keys()230    if sd_model.is_sdxl:231        extend_sdxl(sd_model)232    sd_model.sd_model_hash = sd_model_hash233    sd_model.sd_model_checkpoint = checkpoint_info.filename234    sd_model.sd_checkpoint_info = checkpoint_info235 236    @torch.inference_mode()237    def patched_decode_first_stage(x):238        sample = sd_model.forge_objects.unet.model.model_config.latent_format.process_out(x)239        sample = sd_model.forge_objects.vae.decode(sample).movedim(-1, 1) * 2.0 - 1.0240        return sample.to(x)241 242    @torch.inference_mode()243    def patched_encode_first_stage(x):244        sample = sd_model.forge_objects.vae.encode(x.movedim(1, -1) * 0.5 + 0.5)245        sample = sd_model.forge_objects.unet.model.model_config.latent_format.process_in(sample)246        return sample.to(x)247 248    sd_model.ema_scope = lambda *args, **kwargs: contextlib.nullcontext()249    sd_model.get_first_stage_encoding = lambda x: x250    sd_model.decode_first_stage = patched_decode_first_stage251    sd_model.encode_first_stage = patched_encode_first_stage252    sd_model.clip = sd_model.cond_stage_model253    sd_model.tiling_enabled = False254    timer.record("forge finalize")255 256    sd_model.current_lora_hash = str([])257    return sd_model258