mirror of
https://github.com/lllyasviel/stable-diffusion-webui-forge.git
synced 2026-07-21 21:01:24 +08:00
lint
This commit is contained in:
Haoming
parent
355e7dfcc3
commit
3b84a5d299
@ -1 +1 @@
|
||||
# WIP Backend for Forge
|
||||
<h2 align="center">W.I.P Backend for Forge</h2>
|
||||
|
||||
@ -61,7 +61,4 @@ parser.add_argument("--disable-gpu-warning", action="store_true")
|
||||
args = parser.parse_known_args()[0]
|
||||
|
||||
# Some dynamic args that may be changed by webui rather than cmd flags.
|
||||
dynamic_args = dict(
|
||||
embedding_dir='./embeddings',
|
||||
emphasis_name='original'
|
||||
)
|
||||
dynamic_args = dict(embedding_dir="./embeddings", emphasis_name="original")
|
||||
|
||||
@ -1,10 +1,10 @@
|
||||
import math
|
||||
import torch
|
||||
|
||||
import einops
|
||||
import torch
|
||||
|
||||
from backend.args import args
|
||||
from backend import memory_management
|
||||
|
||||
from backend.args import args
|
||||
|
||||
BROKEN_XFORMERS = False
|
||||
if memory_management.xformers_enabled():
|
||||
@ -42,7 +42,7 @@ def attention_basic(q, k, v, heads, mask=None, attn_precision=None, skip_reshape
|
||||
b, _, dim_head = q.shape
|
||||
dim_head //= heads
|
||||
|
||||
scale = dim_head ** -0.5
|
||||
scale = dim_head**-0.5
|
||||
|
||||
h = heads
|
||||
if skip_reshape:
|
||||
@ -52,26 +52,22 @@ def attention_basic(q, k, v, heads, mask=None, attn_precision=None, skip_reshape
|
||||
)
|
||||
else:
|
||||
q, k, v = map(
|
||||
lambda t: t.unsqueeze(3)
|
||||
.reshape(b, -1, heads, dim_head)
|
||||
.permute(0, 2, 1, 3)
|
||||
.reshape(b * heads, -1, dim_head)
|
||||
.contiguous(),
|
||||
lambda t: t.unsqueeze(3).reshape(b, -1, heads, dim_head).permute(0, 2, 1, 3).reshape(b * heads, -1, dim_head).contiguous(),
|
||||
(q, k, v),
|
||||
)
|
||||
|
||||
if attn_precision == torch.float32:
|
||||
sim = torch.einsum('b i d, b j d -> b i j', q.float(), k.float()) * scale
|
||||
sim = torch.einsum("b i d, b j d -> b i j", q.float(), k.float()) * scale
|
||||
else:
|
||||
sim = torch.einsum('b i d, b j d -> b i j', q, k) * scale
|
||||
sim = torch.einsum("b i d, b j d -> b i j", q, k) * scale
|
||||
|
||||
del q, k
|
||||
|
||||
if exists(mask):
|
||||
if mask.dtype == torch.bool:
|
||||
mask = einops.rearrange(mask, 'b ... -> b (...)')
|
||||
mask = einops.rearrange(mask, "b ... -> b (...)")
|
||||
max_neg_value = -torch.finfo(sim.dtype).max
|
||||
mask = einops.repeat(mask, 'b j -> (b h) () j', h=h)
|
||||
mask = einops.repeat(mask, "b j -> (b h) () j", h=h)
|
||||
sim.masked_fill_(~mask, max_neg_value)
|
||||
else:
|
||||
if len(mask.shape) == 2:
|
||||
@ -82,13 +78,8 @@ def attention_basic(q, k, v, heads, mask=None, attn_precision=None, skip_reshape
|
||||
sim.add_(mask)
|
||||
|
||||
sim = sim.softmax(dim=-1)
|
||||
out = torch.einsum('b i j, b j d -> b i d', sim.to(v.dtype), v)
|
||||
out = (
|
||||
out.unsqueeze(0)
|
||||
.reshape(b, heads, -1, dim_head)
|
||||
.permute(0, 2, 1, 3)
|
||||
.reshape(b, -1, heads * dim_head)
|
||||
)
|
||||
out = torch.einsum("b i j, b j d -> b i d", sim.to(v.dtype), v)
|
||||
out = out.unsqueeze(0).reshape(b, heads, -1, dim_head).permute(0, 2, 1, 3).reshape(b, -1, heads * dim_head)
|
||||
return out
|
||||
|
||||
|
||||
@ -101,7 +92,7 @@ def attention_split(q, k, v, heads, mask=None, attn_precision=None, skip_reshape
|
||||
b, _, dim_head = q.shape
|
||||
dim_head //= heads
|
||||
|
||||
scale = dim_head ** -0.5
|
||||
scale = dim_head**-0.5
|
||||
|
||||
h = heads
|
||||
if skip_reshape:
|
||||
@ -111,11 +102,7 @@ def attention_split(q, k, v, heads, mask=None, attn_precision=None, skip_reshape
|
||||
)
|
||||
else:
|
||||
q, k, v = map(
|
||||
lambda t: t.unsqueeze(3)
|
||||
.reshape(b, -1, heads, dim_head)
|
||||
.permute(0, 2, 1, 3)
|
||||
.reshape(b * heads, -1, dim_head)
|
||||
.contiguous(),
|
||||
lambda t: t.unsqueeze(3).reshape(b, -1, heads, dim_head).permute(0, 2, 1, 3).reshape(b * heads, -1, dim_head).contiguous(),
|
||||
(q, k, v),
|
||||
)
|
||||
|
||||
@ -130,7 +117,7 @@ def attention_split(q, k, v, heads, mask=None, attn_precision=None, skip_reshape
|
||||
element_size = q.element_size()
|
||||
upcast = False
|
||||
|
||||
gb = 1024 ** 3
|
||||
gb = 1024**3
|
||||
tensor_size = q.shape[0] * q.shape[1] * k.shape[1] * element_size
|
||||
modifier = 3
|
||||
mem_required = tensor_size * modifier
|
||||
@ -143,8 +130,7 @@ def attention_split(q, k, v, heads, mask=None, attn_precision=None, skip_reshape
|
||||
|
||||
if steps > 64:
|
||||
max_res = math.floor(math.sqrt(math.sqrt(mem_free_total / 2.5)) / 8) * 64
|
||||
raise RuntimeError(f'Not enough memory, use lower resolution (max approx. {max_res}x{max_res}). '
|
||||
f'Need: {mem_required / 64 / gb:0.1f}GB free, Have:{mem_free_total / gb:0.1f}GB free')
|
||||
raise RuntimeError(f"Not enough memory, use lower resolution (max approx. {max_res}x{max_res}). " f"Need: {mem_required / 64 / gb:0.1f}GB free, Have:{mem_free_total / gb:0.1f}GB free")
|
||||
|
||||
if mask is not None:
|
||||
if len(mask.shape) == 2:
|
||||
@ -162,10 +148,10 @@ def attention_split(q, k, v, heads, mask=None, attn_precision=None, skip_reshape
|
||||
for i in range(0, q.shape[1], slice_size):
|
||||
end = i + slice_size
|
||||
if upcast:
|
||||
with torch.autocast(enabled=False, device_type='cuda'):
|
||||
s1 = torch.einsum('b i d, b j d -> b i j', q[:, i:end].float(), k.float()) * scale
|
||||
with torch.autocast(enabled=False, device_type="cuda"):
|
||||
s1 = torch.einsum("b i d, b j d -> b i j", q[:, i:end].float(), k.float()) * scale
|
||||
else:
|
||||
s1 = torch.einsum('b i d, b j d -> b i j', q[:, i:end], k) * scale
|
||||
s1 = torch.einsum("b i d, b j d -> b i j", q[:, i:end], k) * scale
|
||||
|
||||
if mask is not None:
|
||||
if len(mask.shape) == 2:
|
||||
@ -177,7 +163,7 @@ def attention_split(q, k, v, heads, mask=None, attn_precision=None, skip_reshape
|
||||
del s1
|
||||
first_op_done = True
|
||||
|
||||
r1[:, i:end] = torch.einsum('b i j, b j d -> b i d', s2, v)
|
||||
r1[:, i:end] = torch.einsum("b i j, b j d -> b i d", s2, v)
|
||||
del s2
|
||||
break
|
||||
except memory_management.OOM_EXCEPTION as e:
|
||||
@ -196,12 +182,7 @@ def attention_split(q, k, v, heads, mask=None, attn_precision=None, skip_reshape
|
||||
|
||||
del q, k, v
|
||||
|
||||
r1 = (
|
||||
r1.unsqueeze(0)
|
||||
.reshape(b, heads, -1, dim_head)
|
||||
.permute(0, 2, 1, 3)
|
||||
.reshape(b, -1, heads * dim_head)
|
||||
)
|
||||
r1 = r1.unsqueeze(0).reshape(b, heads, -1, dim_head).permute(0, 2, 1, 3).reshape(b, -1, heads * dim_head)
|
||||
return r1
|
||||
|
||||
|
||||
@ -229,22 +210,15 @@ def attention_xformers(q, k, v, heads, mask=None, attn_precision=None, skip_resh
|
||||
if mask is not None:
|
||||
pad = 8 - q.shape[1] % 8
|
||||
mask_out = torch.empty([q.shape[0], q.shape[1], q.shape[1] + pad], dtype=q.dtype, device=q.device)
|
||||
mask_out[:, :, :mask.shape[-1]] = mask
|
||||
mask = mask_out[:, :, :mask.shape[-1]]
|
||||
mask_out[:, :, : mask.shape[-1]] = mask
|
||||
mask = mask_out[:, :, : mask.shape[-1]]
|
||||
|
||||
out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=mask)
|
||||
|
||||
if skip_reshape:
|
||||
out = (
|
||||
out.unsqueeze(0)
|
||||
.reshape(b, heads, -1, dim_head)
|
||||
.permute(0, 2, 1, 3)
|
||||
.reshape(b, -1, heads * dim_head)
|
||||
)
|
||||
out = out.unsqueeze(0).reshape(b, heads, -1, dim_head).permute(0, 2, 1, 3).reshape(b, -1, heads * dim_head)
|
||||
else:
|
||||
out = (
|
||||
out.reshape(b, -1, heads * dim_head)
|
||||
)
|
||||
out = out.reshape(b, -1, heads * dim_head)
|
||||
|
||||
return out
|
||||
|
||||
@ -261,19 +235,17 @@ def attention_pytorch(q, k, v, heads, mask=None, attn_precision=None, skip_resha
|
||||
)
|
||||
|
||||
out = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=False)
|
||||
out = (
|
||||
out.transpose(1, 2).reshape(b, -1, heads * dim_head)
|
||||
)
|
||||
out = out.transpose(1, 2).reshape(b, -1, heads * dim_head)
|
||||
return out
|
||||
|
||||
|
||||
def slice_attention_single_head_spatial(q, k, v):
|
||||
r1 = torch.zeros_like(k, device=q.device)
|
||||
scale = (int(q.shape[-1]) ** (-0.5))
|
||||
scale = int(q.shape[-1]) ** (-0.5)
|
||||
|
||||
mem_free_total = memory_management.get_free_memory(q.device)
|
||||
|
||||
gb = 1024 ** 3
|
||||
gb = 1024**3
|
||||
tensor_size = q.shape[0] * q.shape[1] * k.shape[2] * q.element_size()
|
||||
modifier = 3 if q.element_size() == 2 else 2.5
|
||||
mem_required = tensor_size * modifier
|
||||
@ -332,8 +304,7 @@ def xformers_attention_single_head_spatial(q, k, v):
|
||||
out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=None)
|
||||
out = out.transpose(1, 2).reshape(B, C, H, W)
|
||||
except NotImplementedError as e:
|
||||
out = slice_attention_single_head_spatial(q.view(B, -1, C), k.view(B, -1, C).transpose(1, 2),
|
||||
v.view(B, -1, C).transpose(1, 2)).reshape(B, C, H, W)
|
||||
out = slice_attention_single_head_spatial(q.view(B, -1, C), k.view(B, -1, C).transpose(1, 2), v.view(B, -1, C).transpose(1, 2)).reshape(B, C, H, W)
|
||||
return out
|
||||
|
||||
|
||||
@ -350,8 +321,7 @@ def pytorch_attention_single_head_spatial(q, k, v):
|
||||
out = out.transpose(2, 3).reshape(B, C, H, W)
|
||||
except memory_management.OOM_EXCEPTION as e:
|
||||
print("scaled_dot_product_attention OOMed: switched to slice attention")
|
||||
out = slice_attention_single_head_spatial(q.view(B, -1, C), k.view(B, -1, C).transpose(1, 2),
|
||||
v.view(B, -1, C).transpose(1, 2)).reshape(B, C, H, W)
|
||||
out = slice_attention_single_head_spatial(q.view(B, -1, C), k.view(B, -1, C).transpose(1, 2), v.view(B, -1, C).transpose(1, 2)).reshape(B, C, H, W)
|
||||
return out
|
||||
|
||||
|
||||
@ -392,9 +362,7 @@ class AttentionProcessorForge:
|
||||
batch_size, channel, height, width = hidden_states.shape
|
||||
hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
|
||||
|
||||
batch_size, sequence_length, _ = (
|
||||
hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
|
||||
)
|
||||
batch_size, sequence_length, _ = hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
|
||||
|
||||
if attention_mask is not None:
|
||||
attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
|
||||
|
||||
@ -1,27 +1,28 @@
|
||||
import os
|
||||
import torch
|
||||
import logging
|
||||
import importlib
|
||||
import logging
|
||||
import os
|
||||
|
||||
import backend.args
|
||||
import huggingface_guess
|
||||
|
||||
import torch
|
||||
from diffusers import DiffusionPipeline
|
||||
from transformers import modeling_utils
|
||||
|
||||
import backend.args
|
||||
from backend import memory_management
|
||||
from backend.utils import read_arbitrary_config, load_torch_file, beautiful_print_gguf_state_dict_statics
|
||||
from backend.state_dict import try_filter_state_dict, load_state_dict
|
||||
from backend.operations import using_forge_operations
|
||||
from backend.nn.vae import IntegratedAutoencoderKL
|
||||
from backend.nn.clip import IntegratedCLIP
|
||||
from backend.nn.unet import IntegratedUNet2DConditionModel
|
||||
|
||||
from backend.diffusion_engine.chroma import Chroma
|
||||
from backend.diffusion_engine.flux import Flux
|
||||
from backend.diffusion_engine.sd15 import StableDiffusion
|
||||
from backend.diffusion_engine.sdxl import StableDiffusionXL, StableDiffusionXLRefiner
|
||||
from backend.diffusion_engine.flux import Flux
|
||||
from backend.diffusion_engine.chroma import Chroma
|
||||
|
||||
from backend.nn.clip import IntegratedCLIP
|
||||
from backend.nn.unet import IntegratedUNet2DConditionModel
|
||||
from backend.nn.vae import IntegratedAutoencoderKL
|
||||
from backend.operations import using_forge_operations
|
||||
from backend.state_dict import load_state_dict, try_filter_state_dict
|
||||
from backend.utils import (
|
||||
beautiful_print_gguf_state_dict_statics,
|
||||
load_torch_file,
|
||||
read_arbitrary_config,
|
||||
)
|
||||
|
||||
possible_models = [StableDiffusion, StableDiffusionXLRefiner, StableDiffusionXL, Chroma, Flux]
|
||||
|
||||
@ -33,34 +34,35 @@ dir_path = os.path.dirname(__file__)
|
||||
def load_huggingface_component(guess, component_name, lib_name, cls_name, repo_path, state_dict):
|
||||
config_path = os.path.join(repo_path, component_name)
|
||||
|
||||
if component_name in ['feature_extractor', 'safety_checker']:
|
||||
if component_name in ["feature_extractor", "safety_checker"]:
|
||||
return None
|
||||
|
||||
if lib_name in ['transformers', 'diffusers']:
|
||||
if component_name in ['scheduler']:
|
||||
if lib_name in ["transformers", "diffusers"]:
|
||||
if component_name in ["scheduler"]:
|
||||
cls = getattr(importlib.import_module(lib_name), cls_name)
|
||||
return cls.from_pretrained(os.path.join(repo_path, component_name))
|
||||
if component_name.startswith('tokenizer'):
|
||||
if component_name.startswith("tokenizer"):
|
||||
cls = getattr(importlib.import_module(lib_name), cls_name)
|
||||
comp = cls.from_pretrained(os.path.join(repo_path, component_name))
|
||||
comp._eventual_warn_about_too_long_sequence = lambda *args, **kwargs: None
|
||||
return comp
|
||||
if cls_name in ['AutoencoderKL']:
|
||||
assert isinstance(state_dict, dict) and len(state_dict) > 16, 'You do not have VAE state dict!'
|
||||
if cls_name in ["AutoencoderKL"]:
|
||||
assert isinstance(state_dict, dict) and len(state_dict) > 16, "You do not have VAE state dict!"
|
||||
|
||||
config = IntegratedAutoencoderKL.load_config(config_path)
|
||||
|
||||
with using_forge_operations(device=memory_management.cpu, dtype=memory_management.vae_dtype()):
|
||||
model = IntegratedAutoencoderKL.from_config(config)
|
||||
|
||||
if 'decoder.up_blocks.0.resnets.0.norm1.weight' in state_dict.keys(): #diffusers format
|
||||
if "decoder.up_blocks.0.resnets.0.norm1.weight" in state_dict.keys(): # diffusers format
|
||||
state_dict = huggingface_guess.diffusers_convert.convert_vae_state_dict(state_dict)
|
||||
load_state_dict(model, state_dict, ignore_start='loss.')
|
||||
load_state_dict(model, state_dict, ignore_start="loss.")
|
||||
return model
|
||||
if component_name.startswith('text_encoder') and cls_name in ['CLIPTextModel', 'CLIPTextModelWithProjection']:
|
||||
assert isinstance(state_dict, dict) and len(state_dict) > 16, 'You do not have CLIP state dict!'
|
||||
if component_name.startswith("text_encoder") and cls_name in ["CLIPTextModel", "CLIPTextModelWithProjection"]:
|
||||
assert isinstance(state_dict, dict) and len(state_dict) > 16, "You do not have CLIP state dict!"
|
||||
|
||||
from transformers import CLIPTextConfig, CLIPTextModel
|
||||
|
||||
config = CLIPTextConfig.from_pretrained(config_path)
|
||||
|
||||
to_args = dict(device=memory_management.cpu, dtype=memory_management.text_encoder_dtype())
|
||||
@ -69,33 +71,30 @@ def load_huggingface_component(guess, component_name, lib_name, cls_name, repo_p
|
||||
with using_forge_operations(**to_args, manual_cast_enabled=True):
|
||||
model = IntegratedCLIP(CLIPTextModel, config, add_text_projection=True).to(**to_args)
|
||||
|
||||
load_state_dict(model, state_dict, ignore_errors=[
|
||||
'transformer.text_projection.weight',
|
||||
'transformer.text_model.embeddings.position_ids',
|
||||
'logit_scale'
|
||||
], log_name=cls_name)
|
||||
load_state_dict(model, state_dict, ignore_errors=["transformer.text_projection.weight", "transformer.text_model.embeddings.position_ids", "logit_scale"], log_name=cls_name)
|
||||
|
||||
return model
|
||||
if cls_name == 'T5EncoderModel':
|
||||
assert isinstance(state_dict, dict) and len(state_dict) > 16, 'You do not have T5 state dict!'
|
||||
if cls_name == "T5EncoderModel":
|
||||
assert isinstance(state_dict, dict) and len(state_dict) > 16, "You do not have T5 state dict!"
|
||||
|
||||
from backend.nn.t5 import IntegratedT5
|
||||
|
||||
config = read_arbitrary_config(config_path)
|
||||
|
||||
storage_dtype = memory_management.text_encoder_dtype()
|
||||
state_dict_dtype = memory_management.state_dict_dtype(state_dict)
|
||||
|
||||
if state_dict_dtype in [torch.float8_e4m3fn, torch.float8_e5m2, 'nf4', 'fp4', 'gguf']:
|
||||
print(f'Using Detected T5 Data Type: {state_dict_dtype}')
|
||||
if state_dict_dtype in [torch.float8_e4m3fn, torch.float8_e5m2, "nf4", "fp4", "gguf"]:
|
||||
print(f"Using Detected T5 Data Type: {state_dict_dtype}")
|
||||
storage_dtype = state_dict_dtype
|
||||
if state_dict_dtype in ['nf4', 'fp4', 'gguf']:
|
||||
print(f'Using pre-quant state dict!')
|
||||
if state_dict_dtype in ['gguf']:
|
||||
if state_dict_dtype in ["nf4", "fp4", "gguf"]:
|
||||
print(f"Using pre-quant state dict!")
|
||||
if state_dict_dtype in ["gguf"]:
|
||||
beautiful_print_gguf_state_dict_statics(state_dict)
|
||||
else:
|
||||
print(f'Using Default T5 Data Type: {storage_dtype}')
|
||||
print(f"Using Default T5 Data Type: {storage_dtype}")
|
||||
|
||||
if storage_dtype in ['nf4', 'fp4', 'gguf']:
|
||||
if storage_dtype in ["nf4", "fp4", "gguf"]:
|
||||
with modeling_utils.no_init_weights():
|
||||
with using_forge_operations(device=memory_management.cpu, dtype=memory_management.text_encoder_dtype(), manual_cast_enabled=False, bnb_dtype=storage_dtype):
|
||||
model = IntegratedT5(config)
|
||||
@ -104,20 +103,22 @@ def load_huggingface_component(guess, component_name, lib_name, cls_name, repo_p
|
||||
with using_forge_operations(device=memory_management.cpu, dtype=storage_dtype, manual_cast_enabled=True):
|
||||
model = IntegratedT5(config)
|
||||
|
||||
load_state_dict(model, state_dict, log_name=cls_name, ignore_errors=['transformer.encoder.embed_tokens.weight', 'logit_scale'])
|
||||
load_state_dict(model, state_dict, log_name=cls_name, ignore_errors=["transformer.encoder.embed_tokens.weight", "logit_scale"])
|
||||
|
||||
return model
|
||||
if cls_name in ['UNet2DConditionModel', 'FluxTransformer2DModel', 'SD3Transformer2DModel', 'ChromaTransformer2DModel']:
|
||||
assert isinstance(state_dict, dict) and len(state_dict) > 16, 'You do not have model state dict!'
|
||||
if cls_name in ["UNet2DConditionModel", "FluxTransformer2DModel", "SD3Transformer2DModel", "ChromaTransformer2DModel"]:
|
||||
assert isinstance(state_dict, dict) and len(state_dict) > 16, "You do not have model state dict!"
|
||||
|
||||
model_loader = None
|
||||
if cls_name == 'UNet2DConditionModel':
|
||||
if cls_name == "UNet2DConditionModel":
|
||||
model_loader = lambda c: IntegratedUNet2DConditionModel.from_config(c)
|
||||
elif cls_name == 'FluxTransformer2DModel':
|
||||
elif cls_name == "FluxTransformer2DModel":
|
||||
from backend.nn.flux import IntegratedFluxTransformer2DModel
|
||||
|
||||
model_loader = lambda c: IntegratedFluxTransformer2DModel(**c)
|
||||
elif cls_name == 'ChromaTransformer2DModel':
|
||||
elif cls_name == "ChromaTransformer2DModel":
|
||||
from backend.nn.chroma import IntegratedChromaTransformer2DModel
|
||||
|
||||
model_loader = lambda c: IntegratedChromaTransformer2DModel(**c)
|
||||
# elif cls_name == 'SD3Transformer2DModel':
|
||||
# from backend.nn.mmditx import MMDiTX
|
||||
@ -129,23 +130,23 @@ def load_huggingface_component(guess, component_name, lib_name, cls_name, repo_p
|
||||
|
||||
storage_dtype = memory_management.unet_dtype(model_params=state_dict_parameters, supported_dtypes=guess.supported_inference_dtypes)
|
||||
|
||||
unet_storage_dtype_overwrite = backend.args.dynamic_args.get('forge_unet_storage_dtype')
|
||||
unet_storage_dtype_overwrite = backend.args.dynamic_args.get("forge_unet_storage_dtype")
|
||||
|
||||
if unet_storage_dtype_overwrite is not None:
|
||||
storage_dtype = unet_storage_dtype_overwrite
|
||||
elif state_dict_dtype in [torch.float8_e4m3fn, torch.float8_e5m2, 'nf4', 'fp4', 'gguf']:
|
||||
print(f'Using Detected UNet Type: {state_dict_dtype}')
|
||||
elif state_dict_dtype in [torch.float8_e4m3fn, torch.float8_e5m2, "nf4", "fp4", "gguf"]:
|
||||
print(f"Using Detected UNet Type: {state_dict_dtype}")
|
||||
storage_dtype = state_dict_dtype
|
||||
if state_dict_dtype in ['nf4', 'fp4', 'gguf']:
|
||||
print(f'Using pre-quant state dict!')
|
||||
if state_dict_dtype in ['gguf']:
|
||||
if state_dict_dtype in ["nf4", "fp4", "gguf"]:
|
||||
print(f"Using pre-quant state dict!")
|
||||
if state_dict_dtype in ["gguf"]:
|
||||
beautiful_print_gguf_state_dict_statics(state_dict)
|
||||
|
||||
load_device = memory_management.get_torch_device()
|
||||
computation_dtype = memory_management.get_computation_dtype(load_device, parameters=state_dict_parameters, supported_dtypes=guess.supported_inference_dtypes)
|
||||
offload_device = memory_management.unet_offload_device()
|
||||
|
||||
if storage_dtype in ['nf4', 'fp4', 'gguf']:
|
||||
if storage_dtype in ["nf4", "fp4", "gguf"]:
|
||||
initial_device = memory_management.unet_inital_load_device(parameters=state_dict_parameters, dtype=computation_dtype)
|
||||
with using_forge_operations(device=initial_device, dtype=computation_dtype, manual_cast_enabled=False, bnb_dtype=storage_dtype):
|
||||
model = model_loader(unet_config)
|
||||
@ -159,7 +160,7 @@ def load_huggingface_component(guess, component_name, lib_name, cls_name, repo_p
|
||||
|
||||
load_state_dict(model, state_dict)
|
||||
|
||||
if hasattr(model, '_internal_dict'):
|
||||
if hasattr(model, "_internal_dict"):
|
||||
model._internal_dict = unet_config
|
||||
else:
|
||||
model.config = unet_config
|
||||
@ -172,7 +173,7 @@ def load_huggingface_component(guess, component_name, lib_name, cls_name, repo_p
|
||||
|
||||
return model
|
||||
|
||||
print(f'Skipped: {component_name} = {lib_name}.{cls_name}')
|
||||
print(f"Skipped: {component_name} = {lib_name}.{cls_name}")
|
||||
return None
|
||||
|
||||
|
||||
@ -180,7 +181,7 @@ def replace_state_dict(sd, asd, guess):
|
||||
vae_key_prefix = guess.vae_key_prefix[0]
|
||||
text_encoder_key_prefix = guess.text_encoder_key_prefix[0]
|
||||
|
||||
if 'enc.blk.0.attn_k.weight' in asd:
|
||||
if "enc.blk.0.attn_k.weight" in asd:
|
||||
wierd_t5_format_from_city96 = {
|
||||
"enc.": "encoder.",
|
||||
".blk.": ".block.",
|
||||
@ -197,7 +198,7 @@ def replace_state_dict(sd, asd, guess):
|
||||
"ffn_gate": "layer.1.DenseReluDense.wi_0",
|
||||
"ffn_norm": "layer.1.layer_norm",
|
||||
}
|
||||
wierd_t5_pre_quant_keys_from_city96 = ['shared.weight']
|
||||
wierd_t5_pre_quant_keys_from_city96 = ["shared.weight"]
|
||||
asd_new = {}
|
||||
for k, v in asd.items():
|
||||
for s, d in wierd_t5_format_from_city96.items():
|
||||
@ -215,7 +216,6 @@ def replace_state_dict(sd, asd, guess):
|
||||
for k, v in asd.items():
|
||||
sd[vae_key_prefix + k] = v
|
||||
|
||||
|
||||
## identify model type
|
||||
flux_test_key = "model.diffusion_model.double_blocks.0.img_attn.norm.key_norm.scale"
|
||||
sd3_test_key = "model.diffusion_model.final_layer.adaLN_modulation.1.bias"
|
||||
@ -229,7 +229,7 @@ def replace_state_dict(sd, asd, guess):
|
||||
case 1024:
|
||||
model_type = "sd2"
|
||||
case 1280:
|
||||
model_type = "xlrf" # sdxl refiner model
|
||||
model_type = "xlrf" # sdxl refiner model
|
||||
case 2048:
|
||||
model_type = "sdxl"
|
||||
elif flux_test_key in sd:
|
||||
@ -239,36 +239,35 @@ def replace_state_dict(sd, asd, guess):
|
||||
|
||||
## prefixes used by various model types for CLIP-L
|
||||
prefix_L = {
|
||||
"-" : None,
|
||||
"sd1" : "cond_stage_model.transformer.",
|
||||
"sd2" : None,
|
||||
"-": None,
|
||||
"sd1": "cond_stage_model.transformer.",
|
||||
"sd2": None,
|
||||
"xlrf": None,
|
||||
"sdxl": "conditioner.embedders.0.transformer.",
|
||||
"flux": "text_encoders.clip_l.transformer.",
|
||||
"sd3" : "text_encoders.clip_l.transformer.",
|
||||
"sd3": "text_encoders.clip_l.transformer.",
|
||||
}
|
||||
## prefixes used by various model types for CLIP-G
|
||||
prefix_G = {
|
||||
"-" : None,
|
||||
"sd1" : None,
|
||||
"sd2" : None,
|
||||
"-": None,
|
||||
"sd1": None,
|
||||
"sd2": None,
|
||||
"xlrf": "conditioner.embedders.0.model.transformer.",
|
||||
"sdxl": "conditioner.embedders.1.model.transformer.",
|
||||
"flux": None,
|
||||
"sd3" : "text_encoders.clip_g.transformer.",
|
||||
"sd3": "text_encoders.clip_g.transformer.",
|
||||
}
|
||||
## prefixes used by various model types for CLIP-H
|
||||
prefix_H = {
|
||||
"-" : None,
|
||||
"sd1" : None,
|
||||
"sd2" : "conditioner.embedders.0.model.",
|
||||
"-": None,
|
||||
"sd1": None,
|
||||
"sd2": "conditioner.embedders.0.model.",
|
||||
"xlrf": None,
|
||||
"sdxl": None,
|
||||
"flux": None,
|
||||
"sd3" : None,
|
||||
"sd3": None,
|
||||
}
|
||||
|
||||
|
||||
## VAE format 0 (extracted from model, could be sd1, sd2, sdxl, sd3).
|
||||
if "first_stage_model.decoder.conv_in.weight" in asd:
|
||||
channels = asd["first_stage_model.decoder.conv_in.weight"].shape[1]
|
||||
@ -282,10 +281,10 @@ def replace_state_dict(sd, asd, guess):
|
||||
sd[k] = v
|
||||
|
||||
## CLIP-H
|
||||
CLIP_H = { # key to identify source model old_prefix
|
||||
'cond_stage_model.model.ln_final.weight' : 'cond_stage_model.model.',
|
||||
# 'text_model.encoder.layers.0.layer_norm1.bias' : 'text_model'. # would need converting
|
||||
}
|
||||
CLIP_H = { # key to identify source model old_prefix
|
||||
"cond_stage_model.model.ln_final.weight": "cond_stage_model.model.",
|
||||
# 'text_model.encoder.layers.0.layer_norm1.bias' : 'text_model'. # would need converting
|
||||
}
|
||||
for CLIP_key in CLIP_H.keys():
|
||||
if CLIP_key in asd and asd[CLIP_key].shape[0] == 1024:
|
||||
new_prefix = prefix_H[model_type]
|
||||
@ -297,36 +296,32 @@ def replace_state_dict(sd, asd, guess):
|
||||
sd[new_k] = v
|
||||
|
||||
## CLIP-G
|
||||
CLIP_G = { # key to identify source model old_prefix
|
||||
'conditioner.embedders.1.model.transformer.resblocks.0.ln_1.bias' : 'conditioner.embedders.1.model.transformer.',
|
||||
'text_encoders.clip_g.transformer.text_model.encoder.layers.0.layer_norm1.bias' : 'text_encoders.clip_g.transformer.',
|
||||
'text_model.encoder.layers.0.layer_norm1.bias' : '',
|
||||
'transformer.resblocks.0.ln_1.bias' : 'transformer.'
|
||||
}
|
||||
CLIP_G = {"conditioner.embedders.1.model.transformer.resblocks.0.ln_1.bias": "conditioner.embedders.1.model.transformer.", "text_encoders.clip_g.transformer.text_model.encoder.layers.0.layer_norm1.bias": "text_encoders.clip_g.transformer.", "text_model.encoder.layers.0.layer_norm1.bias": "", "transformer.resblocks.0.ln_1.bias": "transformer."} # key to identify source model old_prefix
|
||||
for CLIP_key in CLIP_G.keys():
|
||||
if CLIP_key in asd and asd[CLIP_key].shape[0] == 1280:
|
||||
new_prefix = prefix_G[model_type]
|
||||
old_prefix = CLIP_G[CLIP_key]
|
||||
|
||||
if new_prefix is not None:
|
||||
if "resblocks" not in CLIP_key and model_type != "sd3": # need to convert
|
||||
if "resblocks" not in CLIP_key and model_type != "sd3": # need to convert
|
||||
|
||||
def convert_transformers(statedict, prefix_from, prefix_to, number):
|
||||
keys_to_replace = {
|
||||
"{}text_model.embeddings.position_embedding.weight" : "{}positional_embedding",
|
||||
"{}text_model.embeddings.token_embedding.weight" : "{}token_embedding.weight",
|
||||
"{}text_model.final_layer_norm.weight" : "{}ln_final.weight",
|
||||
"{}text_model.final_layer_norm.bias" : "{}ln_final.bias",
|
||||
"text_projection.weight" : "{}text_projection",
|
||||
"{}text_model.embeddings.position_embedding.weight": "{}positional_embedding",
|
||||
"{}text_model.embeddings.token_embedding.weight": "{}token_embedding.weight",
|
||||
"{}text_model.final_layer_norm.weight": "{}ln_final.weight",
|
||||
"{}text_model.final_layer_norm.bias": "{}ln_final.bias",
|
||||
"text_projection.weight": "{}text_projection",
|
||||
}
|
||||
resblock_to_replace = {
|
||||
"layer_norm1" : "ln_1",
|
||||
"layer_norm2" : "ln_2",
|
||||
"mlp.fc1" : "mlp.c_fc",
|
||||
"mlp.fc2" : "mlp.c_proj",
|
||||
"self_attn.out_proj" : "attn.out_proj" ,
|
||||
"layer_norm1": "ln_1",
|
||||
"layer_norm2": "ln_2",
|
||||
"mlp.fc1": "mlp.c_fc",
|
||||
"mlp.fc2": "mlp.c_proj",
|
||||
"self_attn.out_proj": "attn.out_proj",
|
||||
}
|
||||
|
||||
for x in keys_to_replace: # remove trailing 'transformer.' from new prefix
|
||||
for x in keys_to_replace: # remove trailing 'transformer.' from new prefix
|
||||
k = x.format(prefix_from)
|
||||
statedict[keys_to_replace[x].format(prefix_to[:-12])] = statedict.pop(k)
|
||||
|
||||
@ -363,13 +358,7 @@ def replace_state_dict(sd, asd, guess):
|
||||
sd[new_k] = v
|
||||
|
||||
## CLIP-L
|
||||
CLIP_L = { # key to identify source model old_prefix
|
||||
'cond_stage_model.transformer.text_model.encoder.layers.0.layer_norm1.bias' : 'cond_stage_model.transformer.',
|
||||
'conditioner.embedders.0.transformer.text_model.encoder.layers.0.layer_norm1.bias' : 'conditioner.embedders.0.transformer.',
|
||||
'text_encoders.clip_l.transformer.text_model.encoder.layers.0.layer_norm1.bias' : 'text_encoders.clip_l.transformer.',
|
||||
'text_model.encoder.layers.0.layer_norm1.bias' : '',
|
||||
'transformer.resblocks.0.ln_1.bias' : 'transformer.'
|
||||
}
|
||||
CLIP_L = {"cond_stage_model.transformer.text_model.encoder.layers.0.layer_norm1.bias": "cond_stage_model.transformer.", "conditioner.embedders.0.transformer.text_model.encoder.layers.0.layer_norm1.bias": "conditioner.embedders.0.transformer.", "text_encoders.clip_l.transformer.text_model.encoder.layers.0.layer_norm1.bias": "text_encoders.clip_l.transformer.", "text_model.encoder.layers.0.layer_norm1.bias": "", "transformer.resblocks.0.ln_1.bias": "transformer."} # key to identify source model old_prefix
|
||||
|
||||
for CLIP_key in CLIP_L.keys():
|
||||
if CLIP_key in asd and asd[CLIP_key].shape[0] == 768:
|
||||
@ -377,21 +366,22 @@ def replace_state_dict(sd, asd, guess):
|
||||
old_prefix = CLIP_L[CLIP_key]
|
||||
|
||||
if new_prefix is not None:
|
||||
if "resblocks" in CLIP_key: # need to convert
|
||||
if "resblocks" in CLIP_key: # need to convert
|
||||
|
||||
def transformers_convert(statedict, prefix_from, prefix_to, number):
|
||||
keys_to_replace = {
|
||||
"positional_embedding" : "{}text_model.embeddings.position_embedding.weight",
|
||||
"positional_embedding": "{}text_model.embeddings.position_embedding.weight",
|
||||
"token_embedding.weight": "{}text_model.embeddings.token_embedding.weight",
|
||||
"ln_final.weight" : "{}text_model.final_layer_norm.weight",
|
||||
"ln_final.bias" : "{}text_model.final_layer_norm.bias",
|
||||
"text_projection" : "text_projection.weight",
|
||||
"ln_final.weight": "{}text_model.final_layer_norm.weight",
|
||||
"ln_final.bias": "{}text_model.final_layer_norm.bias",
|
||||
"text_projection": "text_projection.weight",
|
||||
}
|
||||
resblock_to_replace = {
|
||||
"ln_1" : "layer_norm1",
|
||||
"ln_2" : "layer_norm2",
|
||||
"mlp.c_fc" : "mlp.fc1",
|
||||
"mlp.c_proj" : "mlp.fc2",
|
||||
"attn.out_proj" : "self_attn.out_proj",
|
||||
"ln_1": "layer_norm1",
|
||||
"ln_2": "layer_norm2",
|
||||
"mlp.c_fc": "mlp.fc1",
|
||||
"mlp.c_proj": "mlp.fc2",
|
||||
"attn.out_proj": "self_attn.out_proj",
|
||||
}
|
||||
|
||||
for k in keys_to_replace:
|
||||
@ -410,7 +400,7 @@ def replace_state_dict(sd, asd, guess):
|
||||
for x in range(3):
|
||||
p = ["self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj"]
|
||||
k_to = "{}text_model.encoder.layers.{}.{}.{}".format(prefix_to, resblock, p[x], y)
|
||||
statedict[k_to] = weights[shape_from*x:shape_from*(x + 1)]
|
||||
statedict[k_to] = weights[shape_from * x : shape_from * (x + 1)]
|
||||
return statedict
|
||||
|
||||
asd = transformers_convert(asd, old_prefix, new_prefix, 12)
|
||||
@ -426,8 +416,7 @@ def replace_state_dict(sd, asd, guess):
|
||||
new_k = k.replace(old_prefix, new_prefix)
|
||||
sd[new_k] = v
|
||||
|
||||
|
||||
if 'encoder.block.0.layer.0.SelfAttention.k.weight' in asd:
|
||||
if "encoder.block.0.layer.0.SelfAttention.k.weight" in asd:
|
||||
keys_to_delete = [k for k in sd if k.startswith(f"{text_encoder_key_prefix}t5xxl.")]
|
||||
for k in keys_to_delete:
|
||||
del sd[k]
|
||||
@ -457,26 +446,23 @@ def split_state_dict(sd, additional_state_dicts: list = None):
|
||||
|
||||
guess.clip_target = guess.clip_target(sd)
|
||||
guess.model_type = guess.model_type(sd)
|
||||
guess.ztsnr = 'ztsnr' in sd
|
||||
guess.ztsnr = "ztsnr" in sd
|
||||
|
||||
sd = guess.process_vae_state_dict(sd)
|
||||
|
||||
state_dict = {
|
||||
guess.unet_target: try_filter_state_dict(sd, guess.unet_key_prefix),
|
||||
guess.vae_target: try_filter_state_dict(sd, guess.vae_key_prefix)
|
||||
}
|
||||
state_dict = {guess.unet_target: try_filter_state_dict(sd, guess.unet_key_prefix), guess.vae_target: try_filter_state_dict(sd, guess.vae_key_prefix)}
|
||||
|
||||
sd = guess.process_clip_state_dict(sd)
|
||||
|
||||
for k, v in guess.clip_target.items():
|
||||
state_dict[v] = try_filter_state_dict(sd, [k + '.'])
|
||||
state_dict[v] = try_filter_state_dict(sd, [k + "."])
|
||||
|
||||
state_dict['ignore'] = sd
|
||||
state_dict["ignore"] = sd
|
||||
|
||||
print_dict = {k: len(v) for k, v in state_dict.items()}
|
||||
print(f'StateDict Keys: {print_dict}')
|
||||
print(f"StateDict Keys: {print_dict}")
|
||||
|
||||
del state_dict['ignore']
|
||||
del state_dict["ignore"]
|
||||
|
||||
return state_dict, guess
|
||||
|
||||
@ -486,11 +472,11 @@ def forge_loader(sd, additional_state_dicts=None):
|
||||
try:
|
||||
state_dicts, estimated_config = split_state_dict(sd, additional_state_dicts=additional_state_dicts)
|
||||
except:
|
||||
raise ValueError('Failed to recognize model type!')
|
||||
raise ValueError("Failed to recognize model type!")
|
||||
|
||||
repo_name = estimated_config.huggingface_repo
|
||||
|
||||
local_path = os.path.join(dir_path, 'huggingface', repo_name)
|
||||
local_path = os.path.join(dir_path, "huggingface", repo_name)
|
||||
config: dict = DiffusionPipeline.load_config(local_path)
|
||||
huggingface_components = {}
|
||||
for component_name, v in config.items():
|
||||
@ -507,44 +493,43 @@ def forge_loader(sd, additional_state_dicts=None):
|
||||
yaml_config_prediction_type = None
|
||||
|
||||
try:
|
||||
import yaml
|
||||
from pathlib import Path
|
||||
config_filename = os.path.splitext(sd)[0] + '.yaml'
|
||||
|
||||
import yaml
|
||||
|
||||
config_filename = os.path.splitext(sd)[0] + ".yaml"
|
||||
if Path(config_filename).is_file():
|
||||
with open(config_filename, 'r') as stream:
|
||||
with open(config_filename, "r") as stream:
|
||||
yaml_config = yaml.safe_load(stream)
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
# Fix Huggingface prediction type using .yaml config or estimated config detection
|
||||
prediction_types = {
|
||||
'EPS': 'epsilon',
|
||||
'V_PREDICTION': 'v_prediction',
|
||||
'EDM': 'edm',
|
||||
"EPS": "epsilon",
|
||||
"V_PREDICTION": "v_prediction",
|
||||
"EDM": "edm",
|
||||
}
|
||||
|
||||
has_prediction_type = 'scheduler' in huggingface_components and hasattr(huggingface_components['scheduler'], 'config') and 'prediction_type' in huggingface_components['scheduler'].config
|
||||
has_prediction_type = "scheduler" in huggingface_components and hasattr(huggingface_components["scheduler"], "config") and "prediction_type" in huggingface_components["scheduler"].config
|
||||
|
||||
if yaml_config is not None:
|
||||
yaml_config_prediction_type: str = (
|
||||
yaml_config.get('model', {}).get('params', {}).get('parameterization', '')
|
||||
or yaml_config.get('model', {}).get('params', {}).get('denoiser_config', {}).get('params', {}).get('scaling_config', {}).get('target', '')
|
||||
)
|
||||
if yaml_config_prediction_type == 'v' or yaml_config_prediction_type.endswith(".VScaling"):
|
||||
yaml_config_prediction_type = 'v_prediction'
|
||||
yaml_config_prediction_type: str = yaml_config.get("model", {}).get("params", {}).get("parameterization", "") or yaml_config.get("model", {}).get("params", {}).get("denoiser_config", {}).get("params", {}).get("scaling_config", {}).get("target", "")
|
||||
if yaml_config_prediction_type == "v" or yaml_config_prediction_type.endswith(".VScaling"):
|
||||
yaml_config_prediction_type = "v_prediction"
|
||||
else:
|
||||
# Use estimated prediction config if no suitable prediction type found
|
||||
yaml_config_prediction_type = ''
|
||||
yaml_config_prediction_type = ""
|
||||
|
||||
if has_prediction_type:
|
||||
if yaml_config_prediction_type:
|
||||
huggingface_components['scheduler'].config.prediction_type = yaml_config_prediction_type
|
||||
huggingface_components["scheduler"].config.prediction_type = yaml_config_prediction_type
|
||||
else:
|
||||
huggingface_components['scheduler'].config.prediction_type = prediction_types.get(estimated_config.model_type.name, huggingface_components['scheduler'].config.prediction_type)
|
||||
huggingface_components["scheduler"].config.prediction_type = prediction_types.get(estimated_config.model_type.name, huggingface_components["scheduler"].config.prediction_type)
|
||||
|
||||
for M in possible_models:
|
||||
if any(isinstance(estimated_config, x) for x in M.matched_guesses):
|
||||
return M(estimated_config=estimated_config, huggingface_components=huggingface_components)
|
||||
|
||||
print('Failed to recognize model type!')
|
||||
print("Failed to recognize model type!")
|
||||
return None
|
||||
|
||||
@ -1,17 +1,17 @@
|
||||
# Cherry-picked some good parts from ComfyUI with some bad parts fixed
|
||||
|
||||
import platform
|
||||
import sys
|
||||
import time
|
||||
from enum import Enum
|
||||
|
||||
import psutil
|
||||
import torch
|
||||
import platform
|
||||
|
||||
from enum import Enum
|
||||
from backend import stream, utils
|
||||
from backend.args import args
|
||||
|
||||
|
||||
cpu = torch.device('cpu')
|
||||
cpu = torch.device("cpu")
|
||||
|
||||
|
||||
class VRAMState(Enum):
|
||||
@ -105,7 +105,7 @@ def get_total_memory(dev=None, torch_total_too=False):
|
||||
if dev is None:
|
||||
dev = get_torch_device()
|
||||
|
||||
if hasattr(dev, 'type') and (dev.type == 'cpu' or dev.type == 'mps'):
|
||||
if hasattr(dev, "type") and (dev.type == "cpu" or dev.type == "mps"):
|
||||
mem_total = psutil.virtual_memory().total
|
||||
mem_total_torch = mem_total
|
||||
else:
|
||||
@ -114,12 +114,12 @@ def get_total_memory(dev=None, torch_total_too=False):
|
||||
mem_total_torch = mem_total
|
||||
elif is_intel_xpu():
|
||||
stats = torch.xpu.memory_stats(dev)
|
||||
mem_reserved = stats['reserved_bytes.all.current']
|
||||
mem_reserved = stats["reserved_bytes.all.current"]
|
||||
mem_total_torch = mem_reserved
|
||||
mem_total = torch.xpu.get_device_properties(dev).total_memory
|
||||
else:
|
||||
stats = torch.cuda.memory_stats(dev)
|
||||
mem_reserved = stats['reserved_bytes.all.current']
|
||||
mem_reserved = stats["reserved_bytes.all.current"]
|
||||
_, mem_total_cuda = torch.cuda.mem_get_info(dev)
|
||||
mem_total_torch = mem_reserved
|
||||
mem_total = mem_total_cuda
|
||||
@ -258,7 +258,7 @@ if PIN_SHARED_MEMORY:
|
||||
|
||||
|
||||
def get_torch_device_name(device):
|
||||
if hasattr(device, 'type'):
|
||||
if hasattr(device, "type"):
|
||||
if device.type == "cuda":
|
||||
try:
|
||||
allocator_backend = torch.cuda.get_allocator_backend()
|
||||
@ -277,12 +277,12 @@ try:
|
||||
torch_device_name = get_torch_device_name(get_torch_device())
|
||||
print("Device: {}".format(torch_device_name))
|
||||
except:
|
||||
torch_device_name = ''
|
||||
torch_device_name = ""
|
||||
print("Could not pick default device.")
|
||||
|
||||
if 'rtx' in torch_device_name.lower():
|
||||
if "rtx" in torch_device_name.lower():
|
||||
if not args.cuda_malloc:
|
||||
print('Hint: your device supports --cuda-malloc for potential speed improvements.')
|
||||
print("Hint: your device supports --cuda-malloc for potential speed improvements.")
|
||||
|
||||
|
||||
current_loaded_models = []
|
||||
@ -310,12 +310,12 @@ def state_dict_parameters(sd):
|
||||
|
||||
def state_dict_dtype(state_dict):
|
||||
for k, v in state_dict.items():
|
||||
if hasattr(v, 'gguf_cls'):
|
||||
return 'gguf'
|
||||
if 'bitsandbytes__nf4' in k:
|
||||
return 'nf4'
|
||||
if 'bitsandbytes__fp4' in k:
|
||||
return 'fp4'
|
||||
if hasattr(v, "gguf_cls"):
|
||||
return "gguf"
|
||||
if "bitsandbytes__nf4" in k:
|
||||
return "nf4"
|
||||
if "bitsandbytes__fp4" in k:
|
||||
return "fp4"
|
||||
|
||||
dtype_counts = {}
|
||||
|
||||
@ -338,11 +338,11 @@ def state_dict_dtype(state_dict):
|
||||
|
||||
|
||||
def bake_gguf_model(model):
|
||||
if getattr(model, 'gguf_baked', False):
|
||||
if getattr(model, "gguf_baked", False):
|
||||
return
|
||||
|
||||
for p in model.parameters():
|
||||
gguf_cls = getattr(p, 'gguf_cls', None)
|
||||
gguf_cls = getattr(p, "gguf_cls", None)
|
||||
if gguf_cls is not None:
|
||||
gguf_cls.bake(p)
|
||||
|
||||
@ -356,7 +356,7 @@ def bake_gguf_model(model):
|
||||
def module_size(module, exclude_device=None, include_device=None, return_split=False):
|
||||
module_mem = 0
|
||||
weight_mem = 0
|
||||
weight_patterns = ['weight']
|
||||
weight_patterns = ["weight"]
|
||||
|
||||
for k, p in module.named_parameters():
|
||||
t = p.data
|
||||
@ -371,7 +371,7 @@ def module_size(module, exclude_device=None, include_device=None, return_split=F
|
||||
|
||||
element_size = t.element_size()
|
||||
|
||||
if getattr(p, 'quant_type', None) in ['fp4', 'nf4']:
|
||||
if getattr(p, "quant_type", None) in ["fp4", "nf4"]:
|
||||
if element_size > 1:
|
||||
# not quanted yet
|
||||
element_size = 0.55 # a bit more than 0.5 because of quant state parameters
|
||||
@ -472,7 +472,7 @@ class LoadedModel:
|
||||
raise e
|
||||
|
||||
if do_not_need_cpu_swap:
|
||||
print('All loaded to GPU.')
|
||||
print("All loaded to GPU.")
|
||||
else:
|
||||
gpu_modules, gpu_modules_only_extras, cpu_modules = build_module_profile(self.real_model, model_gpu_memory_when_using_cpu_swap)
|
||||
pin_memory = PIN_SHARED_MEMORY and is_device_cpu(self.model.offload_device)
|
||||
@ -495,8 +495,8 @@ class LoadedModel:
|
||||
for m in gpu_modules_only_extras:
|
||||
m.prev_parameters_manual_cast = m.parameters_manual_cast
|
||||
m.parameters_manual_cast = True
|
||||
module_move(m, device=self.device, recursive=False, excluded_pattens=['weight'])
|
||||
if hasattr(m, 'weight') and m.weight is not None:
|
||||
module_move(m, device=self.device, recursive=False, excluded_pattens=["weight"])
|
||||
if hasattr(m, "weight") and m.weight is not None:
|
||||
if pin_memory:
|
||||
m.weight = utils.tensor2parameter(m.weight.to(self.model.offload_device).pin_memory())
|
||||
else:
|
||||
@ -504,8 +504,8 @@ class LoadedModel:
|
||||
mem_counter += m.extra_mem
|
||||
swap_counter += m.weight_mem
|
||||
|
||||
swap_flag = 'Shared' if PIN_SHARED_MEMORY else 'CPU'
|
||||
method_flag = 'asynchronous' if stream.should_use_stream() else 'blocked'
|
||||
swap_flag = "Shared" if PIN_SHARED_MEMORY else "CPU"
|
||||
method_flag = "asynchronous" if stream.should_use_stream() else "blocked"
|
||||
print(f"{swap_flag} Swap Loaded ({method_flag} method): {swap_counter / (1024 * 1024):.2f} MB, GPU Loaded: {mem_counter / (1024 * 1024):.2f} MB")
|
||||
|
||||
self.model_accelerated = True
|
||||
@ -596,17 +596,14 @@ def free_memory(memory_required, device, keep_loaded=[], free_all=False):
|
||||
if mem_free_torch > mem_free_total * 0.25:
|
||||
soft_empty_cache()
|
||||
|
||||
print('Done.')
|
||||
print("Done.")
|
||||
return
|
||||
|
||||
|
||||
def compute_model_gpu_memory_when_using_cpu_swap(current_free_mem, inference_memory):
|
||||
maximum_memory_available = current_free_mem - inference_memory
|
||||
|
||||
suggestion = max(
|
||||
maximum_memory_available / 1.3,
|
||||
maximum_memory_available - 1024 * 1024 * 1024 * 1.25
|
||||
)
|
||||
suggestion = max(maximum_memory_available / 1.3, maximum_memory_available - 1024 * 1024 * 1024 * 1.25)
|
||||
|
||||
return int(max(0, suggestion))
|
||||
|
||||
@ -638,7 +635,7 @@ def load_models_gpu(models, memory_required=0, hard_memory_preservation=0):
|
||||
|
||||
moving_time = time.perf_counter() - execution_start_time
|
||||
if moving_time > 0.1:
|
||||
print(f'Memory cleanup has taken {moving_time:.2f} seconds')
|
||||
print(f"Memory cleanup has taken {moving_time:.2f} seconds")
|
||||
|
||||
return
|
||||
|
||||
@ -685,7 +682,7 @@ def load_models_gpu(models, memory_required=0, hard_memory_preservation=0):
|
||||
current_loaded_models.insert(0, loaded_model)
|
||||
|
||||
moving_time = time.perf_counter() - execution_start_time
|
||||
print(f'Moving model(s) has taken {moving_time:.2f} seconds')
|
||||
print(f"Moving model(s) has taken {moving_time:.2f} seconds")
|
||||
|
||||
return
|
||||
|
||||
@ -859,7 +856,7 @@ print(f"VAE dtype preferences: {VAE_DTYPES} -> {vae_dtype()}")
|
||||
|
||||
|
||||
def get_autocast_device(dev):
|
||||
if hasattr(dev, 'type'):
|
||||
if hasattr(dev, "type"):
|
||||
return dev.type
|
||||
return "cuda"
|
||||
|
||||
@ -938,7 +935,7 @@ def cast_to_device(tensor, device, dtype, copy=False):
|
||||
if tensor.dtype == torch.float32 or tensor.dtype == torch.float16:
|
||||
device_supports_cast = True
|
||||
elif tensor.dtype == torch.bfloat16:
|
||||
if hasattr(device, 'type') and device.type.startswith("cuda"):
|
||||
if hasattr(device, "type") and device.type.startswith("cuda"):
|
||||
device_supports_cast = True
|
||||
elif is_intel_xpu():
|
||||
device_supports_cast = True
|
||||
@ -995,7 +992,7 @@ def pytorch_attention_flash_attention():
|
||||
def force_upcast_attention_dtype():
|
||||
upcast = args.force_upcast_attention
|
||||
try:
|
||||
if platform.mac_ver()[0] in ['14.5']: # black image bug on OSX Sonoma 14.5
|
||||
if platform.mac_ver()[0] in ["14.5"]: # black image bug on OSX Sonoma 14.5
|
||||
upcast = True
|
||||
except:
|
||||
pass
|
||||
@ -1010,7 +1007,7 @@ def get_free_memory(dev=None, torch_free_too=False):
|
||||
if dev is None:
|
||||
dev = get_torch_device()
|
||||
|
||||
if hasattr(dev, 'type') and (dev.type == 'cpu' or dev.type == 'mps'):
|
||||
if hasattr(dev, "type") and (dev.type == "cpu" or dev.type == "mps"):
|
||||
mem_free_total = psutil.virtual_memory().available
|
||||
mem_free_torch = mem_free_total
|
||||
else:
|
||||
@ -1019,15 +1016,15 @@ def get_free_memory(dev=None, torch_free_too=False):
|
||||
mem_free_torch = mem_free_total
|
||||
elif is_intel_xpu():
|
||||
stats = torch.xpu.memory_stats(dev)
|
||||
mem_active = stats['active_bytes.all.current']
|
||||
mem_reserved = stats['reserved_bytes.all.current']
|
||||
mem_active = stats["active_bytes.all.current"]
|
||||
mem_reserved = stats["reserved_bytes.all.current"]
|
||||
mem_free_torch = mem_reserved - mem_active
|
||||
mem_free_xpu = torch.xpu.get_device_properties(dev).total_memory - mem_reserved
|
||||
mem_free_total = mem_free_xpu + mem_free_torch
|
||||
else:
|
||||
stats = torch.cuda.memory_stats(dev)
|
||||
mem_active = stats['active_bytes.all.current']
|
||||
mem_reserved = stats['reserved_bytes.all.current']
|
||||
mem_active = stats["active_bytes.all.current"]
|
||||
mem_reserved = stats["reserved_bytes.all.current"]
|
||||
mem_free_cuda, _ = torch.cuda.mem_get_info(dev)
|
||||
mem_free_torch = mem_reserved - mem_active
|
||||
mem_free_total = mem_free_cuda + mem_free_torch
|
||||
@ -1049,22 +1046,22 @@ def mps_mode():
|
||||
|
||||
|
||||
def is_device_type(device, type):
|
||||
if hasattr(device, 'type'):
|
||||
if (device.type == type):
|
||||
if hasattr(device, "type"):
|
||||
if device.type == type:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def is_device_cpu(device):
|
||||
return is_device_type(device, 'cpu')
|
||||
return is_device_type(device, "cpu")
|
||||
|
||||
|
||||
def is_device_mps(device):
|
||||
return is_device_type(device, 'mps')
|
||||
return is_device_type(device, "mps")
|
||||
|
||||
|
||||
def is_device_cuda(device):
|
||||
return is_device_type(device, 'cuda')
|
||||
return is_device_type(device, "cuda")
|
||||
|
||||
|
||||
def should_use_fp16(device=None, model_params=0, prioritize_performance=True, manual_cast=False):
|
||||
@ -1111,7 +1108,7 @@ def should_use_fp16(device=None, model_params=0, prioritize_performance=True, ma
|
||||
if x in props.name.lower():
|
||||
if manual_cast:
|
||||
# For storage dtype
|
||||
free_model_memory = (get_free_memory() * 0.9 - minimum_inference_memory())
|
||||
free_model_memory = get_free_memory() * 0.9 - minimum_inference_memory()
|
||||
if (not prioritize_performance) or model_params * 4 > free_model_memory:
|
||||
return True
|
||||
else:
|
||||
@ -1166,7 +1163,7 @@ def should_use_bf16(device=None, model_params=0, prioritize_performance=True, ma
|
||||
# So in this case bf16 should only be used as storge dtype
|
||||
if manual_cast:
|
||||
# For storage dtype
|
||||
free_model_memory = (get_free_memory() * 0.9 - minimum_inference_memory())
|
||||
free_model_memory = get_free_memory() * 0.9 - minimum_inference_memory()
|
||||
if (not prioritize_performance) or model_params * 4 > free_model_memory:
|
||||
return True
|
||||
|
||||
@ -1178,7 +1175,7 @@ def can_install_bnb():
|
||||
if not torch.cuda.is_available():
|
||||
return False
|
||||
|
||||
cuda_version = tuple(int(x) for x in torch.version.cuda.split('.'))
|
||||
cuda_version = tuple(int(x) for x in torch.version.cuda.split("."))
|
||||
|
||||
if cuda_version >= (11, 7):
|
||||
return True
|
||||
|
||||
@ -1,6 +1,14 @@
|
||||
import torch.nn as nn
|
||||
|
||||
from backend.nn.unet import Downsample, ResBlock, SpatialTransformer, TimestepEmbedSequential, conv_nd, exists, timestep_embedding
|
||||
from backend.nn.unet import (
|
||||
Downsample,
|
||||
ResBlock,
|
||||
SpatialTransformer,
|
||||
TimestepEmbedSequential,
|
||||
conv_nd,
|
||||
exists,
|
||||
timestep_embedding,
|
||||
)
|
||||
|
||||
|
||||
class ControlNet(nn.Module):
|
||||
|
||||
@ -1,40 +1,40 @@
|
||||
# Copyright Forge 2024
|
||||
|
||||
import time
|
||||
import torch
|
||||
import contextlib
|
||||
import time
|
||||
|
||||
from backend import stream, memory_management, utils
|
||||
import torch
|
||||
|
||||
from backend import memory_management, stream, utils
|
||||
from backend.patcher.lora import merge_lora_to_weight
|
||||
|
||||
|
||||
stash = {}
|
||||
|
||||
|
||||
def get_weight_and_bias(layer, weight_args=None, bias_args=None, weight_fn=None, bias_fn=None):
|
||||
scale_weight = getattr(layer, 'scale_weight', None)
|
||||
patches = getattr(layer, 'forge_online_loras', None)
|
||||
scale_weight = getattr(layer, "scale_weight", None)
|
||||
patches = getattr(layer, "forge_online_loras", None)
|
||||
weight_patches, bias_patches = None, None
|
||||
|
||||
if patches is not None:
|
||||
weight_patches = patches.get('weight', None)
|
||||
weight_patches = patches.get("weight", None)
|
||||
|
||||
if patches is not None:
|
||||
bias_patches = patches.get('bias', None)
|
||||
bias_patches = patches.get("bias", None)
|
||||
|
||||
weight = None
|
||||
if layer.weight is not None:
|
||||
weight = layer.weight
|
||||
if weight_fn is not None:
|
||||
if weight_args is not None:
|
||||
fn_device = weight_args.get('device', None)
|
||||
fn_device = weight_args.get("device", None)
|
||||
if fn_device is not None:
|
||||
weight = weight.to(device=fn_device)
|
||||
weight = weight_fn(weight)
|
||||
if weight_args is not None:
|
||||
weight = weight.to(**weight_args)
|
||||
if scale_weight is not None:
|
||||
weight = weight*scale_weight.to(device=weight.device, dtype=weight.dtype)
|
||||
weight = weight * scale_weight.to(device=weight.device, dtype=weight.dtype)
|
||||
if weight_patches is not None:
|
||||
weight = merge_lora_to_weight(patches=weight_patches, weight=weight, key="online weight lora", computation_dtype=weight.dtype)
|
||||
|
||||
@ -43,7 +43,7 @@ def get_weight_and_bias(layer, weight_args=None, bias_args=None, weight_fn=None,
|
||||
bias = layer.bias
|
||||
if bias_fn is not None:
|
||||
if bias_args is not None:
|
||||
fn_device = bias_args.get('device', None)
|
||||
fn_device = bias_args.get("device", None)
|
||||
if fn_device is not None:
|
||||
bias = bias.to(device=fn_device)
|
||||
bias = bias_fn(bias)
|
||||
@ -58,7 +58,7 @@ def weights_manual_cast(layer, x, skip_weight_dtype=False, skip_bias_dtype=False
|
||||
weight, bias, signal = None, None, None
|
||||
non_blocking = True
|
||||
|
||||
if getattr(x.device, 'type', None) == 'mps':
|
||||
if getattr(x.device, "type", None) == "mps":
|
||||
non_blocking = False
|
||||
|
||||
target_dtype = x.dtype
|
||||
@ -135,13 +135,13 @@ class ForgeOperations:
|
||||
self.parameters_manual_cast = current_manual_cast_enabled
|
||||
|
||||
def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs):
|
||||
if hasattr(self, 'dummy'):
|
||||
if prefix + 'weight' in state_dict:
|
||||
self.weight = torch.nn.Parameter(state_dict[prefix + 'weight'].to(self.dummy))
|
||||
if prefix + 'scale_weight' in state_dict:
|
||||
self.scale_weight = torch.nn.Parameter(state_dict[prefix + 'scale_weight'])
|
||||
if prefix + 'bias' in state_dict:
|
||||
self.bias = torch.nn.Parameter(state_dict[prefix + 'bias'].to(self.dummy))
|
||||
if hasattr(self, "dummy"):
|
||||
if prefix + "weight" in state_dict:
|
||||
self.weight = torch.nn.Parameter(state_dict[prefix + "weight"].to(self.dummy))
|
||||
if prefix + "scale_weight" in state_dict:
|
||||
self.scale_weight = torch.nn.Parameter(state_dict[prefix + "scale_weight"])
|
||||
if prefix + "bias" in state_dict:
|
||||
self.bias = torch.nn.Parameter(state_dict[prefix + "bias"].to(self.dummy))
|
||||
del self.dummy
|
||||
else:
|
||||
super()._load_from_state_dict(state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs)
|
||||
@ -158,8 +158,8 @@ class ForgeOperations:
|
||||
class Conv2d(torch.nn.Conv2d):
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
kwargs['device'] = current_device
|
||||
kwargs['dtype'] = current_dtype
|
||||
kwargs["device"] = current_device
|
||||
kwargs["dtype"] = current_dtype
|
||||
super().__init__(*args, **kwargs)
|
||||
self.parameters_manual_cast = current_manual_cast_enabled
|
||||
|
||||
@ -178,8 +178,8 @@ class ForgeOperations:
|
||||
class Conv3d(torch.nn.Conv3d):
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
kwargs['device'] = current_device
|
||||
kwargs['dtype'] = current_dtype
|
||||
kwargs["device"] = current_device
|
||||
kwargs["dtype"] = current_dtype
|
||||
super().__init__(*args, **kwargs)
|
||||
self.parameters_manual_cast = current_manual_cast_enabled
|
||||
|
||||
@ -198,8 +198,8 @@ class ForgeOperations:
|
||||
class Conv1d(torch.nn.Conv1d):
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
kwargs['device'] = current_device
|
||||
kwargs['dtype'] = current_dtype
|
||||
kwargs["device"] = current_device
|
||||
kwargs["dtype"] = current_dtype
|
||||
super().__init__(*args, **kwargs)
|
||||
self.parameters_manual_cast = current_manual_cast_enabled
|
||||
|
||||
@ -218,8 +218,8 @@ class ForgeOperations:
|
||||
class ConvTranspose2d(torch.nn.ConvTranspose2d):
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
kwargs['device'] = current_device
|
||||
kwargs['dtype'] = current_dtype
|
||||
kwargs["device"] = current_device
|
||||
kwargs["dtype"] = current_dtype
|
||||
super().__init__(*args, **kwargs)
|
||||
self.parameters_manual_cast = current_manual_cast_enabled
|
||||
|
||||
@ -243,8 +243,8 @@ class ForgeOperations:
|
||||
class ConvTranspose1d(torch.nn.ConvTranspose1d):
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
kwargs['device'] = current_device
|
||||
kwargs['dtype'] = current_dtype
|
||||
kwargs["device"] = current_device
|
||||
kwargs["dtype"] = current_dtype
|
||||
super().__init__(*args, **kwargs)
|
||||
self.parameters_manual_cast = current_manual_cast_enabled
|
||||
|
||||
@ -268,8 +268,8 @@ class ForgeOperations:
|
||||
class ConvTranspose3d(torch.nn.ConvTranspose3d):
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
kwargs['device'] = current_device
|
||||
kwargs['dtype'] = current_dtype
|
||||
kwargs["device"] = current_device
|
||||
kwargs["dtype"] = current_dtype
|
||||
super().__init__(*args, **kwargs)
|
||||
self.parameters_manual_cast = current_manual_cast_enabled
|
||||
|
||||
@ -293,8 +293,8 @@ class ForgeOperations:
|
||||
class GroupNorm(torch.nn.GroupNorm):
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
kwargs['device'] = current_device
|
||||
kwargs['dtype'] = current_dtype
|
||||
kwargs["device"] = current_device
|
||||
kwargs["dtype"] = current_dtype
|
||||
super().__init__(*args, **kwargs)
|
||||
self.parameters_manual_cast = current_manual_cast_enabled
|
||||
|
||||
@ -312,8 +312,8 @@ class ForgeOperations:
|
||||
class LayerNorm(torch.nn.LayerNorm):
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
kwargs['device'] = current_device
|
||||
kwargs['dtype'] = current_dtype
|
||||
kwargs["device"] = current_device
|
||||
kwargs["dtype"] = current_dtype
|
||||
super().__init__(*args, **kwargs)
|
||||
self.parameters_manual_cast = current_manual_cast_enabled
|
||||
|
||||
@ -331,7 +331,7 @@ class ForgeOperations:
|
||||
class Embedding(torch.nn.Embedding):
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
kwargs['device'] = current_device
|
||||
kwargs["device"] = current_device
|
||||
super().__init__(*args, **kwargs)
|
||||
self.parameters_manual_cast = current_manual_cast_enabled
|
||||
self.bias = None
|
||||
@ -350,7 +350,12 @@ class ForgeOperations:
|
||||
|
||||
|
||||
try:
|
||||
from backend.operations_bnb import ForgeLoader4Bit, ForgeParams4bit, functional_linear_4bits, functional_dequantize_4bit
|
||||
from backend.operations_bnb import (
|
||||
ForgeLoader4Bit,
|
||||
ForgeParams4bit,
|
||||
functional_dequantize_4bit,
|
||||
functional_linear_4bits,
|
||||
)
|
||||
|
||||
class ForgeOperationsBNB4bits(ForgeOperations):
|
||||
class Linear(ForgeLoader4Bit):
|
||||
@ -364,7 +369,7 @@ try:
|
||||
# And it only invokes one time, and most linear does not have bias
|
||||
self.bias = utils.tensor2parameter(self.bias.to(x.dtype))
|
||||
|
||||
if hasattr(self, 'forge_online_loras'):
|
||||
if hasattr(self, "forge_online_loras"):
|
||||
weight, bias, signal = weights_manual_cast(self, x, weight_fn=functional_dequantize_4bit, bias_fn=None, skip_bias_dtype=True)
|
||||
with main_stream_worker(weight, bias, signal):
|
||||
return torch.nn.functional.linear(x, weight, bias)
|
||||
@ -372,7 +377,7 @@ try:
|
||||
if not self.parameters_manual_cast:
|
||||
return functional_linear_4bits(x, self.weight, self.bias)
|
||||
elif not self.weight.bnb_quantized:
|
||||
assert x.device.type == 'cuda', 'BNB Must Use CUDA as Computation Device!'
|
||||
assert x.device.type == "cuda", "BNB Must Use CUDA as Computation Device!"
|
||||
layer_original_device = self.weight.device
|
||||
self.weight = self.weight._quantize(x.device)
|
||||
bias = self.bias.to(x.device) if self.bias is not None else None
|
||||
@ -402,23 +407,23 @@ class ForgeOperationsGGUF(ForgeOperations):
|
||||
self.parameters_manual_cast = current_manual_cast_enabled
|
||||
|
||||
def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs):
|
||||
if hasattr(self, 'dummy'):
|
||||
if hasattr(self, "dummy"):
|
||||
computation_dtype = self.dummy.dtype
|
||||
if computation_dtype not in [torch.float16, torch.bfloat16]:
|
||||
# GGUF cast only supports 16bits otherwise super slow
|
||||
computation_dtype = torch.float16
|
||||
if prefix + 'weight' in state_dict:
|
||||
self.weight = state_dict[prefix + 'weight'].to(device=self.dummy.device)
|
||||
if prefix + "weight" in state_dict:
|
||||
self.weight = state_dict[prefix + "weight"].to(device=self.dummy.device)
|
||||
self.weight.computation_dtype = computation_dtype
|
||||
if prefix + 'bias' in state_dict:
|
||||
self.bias = state_dict[prefix + 'bias'].to(device=self.dummy.device)
|
||||
if prefix + "bias" in state_dict:
|
||||
self.bias = state_dict[prefix + "bias"].to(device=self.dummy.device)
|
||||
self.bias.computation_dtype = computation_dtype
|
||||
del self.dummy
|
||||
else:
|
||||
if prefix + 'weight' in state_dict:
|
||||
self.weight = state_dict[prefix + 'weight']
|
||||
if prefix + 'bias' in state_dict:
|
||||
self.bias = state_dict[prefix + 'bias']
|
||||
if prefix + "weight" in state_dict:
|
||||
self.weight = state_dict[prefix + "weight"]
|
||||
if prefix + "bias" in state_dict:
|
||||
self.bias = state_dict[prefix + "bias"]
|
||||
return
|
||||
|
||||
def _apply(self, fn, recurse=True):
|
||||
@ -430,7 +435,7 @@ class ForgeOperationsGGUF(ForgeOperations):
|
||||
if self.bias is not None and self.bias.dtype != x.dtype:
|
||||
self.bias = utils.tensor2parameter(dequantize_tensor(self.bias).to(x.dtype))
|
||||
|
||||
if self.weight is not None and self.weight.dtype != x.dtype and getattr(self.weight, 'gguf_cls', None) is None:
|
||||
if self.weight is not None and self.weight.dtype != x.dtype and getattr(self.weight, "gguf_cls", None) is None:
|
||||
self.weight = utils.tensor2parameter(self.weight.to(x.dtype))
|
||||
|
||||
weight, bias, signal = weights_manual_cast(self, x, weight_fn=dequantize_tensor, bias_fn=None, skip_bias_dtype=True)
|
||||
@ -445,14 +450,14 @@ def using_forge_operations(operations=None, device=None, dtype=None, manual_cast
|
||||
current_device, current_dtype, current_manual_cast_enabled, current_bnb_dtype = device, dtype, manual_cast_enabled, bnb_dtype
|
||||
|
||||
if operations is None:
|
||||
if bnb_dtype in ['gguf']:
|
||||
if bnb_dtype in ["gguf"]:
|
||||
operations = ForgeOperationsGGUF
|
||||
elif bnb_avaliable and bnb_dtype in ['nf4', 'fp4']:
|
||||
elif bnb_avaliable and bnb_dtype in ["nf4", "fp4"]:
|
||||
operations = ForgeOperationsBNB4bits
|
||||
else:
|
||||
operations = ForgeOperations
|
||||
|
||||
op_names = ['Linear', 'Conv1d', 'Conv2d', 'Conv3d', 'ConvTranspose1d', 'ConvTranspose2d', 'ConvTranspose3d', 'GroupNorm', 'LayerNorm', 'Embedding']
|
||||
op_names = ["Linear", "Conv1d", "Conv2d", "Conv3d", "ConvTranspose1d", "ConvTranspose2d", "ConvTranspose3d", "GroupNorm", "LayerNorm", "Embedding"]
|
||||
backups = {op_name: getattr(torch.nn, op_name) for op_name in op_names}
|
||||
|
||||
try:
|
||||
@ -469,17 +474,14 @@ def using_forge_operations(operations=None, device=None, dtype=None, manual_cast
|
||||
|
||||
def shift_manual_cast(model, enabled):
|
||||
for m in model.modules():
|
||||
if hasattr(m, 'parameters_manual_cast'):
|
||||
if hasattr(m, "parameters_manual_cast"):
|
||||
m.parameters_manual_cast = enabled
|
||||
return
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def automatic_memory_management():
|
||||
memory_management.free_memory(
|
||||
memory_required=3 * 1024 * 1024 * 1024,
|
||||
device=memory_management.get_torch_device()
|
||||
)
|
||||
memory_management.free_memory(memory_required=3 * 1024 * 1024 * 1024, device=memory_management.get_torch_device())
|
||||
|
||||
module_list = []
|
||||
|
||||
@ -511,7 +513,7 @@ def automatic_memory_management():
|
||||
memory_management.soft_empty_cache()
|
||||
end = time.perf_counter()
|
||||
|
||||
print(f'Automatic Memory Management: {len(module_list)} Modules in {(end - start):.2f} seconds.')
|
||||
print(f"Automatic Memory Management: {len(module_list)} Modules in {(end - start):.2f} seconds.")
|
||||
return
|
||||
|
||||
|
||||
@ -519,11 +521,11 @@ class DynamicSwapInstaller:
|
||||
@staticmethod
|
||||
def _install_module(module: torch.nn.Module, target_device: torch.device):
|
||||
original_class = module.__class__
|
||||
module.__dict__['forge_backup_original_class'] = original_class
|
||||
module.__dict__["forge_backup_original_class"] = original_class
|
||||
|
||||
def hacked_get_attr(self, name: str):
|
||||
if '_parameters' in self.__dict__:
|
||||
_parameters = self.__dict__['_parameters']
|
||||
if "_parameters" in self.__dict__:
|
||||
_parameters = self.__dict__["_parameters"]
|
||||
if name in _parameters:
|
||||
p = _parameters[name]
|
||||
if p is None:
|
||||
@ -532,22 +534,26 @@ class DynamicSwapInstaller:
|
||||
return torch.nn.Parameter(p.to(target_device), requires_grad=p.requires_grad)
|
||||
else:
|
||||
return p.to(target_device)
|
||||
if '_buffers' in self.__dict__:
|
||||
_buffers = self.__dict__['_buffers']
|
||||
if "_buffers" in self.__dict__:
|
||||
_buffers = self.__dict__["_buffers"]
|
||||
if name in _buffers:
|
||||
return _buffers[name].to(target_device)
|
||||
return super(original_class, self).__getattr__(name)
|
||||
|
||||
module.__class__ = type('DynamicSwap_' + original_class.__name__, (original_class,), {
|
||||
'__getattr__': hacked_get_attr,
|
||||
})
|
||||
module.__class__ = type(
|
||||
"DynamicSwap_" + original_class.__name__,
|
||||
(original_class,),
|
||||
{
|
||||
"__getattr__": hacked_get_attr,
|
||||
},
|
||||
)
|
||||
|
||||
return
|
||||
|
||||
@staticmethod
|
||||
def _uninstall_module(module: torch.nn.Module):
|
||||
if 'forge_backup_original_class' in module.__dict__:
|
||||
module.__class__ = module.__dict__.pop('forge_backup_original_class')
|
||||
if "forge_backup_original_class" in module.__dict__:
|
||||
module.__class__ = module.__dict__.pop("forge_backup_original_class")
|
||||
return
|
||||
|
||||
@staticmethod
|
||||
@ -561,4 +567,3 @@ class DynamicSwapInstaller:
|
||||
for m in model.modules():
|
||||
DynamicSwapInstaller._uninstall_module(m)
|
||||
return
|
||||
|
||||
|
||||
@ -1,11 +1,11 @@
|
||||
# Copyright Forge 2024
|
||||
|
||||
import torch
|
||||
import bitsandbytes as bnb
|
||||
|
||||
from backend import utils, memory_management
|
||||
from bitsandbytes.nn.modules import Params4bit, QuantState
|
||||
import torch
|
||||
from bitsandbytes.functional import dequantize_4bit
|
||||
from bitsandbytes.nn.modules import Params4bit, QuantState
|
||||
|
||||
from backend import memory_management, utils
|
||||
|
||||
|
||||
def functional_linear_4bits(x, weight, bias):
|
||||
@ -20,12 +20,12 @@ def functional_dequantize_4bit(weight):
|
||||
|
||||
weight_original_device = weight.device
|
||||
|
||||
if weight_original_device.type != 'cuda':
|
||||
if weight_original_device.type != "cuda":
|
||||
weight = weight.cuda()
|
||||
|
||||
weight = dequantize_4bit(weight, quant_state=weight.quant_state, blocksize=weight.blocksize, quant_type=weight.quant_type)
|
||||
|
||||
if weight_original_device.type != 'cuda':
|
||||
if weight_original_device.type != "cuda":
|
||||
weight = weight.to(device=weight_original_device)
|
||||
|
||||
return weight
|
||||
@ -118,26 +118,26 @@ class ForgeLoader4Bit(torch.nn.Module):
|
||||
return
|
||||
|
||||
def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs):
|
||||
quant_state_keys = {k[len(prefix + "weight."):] for k in state_dict.keys() if k.startswith(prefix + "weight.")}
|
||||
quant_state_keys = {k[len(prefix + "weight.") :] for k in state_dict.keys() if k.startswith(prefix + "weight.")}
|
||||
|
||||
if any('bitsandbytes' in k for k in quant_state_keys):
|
||||
if any("bitsandbytes" in k for k in quant_state_keys):
|
||||
quant_state_dict = {k: state_dict[prefix + "weight." + k] for k in quant_state_keys}
|
||||
|
||||
self.weight = ForgeParams4bit.from_prequantized(
|
||||
data=state_dict[prefix + 'weight'],
|
||||
data=state_dict[prefix + "weight"],
|
||||
quantized_stats=quant_state_dict,
|
||||
requires_grad=False,
|
||||
device=self.dummy.device,
|
||||
)
|
||||
|
||||
if prefix + 'bias' in state_dict:
|
||||
self.bias = torch.nn.Parameter(state_dict[prefix + 'bias'].to(self.dummy))
|
||||
if prefix + "bias" in state_dict:
|
||||
self.bias = torch.nn.Parameter(state_dict[prefix + "bias"].to(self.dummy))
|
||||
|
||||
del self.dummy
|
||||
elif hasattr(self, 'dummy'):
|
||||
if prefix + 'weight' in state_dict:
|
||||
elif hasattr(self, "dummy"):
|
||||
if prefix + "weight" in state_dict:
|
||||
self.weight = ForgeParams4bit(
|
||||
state_dict[prefix + 'weight'].to(self.dummy),
|
||||
state_dict[prefix + "weight"].to(self.dummy),
|
||||
requires_grad=False,
|
||||
compress_statistics=False,
|
||||
blocksize=64,
|
||||
@ -145,8 +145,8 @@ class ForgeLoader4Bit(torch.nn.Module):
|
||||
quant_storage=torch.uint8,
|
||||
)
|
||||
|
||||
if prefix + 'bias' in state_dict:
|
||||
self.bias = torch.nn.Parameter(state_dict[prefix + 'bias'].to(self.dummy))
|
||||
if prefix + "bias" in state_dict:
|
||||
self.bias = torch.nn.Parameter(state_dict[prefix + "bias"].to(self.dummy))
|
||||
|
||||
del self.dummy
|
||||
else:
|
||||
@ -154,16 +154,8 @@ class ForgeLoader4Bit(torch.nn.Module):
|
||||
|
||||
def reload_weight(self, weight):
|
||||
weight_original_device = weight.device
|
||||
weight = ForgeParams4bit(
|
||||
weight,
|
||||
requires_grad=False,
|
||||
compress_statistics=self.weight.compress_statistics,
|
||||
blocksize=self.weight.blocksize,
|
||||
quant_type=self.weight.quant_type,
|
||||
quant_storage=self.weight.quant_storage,
|
||||
bnb_quantized=False
|
||||
)
|
||||
if weight_original_device.type == 'cuda':
|
||||
weight = ForgeParams4bit(weight, requires_grad=False, compress_statistics=self.weight.compress_statistics, blocksize=self.weight.blocksize, quant_type=self.weight.quant_type, quant_storage=self.weight.quant_storage, bnb_quantized=False)
|
||||
if weight_original_device.type == "cuda":
|
||||
weight = weight.to(weight_original_device)
|
||||
else:
|
||||
weight = weight.cuda().to(weight_original_device)
|
||||
|
||||
@ -1,7 +1,6 @@
|
||||
import gguf
|
||||
import torch
|
||||
|
||||
|
||||
quants_mapping = {
|
||||
gguf.GGMLQuantizationType.Q2_K: gguf.Q2_K,
|
||||
gguf.GGMLQuantizationType.Q3_K: gguf.Q3_K,
|
||||
@ -60,7 +59,7 @@ def dequantize_tensor(tensor):
|
||||
if tensor is None:
|
||||
return None
|
||||
|
||||
if not hasattr(tensor, 'gguf_cls'):
|
||||
if not hasattr(tensor, "gguf_cls"):
|
||||
return tensor
|
||||
|
||||
gguf_cls = tensor.gguf_cls
|
||||
|
||||
@ -5,7 +5,12 @@ import torch
|
||||
from backend import memory_management, state_dict, utils
|
||||
from backend.misc import image_resize
|
||||
from backend.nn.cnets import cldm, t2i_adapter
|
||||
from backend.operations import ForgeOperations, main_stream_worker, using_forge_operations, weights_manual_cast
|
||||
from backend.operations import (
|
||||
ForgeOperations,
|
||||
main_stream_worker,
|
||||
using_forge_operations,
|
||||
weights_manual_cast,
|
||||
)
|
||||
from backend.patcher.base import ModelPatcher
|
||||
|
||||
|
||||
|
||||
@ -11,7 +11,11 @@ import torch
|
||||
from backend import memory_management, utils
|
||||
from backend.args import args, dynamic_args
|
||||
from backend.operations import cleanup_cache
|
||||
from backend.sampling.condition import Condition, compile_conditions, compile_weighted_conditions
|
||||
from backend.sampling.condition import (
|
||||
Condition,
|
||||
compile_conditions,
|
||||
compile_weighted_conditions,
|
||||
)
|
||||
|
||||
|
||||
def get_area_and_mult(conds, x_in, timestep_in):
|
||||
|
||||
@ -1,6 +1,3 @@
|
||||
import torch
|
||||
|
||||
|
||||
def load_state_dict(model, sd, ignore_errors=[], log_name=None, ignore_start=None):
|
||||
missing, unexpected = model.load_state_dict(sd, strict=False)
|
||||
missing = [x for x in missing if x not in ignore_errors]
|
||||
@ -12,9 +9,9 @@ def load_state_dict(model, sd, ignore_errors=[], log_name=None, ignore_start=Non
|
||||
|
||||
log_name = log_name or type(model).__name__
|
||||
if len(missing) > 0:
|
||||
print(f'{log_name} Missing: {missing}')
|
||||
print(f"{log_name} Missing: {missing}")
|
||||
if len(unexpected) > 0:
|
||||
print(f'{log_name} Unexpected: {unexpected}')
|
||||
print(f"{log_name} Unexpected: {unexpected}")
|
||||
return
|
||||
|
||||
|
||||
@ -22,18 +19,18 @@ def state_dict_has(sd, prefix):
|
||||
return any(x.startswith(prefix) for x in sd.keys())
|
||||
|
||||
|
||||
def filter_state_dict_with_prefix(sd, prefix, new_prefix=''):
|
||||
def filter_state_dict_with_prefix(sd, prefix, new_prefix=""):
|
||||
new_sd = {}
|
||||
|
||||
for k, v in list(sd.items()):
|
||||
if k.startswith(prefix):
|
||||
new_sd[new_prefix + k[len(prefix):]] = v
|
||||
new_sd[new_prefix + k[len(prefix) :]] = v
|
||||
del sd[k]
|
||||
|
||||
return new_sd
|
||||
|
||||
|
||||
def try_filter_state_dict(sd, prefix_list, new_prefix=''):
|
||||
def try_filter_state_dict(sd, prefix_list, new_prefix=""):
|
||||
for prefix in prefix_list:
|
||||
if state_dict_has(sd, prefix):
|
||||
return filter_state_dict_with_prefix(sd, prefix, new_prefix)
|
||||
@ -77,7 +74,7 @@ def transformers_convert(sd, prefix_from, prefix_to, number):
|
||||
for x in range(3):
|
||||
p = ["self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj"]
|
||||
k_to = "{}encoder.layers.{}.{}.{}".format(prefix_to, resblock, p[x], y)
|
||||
sd[k_to] = weights[shape_from*x:shape_from*(x + 1)]
|
||||
sd[k_to] = weights[shape_from * x : shape_from * (x + 1)]
|
||||
return sd
|
||||
|
||||
|
||||
@ -94,7 +91,7 @@ def state_dict_prefix_replace(state_dict, replace_prefix, filter_keys=False):
|
||||
else:
|
||||
out = state_dict
|
||||
for rp in replace_prefix:
|
||||
replace = list(map(lambda a: (a, "{}{}".format(replace_prefix[rp], a[len(rp):])), filter(lambda a: a.startswith(rp), state_dict.keys())))
|
||||
replace = list(map(lambda a: (a, "{}{}".format(replace_prefix[rp], a[len(rp) :])), filter(lambda a: a.startswith(rp), state_dict.keys())))
|
||||
for x in replace:
|
||||
w = state_dict.pop(x[0])
|
||||
out[x[1]] = w
|
||||
|
||||
@ -1,4 +1,5 @@
|
||||
import torch
|
||||
|
||||
from backend.args import args
|
||||
|
||||
|
||||
|
||||
@ -1,19 +1,21 @@
|
||||
import gguf
|
||||
import torch
|
||||
import os
|
||||
import json
|
||||
import os
|
||||
|
||||
import gguf
|
||||
import safetensors.torch
|
||||
import torch
|
||||
|
||||
import backend.misc.checkpoint_pickle
|
||||
from backend.operations_gguf import ParameterGGUF
|
||||
|
||||
|
||||
def read_arbitrary_config(directory):
|
||||
config_path = os.path.join(directory, 'config.json')
|
||||
config_path = os.path.join(directory, "config.json")
|
||||
|
||||
if not os.path.exists(config_path):
|
||||
raise FileNotFoundError(f"No config.json file found in the directory: {directory}")
|
||||
|
||||
with open(config_path, 'rt', encoding='utf-8') as file:
|
||||
with open(config_path, "rt", encoding="utf-8") as file:
|
||||
config_data = json.load(file)
|
||||
|
||||
return config_data
|
||||
@ -31,7 +33,7 @@ def load_torch_file(ckpt, safe_load=False, device=None):
|
||||
sd[str(tensor.name)] = ParameterGGUF(tensor)
|
||||
else:
|
||||
if safe_load:
|
||||
if not 'weights_only' in torch.load.__code__.co_varnames:
|
||||
if not "weights_only" in torch.load.__code__.co_varnames:
|
||||
print("Warning torch.load doesn't support weights_only on this pytorch version, loading unsafely.")
|
||||
safe_load = False
|
||||
if safe_load:
|
||||
@ -147,9 +149,9 @@ def nested_move_to_device(obj, **kwargs):
|
||||
return obj
|
||||
|
||||
|
||||
def get_state_dict_after_quant(model, prefix=''):
|
||||
def get_state_dict_after_quant(model, prefix=""):
|
||||
for m in model.modules():
|
||||
if hasattr(m, 'weight') and hasattr(m.weight, 'bnb_quantized'):
|
||||
if hasattr(m, "weight") and hasattr(m.weight, "bnb_quantized"):
|
||||
if not m.weight.bnb_quantized:
|
||||
original_device = m.weight.device
|
||||
m.cuda()
|
||||
@ -162,14 +164,15 @@ def get_state_dict_after_quant(model, prefix=''):
|
||||
|
||||
def beautiful_print_gguf_state_dict_statics(state_dict):
|
||||
from gguf.constants import GGMLQuantizationType
|
||||
|
||||
type_counts = {}
|
||||
for k, v in state_dict.items():
|
||||
gguf_cls = getattr(v, 'gguf_cls', None)
|
||||
gguf_cls = getattr(v, "gguf_cls", None)
|
||||
if gguf_cls is not None:
|
||||
type_name = gguf_cls.__name__
|
||||
if type_name in type_counts:
|
||||
type_counts[type_name] += 1
|
||||
else:
|
||||
type_counts[type_name] = 1
|
||||
print(f'GGUF state dict: {type_counts}')
|
||||
print(f"GGUF state dict: {type_counts}")
|
||||
return
|
||||
|
||||
Loading…
Reference in New Issue
Block a user