diff --git a/.github/workflows/run_tests.yaml b/.github/workflows/run_tests.yaml
index 0610f4f5436..eccf15e4684 100644
--- a/.github/workflows/run_tests.yaml
+++ b/.github/workflows/run_tests.yaml
@@ -51,6 +51,7 @@ jobs:
--test-server
--do-not-download-clip
--no-half
+ --precision full
--disable-opt-split-attention
--use-cpu all
--api-server-stop
diff --git a/configs/flux1-inference.yaml b/configs/flux1-inference.yaml
new file mode 100644
index 00000000000..f9bbe9073f3
--- /dev/null
+++ b/configs/flux1-inference.yaml
@@ -0,0 +1,4 @@
+model:
+ target: modules.models.flux.FLUX1Inferencer
+ params:
+ state_dict: null
diff --git a/configs/sd_xl_v.yaml b/configs/sd_xl_v.yaml
index c755dc74fda..9fbcbbac8b2 100644
--- a/configs/sd_xl_v.yaml
+++ b/configs/sd_xl_v.yaml
@@ -10,7 +10,7 @@ model:
num_idx: 1000
weighting_config:
- target: sgm.modules.diffusionmodules.denoiser_weighting.EpsWeighting
+ target: sgm.modules.diffusionmodules.denoiser_weighting.VWeighting
scaling_config:
target: sgm.modules.diffusionmodules.denoiser_scaling.VScaling
discretization_config:
@@ -21,7 +21,7 @@ model:
params:
adm_in_channels: 2816
num_classes: sequential
- use_checkpoint: True
+ use_checkpoint: False
in_channels: 4
out_channels: 4
model_channels: 320
diff --git a/extensions-builtin/Lora/network_lora.py b/extensions-builtin/Lora/network_lora.py
index a7a088949ea..2bc6af5d204 100644
--- a/extensions-builtin/Lora/network_lora.py
+++ b/extensions-builtin/Lora/network_lora.py
@@ -2,6 +2,7 @@
import lyco_helpers
import modules.models.sd3.mmdit
+import modules.models.flux.modules.layers
import network
from modules import devices
@@ -37,7 +38,7 @@ def create_module(self, weights, key, none_ok=False):
if weight is None and none_ok:
return None
- is_linear = type(self.sd_module) in [torch.nn.Linear, torch.nn.modules.linear.NonDynamicallyQuantizableLinear, torch.nn.MultiheadAttention, modules.models.sd3.mmdit.QkvLinear]
+ is_linear = type(self.sd_module) in [torch.nn.Linear, torch.nn.modules.linear.NonDynamicallyQuantizableLinear, torch.nn.MultiheadAttention, modules.models.sd3.mmdit.QkvLinear, modules.models.flux.modules.layers.QkvLinear ]
is_conv = type(self.sd_module) in [torch.nn.Conv2d]
if is_linear:
diff --git a/extensions-builtin/Lora/networks.py b/extensions-builtin/Lora/networks.py
index 67f9abe2a37..948fa674033 100644
--- a/extensions-builtin/Lora/networks.py
+++ b/extensions-builtin/Lora/networks.py
@@ -37,7 +37,7 @@
re_digits = re.compile(r"\d+")
-re_x_proj = re.compile(r"(.*)_([qkv]_proj)$")
+re_x_proj = re.compile(r"(.*)_((?:[qkv]|mlp)_proj)$")
re_compiled = {}
suffix_conversion = {
@@ -183,8 +183,12 @@ def load_network(name, network_on_disk):
for key_network, weight in sd.items():
if diffusers_weight_map:
- key_network_without_network_parts, network_name, network_weight = key_network.rsplit(".", 2)
- network_part = network_name + '.' + network_weight
+ if key_network.startswith("lora_unet"):
+ key_network_without_network_parts, _, network_part = key_network.partition(".")
+ key_network_without_network_parts = key_network_without_network_parts.replace("lora_unet", "diffusion_model")
+ else:
+ key_network_without_network_parts, network_name, network_weight = key_network.rsplit(".", 2)
+ network_part = network_name + '.' + network_weight
else:
key_network_without_network_parts, _, network_part = key_network.partition(".")
@@ -373,11 +377,13 @@ def allowed_layer_without_weight(layer):
return False
-def store_weights_backup(weight):
+def store_weights_backup(weight, dtype):
if weight is None:
return None
- return weight.to(devices.cpu, copy=True)
+ if shared.opts.lora_without_backup_weight:
+ return True
+ return weight.to(devices.cpu, dtype=dtype, copy=True)
def restore_weights_backup(obj, field, weight):
@@ -385,16 +391,20 @@ def restore_weights_backup(obj, field, weight):
setattr(obj, field, None)
return
- getattr(obj, field).copy_(weight)
+ old_weight = getattr(obj, field)
+ old_weight.copy_(weight.to(dtype=old_weight.dtype))
-def network_restore_weights_from_backup(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn.GroupNorm, torch.nn.LayerNorm, torch.nn.MultiheadAttention]):
+def network_restore_weights_from_backup(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn.GroupNorm, torch.nn.LayerNorm, torch.nn.MultiheadAttention], cleanup=False):
weights_backup = getattr(self, "network_weights_backup", None)
bias_backup = getattr(self, "network_bias_backup", None)
if weights_backup is None and bias_backup is None:
return
+ if shared.opts.lora_without_backup_weight:
+ return
+
if weights_backup is not None:
if isinstance(self, torch.nn.MultiheadAttention):
restore_weights_backup(self, 'in_proj_weight', weights_backup[0])
@@ -407,55 +417,79 @@ def network_restore_weights_from_backup(self: Union[torch.nn.Conv2d, torch.nn.Li
else:
restore_weights_backup(self, 'bias', bias_backup)
+ if cleanup:
+ if weights_backup is not None:
+ del self.network_weights_backup
+ if bias_backup is not None:
+ del self.network_bias_backup
-def network_apply_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn.GroupNorm, torch.nn.LayerNorm, torch.nn.MultiheadAttention]):
- """
- Applies the currently selected set of networks to the weights of torch layer self.
- If weights already have this particular set of networks applied, does nothing.
- If not, restores original weights from backup and alters weights according to networks.
- """
+def network_backup_weights(self):
network_layer_name = getattr(self, 'network_layer_name', None)
- if network_layer_name is None:
- return
- current_names = getattr(self, "network_current_names", ())
+ _current_names = getattr(self, "network_current_names", ())
wanted_names = tuple((x.name, x.te_multiplier, x.unet_multiplier, x.dyn_dim) for x in loaded_networks)
+ need_backup = False
+ for net in loaded_networks:
+ if network_layer_name in net.modules:
+ need_backup = True
+ break
+ elif network_layer_name + "_q_proj" in net.modules:
+ need_backup = True
+ break
+
+ if not need_backup:
+ return
+
weights_backup = getattr(self, "network_weights_backup", None)
if weights_backup is None and wanted_names != ():
- if current_names != () and not allowed_layer_without_weight(self):
- raise RuntimeError(f"{network_layer_name} - no backup weights found and current weights are not unchanged")
-
if isinstance(self, torch.nn.MultiheadAttention):
- weights_backup = (store_weights_backup(self.in_proj_weight), store_weights_backup(self.out_proj.weight))
+ weights_backup = (store_weights_backup(self.in_proj_weight, self.org_dtype), store_weights_backup(self.out_proj.weight, self.org_dtype))
else:
- weights_backup = store_weights_backup(self.weight)
+ weights_backup = store_weights_backup(self.weight, self.org_dtype)
self.network_weights_backup = weights_backup
bias_backup = getattr(self, "network_bias_backup", None)
if bias_backup is None and wanted_names != ():
if isinstance(self, torch.nn.MultiheadAttention) and self.out_proj.bias is not None:
- bias_backup = store_weights_backup(self.out_proj.bias)
+ bias_backup = store_weights_backup(self.out_proj.bias, self.org_dtype)
elif getattr(self, 'bias', None) is not None:
- bias_backup = store_weights_backup(self.bias)
+ bias_backup = store_weights_backup(self.bias, self.org_dtype)
else:
bias_backup = None
- # Unlike weight which always has value, some modules don't have bias.
- # Only report if bias is not None and current bias are not unchanged.
- if bias_backup is not None and current_names != ():
- raise RuntimeError("no backup bias found and current bias are not unchanged")
-
self.network_bias_backup = bias_backup
- if current_names != wanted_names:
+
+def network_apply_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn.GroupNorm, torch.nn.LayerNorm, torch.nn.MultiheadAttention]):
+ """
+ Applies the currently selected set of networks to the weights of torch layer self.
+ If weights already have this particular set of networks applied, does nothing.
+ If not, restores original weights from backup and alters weights according to networks.
+ """
+
+ network_layer_name = getattr(self, 'network_layer_name', None)
+ if network_layer_name is None:
+ return
+
+ current_names = getattr(self, "network_current_names", ())
+ wanted_names = tuple((x.name, x.te_multiplier, x.unet_multiplier, x.dyn_dim) for x in loaded_networks)
+
+ weights_backup = getattr(self, "network_weights_backup", None)
+ if weights_backup is None and wanted_names != ():
+ network_backup_weights(self)
+ elif current_names != () and current_names != wanted_names and not getattr(self, "weights_restored", False):
network_restore_weights_from_backup(self)
+ if current_names != wanted_names:
+ if hasattr(self, "weights_restored"):
+ self.weights_restored = False
+
for net in loaded_networks:
module = net.modules.get(network_layer_name, None)
- if module is not None and hasattr(self, 'weight') and not isinstance(module, modules.models.sd3.mmdit.QkvLinear):
+ if module is not None and hasattr(self, 'weight') and not all(isinstance(module, linear) for linear in (modules.models.sd3.mmdit.QkvLinear, modules.models.flux.modules.layers.QkvLinear)):
try:
with torch.no_grad():
if getattr(self, 'fp16_weight', None) is None:
@@ -478,6 +512,7 @@ def network_apply_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn
self.bias = torch.nn.Parameter(ex_bias).to(self.weight.dtype)
else:
self.bias.copy_((bias + ex_bias).to(dtype=self.bias.dtype))
+ del weight, bias, updown, ex_bias
except RuntimeError as e:
logging.debug(f"Network {net.name} layer {network_layer_name}: {e}")
extra_network_lora.errors[net.name] = extra_network_lora.errors.get(net.name, 0) + 1
@@ -515,7 +550,9 @@ def network_apply_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn
continue
- if isinstance(self, modules.models.sd3.mmdit.QkvLinear) and module_q and module_k and module_v:
+ module_mlp = net.modules.get(network_layer_name + "_mlp_proj", None)
+
+ if any(isinstance(self, linear) for linear in (modules.models.sd3.mmdit.QkvLinear, modules.models.flux.modules.layers.QkvLinear)) and module_q and module_k and module_v and module_mlp is None and self.weight.shape[0] // 3 == module_q.up_model.weight.shape[0]:
try:
with torch.no_grad():
# Send "real" orig_weight into MHA's lora module
@@ -526,6 +563,31 @@ def network_apply_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn
del qw, kw, vw
updown_qkv = torch.vstack([updown_q, updown_k, updown_v])
self.weight += updown_qkv
+ del updown_qkv
+ del updown_q, updown_k, updown_v
+
+ except RuntimeError as e:
+ logging.debug(f"Network {net.name} layer {network_layer_name}: {e}")
+ extra_network_lora.errors[net.name] = extra_network_lora.errors.get(net.name, 0) + 1
+
+ continue
+
+ if any(isinstance(self, linear) for linear in (modules.models.flux.modules.layers.QkvLinear,)) and module_q and module_k and module_v:
+ try:
+ with torch.no_grad():
+ qw, kw, vw, mlp = torch.tensor_split(self.weight, (3072, 6144, 9216,), 0)
+ updown_q, _ = module_q.calc_updown(qw)
+ updown_k, _ = module_k.calc_updown(kw)
+ updown_v, _ = module_v.calc_updown(vw)
+ if module_mlp is not None:
+ updown_mlp, _ = module_mlp.calc_updown(mlp)
+ else:
+ updown_mlp = torch.zeros(3072 * 4, 3072, dtype=updown_q.dtype, device=updown_q.device)
+ del qw, kw, vw, mlp
+ updown_qkv_mlp = torch.vstack([updown_q, updown_k, updown_v, updown_mlp])
+ self.weight += updown_qkv_mlp
+ del updown_qkv_mlp
+ del updown_q, updown_k, updown_v, updown_mlp
except RuntimeError as e:
logging.debug(f"Network {net.name} layer {network_layer_name}: {e}")
@@ -539,7 +601,12 @@ def network_apply_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn
logging.debug(f"Network {net.name} layer {network_layer_name}: couldn't find supported operation")
extra_network_lora.errors[net.name] = extra_network_lora.errors.get(net.name, 0) + 1
- self.network_current_names = wanted_names
+
+ if shared.opts.lora_without_backup_weight:
+ self.network_weights_backup = None
+ self.network_bias_backup = None
+ else:
+ self.network_current_names = wanted_names
def network_forward(org_module, input, original_forward):
diff --git a/extensions-builtin/Lora/scripts/lora_script.py b/extensions-builtin/Lora/scripts/lora_script.py
index d3ea369ae26..7a23b8d57d4 100644
--- a/extensions-builtin/Lora/scripts/lora_script.py
+++ b/extensions-builtin/Lora/scripts/lora_script.py
@@ -1,15 +1,17 @@
import re
+import torch
import gradio as gr
from fastapi import FastAPI
+import gc
import network
import networks
import lora # noqa:F401
import lora_patches
import extra_networks_lora
import ui_extra_networks_lora
-from modules import script_callbacks, ui_extra_networks, extra_networks, shared
+from modules import script_callbacks, ui_extra_networks, extra_networks, shared, scripts, devices
def unload():
@@ -97,6 +99,82 @@ def network_replacement(m):
d["Prompt"] = re.sub(re_lora, network_replacement, d["Prompt"])
+class ScriptLora(scripts.Script):
+ name = "Lora"
+
+ def title(self):
+ return self.name
+
+ def show(self, is_img2img):
+ return scripts.AlwaysVisible
+
+ def after_extra_networks_activate(self, p, *args, **kwargs):
+ # check modules and setup org_dtype
+ modules = []
+ if shared.sd_model.is_sdxl:
+ for _i, embedder in enumerate(shared.sd_model.conditioner.embedders):
+ if not hasattr(embedder, 'wrapped'):
+ continue
+
+ for _name, module in embedder.wrapped.named_modules():
+ if isinstance(module, (torch.nn.Conv2d, torch.nn.Linear, torch.nn.GroupNorm, torch.nn.LayerNorm, torch.nn.MultiheadAttention)):
+ if hasattr(module, 'weight'):
+ modules.append(module)
+ elif isinstance(module, torch.nn.MultiheadAttention):
+ modules.append(module)
+
+ else:
+ cond_stage_model = getattr(shared.sd_model.cond_stage_model, 'wrapped', shared.sd_model.cond_stage_model)
+
+ for _name, module in cond_stage_model.named_modules():
+ if isinstance(module, (torch.nn.Conv2d, torch.nn.Linear, torch.nn.GroupNorm, torch.nn.LayerNorm, torch.nn.MultiheadAttention)):
+ if hasattr(module, 'weight'):
+ modules.append(module)
+ elif isinstance(module, torch.nn.MultiheadAttention):
+ modules.append(module)
+
+ for _name, module in shared.sd_model.model.named_modules():
+ if isinstance(module, (torch.nn.Conv2d, torch.nn.Linear, torch.nn.GroupNorm, torch.nn.LayerNorm, torch.nn.MultiheadAttention)):
+ if hasattr(module, 'weight'):
+ modules.append(module)
+ elif isinstance(module, torch.nn.MultiheadAttention):
+ modules.append(module)
+
+ print("Total lora modules after_extra_networks_activate() =", len(modules))
+
+ target_dtype = devices.dtype_inference
+ for module in modules:
+ network_layer_name = getattr(module, 'network_layer_name', None)
+ if network_layer_name is None:
+ continue
+
+ if isinstance(module, torch.nn.MultiheadAttention):
+ org_dtype = torch.float32
+ else:
+ org_dtype = None
+ for _name, param in module.named_parameters():
+ if param.dtype != target_dtype:
+ org_dtype = param.dtype
+ break
+
+ # set org_dtype
+ module.org_dtype = org_dtype
+
+ # backup/restore weights
+ current_names = getattr(module, "network_current_names", ())
+ wanted_names = tuple((x.name, x.te_multiplier, x.unet_multiplier, x.dyn_dim) for x in networks.loaded_networks)
+
+ weights_backup = getattr(module, "network_weights_backup", None)
+
+ if current_names == () and current_names != wanted_names and weights_backup is None:
+ networks.network_backup_weights(module)
+ elif current_names != () and current_names != wanted_names:
+ networks.network_restore_weights_from_backup(module, wanted_names == ())
+ module.weights_restored = True
+ if current_names != wanted_names and wanted_names == ():
+ gc.collect()
+
+
script_callbacks.on_infotext_pasted(infotext_pasted)
shared.opts.onchange("lora_in_memory_limit", networks.purge_networks_from_memory)
diff --git a/extensions-builtin/canvas-zoom-and-pan/javascript/zoom.js b/extensions-builtin/canvas-zoom-and-pan/javascript/zoom.js
index 7807f7f6185..b40c860275b 100644
--- a/extensions-builtin/canvas-zoom-and-pan/javascript/zoom.js
+++ b/extensions-builtin/canvas-zoom-and-pan/javascript/zoom.js
@@ -816,7 +816,7 @@ onUiLoaded(async() => {
// Increase or decrease brush size based on scroll direction
adjustBrushSize(elemId, e.deltaY);
}
- });
+ }, {passive: false});
// Handle the move event for pan functionality. Updates the panX and panY variables and applies the new transform to the target element.
function handleMoveKeyDown(e) {
diff --git a/extensions-builtin/hypertile/hypertile.py b/extensions-builtin/hypertile/hypertile.py
index 0f40e2d3925..2da71583871 100644
--- a/extensions-builtin/hypertile/hypertile.py
+++ b/extensions-builtin/hypertile/hypertile.py
@@ -1,7 +1,7 @@
"""
Hypertile module for splitting attention layers in SD-1.5 U-Net and SD-1.5 VAE
Warn: The patch works well only if the input image has a width and height that are multiples of 128
-Original author: @tfernd Github: https://github.com/tfernd/HyperTile
+Original author: @tfernd GitHub: https://github.com/tfernd/HyperTile
"""
from __future__ import annotations
diff --git a/extensions-builtin/postprocessing-for-training/scripts/postprocessing_autosized_crop.py b/extensions-builtin/postprocessing-for-training/scripts/postprocessing_autosized_crop.py
index 7e674989814..c8ea7e0e061 100644
--- a/extensions-builtin/postprocessing-for-training/scripts/postprocessing_autosized_crop.py
+++ b/extensions-builtin/postprocessing-for-training/scripts/postprocessing_autosized_crop.py
@@ -34,14 +34,14 @@ def ui(self):
with ui_components.InputAccordion(False, label="Auto-sized crop") as enable:
gr.Markdown('Each image is center-cropped with an automatically chosen width and height.')
with gr.Row():
- mindim = gr.Slider(minimum=64, maximum=2048, step=8, label="Dimension lower bound", value=384, elem_id="postprocess_multicrop_mindim")
- maxdim = gr.Slider(minimum=64, maximum=2048, step=8, label="Dimension upper bound", value=768, elem_id="postprocess_multicrop_maxdim")
+ mindim = gr.Slider(minimum=64, maximum=2048, step=8, label="Dimension lower bound", value=384, elem_id=self.elem_id_suffix("postprocess_multicrop_mindim"))
+ maxdim = gr.Slider(minimum=64, maximum=2048, step=8, label="Dimension upper bound", value=768, elem_id=self.elem_id_suffix("postprocess_multicrop_maxdim"))
with gr.Row():
- minarea = gr.Slider(minimum=64 * 64, maximum=2048 * 2048, step=1, label="Area lower bound", value=64 * 64, elem_id="postprocess_multicrop_minarea")
- maxarea = gr.Slider(minimum=64 * 64, maximum=2048 * 2048, step=1, label="Area upper bound", value=640 * 640, elem_id="postprocess_multicrop_maxarea")
+ minarea = gr.Slider(minimum=64 * 64, maximum=2048 * 2048, step=1, label="Area lower bound", value=64 * 64, elem_id=self.elem_id_suffix("postprocess_multicrop_minarea"))
+ maxarea = gr.Slider(minimum=64 * 64, maximum=2048 * 2048, step=1, label="Area upper bound", value=640 * 640, elem_id=self.elem_id_suffix("postprocess_multicrop_maxarea"))
with gr.Row():
- objective = gr.Radio(["Maximize area", "Minimize error"], value="Maximize area", label="Resizing objective", elem_id="postprocess_multicrop_objective")
- threshold = gr.Slider(minimum=0, maximum=1, step=0.01, label="Error threshold", value=0.1, elem_id="postprocess_multicrop_threshold")
+ objective = gr.Radio(["Maximize area", "Minimize error"], value="Maximize area", label="Resizing objective", elem_id=self.elem_id_suffix("postprocess_multicrop_objective"))
+ threshold = gr.Slider(minimum=0, maximum=1, step=0.01, label="Error threshold", value=0.1, elem_id=self.elem_id_suffix("postprocess_multicrop_threshold"))
return {
"enable": enable,
diff --git a/extensions-builtin/postprocessing-for-training/scripts/postprocessing_focal_crop.py b/extensions-builtin/postprocessing-for-training/scripts/postprocessing_focal_crop.py
index cff1dbc5470..db835fb6b61 100644
--- a/extensions-builtin/postprocessing-for-training/scripts/postprocessing_focal_crop.py
+++ b/extensions-builtin/postprocessing-for-training/scripts/postprocessing_focal_crop.py
@@ -11,10 +11,10 @@ class ScriptPostprocessingFocalCrop(scripts_postprocessing.ScriptPostprocessing)
def ui(self):
with ui_components.InputAccordion(False, label="Auto focal point crop") as enable:
- face_weight = gr.Slider(label='Focal point face weight', value=0.9, minimum=0.0, maximum=1.0, step=0.05, elem_id="postprocess_focal_crop_face_weight")
- entropy_weight = gr.Slider(label='Focal point entropy weight', value=0.15, minimum=0.0, maximum=1.0, step=0.05, elem_id="postprocess_focal_crop_entropy_weight")
- edges_weight = gr.Slider(label='Focal point edges weight', value=0.5, minimum=0.0, maximum=1.0, step=0.05, elem_id="postprocess_focal_crop_edges_weight")
- debug = gr.Checkbox(label='Create debug image', elem_id="train_process_focal_crop_debug")
+ face_weight = gr.Slider(label='Focal point face weight', value=0.9, minimum=0.0, maximum=1.0, step=0.05, elem_id=self.elem_id_suffix("postprocess_focal_crop_face_weight"))
+ entropy_weight = gr.Slider(label='Focal point entropy weight', value=0.15, minimum=0.0, maximum=1.0, step=0.05, elem_id=self.elem_id_suffix("postprocess_focal_crop_entropy_weight"))
+ edges_weight = gr.Slider(label='Focal point edges weight', value=0.5, minimum=0.0, maximum=1.0, step=0.05, elem_id=self.elem_id_suffix("postprocess_focal_crop_edges_weight"))
+ debug = gr.Checkbox(label='Create debug image', elem_id=self.elem_id_suffix("train_process_focal_crop_debug"))
return {
"enable": enable,
diff --git a/extensions-builtin/postprocessing-for-training/scripts/postprocessing_split_oversized.py b/extensions-builtin/postprocessing-for-training/scripts/postprocessing_split_oversized.py
index 133e199b838..b569a06e373 100644
--- a/extensions-builtin/postprocessing-for-training/scripts/postprocessing_split_oversized.py
+++ b/extensions-builtin/postprocessing-for-training/scripts/postprocessing_split_oversized.py
@@ -35,8 +35,8 @@ class ScriptPostprocessingSplitOversized(scripts_postprocessing.ScriptPostproces
def ui(self):
with ui_components.InputAccordion(False, label="Split oversized images") as enable:
with gr.Row():
- split_threshold = gr.Slider(label='Threshold', value=0.5, minimum=0.0, maximum=1.0, step=0.05, elem_id="postprocess_split_threshold")
- overlap_ratio = gr.Slider(label='Overlap ratio', value=0.2, minimum=0.0, maximum=0.9, step=0.05, elem_id="postprocess_overlap_ratio")
+ split_threshold = gr.Slider(label='Threshold', value=0.5, minimum=0.0, maximum=1.0, step=0.05, elem_id=self.elem_id_suffix("postprocess_split_threshold"))
+ overlap_ratio = gr.Slider(label='Overlap ratio', value=0.2, minimum=0.0, maximum=0.9, step=0.05, elem_id=self.elem_id_suffix("postprocess_overlap_ratio"))
return {
"enable": enable,
diff --git a/html/footer.html b/html/footer.html
index 69b2372c755..2bedddbc19d 100644
--- a/html/footer.html
+++ b/html/footer.html
@@ -1,7 +1,7 @@
API
•
-
Github
+
GitHub
•
Gradio
•
diff --git a/javascript/contextMenus.js b/javascript/contextMenus.js
index e01fd67e80e..661343d524b 100644
--- a/javascript/contextMenus.js
+++ b/javascript/contextMenus.js
@@ -104,7 +104,7 @@ var contextMenuInit = function() {
e.preventDefault();
}
});
- });
+ }, {passive: false});
});
eventListenerApplied = true;
diff --git a/javascript/extraNetworks.js b/javascript/extraNetworks.js
index c5cced97399..8e8adad0bae 100644
--- a/javascript/extraNetworks.js
+++ b/javascript/extraNetworks.js
@@ -201,7 +201,7 @@ function setupExtraNetworks() {
setupExtraNetworksForTab('img2img');
}
-var re_extranet = /<([^:^>]+:[^:]+):[\d.]+>(.*)/;
+var re_extranet = /<([^:^>]+:[^:]+):[\d.]+>(.*)/s;
var re_extranet_g = /<([^:^>]+:[^:]+):[\d.]+>/g;
var re_extranet_neg = /\(([^:^>]+:[\d.]+)\)/;
diff --git a/javascript/progressbar.js b/javascript/progressbar.js
index 23dea64ceda..c81a41e9a47 100644
--- a/javascript/progressbar.js
+++ b/javascript/progressbar.js
@@ -79,11 +79,12 @@ function requestProgress(id_task, progressbarContainer, gallery, atEnd, onProgre
var wakeLock = null;
var requestWakeLock = async function() {
- if (!opts.prevent_screen_sleep_during_generation || wakeLock) return;
+ if (!opts.prevent_screen_sleep_during_generation || wakeLock !== null) return;
try {
wakeLock = await navigator.wakeLock.request('screen');
} catch (err) {
console.error('Wake Lock is not supported.');
+ wakeLock = false;
}
};
diff --git a/javascript/resizeHandle.js b/javascript/resizeHandle.js
index 4aeb14b41f3..0283f21bfe1 100644
--- a/javascript/resizeHandle.js
+++ b/javascript/resizeHandle.js
@@ -124,7 +124,7 @@
} else {
R.screenX = evt.changedTouches[0].screenX;
}
- });
+ }, {passive: false});
});
resizeHandle.addEventListener('dblclick', onDoubleClick);
diff --git a/modules/call_queue.py b/modules/call_queue.py
index 555c35312dd..b20badcaf25 100644
--- a/modules/call_queue.py
+++ b/modules/call_queue.py
@@ -3,7 +3,7 @@
import html
import time
-from modules import shared, progress, errors, devices, fifo_lock, profiling
+from modules import shared, progress, errors, devices, fifo_lock, profiling, manager
queue_lock = fifo_lock.FIFOLock()
@@ -34,7 +34,7 @@ def f(*args, **kwargs):
progress.start_task(id_task)
try:
- res = func(*args, **kwargs)
+ res = manager.task.run_and_wait_result(func, *args, **kwargs)
progress.record_results(id_task, res)
finally:
progress.finish_task(id_task)
diff --git a/modules/dat_model.py b/modules/dat_model.py
index 495d5f4937d..e256a5a3282 100644
--- a/modules/dat_model.py
+++ b/modules/dat_model.py
@@ -1,7 +1,7 @@
import os
from modules import modelloader, errors
-from modules.shared import cmd_opts, opts
+from modules.shared import cmd_opts, opts, hf_endpoint
from modules.upscaler import Upscaler, UpscalerData
from modules.upscaler_utils import upscale_with_model
@@ -49,7 +49,18 @@ def load_model(self, path):
scaler.local_data_path = modelloader.load_file_from_url(
scaler.data_path,
model_dir=self.model_download_path,
+ hash_prefix=scaler.sha256,
)
+
+ if os.path.getsize(scaler.local_data_path) < 200:
+ # Re-download if the file is too small, probably an LFS pointer
+ scaler.local_data_path = modelloader.load_file_from_url(
+ scaler.data_path,
+ model_dir=self.model_download_path,
+ hash_prefix=scaler.sha256,
+ re_download=True,
+ )
+
if not os.path.exists(scaler.local_data_path):
raise FileNotFoundError(f"DAT data missing: {scaler.local_data_path}")
return scaler
@@ -60,20 +71,23 @@ def get_dat_models(scaler):
return [
UpscalerData(
name="DAT x2",
- path="https://github.com/n0kovo/dat_upscaler_models/raw/main/DAT/DAT_x2.pth",
+ path=f"{hf_endpoint}/w-e-w/DAT/resolve/main/experiments/pretrained_models/DAT/DAT_x2.pth",
scale=2,
upscaler=scaler,
+ sha256='7760aa96e4ee77e29d4f89c3a4486200042e019461fdb8aa286f49aa00b89b51',
),
UpscalerData(
name="DAT x3",
- path="https://github.com/n0kovo/dat_upscaler_models/raw/main/DAT/DAT_x3.pth",
+ path=f"{hf_endpoint}/w-e-w/DAT/resolve/main/experiments/pretrained_models/DAT/DAT_x3.pth",
scale=3,
upscaler=scaler,
+ sha256='581973e02c06f90d4eb90acf743ec9604f56f3c2c6f9e1e2c2b38ded1f80d197',
),
UpscalerData(
name="DAT x4",
- path="https://github.com/n0kovo/dat_upscaler_models/raw/main/DAT/DAT_x4.pth",
+ path=f"{hf_endpoint}/w-e-w/DAT/resolve/main/experiments/pretrained_models/DAT/DAT_x4.pth",
scale=4,
upscaler=scaler,
+ sha256='391a6ce69899dff5ea3214557e9d585608254579217169faf3d4c353caff049e',
),
]
diff --git a/modules/devices.py b/modules/devices.py
index ee679141ad7..5b763ec85ec 100644
--- a/modules/devices.py
+++ b/modules/devices.py
@@ -1,5 +1,6 @@
import sys
import contextlib
+from copy import deepcopy
from functools import lru_cache
import torch
@@ -128,6 +129,26 @@ def enable_tf32():
dtype_inference: torch.dtype = torch.float16
unet_needs_upcast = False
+supported_vae_dtypes = [torch.float16, torch.float32]
+
+
+# prepare available dtypes
+if torch.version.cuda:
+ if torch.cuda.is_bf16_supported() and torch.cuda.get_device_properties(torch.cuda.current_device()).major >= 8:
+ supported_vae_dtypes = [torch.bfloat16] + supported_vae_dtypes
+ if has_xpu():
+ supported_vae_dtypes = [torch.bfloat16] + supported_vae_dtypes
+
+
+def supports_non_blocking():
+ if has_mps() or has_xpu():
+ return False
+
+ if npu_specific.has_npu:
+ return False
+
+ return True
+
def cond_cast_unet(input):
if force_fp16:
@@ -146,17 +167,33 @@ def cond_cast_float(input):
torch.nn.MultiheadAttention,
torch.nn.GroupNorm,
torch.nn.LayerNorm,
+ torch.nn.Embedding,
]
-def manual_cast_forward(target_dtype):
+def manual_cast_forward(target_dtype, target_device=None, copy=False):
+ params = dict()
+ if supports_non_blocking():
+ params['non_blocking'] = True
+
+ supported_cast_dtypes = [torch.float16, torch.float32]
+ if torch.cuda.is_bf16_supported():
+ supported_cast_dtypes += [torch.bfloat16]
+
def forward_wrapper(self, *args, **kwargs):
- if any(
- isinstance(arg, torch.Tensor) and arg.dtype != target_dtype
- for arg in args
- ):
- args = [arg.to(target_dtype) if isinstance(arg, torch.Tensor) else arg for arg in args]
- kwargs = {k: v.to(target_dtype) if isinstance(v, torch.Tensor) else v for k, v in kwargs.items()}
+ if target_device is not None:
+ params['device'] = target_device
+ params['dtype'] = target_dtype
+
+ args = list(args)
+ for j in (i for i, arg in enumerate(args) if isinstance(arg, torch.Tensor) and arg.dtype != target_dtype):
+ if args[j].dtype in supported_cast_dtypes:
+ args[j] = args[j].to(**params)
+ args = tuple(args)
+
+ for key in (k for k, v in kwargs.items() if isinstance(v, torch.Tensor) and v.dtype != target_dtype):
+ if kwargs[key].dtype in supported_cast_dtypes:
+ kwargs[key] = kwargs[key].to(**params)
org_dtype = target_dtype
for param in self.parameters():
@@ -164,38 +201,52 @@ def forward_wrapper(self, *args, **kwargs):
org_dtype = param.dtype
break
- if org_dtype != target_dtype:
- self.to(target_dtype)
- result = self.org_forward(*args, **kwargs)
- if org_dtype != target_dtype:
- self.to(org_dtype)
+ if copy and not isinstance(self, torch.nn.Embedding):
+ copied = deepcopy(self)
+ if org_dtype != target_dtype:
+ copied.to(**params)
+
+ result = copied.org_forward(*args, **kwargs)
+ del copied
+ else:
+ if org_dtype != target_dtype:
+ self.to(**params)
+
+ result = self.org_forward(*args, **kwargs)
+
+ if org_dtype != target_dtype:
+ params['dtype'] = org_dtype
+ self.to(**params)
if target_dtype != dtype_inference:
+ params['dtype'] = dtype_inference
if isinstance(result, tuple):
result = tuple(
- i.to(dtype_inference)
+ i.to(**params)
if isinstance(i, torch.Tensor)
else i
for i in result
)
elif isinstance(result, torch.Tensor):
- result = result.to(dtype_inference)
+ result = result.to(**params)
return result
return forward_wrapper
@contextlib.contextmanager
-def manual_cast(target_dtype):
+def manual_cast(target_dtype, target_device=None, copy=None):
applied = False
+
+
for module_type in patch_module_list:
if hasattr(module_type, "org_forward"):
continue
applied = True
org_forward = module_type.forward
if module_type == torch.nn.MultiheadAttention:
- module_type.forward = manual_cast_forward(torch.float32)
+ module_type.forward = manual_cast_forward(torch.float32, target_device, copy)
else:
- module_type.forward = manual_cast_forward(target_dtype)
+ module_type.forward = manual_cast_forward(target_dtype, target_device, copy)
module_type.org_forward = org_forward
try:
yield None
@@ -207,26 +258,37 @@ def manual_cast(target_dtype):
delattr(module_type, "org_forward")
-def autocast(disable=False):
+def autocast(disable=False, current_dtype=None, target_dtype=None, target_device=None, copy=None):
if disable:
return contextlib.nullcontext()
+ copy = copy if copy is not None else shared.opts.lora_without_backup_weight
+
+ if target_dtype is None:
+ target_dtype = dtype
+
if force_fp16:
# No casting during inference if force_fp16 is enabled.
# All tensor dtype conversion happens before inference.
return contextlib.nullcontext()
- if fp8 and device==cpu:
+ if fp8 and target_device==cpu:
return torch.autocast("cpu", dtype=torch.bfloat16, enabled=True)
if fp8 and dtype_inference == torch.float32:
- return manual_cast(dtype)
+ return manual_cast(target_dtype, target_device, copy=copy)
+
+ if target_dtype != dtype_inference or copy:
+ return manual_cast(target_dtype, target_device, copy=copy)
+
+ if current_dtype is not None and current_dtype != target_dtype:
+ return manual_cast(target_dtype, target_device, copy=copy)
- if dtype == torch.float32 or dtype_inference == torch.float32:
+ if target_dtype == torch.float32 or dtype_inference == torch.float32:
return contextlib.nullcontext()
if has_xpu() or has_mps() or cuda_no_autocast():
- return manual_cast(dtype)
+ return manual_cast(target_dtype, target_device)
return torch.autocast("cuda")
diff --git a/modules/launch_utils.py b/modules/launch_utils.py
index 20c7dc127a7..5c868747e7a 100644
--- a/modules/launch_utils.py
+++ b/modules/launch_utils.py
@@ -463,11 +463,17 @@ def configure_for_tests():
def start():
print(f"Launching {'API server' if '--nowebui' in sys.argv else 'Web UI'} with arguments: {shlex.join(sys.argv[1:])}")
import webui
+
+ from modules import manager
+
if '--nowebui' in sys.argv:
webui.api_only()
else:
webui.webui()
+ manager.task.main_loop()
+ return
+
def dump_sysinfo():
from modules import sysinfo
diff --git a/modules/lowvram.py b/modules/lowvram.py
index 6728c337b64..9914a06c6dd 100644
--- a/modules/lowvram.py
+++ b/modules/lowvram.py
@@ -53,6 +53,7 @@ def send_me_to_gpu(module, _):
if module_in_gpu is not None:
module_in_gpu.to(cpu)
+ devices.torch_gc()
module.to(devices.device)
module_in_gpu = module
diff --git a/modules/manager.py b/modules/manager.py
new file mode 100644
index 00000000000..34c67c6b3cc
--- /dev/null
+++ b/modules/manager.py
@@ -0,0 +1,83 @@
+#
+# based on forge's work from https://github.com/lllyasviel/stable-diffusion-webui-forge/blob/main/modules_forge/main_thread.py
+#
+# Original author comment:
+# This file is the main thread that handles all gradio calls for major t2i or i2i processing.
+# Other gradio calls (like those from extensions) are not influenced.
+# By using one single thread to process all major calls, model moving is significantly faster.
+#
+# 2024/09/28 classified,
+
+import random
+import string
+import threading
+import time
+
+from collections import OrderedDict
+
+
+class Task:
+ def __init__(self, **kwargs):
+ self.__dict__.update(kwargs)
+
+
+class TaskManager:
+ last_exception = None
+ pending_tasks = []
+ finished_tasks = OrderedDict()
+ lock = None
+ running = False
+
+ def __init__(self):
+ self.lock = threading.Lock()
+
+ def work(self, task):
+ try:
+ task.result = task.func(*task.args, **task.kwargs)
+ except Exception as e:
+ task.exception = e
+ self.last_exception = e
+
+
+ def stop(self):
+ self.running = False
+
+
+ def main_loop(self):
+ self.running = True
+ while self.running:
+ time.sleep(0.01)
+ if len(self.pending_tasks) > 0:
+ with self.lock:
+ task = self.pending_tasks.pop(0)
+
+ self.work(task)
+
+ self.finished_tasks[task.task_id] = task
+
+
+ def push_task(self, func, *args, **kwargs):
+ if args and type(args[0]) == str and args[0].startswith("task(") and args[0].endswith(")"):
+ task_id = args[0]
+ else:
+ task_id = ''.join(random.choices(string.ascii_uppercase + string.digits, k=7))
+ task = Task(task_id=task_id, func=func, args=args, kwargs=kwargs, result=None, exception=None)
+ self.pending_tasks.append(task)
+
+ return task.task_id
+
+
+ def run_and_wait_result(self, func, *args, **kwargs):
+ current_id = self.push_task(func, *args, **kwargs)
+
+ while True:
+ time.sleep(0.01)
+ if current_id in self.finished_tasks:
+ finished = self.finished_tasks.pop(current_id)
+ if finished.exception is not None:
+ raise finished.exception
+
+ return finished.result
+
+
+task = TaskManager()
diff --git a/modules/modelloader.py b/modules/modelloader.py
index 36e7415af43..f5a2ff79c30 100644
--- a/modules/modelloader.py
+++ b/modules/modelloader.py
@@ -10,6 +10,7 @@
from modules import shared
from modules.upscaler import Upscaler, UpscalerLanczos, UpscalerNearest, UpscalerNone
+from modules.util import load_file_from_url # noqa, backwards compatibility
if TYPE_CHECKING:
import spandrel
@@ -17,30 +18,6 @@
logger = logging.getLogger(__name__)
-def load_file_from_url(
- url: str,
- *,
- model_dir: str,
- progress: bool = True,
- file_name: str | None = None,
- hash_prefix: str | None = None,
-) -> str:
- """Download a file from `url` into `model_dir`, using the file present if possible.
-
- Returns the path to the downloaded file.
- """
- os.makedirs(model_dir, exist_ok=True)
- if not file_name:
- parts = urlparse(url)
- file_name = os.path.basename(parts.path)
- cached_file = os.path.abspath(os.path.join(model_dir, file_name))
- if not os.path.exists(cached_file):
- print(f'Downloading: "{url}" to {cached_file}\n')
- from torch.hub import download_url_to_file
- download_url_to_file(url, cached_file, progress=progress, hash_prefix=hash_prefix)
- return cached_file
-
-
def load_models(model_path: str, model_url: str = None, command_path: str = None, ext_filter=None, download_name=None, ext_blacklist=None, hash_prefix=None) -> list:
"""
A one-and done loader to try finding the desired models in specified directories.
diff --git a/modules/models/flux/__init__.py b/modules/models/flux/__init__.py
new file mode 100644
index 00000000000..1cc52a00bef
--- /dev/null
+++ b/modules/models/flux/__init__.py
@@ -0,0 +1,5 @@
+from .flux import FLUX1Inferencer
+
+__all__ = [
+ "FLUX1Inferencer",
+]
diff --git a/modules/models/flux/flux.py b/modules/models/flux/flux.py
new file mode 100644
index 00000000000..a7370af253b
--- /dev/null
+++ b/modules/models/flux/flux.py
@@ -0,0 +1,360 @@
+import contextlib
+
+import os
+import safetensors
+import torch
+import math
+
+import k_diffusion
+from transformers import CLIPTokenizer
+
+from modules import shared, devices, modelloader, sd_hijack_clip
+
+from modules.models.sd3.sd3_impls import SDVAE
+from modules.models.sd3.sd3_cond import CLIPL_CONFIG, T5_CONFIG, CLIPL_URL, T5_URL, SafetensorsMapping, Sd3T5
+from modules.models.sd3.other_impls import SDClipModel, T5XXLModel, SDTokenizer, T5XXLTokenizer
+from PIL import Image
+
+from .model import Flux
+
+
+class FluxTokenizer:
+ def __init__(self):
+ clip_tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14")
+ self.clip_l = SDTokenizer(tokenizer=clip_tokenizer)
+ self.t5xxl = T5XXLTokenizer()
+
+ def tokenize_with_weights(self, text:str):
+ out = {}
+ out["l"] = self.clip_l.tokenize_with_weights(text)
+ out["t5xxl"] = self.t5xxl.tokenize_with_weights(text)
+ return out
+
+
+class Flux1ClipL(sd_hijack_clip.TextConditionalModel):
+ def __init__(self, clip_l):
+ super().__init__()
+
+ self.clip_l = clip_l
+
+ self.tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14")
+
+ empty = self.tokenizer('')["input_ids"]
+ self.id_start = empty[0]
+ self.id_end = empty[1]
+ self.id_pad = empty[1]
+
+ self.return_pooled = True
+
+ def tokenize(self, texts):
+ return self.tokenizer(texts, truncation=False, add_special_tokens=False)["input_ids"]
+
+ def encode_with_transformers(self, tokens):
+ l_out, l_pooled = self.clip_l(tokens)
+ l_out = torch.cat([l_out], dim=-1)
+ l_out = torch.nn.functional.pad(l_out, (0, 4096 - l_out.shape[-1]))
+
+ vector_out = torch.cat([l_pooled], dim=-1)
+
+ l_out.pooled = vector_out
+
+ return l_out
+
+ def encode_embedding_init_text(self, init_text, nvpt):
+ return torch.zeros((nvpt, 768+1280), device=devices.device) # XXX
+
+
+
+class FluxCond(torch.nn.Module):
+ def __init__(self, *args, **kwargs):
+ super().__init__(*args, **kwargs)
+
+ self.tokenizer = FluxTokenizer()
+
+ with torch.no_grad():
+ self.clip_l = SDClipModel(layer="hidden", layer_idx=-2, device="cpu", dtype=devices.dtype_inference, layer_norm_hidden_state=False, return_projected_pooled=False, textmodel_json_config=CLIPL_CONFIG)
+
+ if shared.opts.flux_enable_t5:
+ self.t5xxl = T5XXLModel(T5_CONFIG, device="cpu", dtype=devices.dtype_inference)
+ else:
+ self.t5xxl = None
+
+ self.model_l = Flux1ClipL(self.clip_l)
+ self.model_t5 = Sd3T5(self.t5xxl)
+
+ def forward(self, prompts: list[str]):
+ with devices.without_autocast():
+ l_out, vector_out = self.model_l(prompts)
+ t5_out = self.model_t5(prompts, token_count=l_out.shape[1])
+ lt_out = torch.cat([l_out, t5_out], dim=-2)
+
+ return {
+ 'crossattn': lt_out,
+ 'vector': vector_out,
+ }
+
+ def before_load_weights(self, state_dict):
+ clip_path = os.path.join(shared.models_path, "CLIP")
+
+ if 'text_encoders.clip_l.transformer.text_model.embeddings.position_embedding.weight' not in state_dict:
+ clip_l_file = modelloader.load_file_from_url(CLIPL_URL, model_dir=clip_path, file_name="clip_l.safetensors")
+ with safetensors.safe_open(clip_l_file, framework="pt") as file:
+ self.clip_l.transformer.load_state_dict(SafetensorsMapping(file), strict=False)
+
+ if self.t5xxl and 'text_encoders.t5xxl.transformer.encoder.block.0.layer.0.SelfAttention.k.weight' not in state_dict:
+ t5_file = modelloader.load_file_from_url(T5_URL, model_dir=clip_path, file_name="t5xxl_fp8_e4m3fn.safetensors")
+ with safetensors.safe_open(t5_file, framework="pt") as file:
+ self.t5xxl.transformer.load_state_dict(SafetensorsMapping(file), strict=False)
+
+ def encode_embedding_init_text(self, init_text, nvpt):
+ return self.model_l.encode_embedding_init_text(init_text, nvpt)
+
+ def tokenize(self, texts):
+ return self.model_l.tokenize(texts)
+
+ def medvram_modules(self):
+ return [self.clip_l, self.t5xxl]
+
+ def get_token_count(self, text):
+ _, token_count = self.model_l.process_texts([text])
+
+ return token_count
+
+ def get_target_prompt_token_count(self, token_count):
+ return self.model_l.get_target_prompt_token_count(token_count)
+
+def flux_time_shift(mu: float, sigma: float, t):
+ return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma)
+
+class ModelSamplingFlux(torch.nn.Module):
+ def __init__(self, shift=1.15):
+ super().__init__()
+
+ self.set_parameters(shift=shift)
+
+ def set_parameters(self, shift=1.15, timesteps=10000):
+ self.shift = shift
+ ts = self.sigma((torch.arange(1, timesteps + 1, 1) / timesteps))
+ self.register_buffer('sigmas', ts)
+
+ @property
+ def sigma_min(self):
+ return self.sigmas[0]
+
+ @property
+ def sigma_max(self):
+ return self.sigmas[-1]
+
+ def timestep(self, sigma):
+ return sigma
+
+ def sigma(self, timestep):
+ return flux_time_shift(self.shift, 1.0, timestep)
+
+ def percent_to_sigma(self, percent):
+ if percent <= 0.0:
+ return 1.0
+ if percent >= 1.0:
+ return 0.0
+ return 1.0 - percent
+
+ def calculate_denoised(self, sigma, model_output, model_input):
+ sigma = sigma.view(sigma.shape[:1] + (1,) * (model_output.ndim - 1))
+ return model_input - model_output * sigma
+
+
+class BaseModel(torch.nn.Module):
+ """Wrapper around the core FLUX model"""
+ def __init__(self, shift=1.15, device=None, dtype=torch.float16, state_dict=None, prefix="", **kwargs):
+ super().__init__()
+
+ self.diffusion_model = Flux(device=device, dtype=dtype, **kwargs)
+ self.model_sampling = ModelSamplingFlux(shift=shift)
+ self.depth = kwargs['depth']
+ self.depth_single_block = kwargs['depth_single_blocks']
+
+ def apply_model(self, x, sigma, c_crossattn=None, y=None):
+ dtype = self.get_dtype()
+ timestep = self.model_sampling.timestep(sigma).float()
+ guidance = torch.FloatTensor([3.5]).to(device=devices.device, dtype=torch.float32)
+ model_output = self.diffusion_model(x.to(dtype), timestep, context=c_crossattn.to(dtype), y=y.to(dtype), guidance=guidance).to(x.dtype)
+ return self.model_sampling.calculate_denoised(sigma, model_output, x)
+
+ def forward(self, *args, **kwargs):
+ return self.apply_model(*args, **kwargs)
+
+ def get_dtype(self):
+ return self.diffusion_model.dtype
+
+
+class FLUX1LatentFormat:
+ """Latents are slightly shifted from center - this class must be called after VAE Decode to correct for the shift"""
+ def __init__(self, scale_factor=0.3611, shift_factor=0.1159):
+ self.scale_factor = scale_factor
+ self.shift_factor = shift_factor
+
+ def process_in(self, latent):
+ return (latent - self.shift_factor) * self.scale_factor
+
+ def process_out(self, latent):
+ return (latent / self.scale_factor) + self.shift_factor
+
+ def decode_latent_to_preview(self, x0):
+ """Quick RGB approximate preview of sd3 latents"""
+ factors = torch.tensor([
+ [-0.0404, 0.0159, 0.0609], [ 0.0043, 0.0298, 0.0850],
+ [ 0.0328, -0.0749, -0.0503], [-0.0245, 0.0085, 0.0549],
+ [ 0.0966, 0.0894, 0.0530], [ 0.0035, 0.0399, 0.0123],
+ [ 0.0583, 0.1184, 0.1262], [-0.0191, -0.0206, -0.0306],
+ [-0.0324, 0.0055, 0.1001], [ 0.0955, 0.0659, -0.0545],
+ [-0.0504, 0.0231, -0.0013], [ 0.0500, -0.0008, -0.0088],
+ [ 0.0982, 0.0941, 0.0976], [-0.1233, -0.0280, -0.0897],
+ [-0.0005, -0.0530, -0.0020], [-0.1273, -0.0932, -0.0680],
+ ], device="cpu")
+ latent_image = x0[0].permute(1, 2, 0).cpu() @ factors
+
+ latents_ubyte = (((latent_image + 1) / 2)
+ .clamp(0, 1) # change scale from -1..1 to 0..1
+ .mul(0xFF) # to 0..255
+ .byte()).cpu()
+
+ return Image.fromarray(latents_ubyte.numpy())
+
+
+class FLUX1Denoiser(k_diffusion.external.DiscreteSchedule):
+ def __init__(self, inner_model, sigmas):
+ super().__init__(sigmas, quantize=shared.opts.enable_quantization)
+ self.inner_model = inner_model
+
+ def forward(self, input, sigma, **kwargs):
+ return self.inner_model.apply_model(input, sigma, **kwargs)
+
+
+class FLUX1Inferencer(torch.nn.Module):
+ def __init__(self, state_dict, use_ema=False):
+ super().__init__()
+
+ params = dict(
+ image_model="flux",
+ in_channels=16,
+ vec_in_dim=768,
+ context_in_dim=4096,
+ hidden_size=3072,
+ mlp_ratio=4.0,
+ num_heads=24,
+ depth=19,
+ depth_single_blocks=38,
+ axes_dim=[16, 56, 56],
+ theta=10000,
+ qkv_bias=True,
+ guidance_embed=True,
+ )
+
+ # detect model_prefix
+ diffusion_model_prefix = "model.diffusion_model."
+ if "model.diffusion_model.double_blocks.0.img_attn.norm.key_norm.scale" in state_dict:
+ diffusion_model_prefix = "model.diffusion_model."
+ elif "double_blocks.0.img_attn.norm.key_norm.scale" in state_dict:
+ diffusion_model_prefix = ""
+
+ shift=1.15
+ # check guidance_in to detect Flux schnell
+ if f"{diffusion_model_prefix}guidance_in.in_layer.weight" not in state_dict:
+ print("Flux schnell detected")
+ params.update(dict(guidance_embed=False,))
+ shift=1.0
+
+ with torch.no_grad():
+ self.model = BaseModel(shift=shift, state_dict=state_dict, prefix=diffusion_model_prefix, device="cpu", dtype=devices.dtype_inference, **params)
+ self.first_stage_model = SDVAE(device="cpu", dtype=devices.dtype_vae)
+ self.first_stage_model.dtype = devices.dtype_vae
+ self.vae = self.first_stage_model # real vae
+
+ self.alphas_cumprod = 1 / (self.model.model_sampling.sigmas ** 2 + 1)
+
+ self.text_encoders = FluxCond()
+ self.cond_stage_key = 'txt'
+
+ self.parameterization = "eps"
+ self.model.conditioning_key = "crossattn"
+
+ self.latent_format = FLUX1LatentFormat()
+ self.latent_channels = 16
+
+ @property
+ def cond_stage_model(self):
+ return self.text_encoders
+
+ def before_load_weights(self, state_dict):
+ self.cond_stage_model.before_load_weights(state_dict)
+
+ def ema_scope(self):
+ return contextlib.nullcontext()
+
+ def get_learned_conditioning(self, batch: list[str]):
+ return self.cond_stage_model(batch)
+
+ def apply_model(self, x, t, cond):
+ return self.model(x, t, c_crossattn=cond['crossattn'], y=cond['vector'])
+
+ def decode_first_stage(self, latent):
+ latent = self.latent_format.process_out(latent)
+ x = self.first_stage_model.decode(latent)
+ if x.dtype == torch.float16:
+ x = torch.nan_to_num(x, nan=0.0, posinf=65504, neginf=-65504)
+ return x
+
+ def encode_first_stage(self, image):
+ latent = self.first_stage_model.encode(image)
+ return self.latent_format.process_in(latent)
+
+ def get_first_stage_encoding(self, x):
+ return x
+
+ def create_denoiser(self):
+ return FLUX1Denoiser(self, self.model.model_sampling.sigmas)
+
+ def medvram_fields(self):
+ return [
+ (self, 'first_stage_model'),
+ (self, 'text_encoders'),
+ (self, 'model'),
+ ]
+
+ def add_noise_to_latent(self, x, noise, amount):
+ return x * (1 - amount) + noise * amount
+
+ def fix_dimensions(self, width, height):
+ return width // 16 * 16, height // 16 * 16
+
+ def diffusers_weight_mapping(self):
+ # https://github.com/kohya-ss/sd-scripts/blob/a61cf73a5cb5209c3f4d1a3688dd276a4dfd1ecb/networks/convert_flux_lora.py
+ # please see also https://github.com/huggingface/diffusers/blob/main/src/diffusers/loaders/lora_conversion_utils.py
+ for i in range(self.model.depth):
+ yield f"transformer.transformer_blocks.{i}.attn.add_k_proj", f"diffusion_model_double_blocks_{i}_txt_attn_qkv_k_proj"
+ yield f"transformer.transformer_blocks.{i}.attn.add_q_proj", f"diffusion_model_double_blocks_{i}_txt_attn_qkv_q_proj"
+ yield f"transformer.transformer_blocks.{i}.attn.add_v_proj", f"diffusion_model_double_blocks_{i}_txt_attn_qkv_v_proj"
+
+ yield f"transformer.transformer_blocks.{i}.attn.to_add_out", f"diffusion_model_double_blocks_{i}_txt_attn_proj"
+
+ yield f"transformer.transformer_blocks.{i}.attn.to_k", f"diffusion_model_double_blocks_{i}_img_attn_qkv_k_proj"
+ yield f"transformer.transformer_blocks.{i}.attn.to_q", f"diffusion_model_double_blocks_{i}_img_attn_qkv_q_proj"
+ yield f"transformer.transformer_blocks.{i}.attn.to_v", f"diffusion_model_double_blocks_{i}_img_attn_qkv_v_proj"
+
+ yield f"transformer.transformer_blocks.{i}.attn.to_out.0", f"diffusion_model_double_blocks_{i}_img_attn_proj"
+
+ yield f"transformer.transformer_blocks.{i}.ff.net.0.proj", f"diffusion_model_double_blocks_{i}_img_mlp_0"
+ yield f"transformer.transformer_blocks.{i}.ff.net.2", f"diffusion_model_double_blocks_{i}_img_mlp_2"
+ yield f"transformer.transformer_blocks.{i}.ff_context.net.0.proj", f"diffusion_model_double_blocks_{i}_txt_mlp_0"
+ yield f"transformer.transformer_blocks.{i}.ff_context.net.2", f"diffusion_model_double_blocks_{i}_txt_mlp_2"
+ yield f"transformer.transformer_blocks.{i}.norm1.linear", f"diffusion_model_double_blocks_{i}_img_mod_lin"
+ yield f"transformer.transformer_blocks.{i}.norm1_context.linear", f"diffusion_model_double_blocks_{i}_txt_mod_lin"
+
+ for i in range(self.model.depth_single_block):
+ yield f"transformer.single_transformer_blocks.{i}.attn.to_q", f"diffusion_model_single_blocks_{i}_linear1_q_proj"
+ yield f"transformer.single_transformer_blocks.{i}.attn.to_k", f"diffusion_model_single_blocks_{i}_linear1_k_proj"
+ yield f"transformer.single_transformer_blocks.{i}.attn.to_v", f"diffusion_model_single_blocks_{i}_linear1_v_proj"
+ yield f"transformer.single_transformer_blocks.{i}.proj_mlp", f"diffusion_model_single_blocks_{i}_linear1_mlp_proj"
+
+ yield f"transformer.single_transformer_blocks.{i}.proj_out", f"diffusion_model_single_blocks_{i}_linear2"
+ yield f"transformer.single_transformer_blocks.{i}.norm.linear", f"diffusion_model_single_blocks_{i}_modulation_lin"
diff --git a/modules/models/flux/math.py b/modules/models/flux/math.py
new file mode 100644
index 00000000000..4ad99818e2c
--- /dev/null
+++ b/modules/models/flux/math.py
@@ -0,0 +1,30 @@
+import torch
+from einops import rearrange
+from torch import Tensor
+
+
+def attention(q: Tensor, k: Tensor, v: Tensor, pe: Tensor) -> Tensor:
+ q, k = apply_rope(q, k, pe)
+
+ x = torch.nn.functional.scaled_dot_product_attention(q, k, v)
+ x = rearrange(x, "B H L D -> B L (H D)")
+
+ return x
+
+
+def rope(pos: Tensor, dim: int, theta: int) -> Tensor:
+ assert dim % 2 == 0
+ scale = torch.arange(0, dim, 2, dtype=torch.float64, device=pos.device) / dim
+ omega = 1.0 / (theta**scale)
+ out = torch.einsum("...n,d->...nd", pos, omega)
+ out = torch.stack([torch.cos(out), -torch.sin(out), torch.sin(out), torch.cos(out)], dim=-1)
+ out = rearrange(out, "b n d (i j) -> b n d i j", i=2, j=2)
+ return out.float()
+
+
+def apply_rope(xq: Tensor, xk: Tensor, freqs_cis: Tensor):
+ xq_ = xq.float().reshape(*xq.shape[:-1], -1, 1, 2)
+ xk_ = xk.float().reshape(*xk.shape[:-1], -1, 1, 2)
+ xq_out = freqs_cis[..., 0] * xq_[..., 0] + freqs_cis[..., 1] * xq_[..., 1]
+ xk_out = freqs_cis[..., 0] * xk_[..., 0] + freqs_cis[..., 1] * xk_[..., 1]
+ return xq_out.reshape(*xq.shape).type_as(xq), xk_out.reshape(*xk.shape).type_as(xk)
diff --git a/modules/models/flux/model.py b/modules/models/flux/model.py
new file mode 100644
index 00000000000..c873978277d
--- /dev/null
+++ b/modules/models/flux/model.py
@@ -0,0 +1,147 @@
+# original code from https://github.com/black-forest-labs/flux
+#
+from dataclasses import dataclass
+
+import torch
+
+from einops import rearrange, repeat
+from torch import Tensor, nn
+
+from .modules.layers import (DoubleStreamBlock, EmbedND, LastLayer,
+ MLPEmbedder, SingleStreamBlock,
+ timestep_embedding)
+
+
+@dataclass
+class FluxParams:
+ in_channels: int
+ vec_in_dim: int
+ context_in_dim: int
+ hidden_size: int
+ mlp_ratio: float
+ num_heads: int
+ depth: int
+ depth_single_blocks: int
+ axes_dim: list
+ theta: int
+ qkv_bias: bool
+ guidance_embed: bool
+
+
+class Flux(nn.Module):
+ """
+ Transformer model for flow matching on sequences.
+ """
+
+ def __init__(self, image_model=None, final_layer=True, dtype=None, device=None, **kwargs):
+ super().__init__()
+
+ self.dtype = dtype
+ params = FluxParams(**kwargs)
+ self.params = params
+ self.in_channels = params.in_channels * 2 * 2
+ self.out_channels = self.in_channels
+ if params.hidden_size % params.num_heads != 0:
+ raise ValueError(
+ f"Hidden size {params.hidden_size} must be divisible by num_heads {params.num_heads}"
+ )
+ pe_dim = params.hidden_size // params.num_heads
+ if sum(params.axes_dim) != pe_dim:
+ raise ValueError(f"Got {params.axes_dim} but expected positional dim {pe_dim}")
+ self.hidden_size = params.hidden_size
+ self.num_heads = params.num_heads
+ self.pe_embedder = EmbedND(dim=pe_dim, theta=params.theta, axes_dim=params.axes_dim)
+ self.img_in = nn.Linear(self.in_channels, self.hidden_size, bias=True, dtype=dtype, device=device)
+ self.time_in = MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size, dtype=dtype, device=device)
+ self.vector_in = MLPEmbedder(params.vec_in_dim, self.hidden_size, dtype=dtype, device=device)
+ self.guidance_in = (
+ MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size, dtype=dtype, device=device) if params.guidance_embed else nn.Identity()
+ )
+ self.txt_in = nn.Linear(params.context_in_dim, self.hidden_size, dtype=dtype, device=device)
+
+ self.double_blocks = nn.ModuleList(
+ [
+ DoubleStreamBlock(
+ self.hidden_size,
+ self.num_heads,
+ mlp_ratio=params.mlp_ratio,
+ qkv_bias=params.qkv_bias,
+ dtype=dtype, device=device,
+ )
+ for _ in range(params.depth)
+ ]
+ )
+
+ self.single_blocks = nn.ModuleList(
+ [
+ SingleStreamBlock(self.hidden_size, self.num_heads, mlp_ratio=params.mlp_ratio, dtype=dtype, device=device)
+ for _ in range(params.depth_single_blocks)
+ ]
+ )
+
+ if final_layer:
+ self.final_layer = LastLayer(self.hidden_size, 1, self.out_channels, dtype=dtype, device=device)
+
+ def forward_orig(
+ self,
+ img: Tensor,
+ img_ids: Tensor,
+ txt: Tensor,
+ txt_ids: Tensor,
+ timesteps: Tensor,
+ y: Tensor,
+ guidance: Tensor = None,
+ ) -> Tensor:
+ if img.ndim != 3 or txt.ndim != 3:
+ raise ValueError("Input img and txt tensors must have 3 dimensions.")
+
+ # running on sequences img
+ img = self.img_in(img)
+ vec = self.time_in(timestep_embedding(timesteps, 256).to(img.dtype))
+ if self.params.guidance_embed:
+ if guidance is None:
+ raise ValueError("Didn't get guidance strength for guidance distilled model.")
+ vec = vec + self.guidance_in(timestep_embedding(guidance, 256).to(img.dtype))
+ vec = vec + self.vector_in(y)
+ txt = self.txt_in(txt)
+
+ ids = torch.cat((txt_ids, img_ids), dim=1)
+ pe = self.pe_embedder(ids)
+
+ for block in self.double_blocks:
+ img, txt = block(img=img, txt=txt, vec=vec, pe=pe)
+
+ img = torch.cat((txt, img), 1)
+ for block in self.single_blocks:
+ img = block(img, vec=vec, pe=pe)
+ img = img[:, txt.shape[1] :, ...]
+
+ if self.final_layer:
+ img = self.final_layer(img, vec) # (N, T, patch_size ** 2 * out_channels)
+ return img
+
+ def forward(self, x, timestep, context, y, guidance, **kwargs):
+ # from comfy/ldm/common_dit.py
+ def pad_to_patch_size(img, patch_size=(2, 2), padding_mode="circular"):
+ if padding_mode == "circular" and torch.jit.is_tracing() or torch.jit.is_scripting():
+ padding_mode = "reflect"
+ pad_h = (patch_size[0] - img.shape[-2] % patch_size[0]) % patch_size[0]
+ pad_w = (patch_size[1] - img.shape[-1] % patch_size[1]) % patch_size[1]
+ return torch.nn.functional.pad(img, (0, pad_w, 0, pad_h), mode=padding_mode)
+
+ bs, c, h, w = x.shape
+ patch_size = 2
+ x = pad_to_patch_size(x, (patch_size, patch_size))
+
+ img = rearrange(x, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=patch_size, pw=patch_size)
+
+ h_len = ((h + (patch_size // 2)) // patch_size)
+ w_len = ((w + (patch_size // 2)) // patch_size)
+ img_ids = torch.zeros((h_len, w_len, 3), device=x.device, dtype=x.dtype)
+ img_ids[..., 1] = img_ids[..., 1] + torch.linspace(0, h_len - 1, steps=h_len, device=x.device, dtype=x.dtype)[:, None]
+ img_ids[..., 2] = img_ids[..., 2] + torch.linspace(0, w_len - 1, steps=w_len, device=x.device, dtype=x.dtype)[None, :]
+ img_ids = repeat(img_ids, "h w c -> b (h w) c", b=bs)
+
+ txt_ids = torch.zeros((bs, context.shape[1], 3), device=x.device, dtype=x.dtype)
+ out = self.forward_orig(img, img_ids, context, txt_ids, timestep, y, guidance)
+ return rearrange(out, "b (h w) (c ph pw) -> b c (h ph) (w pw)", h=h_len, w=w_len, ph=2, pw=2)[:,:,:h,:w]
diff --git a/modules/models/flux/modules/layers.py b/modules/models/flux/modules/layers.py
new file mode 100644
index 00000000000..2202f5dcb35
--- /dev/null
+++ b/modules/models/flux/modules/layers.py
@@ -0,0 +1,265 @@
+import math
+from dataclasses import dataclass
+
+import torch
+from einops import rearrange
+from torch import Tensor, nn
+
+from ..math import attention, rope
+
+
+class EmbedND(nn.Module):
+ def __init__(self, dim: int, theta: int, axes_dim: list):
+ super().__init__()
+ self.dim = dim
+ self.theta = theta
+ self.axes_dim = axes_dim
+
+ def forward(self, ids: Tensor) -> Tensor:
+ n_axes = ids.shape[-1]
+ emb = torch.cat(
+ [rope(ids[..., i], self.axes_dim[i], self.theta) for i in range(n_axes)],
+ dim=-3,
+ )
+
+ return emb.unsqueeze(1)
+
+
+def timestep_embedding(t: Tensor, dim, max_period=10000, time_factor: float = 1000.0):
+ """
+ Create sinusoidal timestep embeddings.
+ :param t: a 1-D Tensor of N indices, one per batch element.
+ These may be fractional.
+ :param dim: the dimension of the output.
+ :param max_period: controls the minimum frequency of the embeddings.
+ :return: an (N, D) Tensor of positional embeddings.
+ """
+ t = time_factor * t
+ half = dim // 2
+ freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32, device=t.device) / half)
+
+ args = t[:, None].float() * freqs[None]
+ embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+ if dim % 2:
+ embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+ if torch.is_floating_point(t):
+ embedding = embedding.to(t)
+ return embedding
+
+
+class MLPEmbedder(nn.Module):
+ def __init__(self, in_dim: int, hidden_dim: int, dtype=None, device=None):
+ super().__init__()
+ self.in_layer = nn.Linear(in_dim, hidden_dim, bias=True, dtype=dtype, device=device)
+ self.silu = nn.SiLU()
+ self.out_layer = nn.Linear(hidden_dim, hidden_dim, bias=True, dtype=dtype, device=device)
+
+ def forward(self, x: Tensor) -> Tensor:
+ return self.out_layer(self.silu(self.in_layer(x)))
+
+
+class RMSNorm(torch.nn.Module):
+ def __init__(self, dim: int, dtype=None, device=None):
+ super().__init__()
+ self.scale = nn.Parameter(torch.empty((dim), dtype=dtype, device=device))
+
+ def forward(self, x: Tensor):
+ x_dtype = x.dtype
+ x = x.float()
+ rrms = torch.rsqrt(torch.mean(x**2, dim=-1, keepdim=True) + 1e-6)
+ return (x * rrms).to(dtype=x_dtype) * self.scale
+
+
+class QKNorm(torch.nn.Module):
+ def __init__(self, dim: int, dtype=None, device=None):
+ super().__init__()
+ self.query_norm = RMSNorm(dim, dtype=dtype, device=device)
+ self.key_norm = RMSNorm(dim, dtype=dtype, device=device)
+
+ def forward(self, q: Tensor, k: Tensor, v: Tensor) -> tuple:
+ q = self.query_norm(q)
+ k = self.key_norm(k)
+ return q.to(v), k.to(v)
+
+
+class QkvLinear(torch.nn.Linear):
+ pass
+
+
+class SelfAttention(nn.Module):
+ def __init__(self, dim: int, num_heads: int = 8, qkv_bias: bool = False, dtype=None, device=None):
+ super().__init__()
+ self.num_heads = num_heads
+ head_dim = dim // num_heads
+
+ self.qkv = QkvLinear(dim, dim * 3, bias=qkv_bias, dtype=dtype, device=device)
+ self.norm = QKNorm(head_dim, dtype=dtype, device=device)
+ self.proj = nn.Linear(dim, dim, dtype=dtype, device=device)
+
+ def forward(self, x: Tensor, pe: Tensor) -> Tensor:
+ qkv = self.qkv(x)
+ q, k, v = rearrange(qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads)
+ q, k = self.norm(q, k, v)
+ x = attention(q, k, v, pe=pe)
+ x = self.proj(x)
+ return x
+
+
+@dataclass
+class ModulationOut:
+ shift: Tensor
+ scale: Tensor
+ gate: Tensor
+
+
+class Modulation(nn.Module):
+ def __init__(self, dim: int, double: bool, dtype=None, device=None):
+ super().__init__()
+ self.is_double = double
+ self.multiplier = 6 if double else 3
+ self.lin = nn.Linear(dim, self.multiplier * dim, bias=True, dtype=dtype, device=device)
+
+ def forward(self, vec: Tensor) -> tuple:
+ out = self.lin(nn.functional.silu(vec))[:, None, :].chunk(self.multiplier, dim=-1)
+
+ return (
+ ModulationOut(*out[:3]),
+ ModulationOut(*out[3:]) if self.is_double else None,
+ )
+
+
+class DoubleStreamBlock(nn.Module):
+ def __init__(self, hidden_size: int, num_heads: int, mlp_ratio: float, qkv_bias: bool = False, dtype=None, device=None):
+ super().__init__()
+
+ mlp_hidden_dim = int(hidden_size * mlp_ratio)
+ self.num_heads = num_heads
+ self.hidden_size = hidden_size
+ self.img_mod = Modulation(hidden_size, double=True, dtype=dtype, device=device)
+ self.img_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, dtype=dtype, device=device)
+ self.img_attn = SelfAttention(dim=hidden_size, num_heads=num_heads, qkv_bias=qkv_bias, dtype=dtype, device=device)
+
+ self.img_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, dtype=dtype, device=device)
+ self.img_mlp = nn.Sequential(
+ nn.Linear(hidden_size, mlp_hidden_dim, bias=True, dtype=dtype, device=device),
+ nn.GELU(approximate="tanh"),
+ nn.Linear(mlp_hidden_dim, hidden_size, bias=True, dtype=dtype, device=device),
+ )
+
+ self.txt_mod = Modulation(hidden_size, double=True, dtype=dtype, device=device)
+ self.txt_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, dtype=dtype, device=device)
+ self.txt_attn = SelfAttention(dim=hidden_size, num_heads=num_heads, qkv_bias=qkv_bias, dtype=dtype, device=device)
+
+ self.txt_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, dtype=dtype, device=device)
+ self.txt_mlp = nn.Sequential(
+ nn.Linear(hidden_size, mlp_hidden_dim, bias=True, dtype=dtype, device=device),
+ nn.GELU(approximate="tanh"),
+ nn.Linear(mlp_hidden_dim, hidden_size, bias=True, dtype=dtype, device=device),
+ )
+
+ def forward(self, img: Tensor, txt: Tensor, vec: Tensor, pe: Tensor):
+ img_mod1, img_mod2 = self.img_mod(vec)
+ txt_mod1, txt_mod2 = self.txt_mod(vec)
+
+ # prepare image for attention
+ img_modulated = self.img_norm1(img)
+ img_modulated = (1 + img_mod1.scale) * img_modulated + img_mod1.shift
+ img_qkv = self.img_attn.qkv(img_modulated)
+ img_q, img_k, img_v = rearrange(img_qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads)
+ img_q, img_k = self.img_attn.norm(img_q, img_k, img_v)
+
+ # prepare txt for attention
+ txt_modulated = self.txt_norm1(txt)
+ txt_modulated = (1 + txt_mod1.scale) * txt_modulated + txt_mod1.shift
+ txt_qkv = self.txt_attn.qkv(txt_modulated)
+ txt_q, txt_k, txt_v = rearrange(txt_qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads)
+ txt_q, txt_k = self.txt_attn.norm(txt_q, txt_k, txt_v)
+
+ # run actual attention
+ q = torch.cat((txt_q, img_q), dim=2)
+ k = torch.cat((txt_k, img_k), dim=2)
+ v = torch.cat((txt_v, img_v), dim=2)
+
+ attn = attention(q, k, v, pe=pe)
+ txt_attn, img_attn = attn[:, : txt.shape[1]], attn[:, txt.shape[1] :]
+
+ # calculate the img bloks
+ img = img + img_mod1.gate * self.img_attn.proj(img_attn)
+ img = img + img_mod2.gate * self.img_mlp((1 + img_mod2.scale) * self.img_norm2(img) + img_mod2.shift)
+
+ # calculate the txt bloks
+ txt = txt + txt_mod1.gate * self.txt_attn.proj(txt_attn)
+ txt = txt + txt_mod2.gate * self.txt_mlp((1 + txt_mod2.scale) * self.txt_norm2(txt) + txt_mod2.shift)
+
+ if txt.dtype == torch.float16:
+ txt = torch.nan_to_num(txt, nan=0.0, posinf=65504, neginf=-65504)
+
+
+ return img, txt
+
+
+class SingleStreamBlock(nn.Module):
+ """
+ A DiT block with parallel linear layers as described in
+ https://arxiv.org/abs/2302.05442 and adapted modulation interface.
+ """
+
+ def __init__(
+ self,
+ hidden_size: int,
+ num_heads: int,
+ mlp_ratio: float = 4.0,
+ qk_scale: float = None,
+ dtype=None,
+ device=None,
+ ):
+ super().__init__()
+ self.hidden_dim = hidden_size
+ self.num_heads = num_heads
+ head_dim = hidden_size // num_heads
+ self.scale = qk_scale or head_dim**-0.5
+
+ self.mlp_hidden_dim = int(hidden_size * mlp_ratio)
+ # qkv and mlp_in
+ self.linear1 = QkvLinear(hidden_size, hidden_size * 3 + self.mlp_hidden_dim, dtype=dtype, device=device)
+ # proj and mlp_out
+ self.linear2 = nn.Linear(hidden_size + self.mlp_hidden_dim, hidden_size, dtype=dtype, device=device)
+
+ self.norm = QKNorm(head_dim, dtype=dtype, device=device)
+
+ self.hidden_size = hidden_size
+ self.pre_norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, dtype=dtype, device=device)
+
+ self.mlp_act = nn.GELU(approximate="tanh")
+ self.modulation = Modulation(hidden_size, double=False, dtype=dtype, device=device)
+
+ def forward(self, x: Tensor, vec: Tensor, pe: Tensor) -> Tensor:
+ mod, _ = self.modulation(vec)
+ x_mod = (1 + mod.scale) * self.pre_norm(x) + mod.shift
+ qkv, mlp = torch.split(self.linear1(x_mod), [3 * self.hidden_size, self.mlp_hidden_dim], dim=-1)
+
+ q, k, v = rearrange(qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads)
+ q, k = self.norm(q, k, v)
+
+ # compute attention
+ attn = attention(q, k, v, pe=pe)
+ # compute activation in mlp stream, cat again and run second linear layer
+ output = self.linear2(torch.cat((attn, self.mlp_act(mlp)), 2))
+ x += mod.gate * output
+ if x.dtype == torch.float16:
+ x = torch.nan_to_num(x, nan=0.0, posinf=65504, neginf=-65504)
+ return x
+
+
+class LastLayer(nn.Module):
+ def __init__(self, hidden_size: int, patch_size: int, out_channels: int, dtype=None, device=None):
+ super().__init__()
+ self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, dtype=dtype, device=device)
+ self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels, bias=True, dtype=dtype, device=device)
+ self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size, bias=True, dtype=dtype, device=device))
+
+ def forward(self, x: Tensor, vec: Tensor) -> Tensor:
+ shift, scale = self.adaLN_modulation(vec).chunk(2, dim=1)
+ x = (1 + scale[:, None, :]) * self.norm_final(x) + shift[:, None, :]
+ x = self.linear(x)
+ return x
diff --git a/modules/models/flux/util.py b/modules/models/flux/util.py
new file mode 100644
index 00000000000..9303eb7cffe
--- /dev/null
+++ b/modules/models/flux/util.py
@@ -0,0 +1,201 @@
+import os
+from dataclasses import dataclass
+
+import torch
+from einops import rearrange
+from huggingface_hub import hf_hub_download
+from imwatermark import WatermarkEncoder
+from safetensors.torch import load_file as load_sft
+
+from .model import Flux, FluxParams
+from .modules.autoencoder import AutoEncoder, AutoEncoderParams
+from .modules.conditioner import HFEmbedder
+
+
+@dataclass
+class ModelSpec:
+ params: FluxParams
+ ae_params: AutoEncoderParams
+ ckpt_path: str | None
+ ae_path: str | None
+ repo_id: str | None
+ repo_flow: str | None
+ repo_ae: str | None
+
+
+configs = {
+ "flux-dev": ModelSpec(
+ repo_id="black-forest-labs/FLUX.1-dev",
+ repo_flow="flux1-dev.safetensors",
+ repo_ae="ae.safetensors",
+ ckpt_path=os.getenv("FLUX_DEV"),
+ params=FluxParams(
+ in_channels=64,
+ vec_in_dim=768,
+ context_in_dim=4096,
+ hidden_size=3072,
+ mlp_ratio=4.0,
+ num_heads=24,
+ depth=19,
+ depth_single_blocks=38,
+ axes_dim=[16, 56, 56],
+ theta=10_000,
+ qkv_bias=True,
+ guidance_embed=True,
+ ),
+ ae_path=os.getenv("AE"),
+ ae_params=AutoEncoderParams(
+ resolution=256,
+ in_channels=3,
+ ch=128,
+ out_ch=3,
+ ch_mult=[1, 2, 4, 4],
+ num_res_blocks=2,
+ z_channels=16,
+ scale_factor=0.3611,
+ shift_factor=0.1159,
+ ),
+ ),
+ "flux-schnell": ModelSpec(
+ repo_id="black-forest-labs/FLUX.1-schnell",
+ repo_flow="flux1-schnell.safetensors",
+ repo_ae="ae.safetensors",
+ ckpt_path=os.getenv("FLUX_SCHNELL"),
+ params=FluxParams(
+ in_channels=64,
+ vec_in_dim=768,
+ context_in_dim=4096,
+ hidden_size=3072,
+ mlp_ratio=4.0,
+ num_heads=24,
+ depth=19,
+ depth_single_blocks=38,
+ axes_dim=[16, 56, 56],
+ theta=10_000,
+ qkv_bias=True,
+ guidance_embed=False,
+ ),
+ ae_path=os.getenv("AE"),
+ ae_params=AutoEncoderParams(
+ resolution=256,
+ in_channels=3,
+ ch=128,
+ out_ch=3,
+ ch_mult=[1, 2, 4, 4],
+ num_res_blocks=2,
+ z_channels=16,
+ scale_factor=0.3611,
+ shift_factor=0.1159,
+ ),
+ ),
+}
+
+
+def print_load_warning(missing: list[str], unexpected: list[str]) -> None:
+ if len(missing) > 0 and len(unexpected) > 0:
+ print(f"Got {len(missing)} missing keys:\n\t" + "\n\t".join(missing))
+ print("\n" + "-" * 79 + "\n")
+ print(f"Got {len(unexpected)} unexpected keys:\n\t" + "\n\t".join(unexpected))
+ elif len(missing) > 0:
+ print(f"Got {len(missing)} missing keys:\n\t" + "\n\t".join(missing))
+ elif len(unexpected) > 0:
+ print(f"Got {len(unexpected)} unexpected keys:\n\t" + "\n\t".join(unexpected))
+
+
+def load_flow_model(name: str, device: str | torch.device = "cuda", hf_download: bool = True):
+ # Loading Flux
+ print("Init model")
+ ckpt_path = configs[name].ckpt_path
+ if (
+ ckpt_path is None
+ and configs[name].repo_id is not None
+ and configs[name].repo_flow is not None
+ and hf_download
+ ):
+ ckpt_path = hf_hub_download(configs[name].repo_id, configs[name].repo_flow)
+
+ with torch.device("meta" if ckpt_path is not None else device):
+ model = Flux(configs[name].params).to(torch.bfloat16)
+
+ if ckpt_path is not None:
+ print("Loading checkpoint")
+ # load_sft doesn't support torch.device
+ sd = load_sft(ckpt_path, device=str(device))
+ missing, unexpected = model.load_state_dict(sd, strict=False, assign=True)
+ print_load_warning(missing, unexpected)
+ return model
+
+
+def load_t5(device: str | torch.device = "cuda", max_length: int = 512) -> HFEmbedder:
+ # max length 64, 128, 256 and 512 should work (if your sequence is short enough)
+ return HFEmbedder("google/t5-v1_1-xxl", max_length=max_length, torch_dtype=torch.bfloat16).to(device)
+
+
+def load_clip(device: str | torch.device = "cuda") -> HFEmbedder:
+ return HFEmbedder("openai/clip-vit-large-patch14", max_length=77, torch_dtype=torch.bfloat16).to(device)
+
+
+def load_ae(name: str, device: str | torch.device = "cuda", hf_download: bool = True) -> AutoEncoder:
+ ckpt_path = configs[name].ae_path
+ if (
+ ckpt_path is None
+ and configs[name].repo_id is not None
+ and configs[name].repo_ae is not None
+ and hf_download
+ ):
+ ckpt_path = hf_hub_download(configs[name].repo_id, configs[name].repo_ae)
+
+ # Loading the autoencoder
+ print("Init AE")
+ with torch.device("meta" if ckpt_path is not None else device):
+ ae = AutoEncoder(configs[name].ae_params)
+
+ if ckpt_path is not None:
+ sd = load_sft(ckpt_path, device=str(device))
+ missing, unexpected = ae.load_state_dict(sd, strict=False, assign=True)
+ print_load_warning(missing, unexpected)
+ return ae
+
+
+class WatermarkEmbedder:
+ def __init__(self, watermark):
+ self.watermark = watermark
+ self.num_bits = len(WATERMARK_BITS)
+ self.encoder = WatermarkEncoder()
+ self.encoder.set_watermark("bits", self.watermark)
+
+ def __call__(self, image: torch.Tensor) -> torch.Tensor:
+ """
+ Adds a predefined watermark to the input image
+
+ Args:
+ image: ([N,] B, RGB, H, W) in range [-1, 1]
+
+ Returns:
+ same as input but watermarked
+ """
+ image = 0.5 * image + 0.5
+ squeeze = len(image.shape) == 4
+ if squeeze:
+ image = image[None, ...]
+ n = image.shape[0]
+ image_np = rearrange((255 * image).detach().cpu(), "n b c h w -> (n b) h w c").numpy()[:, :, :, ::-1]
+ # torch (b, c, h, w) in [0, 1] -> numpy (b, h, w, c) [0, 255]
+ # watermarking libary expects input as cv2 BGR format
+ for k in range(image_np.shape[0]):
+ image_np[k] = self.encoder.encode(image_np[k], "dwtDct")
+ image = torch.from_numpy(rearrange(image_np[:, :, :, ::-1], "(n b) h w c -> n b c h w", n=n)).to(
+ image.device
+ )
+ image = torch.clamp(image / 255, min=0.0, max=1.0)
+ if squeeze:
+ image = image[0]
+ image = 2 * image - 1
+ return image
+
+
+# A fixed 48-bit message that was chosen at random
+WATERMARK_MESSAGE = 0b001010101111111010000111100111001111010100101110
+# bin(x)[2:] gives bits of x as str, use int to convert them to 0/1
+WATERMARK_BITS = [int(bit) for bit in bin(WATERMARK_MESSAGE)[2:]]
+embed_watermark = WatermarkEmbedder(WATERMARK_BITS)
diff --git a/modules/models/sd3/other_impls.py b/modules/models/sd3/other_impls.py
index 78c1dc68758..4524fa019b7 100644
--- a/modules/models/sd3/other_impls.py
+++ b/modules/models/sd3/other_impls.py
@@ -24,6 +24,11 @@ class AutocastLinear(nn.Linear):
def forward(self, x):
return torch.nn.functional.linear(x, self.weight.to(x.dtype), self.bias.to(x.dtype) if self.bias is not None else None)
+class AutocastLayerNorm(nn.LayerNorm):
+ def forward(self, x):
+ return torch.nn.functional.layer_norm(
+ x, self.normalized_shape, self.weight.to(x.dtype), self.bias.to(x.dtype) if self.bias is not None else None, self.eps)
+
def attention(q, k, v, heads, mask=None):
"""Convenience wrapper around a basic attention operation"""
@@ -41,9 +46,9 @@ def __init__(self, in_features, hidden_features=None, out_features=None, act_lay
out_features = out_features or in_features
hidden_features = hidden_features or in_features
- self.fc1 = nn.Linear(in_features, hidden_features, bias=bias, dtype=dtype, device=device)
+ self.fc1 = AutocastLinear(in_features, hidden_features, bias=bias, dtype=dtype, device=device)
self.act = act_layer
- self.fc2 = nn.Linear(hidden_features, out_features, bias=bias, dtype=dtype, device=device)
+ self.fc2 = AutocastLinear(hidden_features, out_features, bias=bias, dtype=dtype, device=device)
def forward(self, x):
x = self.fc1(x)
@@ -61,10 +66,10 @@ class CLIPAttention(torch.nn.Module):
def __init__(self, embed_dim, heads, dtype, device):
super().__init__()
self.heads = heads
- self.q_proj = nn.Linear(embed_dim, embed_dim, bias=True, dtype=dtype, device=device)
- self.k_proj = nn.Linear(embed_dim, embed_dim, bias=True, dtype=dtype, device=device)
- self.v_proj = nn.Linear(embed_dim, embed_dim, bias=True, dtype=dtype, device=device)
- self.out_proj = nn.Linear(embed_dim, embed_dim, bias=True, dtype=dtype, device=device)
+ self.q_proj = AutocastLinear(embed_dim, embed_dim, bias=True, dtype=dtype, device=device)
+ self.k_proj = AutocastLinear(embed_dim, embed_dim, bias=True, dtype=dtype, device=device)
+ self.v_proj = AutocastLinear(embed_dim, embed_dim, bias=True, dtype=dtype, device=device)
+ self.out_proj = AutocastLinear(embed_dim, embed_dim, bias=True, dtype=dtype, device=device)
def forward(self, x, mask=None):
q = self.q_proj(x)
@@ -82,9 +87,11 @@ def forward(self, x, mask=None):
class CLIPLayer(torch.nn.Module):
def __init__(self, embed_dim, heads, intermediate_size, intermediate_activation, dtype, device):
super().__init__()
- self.layer_norm1 = nn.LayerNorm(embed_dim, dtype=dtype, device=device)
+ #self.layer_norm1 = nn.LayerNorm(embed_dim, dtype=dtype, device=device)
+ self.layer_norm1 = AutocastLayerNorm(embed_dim, dtype=dtype, device=device)
self.self_attn = CLIPAttention(embed_dim, heads, dtype, device)
- self.layer_norm2 = nn.LayerNorm(embed_dim, dtype=dtype, device=device)
+ self.layer_norm2 = AutocastLayerNorm(embed_dim, dtype=dtype, device=device)
+ #self.layer_norm2 = nn.LayerNorm(embed_dim, dtype=dtype, device=device)
#self.mlp = CLIPMLP(embed_dim, intermediate_size, intermediate_activation, dtype, device)
self.mlp = Mlp(embed_dim, intermediate_size, embed_dim, act_layer=ACTIVATIONS[intermediate_activation], dtype=dtype, device=device)
@@ -131,7 +138,7 @@ def __init__(self, config_dict, dtype, device):
super().__init__()
self.embeddings = CLIPEmbeddings(embed_dim, dtype=torch.float32, device=device, textual_inversion_key=config_dict.get('textual_inversion_key', 'clip_l'))
self.encoder = CLIPEncoder(num_layers, embed_dim, heads, intermediate_size, intermediate_activation, dtype, device)
- self.final_layer_norm = nn.LayerNorm(embed_dim, dtype=dtype, device=device)
+ self.final_layer_norm = AutocastLayerNorm(embed_dim, dtype=dtype, device=device)
def forward(self, input_tokens, intermediate_output=None, final_layer_norm_intermediate=True):
x = self.embeddings(input_tokens)
@@ -150,7 +157,7 @@ def __init__(self, config_dict, dtype, device):
self.num_layers = config_dict["num_hidden_layers"]
self.text_model = CLIPTextModel_(config_dict, dtype, device)
embed_dim = config_dict["hidden_size"]
- self.text_projection = nn.Linear(embed_dim, embed_dim, bias=False, dtype=dtype, device=device)
+ self.text_projection = AutocastLinear(embed_dim, embed_dim, bias=False, dtype=dtype, device=device)
self.text_projection.weight.copy_(torch.eye(embed_dim))
self.dtype = dtype
@@ -370,7 +377,7 @@ def __init__(self, model_dim, inner_dim, num_heads, relative_attention_bias, dty
if relative_attention_bias:
self.relative_attention_num_buckets = 32
self.relative_attention_max_distance = 128
- self.relative_attention_bias = torch.nn.Embedding(self.relative_attention_num_buckets, self.num_heads, device=device)
+ self.relative_attention_bias = torch.nn.Embedding(self.relative_attention_num_buckets, self.num_heads, device=device, dtype=torch.float32)
@staticmethod
def _relative_position_bucket(relative_position, bidirectional=True, num_buckets=32, max_distance=128):
@@ -442,7 +449,7 @@ def forward(self, x, past_bias=None):
else:
mask = None
- out = attention(q, k * ((k.shape[-1] / self.num_heads) ** 0.5), v, self.num_heads, mask.to(x.dtype) if mask is not None else None)
+ out = attention(q, k * ((k.shape[-1] / self.num_heads) ** 0.5), v, self.num_heads, mask.to(q.dtype) if mask is not None else None)
return self.o(out), past_bias
@@ -475,19 +482,21 @@ def forward(self, x, past_bias=None):
class T5Stack(torch.nn.Module):
def __init__(self, num_layers, model_dim, inner_dim, ff_dim, num_heads, vocab_size, dtype, device):
super().__init__()
- self.embed_tokens = torch.nn.Embedding(vocab_size, model_dim, device=device)
+ #self.embed_tokens = torch.nn.Embedding(vocab_size, model_dim, device=device, dtype=torch.float32)
self.block = torch.nn.ModuleList([T5Block(model_dim, inner_dim, ff_dim, num_heads, relative_attention_bias=(i == 0), dtype=dtype, device=device) for i in range(num_layers)])
self.final_layer_norm = T5LayerNorm(model_dim, dtype=dtype, device=device)
- def forward(self, input_ids, intermediate_output=None, final_layer_norm_intermediate=True):
+ def forward(self, x, intermediate_output=None, final_layer_norm_intermediate=True):
intermediate = None
- x = self.embed_tokens(input_ids).to(torch.float32) # needs float32 or else T5 returns all zeroes
+ #x = self.embed_tokens(input_ids).to(torch.float32) # needs float32 or else T5 returns all zeroes
+ # some T5XXL do not embed_token. use shared token instead like comfy
past_bias = None
for i, layer in enumerate(self.block):
x, past_bias = layer(x, past_bias)
if i == intermediate_output:
intermediate = x.clone()
x = self.final_layer_norm(x)
+ x = torch.nan_to_num(x)
if intermediate is not None and final_layer_norm_intermediate:
intermediate = self.final_layer_norm(intermediate)
return x, intermediate
@@ -498,13 +507,18 @@ def __init__(self, config_dict, dtype, device):
super().__init__()
self.num_layers = config_dict["num_layers"]
self.encoder = T5Stack(self.num_layers, config_dict["d_model"], config_dict["d_model"], config_dict["d_ff"], config_dict["num_heads"], config_dict["vocab_size"], dtype, device)
+ self.shared = torch.nn.Embedding(config_dict["vocab_size"], config_dict["d_model"], device=device, dtype=torch.float32)
self.dtype = dtype
def get_input_embeddings(self):
- return self.encoder.embed_tokens
+ #return self.encoder.embed_tokens
+ return self.shared
def set_input_embeddings(self, embeddings):
- self.encoder.embed_tokens = embeddings
+ #self.encoder.embed_tokens = embeddings
+ self.shared = embeddings
- def forward(self, *args, **kwargs):
- return self.encoder(*args, **kwargs)
+ def forward(self, input_ids, *args, **kwargs):
+ x = self.shared(input_ids).float()
+ x = torch.nan_to_num(x)
+ return self.encoder(x, *args, **kwargs)
diff --git a/modules/models/sd3/sd3_cond.py b/modules/models/sd3/sd3_cond.py
index 325c512d594..fc0232325ff 100644
--- a/modules/models/sd3/sd3_cond.py
+++ b/modules/models/sd3/sd3_cond.py
@@ -24,7 +24,7 @@ def __getitem__(self, key):
return self.file.get_tensor(key)
-CLIPL_URL = "https://huggingface.co/AUTOMATIC/stable-diffusion-3-medium-text-encoders/resolve/main/clip_l.safetensors"
+CLIPL_URL = f"{shared.hf_endpoint}/AUTOMATIC/stable-diffusion-3-medium-text-encoders/resolve/main/clip_l.safetensors"
CLIPL_CONFIG = {
"hidden_act": "quick_gelu",
"hidden_size": 768,
@@ -33,7 +33,7 @@ def __getitem__(self, key):
"num_hidden_layers": 12,
}
-CLIPG_URL = "https://huggingface.co/AUTOMATIC/stable-diffusion-3-medium-text-encoders/resolve/main/clip_g.safetensors"
+CLIPG_URL = f"{shared.hf_endpoint}/AUTOMATIC/stable-diffusion-3-medium-text-encoders/resolve/main/clip_g.safetensors"
CLIPG_CONFIG = {
"hidden_act": "gelu",
"hidden_size": 1280,
@@ -43,7 +43,7 @@ def __getitem__(self, key):
"textual_inversion_key": "clip_g",
}
-T5_URL = "https://huggingface.co/AUTOMATIC/stable-diffusion-3-medium-text-encoders/resolve/main/t5xxl_fp16.safetensors"
+T5_URL = f"{shared.hf_endpoint}/AUTOMATIC/stable-diffusion-3-medium-text-encoders/resolve/main/t5xxl_fp16_e4m3fn.safetensors"
T5_CONFIG = {
"d_ff": 10240,
"d_model": 4096,
@@ -140,7 +140,7 @@ def tokenize_line(self, line, *, target_token_count=None):
return tokens, multipliers
def forward(self, texts, *, token_count):
- if not self.t5xxl or not shared.opts.sd3_enable_t5:
+ if not self.t5xxl:
return torch.zeros((len(texts), token_count, 4096), device=devices.device, dtype=devices.dtype)
tokens_batch = []
@@ -164,11 +164,11 @@ def __init__(self, *args, **kwargs):
self.tokenizer = SD3Tokenizer()
with torch.no_grad():
- self.clip_g = SDXLClipG(CLIPG_CONFIG, device="cpu", dtype=devices.dtype)
- self.clip_l = SDClipModel(layer="hidden", layer_idx=-2, device="cpu", dtype=devices.dtype, layer_norm_hidden_state=False, return_projected_pooled=False, textmodel_json_config=CLIPL_CONFIG)
+ self.clip_g = SDXLClipG(CLIPG_CONFIG, device="cpu", dtype=devices.dtype_inference)
+ self.clip_l = SDClipModel(layer="hidden", layer_idx=-2, device="cpu", dtype=devices.dtype_inference, layer_norm_hidden_state=False, return_projected_pooled=False, textmodel_json_config=CLIPL_CONFIG)
if shared.opts.sd3_enable_t5:
- self.t5xxl = T5XXLModel(T5_CONFIG, device="cpu", dtype=devices.dtype)
+ self.t5xxl = T5XXLModel(T5_CONFIG, device="cpu", dtype=devices.dtype_inference)
else:
self.t5xxl = None
@@ -199,8 +199,8 @@ def before_load_weights(self, state_dict):
with safetensors.safe_open(clip_l_file, framework="pt") as file:
self.clip_l.transformer.load_state_dict(SafetensorsMapping(file), strict=False)
- if self.t5xxl and 'text_encoders.t5xxl.transformer.encoder.embed_tokens.weight' not in state_dict:
- t5_file = modelloader.load_file_from_url(T5_URL, model_dir=clip_path, file_name="t5xxl_fp16.safetensors")
+ if self.t5xxl and 'text_encoders.t5xxl.transformer.encoder.block.0.layer.0.SelfAttention.k.weight' not in state_dict:
+ t5_file = modelloader.load_file_from_url(T5_URL, model_dir=clip_path, file_name="t5xxl_fp8_e4m3fn.safetensors")
with safetensors.safe_open(t5_file, framework="pt") as file:
self.t5xxl.transformer.load_state_dict(SafetensorsMapping(file), strict=False)
diff --git a/modules/processing.py b/modules/processing.py
index 7535b56e18c..ed140983d9e 100644
--- a/modules/processing.py
+++ b/modules/processing.py
@@ -484,7 +484,7 @@ def get_conds_with_caching(self, function, required_prompts, steps, caches, extr
cache = caches[0]
- with devices.autocast():
+ with devices.autocast(target_dtype=devices.dtype_inference):
cache[1] = function(shared.sd_model, required_prompts, steps, hires_steps, shared.opts.use_old_scheduling)
cache[0] = cached_params
@@ -984,7 +984,7 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
sd_models.apply_alpha_schedule_override(p.sd_model, p)
- with devices.without_autocast() if devices.unet_needs_upcast else devices.autocast():
+ with devices.without_autocast() if devices.unet_needs_upcast else devices.autocast(target_dtype=devices.dtype_inference, current_dtype=devices.dtype_unet):
samples_ddim = p.sample(conditioning=p.c, unconditional_conditioning=p.uc, seeds=p.seeds, subseeds=p.subseeds, subseed_strength=p.subseed_strength, prompts=p.prompts)
if p.scripts is not None:
@@ -1147,6 +1147,11 @@ def infotext(index=0, use_main_prompt=False):
if p.scripts is not None:
p.scripts.postprocess(p, res)
+
+ if lowvram.is_enabled(shared.sd_model):
+ # for interrupted case
+ lowvram.send_everything_to_cpu()
+
return res
@@ -1259,7 +1264,10 @@ def init(self, all_prompts, all_seeds, all_subseeds):
if self.hr_checkpoint_info is None:
raise Exception(f'Could not find checkpoint with name {self.hr_checkpoint_name}')
- self.extra_generation_params["Hires checkpoint"] = self.hr_checkpoint_info.short_title
+ if shared.sd_model.sd_checkpoint_info == self.hr_checkpoint_info:
+ self.hr_checkpoint_info = None
+ else:
+ self.extra_generation_params["Hires checkpoint"] = self.hr_checkpoint_info.short_title
if self.hr_sampler_name is not None and self.hr_sampler_name != self.sampler_name:
self.extra_generation_params["Hires sampler"] = self.hr_sampler_name
@@ -1436,7 +1444,7 @@ def save_intermediate(image, index):
with devices.autocast():
extra_networks.activate(self, self.hr_extra_network_data)
- with devices.autocast():
+ with devices.autocast(target_dtype=devices.dtype_inference):
self.calculate_hr_conds()
sd_models.apply_token_merging(self.sd_model, self.get_token_merging_ratio(for_hr=True))
diff --git a/modules/scripts_auto_postprocessing.py b/modules/scripts_auto_postprocessing.py
index d63078de50e..cf981625034 100644
--- a/modules/scripts_auto_postprocessing.py
+++ b/modules/scripts_auto_postprocessing.py
@@ -13,6 +13,7 @@ def show(self, is_img2img):
return scripts.AlwaysVisible
def ui(self, is_img2img):
+ self.script.tab_name = '_img2img' if is_img2img else '_txt2img'
self.postprocessing_controls = self.script.ui()
return self.postprocessing_controls.values()
@@ -33,7 +34,7 @@ def create_auto_preprocessing_script_data():
for name in shared.opts.postprocessing_enable_in_main_ui:
script = next(iter([x for x in scripts.postprocessing_scripts_data if x.script_class.name == name]), None)
- if script is None:
+ if script is None or script.script_class.extra_only:
continue
constructor = lambda s=script: ScriptPostprocessingForMainUI(s.script_class())
diff --git a/modules/scripts_postprocessing.py b/modules/scripts_postprocessing.py
index 4b3b7afda1c..70270c658d9 100644
--- a/modules/scripts_postprocessing.py
+++ b/modules/scripts_postprocessing.py
@@ -1,3 +1,4 @@
+import re
import dataclasses
import os
import gradio as gr
@@ -59,6 +60,10 @@ class ScriptPostprocessing:
args_from = None
args_to = None
+ # define if the script should be used only in extras or main UI
+ extra_only = None
+ main_ui_only = None
+
order = 1000
"""scripts will be ordred by this value in postprocessing UI"""
@@ -97,6 +102,31 @@ def process_firstpass(self, pp: PostprocessedImage, **args):
def image_changed(self):
pass
+ tab_name = '' # used by ScriptPostprocessingForMainUI
+ replace_pattern = re.compile(r'\s')
+ rm_pattern = re.compile(r'[^a-z_0-9]')
+
+ def elem_id(self, item_id):
+ """
+ Helper function to generate id for a HTML element
+ constructs final id out of script name and user-supplied item_id
+ 'script_extras_{self.name.lower()}_{item_id}'
+ {tab_name} will append to the end of the id if set
+ tab_name will be set to '_img2img' or '_txt2img' if use by ScriptPostprocessingForMainUI
+
+ Extensions should use this function to generate element IDs
+ """
+ return self.elem_id_suffix(f'extras_{self.name.lower()}_{item_id}')
+
+ def elem_id_suffix(self, base_id):
+ """
+ Append tab_name to the base_id
+
+ Extensions that already have specific there element IDs and wish to keep their IDs the same when possible should use this function
+ """
+ base_id = self.rm_pattern.sub('', self.replace_pattern.sub('_', base_id))
+ return f'{base_id}{self.tab_name}'
+
def wrap_call(func, filename, funcname, *args, default=None, **kwargs):
try:
@@ -119,10 +149,6 @@ def initialize_scripts(self, scripts_data):
for script_data in scripts_data:
script: ScriptPostprocessing = script_data.script_class()
script.filename = script_data.path
-
- if script.name == "Simple Upscale":
- continue
-
self.scripts.append(script)
def create_script_ui(self, script, inputs):
@@ -152,7 +178,7 @@ def script_score(name):
return len(self.scripts)
- filtered_scripts = [script for script in self.scripts if script.name not in scripts_filter_out]
+ filtered_scripts = [script for script in self.scripts if script.name not in scripts_filter_out and not script.main_ui_only]
script_scores = {script.name: (script_score(script.name), script.order, script.name, original_index) for original_index, script in enumerate(filtered_scripts)}
return sorted(filtered_scripts, key=lambda x: script_scores[x.name])
diff --git a/modules/sd_disable_initialization.py b/modules/sd_disable_initialization.py
index 273a7edd8b4..0261db08f18 100644
--- a/modules/sd_disable_initialization.py
+++ b/modules/sd_disable_initialization.py
@@ -76,7 +76,7 @@ def transformers_modeling_utils_load_pretrained_model(*args, **kwargs):
def transformers_utils_hub_get_file_from_cache(original, url, *args, **kwargs):
# this file is always 404, prevent making request
- if url == 'https://huggingface.co/openai/clip-vit-large-patch14/resolve/main/added_tokens.json' or url == 'openai/clip-vit-large-patch14' and args[0] == 'added_tokens.json':
+ if url == f'{shared.hf_endpoint}/openai/clip-vit-large-patch14/resolve/main/added_tokens.json' or url == 'openai/clip-vit-large-patch14' and args[0] == 'added_tokens.json':
return None
try:
@@ -160,7 +160,7 @@ def __init__(self, state_dict, device, weight_dtype_conversion=None):
self.state_dict = state_dict
self.device = device
self.weight_dtype_conversion = weight_dtype_conversion or {}
- self.default_dtype = self.weight_dtype_conversion.get('')
+ self.default_dtype = self.weight_dtype_conversion.get('', None)
def get_weight_dtype(self, key):
key_first_term, _ = key.split('.', 1)
@@ -176,6 +176,11 @@ def __enter__(self):
def load_from_state_dict(original, module, state_dict, prefix, *args, **kwargs):
used_param_keys = []
+ if isinstance(module, (torch.nn.Linear, torch.nn.Conv2d, torch.nn.GroupNorm,)):
+ # HACK add assign=True to local_metadata for some cases
+ args[0]['assign_to_params_buffers'] = True
+
+
for name, param in module._parameters.items():
if param is None:
continue
@@ -183,12 +188,16 @@ def load_from_state_dict(original, module, state_dict, prefix, *args, **kwargs):
key = prefix + name
sd_param = sd.pop(key, None)
if sd_param is not None:
- state_dict[key] = sd_param.to(dtype=self.get_weight_dtype(key))
+ dtype = self.get_weight_dtype(key)
+ if dtype is None:
+ state_dict[key] = sd_param
+ else:
+ state_dict[key] = sd_param.to(dtype=dtype)
used_param_keys.append(key)
if param.is_meta:
dtype = sd_param.dtype if sd_param is not None else param.dtype
- module._parameters[name] = torch.nn.parameter.Parameter(torch.zeros_like(param, device=device, dtype=dtype), requires_grad=param.requires_grad)
+ module._parameters[name] = torch.nn.parameter.Parameter(torch.empty_like(param, device=device, dtype=dtype), requires_grad=param.requires_grad)
for name in module._buffers:
key = prefix + name
diff --git a/modules/sd_hijack_unet.py b/modules/sd_hijack_unet.py
index b4f03b138a4..842030be849 100644
--- a/modules/sd_hijack_unet.py
+++ b/modules/sd_hijack_unet.py
@@ -42,12 +42,12 @@ def apply_model(orig_func, self, x_noisy, t, cond, **kwargs):
if isinstance(cond, dict):
for y in cond.keys():
if isinstance(cond[y], list):
- cond[y] = [x.to(devices.dtype_unet) if isinstance(x, torch.Tensor) else x for x in cond[y]]
+ cond[y] = [x.to(devices.dtype_inference) if isinstance(x, torch.Tensor) else x for x in cond[y]]
else:
- cond[y] = cond[y].to(devices.dtype_unet) if isinstance(cond[y], torch.Tensor) else cond[y]
+ cond[y] = cond[y].to(devices.dtype_inference) if isinstance(cond[y], torch.Tensor) else cond[y]
with devices.autocast():
- result = orig_func(self, x_noisy.to(devices.dtype_unet), t.to(devices.dtype_unet), cond, **kwargs)
+ result = orig_func(self, x_noisy.to(devices.dtype_inference), t.to(devices.dtype_inference), cond, **kwargs)
if devices.unet_needs_upcast:
return result.float()
else:
@@ -107,7 +107,7 @@ def __init__(self, *args, **kwargs):
torch.nn.GELU.__init__(self, *args, **kwargs)
def forward(self, x):
if devices.unet_needs_upcast:
- return torch.nn.GELU.forward(self.float(), x.float()).to(devices.dtype_unet)
+ return torch.nn.GELU.forward(self.float(), x.float()).to(devices.dtype_inference)
else:
return torch.nn.GELU.forward(self, x)
@@ -125,11 +125,11 @@ def hijack_ddpm_edit():
CondFunc('ldm.models.diffusion.ddpm.LatentDiffusion.apply_model', apply_model, unet_needs_upcast)
CondFunc('ldm.modules.diffusionmodules.openaimodel.timestep_embedding', timestep_embedding)
CondFunc('ldm.modules.attention.SpatialTransformer.forward', spatial_transformer_forward)
-CondFunc('ldm.modules.diffusionmodules.openaimodel.timestep_embedding', lambda orig_func, timesteps, *args, **kwargs: orig_func(timesteps, *args, **kwargs).to(torch.float32 if timesteps.dtype == torch.int64 else devices.dtype_unet), unet_needs_upcast)
+CondFunc('ldm.modules.diffusionmodules.openaimodel.timestep_embedding', lambda orig_func, timesteps, *args, **kwargs: orig_func(timesteps, *args, **kwargs).to(torch.float32 if timesteps.dtype == torch.int64 else devices.dtype_inference), unet_needs_upcast)
if version.parse(torch.__version__) <= version.parse("1.13.2") or torch.cuda.is_available():
CondFunc('ldm.modules.diffusionmodules.util.GroupNorm32.forward', lambda orig_func, self, *args, **kwargs: orig_func(self.float(), *args, **kwargs), unet_needs_upcast)
- CondFunc('ldm.modules.attention.GEGLU.forward', lambda orig_func, self, x: orig_func(self.float(), x.float()).to(devices.dtype_unet), unet_needs_upcast)
+ CondFunc('ldm.modules.attention.GEGLU.forward', lambda orig_func, self, x: orig_func(self.float(), x.float()).to(devices.dtype_inference), unet_needs_upcast)
CondFunc('open_clip.transformer.ResidualAttentionBlock.__init__', lambda orig_func, *args, **kwargs: kwargs.update({'act_layer': GELUHijack}) and False or orig_func(*args, **kwargs), lambda _, *args, **kwargs: kwargs.get('act_layer') is None or kwargs['act_layer'] == torch.nn.GELU)
first_stage_cond = lambda _, self, *args, **kwargs: devices.unet_needs_upcast and self.model.diffusion_model.dtype == torch.float16
@@ -146,7 +146,7 @@ def timestep_embedding_cast_result(orig_func, timesteps, *args, **kwargs):
if devices.unet_needs_upcast and timesteps.dtype == torch.int64:
dtype = torch.float32
else:
- dtype = devices.dtype_unet
+ dtype = devices.dtype_inference
return orig_func(timesteps, *args, **kwargs).to(dtype=dtype)
diff --git a/modules/sd_models.py b/modules/sd_models.py
index 1c7d370e97b..da9b7c6f4ff 100644
--- a/modules/sd_models.py
+++ b/modules/sd_models.py
@@ -33,6 +33,7 @@ class ModelType(enum.Enum):
SDXL = 3
SSD = 4
SD3 = 5
+ FLUX1 = 6
def replace_key(d, key, new_key, value):
@@ -281,6 +282,30 @@ def get_state_dict_from_checkpoint(pl_sd):
return pl_sd
+def fix_unet_prefix(state_dict):
+ known_prefixes = ("model.diffusion_model.", "first_stage_model.", "cond_stage_model.", "conditioner", "vae.", "text_encoders.")
+
+ for k in state_dict.keys():
+ found = [prefix for prefix in known_prefixes if k.startswith(prefix)]
+ if len(found) > 0:
+ return state_dict
+
+ # no known prefix found.
+ # in this case, this is a unet only state_dict
+ known_keys = (
+ "input_blocks.0.0.weight", # SD1.5, SD2, SDXL
+ "joint_blocks.0.context_block.adaLN_modulation.1.weight", # SD3
+ "double_blocks.0.img_attn.proj.weight", # FLUX
+ )
+
+ if any(key in state_dict for key in known_keys):
+ state_dict = {f"model.diffusion_model.{k}": v for k, v in state_dict.items()}
+ print("Fixed state_dict keys...")
+ return state_dict
+
+ return state_dict
+
+
def read_metadata_from_safetensors(filename):
import json
@@ -342,6 +367,7 @@ def get_checkpoint_state_dict(checkpoint_info: CheckpointInfo, timer):
print(f"Loading weights [{sd_model_hash}] from {checkpoint_info.filename}")
res = read_state_dict(checkpoint_info.filename)
+ res = fix_unet_prefix(res)
timer.record("load weights from disk")
return res
@@ -369,7 +395,7 @@ def check_fp8(model):
enable_fp8 = False
elif shared.opts.fp8_storage == "Enable":
enable_fp8 = True
- elif getattr(model, "is_sdxl", False) and shared.opts.fp8_storage == "Enable for SDXL":
+ elif any(getattr(model, attr, False) for attr in ("is_sdxl", "is_flux1")) and shared.opts.fp8_storage == "Enable for SDXL":
enable_fp8 = True
else:
enable_fp8 = False
@@ -382,10 +408,14 @@ def set_model_type(model, state_dict):
model.is_sdxl = False
model.is_ssd = False
model.is_sd3 = False
+ model.is_flux1 = False
if "model.diffusion_model.x_embedder.proj.weight" in state_dict:
model.is_sd3 = True
model.model_type = ModelType.SD3
+ elif "model.diffusion_model.double_blocks.0.img_attn.norm.key_norm.scale" in state_dict:
+ model.is_flux1 = True
+ model.model_type = ModelType.FLUX1
elif hasattr(model, 'conditioner'):
model.is_sdxl = True
@@ -407,6 +437,82 @@ def set_model_fields(model):
model.latent_channels = 4
+def get_state_dict_dtype(state_dict):
+ # detect dtypes of state_dict
+ state_dict_dtype = {}
+
+ known_prefixes = ("model.diffusion_model.", "first_stage_model.", "cond_stage_model.", "conditioner", "vae.", "text_encoders.")
+
+ for k in state_dict.keys():
+ found = [prefix for prefix in known_prefixes if k.startswith(prefix)]
+ if len(found) > 0:
+ prefix = found[0]
+ dtype = state_dict[k].dtype
+ dtypes = state_dict_dtype.get(prefix, {})
+ if dtype in dtypes:
+ dtypes[dtype] += 1
+ else:
+ dtypes[dtype] = 1
+ state_dict_dtype[prefix] = dtypes
+
+ for prefix in state_dict_dtype:
+ dtypes = state_dict_dtype[prefix]
+ # sort by count
+ state_dict_dtype[prefix] = dict(sorted(dtypes.items(), key=lambda item: item[1], reverse=True))
+
+ print("Detected dtypes:", state_dict_dtype)
+ return state_dict_dtype
+
+
+def get_loadable_dtype(prefix="model.diffusion_model.", state_dict=None, state_dict_dtype=None):
+ if state_dict is not None:
+ state_dict_dtype = get_state_dict_dtype(state_dict)
+
+ # get the first dtype
+ if prefix in state_dict_dtype:
+ return list(state_dict_dtype[prefix])[0]
+ return None
+
+
+def get_vae_dtype(state_dict=None, state_dict_dtype=None):
+ if state_dict is not None:
+ state_dict_dtype = get_state_dict_dtype(state_dict)
+
+ if state_dict_dtype is None:
+ raise ValueError("fail to get vae dtype")
+
+
+ vae_prefixes = [prefix for prefix in ("vae.", "first_stage_model.") if prefix in state_dict_dtype]
+
+ if len(vae_prefixes) > 0:
+ vae_prefix = vae_prefixes[0]
+ for dtype in state_dict_dtype[vae_prefix]:
+ if state_dict_dtype[vae_prefix][dtype] > 240 and dtype in (torch.float16, torch.float32, torch.bfloat16):
+ # vae items: 248 for SD1, SDXL 245 for flux
+ return dtype
+
+ return None
+
+
+def fix_position_ids(state_dict, force=False):
+ # for SD1.5 or some SDXL with position_ids
+ for prefix in ("cond_stage_models.", "conditioner.embedders.0."):
+ position_id_key = f"{prefix}transformer.text_model.embeddings.position_ids"
+ if position_id_key in state_dict:
+ original = state_dict[position_id_key]
+ if original.dtype == torch.int64:
+ return
+
+ if force:
+ # regenerate
+ fixed = torch.tensor([list(range(77))], dtype=torch.int64, device=original.device)
+ else:
+ fixed = state_dict[position_id_key].to(torch.int64)
+ print(f"Warning: Fixed position_ids dtype from {original.dtype} to {fixed.dtype}")
+
+ state_dict[position_id_key] = fixed
+
+
def load_model_weights(model, checkpoint_info: CheckpointInfo, state_dict, timer):
sd_model_hash = checkpoint_info.calculate_shorthash()
timer.record("calculate hash")
@@ -428,6 +534,9 @@ def load_model_weights(model, checkpoint_info: CheckpointInfo, state_dict, timer
else:
model.ztsnr = False
+ fix_position_ids(state_dict)
+
+
if model.is_sdxl:
sd_models_xl.extend_sdxl(model)
@@ -441,6 +550,9 @@ def load_model_weights(model, checkpoint_info: CheckpointInfo, state_dict, timer
if hasattr(model, "before_load_weights"):
model.before_load_weights(state_dict)
+ # get all dtypes of state_dict
+ state_dict_dtype = get_state_dict_dtype(state_dict)
+
model.load_state_dict(state_dict, strict=False)
timer.record("apply weights to model")
@@ -466,7 +578,13 @@ def load_model_weights(model, checkpoint_info: CheckpointInfo, state_dict, timer
model.to(memory_format=torch.channels_last)
timer.record("apply channels_last")
- if shared.cmd_opts.no_half:
+ # check dtype of vae
+ dtype_vae = get_vae_dtype(state_dict_dtype=state_dict_dtype)
+ found_unet_dtype = get_loadable_dtype("model.diffusion_model.", state_dict_dtype=state_dict_dtype)
+ unet_has_float = found_unet_dtype in (torch.float16, torch.float32, torch.bfloat16)
+
+ if (found_unet_dtype is None or unet_has_float) and shared.cmd_opts.no_half:
+ # unet type is not detected or unet has float dtypes
model.float()
model.alphas_cumprod_original = model.alphas_cumprod
devices.dtype_unet = torch.float32
@@ -476,8 +594,11 @@ def load_model_weights(model, checkpoint_info: CheckpointInfo, state_dict, timer
vae = model.first_stage_model
depth_model = getattr(model, 'depth_model', None)
+ if dtype_vae == torch.bfloat16 and dtype_vae in devices.supported_vae_dtypes:
+ # preserve bfloat16 if it supported
+ model.first_stage_model = None
# with --no-half-vae, remove VAE from model when doing half() to prevent its weights from being converted to float16
- if shared.cmd_opts.no_half_vae:
+ elif shared.cmd_opts.no_half_vae:
model.first_stage_model = None
# with --upcast-sampling, don't convert the depth model weights to float16
if shared.cmd_opts.upcast_sampling and depth_model:
@@ -485,15 +606,28 @@ def load_model_weights(model, checkpoint_info: CheckpointInfo, state_dict, timer
alphas_cumprod = model.alphas_cumprod
model.alphas_cumprod = None
- model.half()
+
+
+ if found_unet_dtype in (torch.float16, torch.float32, torch.bfloat16):
+ model.half()
+ elif found_unet_dtype in (torch.float8_e4m3fn, torch.float8_e5m2):
+ pass
+ else:
+ print("Fail to get a vaild UNet dtype. ignore...")
+
model.alphas_cumprod = alphas_cumprod
model.alphas_cumprod_original = alphas_cumprod
model.first_stage_model = vae
if depth_model:
model.depth_model = depth_model
- devices.dtype_unet = torch.float16
- timer.record("apply half()")
+ if found_unet_dtype in (torch.float16, torch.float32):
+ devices.dtype_unet = torch.float16
+ timer.record("apply half()")
+ else:
+ print(f"load Unet {found_unet_dtype} as is ...")
+ devices.dtype_unet = found_unet_dtype if found_unet_dtype else torch.float16
+ timer.record("load UNet")
apply_alpha_schedule_override(model)
@@ -503,10 +637,18 @@ def load_model_weights(model, checkpoint_info: CheckpointInfo, state_dict, timer
if hasattr(module, 'fp16_bias'):
del module.fp16_bias
- if check_fp8(model):
+ if found_unet_dtype not in (torch.float8_e4m3fn,torch.float8_e5m2) and check_fp8(model):
devices.fp8 = True
+
+ # do not convert vae, text_encoders.clip_l, clip_g, t5xxl
first_stage = model.first_stage_model
model.first_stage_model = None
+ vae = getattr(model, 'vae', None)
+ if vae is not None:
+ model.vae = None
+ text_encoders = getattr(model, 'text_encoders', None)
+ if text_encoders is not None:
+ model.text_encoders = None
for module in model.modules():
if isinstance(module, (torch.nn.Conv2d, torch.nn.Linear)):
if shared.opts.cache_fp16_weight:
@@ -514,6 +656,10 @@ def load_model_weights(model, checkpoint_info: CheckpointInfo, state_dict, timer
if module.bias is not None:
module.fp16_bias = module.bias.data.clone().cpu().half()
module.to(torch.float8_e4m3fn)
+ if text_encoders is not None:
+ model.text_encoders = text_encoders
+ if vae is not None:
+ model.vae = vae
model.first_stage_model = first_stage
timer.record("apply fp8")
else:
@@ -521,8 +667,16 @@ def load_model_weights(model, checkpoint_info: CheckpointInfo, state_dict, timer
devices.unet_needs_upcast = shared.cmd_opts.upcast_sampling and devices.dtype == torch.float16 and devices.dtype_unet == torch.float16
- model.first_stage_model.to(devices.dtype_vae)
- timer.record("apply dtype to VAE")
+ # check supported vae dtype
+ dtype_vae = get_vae_dtype(state_dict_dtype=state_dict_dtype)
+ if dtype_vae == torch.bfloat16 and dtype_vae in devices.supported_vae_dtypes:
+ devices.dtype_vae = torch.bfloat16
+ print(f"VAE dtype {dtype_vae} detected. load as is.")
+ else:
+ # use default devices.dtype_vae
+ model.first_stage_model.to(devices.dtype_vae)
+ print(f"Use VAE dtype {devices.dtype_vae}")
+ timer.record("apply dtype to VAE")
# clean up cache if limit is reached
while len(checkpoints_loaded) > shared.opts.sd_checkpoint_cache:
@@ -675,6 +829,9 @@ def apply_alpha_schedule_override(sd_model, p=None):
sd2_clip_weight = 'cond_stage_model.model.transformer.resblocks.0.attn.in_proj_weight'
sdxl_clip_weight = 'conditioner.embedders.1.model.ln_final.weight'
sdxl_refiner_clip_weight = 'conditioner.embedders.0.model.ln_final.weight'
+clip_l_clip_weight = 'text_encoders.clip_l.transformer.text_model.final_layer_norm.weight'
+clip_g_clip_weight = 'text_encoders.clip_g.transformer.text_model.final_layer_norm.weight'
+t5xxl_clip_weight = 'text_encoders.t5xxl.transformer.encoder.final_layer_norm.weight'
class SdModelData:
@@ -807,7 +964,7 @@ def load_model(checkpoint_info=None, already_loaded_state_dict=None, checkpoint_
if not checkpoint_config:
checkpoint_config = sd_models_config.find_checkpoint_config(state_dict, checkpoint_info)
- clip_is_included_into_sd = any(x for x in [sd1_clip_weight, sd2_clip_weight, sdxl_clip_weight, sdxl_refiner_clip_weight] if x in state_dict)
+ clip_is_included_into_sd = any(x for x in [sd1_clip_weight, sd2_clip_weight, sdxl_clip_weight, sdxl_refiner_clip_weight, clip_l_clip_weight, clip_g_clip_weight ] if x in state_dict)
timer.record("find config")
@@ -818,6 +975,18 @@ def load_model(checkpoint_info=None, already_loaded_state_dict=None, checkpoint_
print(f"Creating model from config: {checkpoint_config}")
+ # get all dtypes of state_dict
+ state_dict_dtype = get_state_dict_dtype(state_dict)
+
+ # check loadable unet dtype before loading
+ loadable_unet_dtype = get_loadable_dtype("model.diffusion_model.", state_dict_dtype=state_dict_dtype)
+
+ # check dtype of vae
+ dtype_vae = get_vae_dtype(state_dict_dtype=state_dict_dtype)
+ if dtype_vae == torch.bfloat16 and dtype_vae in devices.supported_vae_dtypes:
+ devices.dtype_vae = torch.bfloat16
+ print(f"VAE dtype {dtype_vae} detected.")
+
sd_model = None
try:
with sd_disable_initialization.DisableInitialization(disable_clip=clip_is_included_into_sd or shared.cmd_opts.do_not_download_clip):
@@ -842,8 +1011,10 @@ def load_model(checkpoint_info=None, already_loaded_state_dict=None, checkpoint_
else:
weight_dtype_conversion = {
'first_stage_model': None,
+ 'text_encoders': None,
+ 'vae': None,
'alphas_cumprod': None,
- '': torch.float16,
+ '': torch.float16 if loadable_unet_dtype in (torch.float16, torch.float32, torch.bfloat16) else None,
}
with sd_disable_initialization.LoadStateDictOnMeta(state_dict, device=model_target_device(sd_model), weight_dtype_conversion=weight_dtype_conversion):
@@ -870,7 +1041,7 @@ def load_model(checkpoint_info=None, already_loaded_state_dict=None, checkpoint_
timer.record("scripts callbacks")
- with devices.autocast(), torch.no_grad():
+ with devices.autocast(target_dtype=devices.dtype_inference), torch.no_grad():
sd_model.cond_stage_model_empty_prompt = get_empty_cond(sd_model)
timer.record("calculate empty prompt")
diff --git a/modules/sd_models_config.py b/modules/sd_models_config.py
index 3c1e4a1518f..4251062c8cc 100644
--- a/modules/sd_models_config.py
+++ b/modules/sd_models_config.py
@@ -25,6 +25,7 @@
config_alt_diffusion = os.path.join(sd_configs_path, "alt-diffusion-inference.yaml")
config_alt_diffusion_m18 = os.path.join(sd_configs_path, "alt-diffusion-m18-inference.yaml")
config_sd3 = os.path.join(sd_configs_path, "sd3-inference.yaml")
+config_flux1 = os.path.join(sd_configs_path, "flux1-inference.yaml")
def is_using_v_parameterization_for_sd2(state_dict):
@@ -78,6 +79,9 @@ def guess_model_config_from_state_dict(sd, filename):
if "model.diffusion_model.x_embedder.proj.weight" in sd:
return config_sd3
+ if "model.diffusion_model.double_blocks.0.img_attn.norm.key_norm.scale" in sd:
+ return config_flux1
+
if sd.get('conditioner.embedders.1.model.ln_final.weight', None) is not None:
if diffusion_model_input.shape[1] == 9:
return config_sdxl_inpainting
diff --git a/modules/sd_models_types.py b/modules/sd_models_types.py
index 2fce2777b2f..867f8b6e20a 100644
--- a/modules/sd_models_types.py
+++ b/modules/sd_models_types.py
@@ -36,5 +36,8 @@ class WebuiSdModel(LatentDiffusion):
is_sd3: bool
"""True if the model's architecture is SD 3"""
+ is_flux1: bool
+ """True if the model's architecture is FLUX 1"""
+
latent_channels: int
"""number of layer in latent image representation; will be 16 in SD3 and 4 in other version"""
diff --git a/modules/sd_models_xl.py b/modules/sd_models_xl.py
index 1242a59369f..5abb75d1f02 100644
--- a/modules/sd_models_xl.py
+++ b/modules/sd_models_xl.py
@@ -18,7 +18,7 @@ def get_learned_conditioning(self: sgm.models.diffusion.DiffusionEngine, batch:
is_negative_prompt = getattr(batch, 'is_negative_prompt', False)
aesthetic_score = shared.opts.sdxl_refiner_low_aesthetic_score if is_negative_prompt else shared.opts.sdxl_refiner_high_aesthetic_score
- devices_args = dict(device=devices.device, dtype=devices.dtype)
+ devices_args = dict(device=devices.device, dtype=devices.dtype_inference)
sdxl_conds = {
"txt": batch,
diff --git a/modules/sd_samplers_common.py b/modules/sd_samplers_common.py
index c060cccb24b..b312c41d8e6 100644
--- a/modules/sd_samplers_common.py
+++ b/modules/sd_samplers_common.py
@@ -64,7 +64,7 @@ def single_sample_to_image(sample, approximation=None):
x_sample = samples_to_images_tensor(sample.unsqueeze(0), approximation)[0] * 0.5 + 0.5
x_sample = torch.clamp(x_sample, min=0.0, max=1.0)
- x_sample = 255. * np.moveaxis(x_sample.cpu().numpy(), 0, 2)
+ x_sample = 255. * np.moveaxis(x_sample.to(dtype=devices.dtype).cpu().numpy(), 0, 2)
x_sample = x_sample.astype(np.uint8)
return Image.fromarray(x_sample)
diff --git a/modules/sd_vae.py b/modules/sd_vae.py
index 43687e48dcf..6ae038333d4 100644
--- a/modules/sd_vae.py
+++ b/modules/sd_vae.py
@@ -197,47 +197,58 @@ def load_vae(model, vae_file=None, vae_source="from unknown source"):
cache_enabled = shared.opts.sd_vae_checkpoint_cache > 0
+ loaded = False
if vae_file:
if cache_enabled and vae_file in checkpoints_loaded:
# use vae checkpoint cache
print(f"Loading VAE weights {vae_source}: cached {get_filename(vae_file)}")
store_base_vae(model)
- _load_vae_dict(model, checkpoints_loaded[vae_file])
+ loaded = _load_vae_dict(model, checkpoints_loaded[vae_file])
else:
assert os.path.isfile(vae_file), f"VAE {vae_source} doesn't exist: {vae_file}"
print(f"Loading VAE weights {vae_source}: {vae_file}")
store_base_vae(model)
vae_dict_1 = load_vae_dict(vae_file, map_location=shared.weight_load_location)
- _load_vae_dict(model, vae_dict_1)
+ loaded = _load_vae_dict(model, vae_dict_1)
- if cache_enabled:
+ if loaded and cache_enabled:
# cache newly loaded vae
checkpoints_loaded[vae_file] = vae_dict_1.copy()
# clean up cache if limit is reached
- if cache_enabled:
+ if loaded and cache_enabled:
while len(checkpoints_loaded) > shared.opts.sd_vae_checkpoint_cache + 1: # we need to count the current model
checkpoints_loaded.popitem(last=False) # LRU
# If vae used is not in dict, update it
# It will be removed on refresh though
vae_opt = get_filename(vae_file)
- if vae_opt not in vae_dict:
+ if loaded and vae_opt not in vae_dict:
vae_dict[vae_opt] = vae_file
elif loaded_vae_file:
restore_base_vae(model)
+ loaded = True
- loaded_vae_file = vae_file
+ if loaded:
+ loaded_vae_file = vae_file
model.base_vae = base_vae
model.loaded_vae_file = loaded_vae_file
+ return loaded
# don't call this from outside
def _load_vae_dict(model, vae_dict_1):
+ conv_out = model.first_stage_model.state_dict().get("encoder.conv_out.weight")
+ # check shape of "encoder.conv_out.weight". SD1.5/SDXL: [8, 512, 3, 3], FLUX/SD3: [32, 512, 3, 3]
+ if conv_out.shape != vae_dict_1["encoder.conv_out.weight"].shape:
+ print("Failed to load VAE. Size mismatched!")
+ return False
+
model.first_stage_model.load_state_dict(vae_dict_1)
model.first_stage_model.to(devices.dtype_vae)
+ return True
def clear_loaded_vae():
@@ -270,7 +281,7 @@ def reload_vae_weights(sd_model=None, vae_file=unspecified):
sd_hijack.model_hijack.undo_hijack(sd_model)
- load_vae(sd_model, vae_file, vae_source)
+ loaded = load_vae(sd_model, vae_file, vae_source)
sd_hijack.model_hijack.hijack(sd_model)
@@ -279,5 +290,6 @@ def reload_vae_weights(sd_model=None, vae_file=unspecified):
script_callbacks.model_loaded_callback(sd_model)
- print("VAE weights loaded.")
+ if loaded:
+ print("VAE weights loaded.")
return sd_model
diff --git a/modules/sd_vae_approx.py b/modules/sd_vae_approx.py
index c5dda7431f1..48ffd672d5c 100644
--- a/modules/sd_vae_approx.py
+++ b/modules/sd_vae_approx.py
@@ -44,6 +44,8 @@ def model():
model_name = "vaeapprox-sd3.pt"
elif shared.sd_model.is_sdxl:
model_name = "vaeapprox-sdxl.pt"
+ elif shared.sd_model.is_flux1:
+ model_name = "vaeapprox-sd3.pt"
else:
model_name = "model.pt"
@@ -81,6 +83,18 @@ def cheap_approximation(sample):
[ 0.0815, 0.0846, 0.1207], [-0.0120, -0.0055, -0.0867],
[-0.0749, -0.0634, -0.0456], [-0.1418, -0.1457, -0.1259],
]
+ elif shared.sd_model.is_flux1:
+ coeffs = [
+ # from comfy
+ [-0.0404, 0.0159, 0.0609], [ 0.0043, 0.0298, 0.0850],
+ [ 0.0328, -0.0749, -0.0503], [-0.0245, 0.0085, 0.0549],
+ [ 0.0966, 0.0894, 0.0530], [ 0.0035, 0.0399, 0.0123],
+ [ 0.0583, 0.1184, 0.1262], [-0.0191, -0.0206, -0.0306],
+ [-0.0324, 0.0055, 0.1001], [ 0.0955, 0.0659, -0.0545],
+ [-0.0504, 0.0231, -0.0013], [ 0.0500, -0.0008, -0.0088],
+ [ 0.0982, 0.0941, 0.0976], [-0.1233, -0.0280, -0.0897],
+ [-0.0005, -0.0530, -0.0020], [-0.1273, -0.0932, -0.0680],
+ ]
elif shared.sd_model.is_sdxl:
coeffs = [
[ 0.3448, 0.4168, 0.4395],
diff --git a/modules/sd_vae_taesd.py b/modules/sd_vae_taesd.py
index d06253d2a88..76771e95e55 100644
--- a/modules/sd_vae_taesd.py
+++ b/modules/sd_vae_taesd.py
@@ -63,7 +63,7 @@ def __init__(self, decoder_path="taesd_decoder.pth", latent_channels=None):
super().__init__()
if latent_channels is None:
- latent_channels = 16 if "taesd3" in str(decoder_path) else 4
+ latent_channels = 16 if any(typ in str(decoder_path) for typ in ("taesd3", "taef1")) else 4
self.decoder = decoder(latent_channels)
self.decoder.load_state_dict(
@@ -79,7 +79,7 @@ def __init__(self, encoder_path="taesd_encoder.pth", latent_channels=None):
super().__init__()
if latent_channels is None:
- latent_channels = 16 if "taesd3" in str(encoder_path) else 4
+ latent_channels = 16 if any(typ in str(encoder_path) for typ in ("taesd3", "taef1")) else 4
self.encoder = encoder(latent_channels)
self.encoder.load_state_dict(
@@ -97,6 +97,8 @@ def download_model(model_path, model_url):
def decoder_model():
if shared.sd_model.is_sd3:
model_name = "taesd3_decoder.pth"
+ elif shared.sd_model.is_flux1:
+ model_name = "taef1_decoder.pth"
elif shared.sd_model.is_sdxl:
model_name = "taesdxl_decoder.pth"
else:
@@ -122,6 +124,8 @@ def decoder_model():
def encoder_model():
if shared.sd_model.is_sd3:
model_name = "taesd3_encoder.pth"
+ elif shared.sd_model.is_flux1:
+ model_name = "taef1_encoder.pth"
elif shared.sd_model.is_sdxl:
model_name = "taesdxl_encoder.pth"
else:
diff --git a/modules/shared_init.py b/modules/shared_init.py
index a6ad0433d6f..2d58c637450 100644
--- a/modules/shared_init.py
+++ b/modules/shared_init.py
@@ -29,7 +29,7 @@ def initialize():
devices.dtype = torch.float32 if cmd_opts.no_half else torch.float16
devices.dtype_vae = torch.float32 if cmd_opts.no_half or cmd_opts.no_half_vae else torch.float16
- devices.dtype_inference = torch.float32 if cmd_opts.precision == 'full' else devices.dtype
+ devices.dtype_inference = torch.float32 if cmd_opts.precision == 'full' else devices.dtype_inference
if cmd_opts.precision == "half":
msg = "--no-half and --no-half-vae conflict with --precision half"
diff --git a/modules/shared_items.py b/modules/shared_items.py
index 11f10b3f7b1..3aaf0649028 100644
--- a/modules/shared_items.py
+++ b/modules/shared_items.py
@@ -16,10 +16,12 @@ def dat_models_names():
return [x.name for x in modules.dat_model.get_dat_models(None)]
-def postprocessing_scripts():
+def postprocessing_scripts(filter_out_extra_only=False, filter_out_main_ui_only=False):
import modules.scripts
-
- return modules.scripts.scripts_postproc.scripts
+ return list(filter(
+ lambda s: (not filter_out_extra_only or not s.extra_only) and (not filter_out_main_ui_only or not s.main_ui_only),
+ modules.scripts.scripts_postproc.scripts,
+ ))
def sd_vae_items():
diff --git a/modules/shared_options.py b/modules/shared_options.py
index efede7067f2..0cb10acfd66 100644
--- a/modules/shared_options.py
+++ b/modules/shared_options.py
@@ -195,6 +195,9 @@
options_templates.update(options_section(('sd3', "Stable Diffusion 3", "sd"), {
"sd3_enable_t5": OptionInfo(False, "Enable T5").info("load T5 text encoder; increases VRAM use by a lot, potentially improving quality of generation; requires model reload to apply"),
}))
+options_templates.update(options_section(('flux', "Stable Diffusion FLUX", "sd"), {
+ "flux_enable_t5": OptionInfo(False, "Enable T5").info("load T5 text encoder; increases VRAM use by a lot, potentially improving quality of generation; requires model reload to apply"),
+}))
options_templates.update(options_section(('vae', "VAE", "sd"), {
"sd_vae_explanation": OptionHTML("""
@@ -242,6 +245,7 @@
"batch_cond_uncond": OptionInfo(True, "Batch cond/uncond").info("do both conditional and unconditional denoising in one batch; uses a bit more VRAM during sampling, but improves speed; previously this was controlled by --always-batch-cond-uncond commandline argument"),
"fp8_storage": OptionInfo("Disable", "FP8 weight", gr.Radio, {"choices": ["Disable", "Enable for SDXL", "Enable"]}).info("Use FP8 to store Linear/Conv layers' weight. Require pytorch>=2.1.0."),
"cache_fp16_weight": OptionInfo(False, "Cache FP16 weight for LoRA").info("Cache fp16 weight when enabling FP8, will increase the quality of LoRA. Use more system ram."),
+ "lora_without_backup_weight": OptionInfo(False, "LoRA without backup weights").info("LoRA without backup weights to save RAM."),
}))
options_templates.update(options_section(('compatibility', "Compatibility", "sd"), {
@@ -291,6 +295,7 @@
"textual_inversion_print_at_load": OptionInfo(False, "Print a list of Textual Inversion embeddings when loading model"),
"textual_inversion_add_hashes_to_infotext": OptionInfo(True, "Add Textual Inversion hashes to infotext"),
"sd_hypernetwork": OptionInfo("None", "Add hypernetwork to prompt", gr.Dropdown, lambda: {"choices": ["None", *shared.hypernetworks]}, refresh=shared_items.reload_hypernetworks),
+ "textual_inversion_image_embedding_data_cache": OptionInfo(False, 'Cache the data of image embeddings').info('potentially increase TI load time at the cost some disk space'),
}))
options_templates.update(options_section(('ui_prompt_editing', "Prompt editing", "ui"), {
@@ -410,9 +415,9 @@
}))
options_templates.update(options_section(('postprocessing', "Postprocessing", "postprocessing"), {
- 'postprocessing_enable_in_main_ui': OptionInfo([], "Enable postprocessing operations in txt2img and img2img tabs", ui_components.DropdownMulti, lambda: {"choices": [x.name for x in shared_items.postprocessing_scripts()]}),
- 'postprocessing_disable_in_extras': OptionInfo([], "Disable postprocessing operations in extras tab", ui_components.DropdownMulti, lambda: {"choices": [x.name for x in shared_items.postprocessing_scripts()]}),
- 'postprocessing_operation_order': OptionInfo([], "Postprocessing operation order", ui_components.DropdownMulti, lambda: {"choices": [x.name for x in shared_items.postprocessing_scripts()]}),
+ 'postprocessing_enable_in_main_ui': OptionInfo([], "Enable postprocessing operations in txt2img and img2img tabs", ui_components.DropdownMulti, lambda: {"choices": [x.name for x in shared_items.postprocessing_scripts(filter_out_extra_only=True)]}),
+ 'postprocessing_disable_in_extras': OptionInfo([], "Disable postprocessing operations in extras tab", ui_components.DropdownMulti, lambda: {"choices": [x.name for x in shared_items.postprocessing_scripts(filter_out_main_ui_only=True)]}),
+ 'postprocessing_operation_order': OptionInfo([], "Postprocessing operation order", ui_components.DropdownMulti, lambda: {"choices": [x.name for x in shared_items.postprocessing_scripts(filter_out_main_ui_only=True)]}),
'upscaling_max_images_in_cache': OptionInfo(5, "Maximum number of images in upscaling cache", gr.Slider, {"minimum": 0, "maximum": 10, "step": 1}),
'postprocessing_existing_caption_action': OptionInfo("Ignore", "Action for existing captions", gr.Radio, {"choices": ["Ignore", "Keep", "Prepend", "Append"]}).info("when generating captions using postprocessing; Ignore = use generated; Keep = use original; Prepend/Append = combine both"),
}))
diff --git a/modules/textual_inversion/textual_inversion.py b/modules/textual_inversion/textual_inversion.py
index dc7833e9394..f209b883450 100644
--- a/modules/textual_inversion/textual_inversion.py
+++ b/modules/textual_inversion/textual_inversion.py
@@ -12,7 +12,7 @@
import numpy as np
from PIL import Image, PngImagePlugin
-from modules import shared, devices, sd_hijack, sd_models, images, sd_samplers, sd_hijack_checkpoint, errors, hashes
+from modules import shared, devices, sd_hijack, sd_models, images, sd_samplers, sd_hijack_checkpoint, errors, hashes, cache
import modules.textual_inversion.dataset
from modules.textual_inversion.learn_schedule import LearnRateScheduler
@@ -116,6 +116,7 @@ def __init__(self):
self.expected_shape = -1
self.embedding_dirs = {}
self.previously_displayed_embeddings = ()
+ self.image_embedding_cache = cache.cache('image-embedding')
def add_embedding_dir(self, path):
self.embedding_dirs[path] = DirWithTextualInversionEmbeddings(path)
@@ -154,6 +155,31 @@ def get_expected_shape(self):
vec = shared.sd_model.cond_stage_model.encode_embedding_init_text(",", 1)
return vec.shape[1]
+ def read_embedding_from_image(self, path, name):
+ try:
+ ondisk_mtime = os.path.getmtime(path)
+
+ if (cache_embedding := self.image_embedding_cache.get(path)) and ondisk_mtime == cache_embedding.get('mtime', 0):
+ # cache will only be used if the file has not been modified time matches
+ return cache_embedding.get('data', None), cache_embedding.get('name', None)
+
+ embed_image = Image.open(path)
+ if hasattr(embed_image, 'text') and 'sd-ti-embedding' in embed_image.text:
+ data = embedding_from_b64(embed_image.text['sd-ti-embedding'])
+ name = data.get('name', name)
+ elif data := extract_image_data_embed(embed_image):
+ name = data.get('name', name)
+
+ if data is None or shared.opts.textual_inversion_image_embedding_data_cache:
+ # data of image embeddings only will be cached if the option textual_inversion_image_embedding_data_cache is enabled
+ # results of images that are not embeddings will allways be cached to reduce unnecessary future disk reads
+ self.image_embedding_cache[path] = {'data': data, 'name': None if data is None else name, 'mtime': ondisk_mtime}
+
+ return data, name
+ except Exception:
+ errors.report(f"Error loading embedding {path}", exc_info=True)
+ return None, None
+
def load_from_file(self, path, filename):
name, ext = os.path.splitext(filename)
ext = ext.upper()
@@ -163,17 +189,10 @@ def load_from_file(self, path, filename):
if second_ext.upper() == '.PREVIEW':
return
- embed_image = Image.open(path)
- if hasattr(embed_image, 'text') and 'sd-ti-embedding' in embed_image.text:
- data = embedding_from_b64(embed_image.text['sd-ti-embedding'])
- name = data.get('name', name)
- else:
- data = extract_image_data_embed(embed_image)
- if data:
- name = data.get('name', name)
- else:
- # if data is None, means this is not an embedding, just a preview image
- return
+ data, name = self.read_embedding_from_image(path, name)
+ if data is None:
+ return
+
elif ext in ['.BIN', '.PT']:
data = torch.load(path, map_location="cpu")
elif ext in ['.SAFETENSORS']:
@@ -191,7 +210,6 @@ def load_from_file(self, path, filename):
else:
print(f"Unable to load Textual inversion embedding due to data issue: '{name}'.")
-
def load_from_dir(self, embdir):
if not os.path.isdir(embdir.path):
return
diff --git a/modules/ui_components.py b/modules/ui_components.py
index 9cf67722a3d..3e3ea54bafe 100644
--- a/modules/ui_components.py
+++ b/modules/ui_components.py
@@ -91,6 +91,7 @@ class InputAccordion(gr.Checkbox):
Actually just a hidden checkbox, but creates an accordion that follows and is followed by the state of the checkbox.
"""
+ accordion_id_set = set()
global_index = 0
def __init__(self, value, **kwargs):
@@ -99,6 +100,18 @@ def __init__(self, value, **kwargs):
self.accordion_id = f"input-accordion-{InputAccordion.global_index}"
InputAccordion.global_index += 1
+ if not InputAccordion.accordion_id_set:
+ from modules import script_callbacks
+ script_callbacks.on_script_unloaded(InputAccordion.reset)
+
+ if self.accordion_id in InputAccordion.accordion_id_set:
+ count = 1
+ while (unique_id := f'{self.accordion_id}-{count}') in InputAccordion.accordion_id_set:
+ count += 1
+ self.accordion_id = unique_id
+
+ InputAccordion.accordion_id_set.add(self.accordion_id)
+
kwargs_checkbox = {
**kwargs,
"elem_id": f"{self.accordion_id}-checkbox",
@@ -143,3 +156,7 @@ def __exit__(self, exc_type, exc_val, exc_tb):
def get_block_name(self):
return "checkbox"
+ @classmethod
+ def reset(cls):
+ cls.global_index = 0
+ cls.accordion_id_set.clear()
diff --git a/modules/ui_extra_networks.py b/modules/ui_extra_networks.py
index 6e9ec164552..1f19bd36d12 100644
--- a/modules/ui_extra_networks.py
+++ b/modules/ui_extra_networks.py
@@ -177,10 +177,8 @@ def add_pages_to_demo(app):
app.add_api_route("/sd_extra_networks/get-single-card", get_single_card, methods=["GET"])
-def quote_js(s):
- s = s.replace('\\', '\\\\')
- s = s.replace('"', '\\"')
- return f'"{s}"'
+def quote_js(s: str):
+ return json.dumps(s, ensure_ascii=False)
class ExtraNetworksPage:
diff --git a/modules/ui_loadsave.py b/modules/ui_loadsave.py
index 0cc1ab82af4..95a776d947e 100644
--- a/modules/ui_loadsave.py
+++ b/modules/ui_loadsave.py
@@ -176,7 +176,7 @@ def iter_changes(self, current_ui_settings, values):
if new_value == old_value:
continue
- if old_value is None and new_value == '' or new_value == []:
+ if old_value is None and (new_value == '' or new_value == []):
continue
yield path, old_value, new_value
diff --git a/modules/upscaler.py b/modules/upscaler.py
index 507881fede2..12ab3547cf6 100644
--- a/modules/upscaler.py
+++ b/modules/upscaler.py
@@ -93,13 +93,14 @@ class UpscalerData:
scaler: Upscaler = None
model: None
- def __init__(self, name: str, path: str, upscaler: Upscaler = None, scale: int = 4, model=None):
+ def __init__(self, name: str, path: str, upscaler: Upscaler = None, scale: int = 4, model=None, sha256: str = None):
self.name = name
self.data_path = path
self.local_data_path = path
self.scaler = upscaler
self.scale = scale
self.model = model
+ self.sha256 = sha256
def __repr__(self):
return f"
"
diff --git a/modules/util.py b/modules/util.py
index 7911b0db72c..baeba2fa271 100644
--- a/modules/util.py
+++ b/modules/util.py
@@ -211,3 +211,80 @@ def open_folder(path):
subprocess.Popen(["explorer.exe", subprocess.check_output(["wslpath", "-w", path])])
else:
subprocess.Popen(["xdg-open", path])
+
+
+def load_file_from_url(
+ url: str,
+ *,
+ model_dir: str,
+ progress: bool = True,
+ file_name: str | None = None,
+ hash_prefix: str | None = None,
+ re_download: bool = False,
+) -> str:
+ """Download a file from `url` into `model_dir`, using the file present if possible.
+ Returns the path to the downloaded file.
+
+ file_name: if specified, it will be used as the filename, otherwise the filename will be extracted from the url.
+ file is downloaded to {file_name}.tmp then moved to the final location after download is complete.
+ hash_prefix: sha256 hex string, if provided, the hash of the downloaded file will be checked against this prefix.
+ if the hash does not match, the temporary file is deleted and a ValueError is raised.
+ re_download: forcibly re-download the file even if it already exists.
+ """
+ from urllib.parse import urlparse
+ import requests
+ try:
+ from tqdm import tqdm
+ except ImportError:
+ class tqdm:
+ def __init__(self, *args, **kwargs):
+ pass
+
+ def update(self, n=1, *args, **kwargs):
+ pass
+
+ def __enter__(self):
+ return self
+
+ def __exit__(self, exc_type, exc_val, exc_tb):
+ pass
+
+ if not file_name:
+ parts = urlparse(url)
+ file_name = os.path.basename(parts.path)
+
+ cached_file = os.path.abspath(os.path.join(model_dir, file_name))
+
+ if re_download or not os.path.exists(cached_file):
+ os.makedirs(model_dir, exist_ok=True)
+ temp_file = os.path.join(model_dir, f"{file_name}.tmp")
+ print(f'\nDownloading: "{url}" to {cached_file}')
+ response = requests.get(url, stream=True)
+ response.raise_for_status()
+ total_size = int(response.headers.get('content-length', 0))
+ with tqdm(total=total_size, unit='B', unit_scale=True, desc=file_name, disable=not progress) as progress_bar:
+ with open(temp_file, 'wb') as file:
+ for chunk in response.iter_content(chunk_size=1024):
+ if chunk:
+ file.write(chunk)
+ progress_bar.update(len(chunk))
+
+ if hash_prefix and not compare_sha256(temp_file, hash_prefix):
+ print(f"Hash mismatch for {temp_file}. Deleting the temporary file.")
+ os.remove(temp_file)
+ raise ValueError(f"File hash does not match the expected hash prefix {hash_prefix}!")
+
+ os.rename(temp_file, cached_file)
+ return cached_file
+
+
+def compare_sha256(file_path: str, hash_prefix: str) -> bool:
+ """Check if the SHA256 hash of the file matches the given prefix."""
+ import hashlib
+ hash_sha256 = hashlib.sha256()
+ blksize = 1024 * 1024
+
+ with open(file_path, "rb") as f:
+ for chunk in iter(lambda: f.read(blksize), b""):
+ hash_sha256.update(chunk)
+ return hash_sha256.hexdigest().startswith(hash_prefix.strip().lower())
diff --git a/scripts/postprocessing_codeformer.py b/scripts/postprocessing_codeformer.py
index 53a0cc44cde..f86e99a04db 100644
--- a/scripts/postprocessing_codeformer.py
+++ b/scripts/postprocessing_codeformer.py
@@ -12,8 +12,8 @@ class ScriptPostprocessingCodeFormer(scripts_postprocessing.ScriptPostprocessing
def ui(self):
with ui_components.InputAccordion(False, label="CodeFormer") as enable:
with gr.Row():
- codeformer_visibility = gr.Slider(minimum=0.0, maximum=1.0, step=0.001, label="Visibility", value=1.0, elem_id="extras_codeformer_visibility")
- codeformer_weight = gr.Slider(minimum=0.0, maximum=1.0, step=0.001, label="Weight (0 = maximum effect, 1 = minimum effect)", value=0, elem_id="extras_codeformer_weight")
+ codeformer_visibility = gr.Slider(minimum=0.0, maximum=1.0, step=0.001, label="Visibility", value=1.0, elem_id=self.elem_id_suffix("extras_codeformer_visibility"))
+ codeformer_weight = gr.Slider(minimum=0.0, maximum=1.0, step=0.001, label="Weight (0 = maximum effect, 1 = minimum effect)", value=0, elem_id=self.elem_id_suffix("extras_codeformer_weight"))
return {
"enable": enable,
diff --git a/scripts/postprocessing_gfpgan.py b/scripts/postprocessing_gfpgan.py
index 57e3623995c..3a130fd6336 100644
--- a/scripts/postprocessing_gfpgan.py
+++ b/scripts/postprocessing_gfpgan.py
@@ -11,7 +11,7 @@ class ScriptPostprocessingGfpGan(scripts_postprocessing.ScriptPostprocessing):
def ui(self):
with ui_components.InputAccordion(False, label="GFPGAN") as enable:
- gfpgan_visibility = gr.Slider(minimum=0.0, maximum=1.0, step=0.001, label="Visibility", value=1.0, elem_id="extras_gfpgan_visibility")
+ gfpgan_visibility = gr.Slider(minimum=0.0, maximum=1.0, step=0.001, label="Visibility", value=1.0, elem_id=self.elem_id_suffix("extras_gfpgan_visibility"))
return {
"enable": enable,
diff --git a/scripts/postprocessing_upscale.py b/scripts/postprocessing_upscale.py
index 2409fd2073e..838872bf1b9 100644
--- a/scripts/postprocessing_upscale.py
+++ b/scripts/postprocessing_upscale.py
@@ -30,31 +30,31 @@ class ScriptPostprocessingUpscale(scripts_postprocessing.ScriptPostprocessing):
def ui(self):
selected_tab = gr.Number(value=0, visible=False)
- with InputAccordion(True, label="Upscale", elem_id="extras_upscale") as upscale_enabled:
+ with InputAccordion(True, label="Upscale", elem_id=self.elem_id_suffix("extras_upscale")) as upscale_enabled:
with FormRow():
- extras_upscaler_1 = gr.Dropdown(label='Upscaler 1', elem_id="extras_upscaler_1", choices=[x.name for x in shared.sd_upscalers], value=shared.sd_upscalers[0].name)
+ extras_upscaler_1 = gr.Dropdown(label='Upscaler 1', elem_id=self.elem_id_suffix("extras_upscaler_1"), choices=[x.name for x in shared.sd_upscalers], value=shared.sd_upscalers[0].name)
with FormRow():
- extras_upscaler_2 = gr.Dropdown(label='Upscaler 2', elem_id="extras_upscaler_2", choices=[x.name for x in shared.sd_upscalers], value=shared.sd_upscalers[0].name)
- extras_upscaler_2_visibility = gr.Slider(minimum=0.0, maximum=1.0, step=0.001, label="Upscaler 2 visibility", value=0.0, elem_id="extras_upscaler_2_visibility")
+ extras_upscaler_2 = gr.Dropdown(label='Upscaler 2', elem_id=self.elem_id_suffix("extras_upscaler_2"), choices=[x.name for x in shared.sd_upscalers], value=shared.sd_upscalers[0].name)
+ extras_upscaler_2_visibility = gr.Slider(minimum=0.0, maximum=1.0, step=0.001, label="Upscaler 2 visibility", value=0.0, elem_id=self.elem_id_suffix("extras_upscaler_2_visibility"))
with FormRow():
- with gr.Tabs(elem_id="extras_resize_mode"):
- with gr.TabItem('Scale by', elem_id="extras_scale_by_tab") as tab_scale_by:
+ with gr.Tabs(elem_id=self.elem_id_suffix("extras_resize_mode")):
+ with gr.TabItem('Scale by', elem_id=self.elem_id_suffix("extras_scale_by_tab")) as tab_scale_by:
with gr.Row():
with gr.Column(scale=4):
- upscaling_resize = gr.Slider(minimum=1.0, maximum=8.0, step=0.05, label="Resize", value=4, elem_id="extras_upscaling_resize")
+ upscaling_resize = gr.Slider(minimum=1.0, maximum=8.0, step=0.05, label="Resize", value=4, elem_id=self.elem_id_suffix("extras_upscaling_resize"))
with gr.Column(scale=1, min_width=160):
- max_side_length = gr.Number(label="Max side length", value=0, elem_id="extras_upscale_max_side_length", tooltip="If any of two sides of the image ends up larger than specified, will downscale it to fit. 0 = no limit.", min_width=160, step=8, minimum=0)
+ max_side_length = gr.Number(label="Max side length", value=0, elem_id=self.elem_id_suffix("extras_upscale_max_side_length"), tooltip="If any of two sides of the image ends up larger than specified, will downscale it to fit. 0 = no limit.", min_width=160, step=8, minimum=0)
- with gr.TabItem('Scale to', elem_id="extras_scale_to_tab") as tab_scale_to:
+ with gr.TabItem('Scale to', elem_id=self.elem_id_suffix("extras_scale_to_tab")) as tab_scale_to:
with FormRow():
- with gr.Column(elem_id="upscaling_column_size", scale=4):
- upscaling_resize_w = gr.Slider(minimum=64, maximum=8192, step=8, label="Width", value=512, elem_id="extras_upscaling_resize_w")
- upscaling_resize_h = gr.Slider(minimum=64, maximum=8192, step=8, label="Height", value=512, elem_id="extras_upscaling_resize_h")
- with gr.Column(elem_id="upscaling_dimensions_row", scale=1, elem_classes="dimensions-tools"):
- upscaling_res_switch_btn = ToolButton(value=switch_values_symbol, elem_id="upscaling_res_switch_btn", tooltip="Switch width/height")
- upscaling_crop = gr.Checkbox(label='Crop to fit', value=True, elem_id="extras_upscaling_crop")
+ with gr.Column(elem_id=self.elem_id_suffix("upscaling_column_size"), scale=4):
+ upscaling_resize_w = gr.Slider(minimum=64, maximum=8192, step=8, label="Width", value=512, elem_id=self.elem_id_suffix("extras_upscaling_resize_w"))
+ upscaling_resize_h = gr.Slider(minimum=64, maximum=8192, step=8, label="Height", value=512, elem_id=self.elem_id_suffix("extras_upscaling_resize_h"))
+ with gr.Column(elem_id=self.elem_id_suffix("upscaling_dimensions_row"), scale=1, elem_classes="dimensions-tools"):
+ upscaling_res_switch_btn = ToolButton(value=switch_values_symbol, elem_id=self.elem_id_suffix("upscaling_res_switch_btn"), tooltip="Switch width/height")
+ upscaling_crop = gr.Checkbox(label='Crop to fit', value=True, elem_id=self.elem_id_suffix("extras_upscaling_crop"))
def on_selected_upscale_method(upscale_method):
if not shared.opts.set_scale_by_when_changing_upscaler:
@@ -169,6 +169,7 @@ def image_changed(self):
class ScriptPostprocessingUpscaleSimple(ScriptPostprocessingUpscale):
name = "Simple Upscale"
order = 900
+ main_ui_only = True
def ui(self):
with FormRow():
diff --git a/scripts/xyz_grid.py b/scripts/xyz_grid.py
index c60dd6dda2f..5778d5f4c8d 100644
--- a/scripts/xyz_grid.py
+++ b/scripts/xyz_grid.py
@@ -20,7 +20,7 @@
import modules.sd_vae
import re
-from modules.ui_components import ToolButton
+from modules.ui_components import ToolButton, InputAccordion
fill_values_symbol = "\U0001f4d2" # 📒
@@ -285,7 +285,7 @@ def __init__(self, *args, **kwargs):
]
-def draw_xyz_grid(p, xs, ys, zs, x_labels, y_labels, z_labels, cell, draw_legend, include_lone_images, include_sub_grids, first_axes_processed, second_axes_processed, margin_size):
+def draw_xyz_grid(p, xs, ys, zs, x_labels, y_labels, z_labels, cell, draw_legend, include_lone_images, include_sub_grids, first_axes_processed, second_axes_processed, margin_size, draw_grid):
hor_texts = [[images.GridAnnotation(x)] for x in x_labels]
ver_texts = [[images.GridAnnotation(y)] for y in y_labels]
title_texts = [[images.GridAnnotation(z)] for z in z_labels]
@@ -370,29 +370,30 @@ def index(ix, iy, iz):
print("Unexpected error: draw_xyz_grid failed to return even a single processed image")
return Processed(p, [])
- z_count = len(zs)
+ if draw_grid:
+ z_count = len(zs)
- for i in range(z_count):
- start_index = (i * len(xs) * len(ys)) + i
- end_index = start_index + len(xs) * len(ys)
- grid = images.image_grid(processed_result.images[start_index:end_index], rows=len(ys))
+ for i in range(z_count):
+ start_index = (i * len(xs) * len(ys)) + i
+ end_index = start_index + len(xs) * len(ys)
+ grid = images.image_grid(processed_result.images[start_index:end_index], rows=len(ys))
+ if draw_legend:
+ grid_max_w, grid_max_h = map(max, zip(*(img.size for img in processed_result.images[start_index:end_index])))
+ grid = images.draw_grid_annotations(grid, grid_max_w, grid_max_h, hor_texts, ver_texts, margin_size)
+ processed_result.images.insert(i, grid)
+ processed_result.all_prompts.insert(i, processed_result.all_prompts[start_index])
+ processed_result.all_seeds.insert(i, processed_result.all_seeds[start_index])
+ processed_result.infotexts.insert(i, processed_result.infotexts[start_index])
+
+ z_grid = images.image_grid(processed_result.images[:z_count], rows=1)
+ z_sub_grid_max_w, z_sub_grid_max_h = map(max, zip(*(img.size for img in processed_result.images[:z_count])))
if draw_legend:
- grid_max_w, grid_max_h = map(max, zip(*(img.size for img in processed_result.images[start_index:end_index])))
- grid = images.draw_grid_annotations(grid, grid_max_w, grid_max_h, hor_texts, ver_texts, margin_size)
- processed_result.images.insert(i, grid)
- processed_result.all_prompts.insert(i, processed_result.all_prompts[start_index])
- processed_result.all_seeds.insert(i, processed_result.all_seeds[start_index])
- processed_result.infotexts.insert(i, processed_result.infotexts[start_index])
-
- z_grid = images.image_grid(processed_result.images[:z_count], rows=1)
- z_sub_grid_max_w, z_sub_grid_max_h = map(max, zip(*(img.size for img in processed_result.images[:z_count])))
- if draw_legend:
- z_grid = images.draw_grid_annotations(z_grid, z_sub_grid_max_w, z_sub_grid_max_h, title_texts, [[images.GridAnnotation()]])
- processed_result.images.insert(0, z_grid)
- # TODO: Deeper aspects of the program rely on grid info being misaligned between metadata arrays, which is not ideal.
- # processed_result.all_prompts.insert(0, processed_result.all_prompts[0])
- # processed_result.all_seeds.insert(0, processed_result.all_seeds[0])
- processed_result.infotexts.insert(0, processed_result.infotexts[0])
+ z_grid = images.draw_grid_annotations(z_grid, z_sub_grid_max_w, z_sub_grid_max_h, title_texts, [[images.GridAnnotation()]])
+ processed_result.images.insert(0, z_grid)
+ # TODO: Deeper aspects of the program rely on grid info being misaligned between metadata arrays, which is not ideal.
+ # processed_result.all_prompts.insert(0, processed_result.all_prompts[0])
+ # processed_result.all_seeds.insert(0, processed_result.all_seeds[0])
+ processed_result.infotexts.insert(0, processed_result.infotexts[0])
return processed_result
@@ -442,7 +443,6 @@ def ui(self, is_img2img):
with gr.Row(variant="compact", elem_id="axis_options"):
with gr.Column():
- draw_legend = gr.Checkbox(label='Draw legend', value=True, elem_id=self.elem_id("draw_legend"))
no_fixed_seeds = gr.Checkbox(label='Keep -1 for seeds', value=False, elem_id=self.elem_id("no_fixed_seeds"))
with gr.Row():
vary_seeds_x = gr.Checkbox(label='Vary seeds for X', value=False, min_width=80, elem_id=self.elem_id("vary_seeds_x"), tooltip="Use different seeds for images along X axis.")
@@ -450,9 +450,12 @@ def ui(self, is_img2img):
vary_seeds_z = gr.Checkbox(label='Vary seeds for Z', value=False, min_width=80, elem_id=self.elem_id("vary_seeds_z"), tooltip="Use different seeds for images along Z axis.")
with gr.Column():
include_lone_images = gr.Checkbox(label='Include Sub Images', value=False, elem_id=self.elem_id("include_lone_images"))
- include_sub_grids = gr.Checkbox(label='Include Sub Grids', value=False, elem_id=self.elem_id("include_sub_grids"))
csv_mode = gr.Checkbox(label='Use text inputs instead of dropdowns', value=False, elem_id=self.elem_id("csv_mode"))
- with gr.Column():
+
+ with InputAccordion(True, label='Draw grid', elem_id=self.elem_id('draw_grid')) as draw_grid:
+ with gr.Row():
+ include_sub_grids = gr.Checkbox(label='Include Sub Grids', value=False, elem_id=self.elem_id("include_sub_grids"))
+ draw_legend = gr.Checkbox(label='Draw legend', value=True, elem_id=self.elem_id("draw_legend"))
margin_size = gr.Slider(label="Grid margins (px)", minimum=0, maximum=500, value=0, step=2, elem_id=self.elem_id("margin_size"))
with gr.Row(variant="compact", elem_id="swap_axes"):
@@ -534,9 +537,9 @@ def get_dropdown_update_from_params(axis, params):
(z_values_dropdown, lambda params: get_dropdown_update_from_params("Z", params)),
)
- return [x_type, x_values, x_values_dropdown, y_type, y_values, y_values_dropdown, z_type, z_values, z_values_dropdown, draw_legend, include_lone_images, include_sub_grids, no_fixed_seeds, vary_seeds_x, vary_seeds_y, vary_seeds_z, margin_size, csv_mode]
+ return [x_type, x_values, x_values_dropdown, y_type, y_values, y_values_dropdown, z_type, z_values, z_values_dropdown, draw_legend, include_lone_images, include_sub_grids, no_fixed_seeds, vary_seeds_x, vary_seeds_y, vary_seeds_z, margin_size, csv_mode, draw_grid]
- def run(self, p, x_type, x_values, x_values_dropdown, y_type, y_values, y_values_dropdown, z_type, z_values, z_values_dropdown, draw_legend, include_lone_images, include_sub_grids, no_fixed_seeds, vary_seeds_x, vary_seeds_y, vary_seeds_z, margin_size, csv_mode):
+ def run(self, p, x_type, x_values, x_values_dropdown, y_type, y_values, y_values_dropdown, z_type, z_values, z_values_dropdown, draw_legend, include_lone_images, include_sub_grids, no_fixed_seeds, vary_seeds_x, vary_seeds_y, vary_seeds_z, margin_size, csv_mode, draw_grid):
x_type, y_type, z_type = x_type or 0, y_type or 0, z_type or 0 # if axle type is None set to 0
if not no_fixed_seeds:
@@ -781,7 +784,8 @@ def cell(x, y, z, ix, iy, iz):
include_sub_grids=include_sub_grids,
first_axes_processed=first_axes_processed,
second_axes_processed=second_axes_processed,
- margin_size=margin_size
+ margin_size=margin_size,
+ draw_grid=draw_grid,
)
if not processed.images:
@@ -790,14 +794,15 @@ def cell(x, y, z, ix, iy, iz):
z_count = len(zs)
- # Set the grid infotexts to the real ones with extra_generation_params (1 main grid + z_count sub-grids)
- processed.infotexts[:1 + z_count] = grid_infotext[:1 + z_count]
+ if draw_grid:
+ # Set the grid infotexts to the real ones with extra_generation_params (1 main grid + z_count sub-grids)
+ processed.infotexts[:1 + z_count] = grid_infotext[:1 + z_count]
if not include_lone_images:
# Don't need sub-images anymore, drop from list:
- processed.images = processed.images[:z_count + 1]
+ processed.images = processed.images[:z_count + 1] if draw_grid else []
- if opts.grid_save:
+ if draw_grid and opts.grid_save:
# Auto-save main and sub-grids:
grid_count = z_count + 1 if z_count > 1 else 1
for g in range(grid_count):
@@ -807,7 +812,7 @@ def cell(x, y, z, ix, iy, iz):
if not include_sub_grids: # if not include_sub_grids then skip saving after the first grid
break
- if not include_sub_grids:
+ if draw_grid and not include_sub_grids:
# Done with sub-grids, drop all related information:
for _ in range(z_count):
del processed.images[1]
diff --git a/webui.bat b/webui.bat
index 7b162ce27cc..f29a40b36e9 100644
--- a/webui.bat
+++ b/webui.bat
@@ -4,7 +4,16 @@ if exist webui.settings.bat (
call webui.settings.bat
)
-if not defined PYTHON (set PYTHON=python)
+if not defined PYTHON (
+ for /f "delims=" %%A in ('where python ^| findstr /n . ^| findstr ^^1:') do (
+ if /i "%%~xA" == ".exe" (
+ set PYTHON=python
+ ) else (
+ set PYTHON=call python
+ )
+ )
+)
+
if defined GIT (set "GIT_PYTHON_GIT_EXECUTABLE=%GIT%")
if not defined VENV_DIR (set "VENV_DIR=%~dp0%venv")
diff --git a/webui.py b/webui.py
index 2c417168aa6..398d8355023 100644
--- a/webui.py
+++ b/webui.py
@@ -6,6 +6,8 @@
from modules import timer
from modules import initialize_util
from modules import initialize
+from modules import manager
+from threading import Thread
startup_timer = timer.startup_timer
startup_timer.record("launcher")
@@ -14,6 +16,8 @@
initialize.check_versions()
+initialize.initialize()
+
def create_api(app):
from modules.api.api import Api
@@ -23,12 +27,10 @@ def create_api(app):
return api
-def api_only():
+def _api_only():
from fastapi import FastAPI
from modules.shared_cmd_options import cmd_opts
- initialize.initialize()
-
app = FastAPI()
initialize_util.setup_middleware(app)
api = create_api(app)
@@ -45,14 +47,53 @@ def api_only():
)
-def webui():
+def warning_if_invalid_install_dir():
+ """
+ Shows a warning if the webui is installed under a path that contains a leading dot in any of its parent directories.
+
+ Gradio '/file=' route will block access to files that have a leading dot in the path segments.
+ We use this route to serve files such as JavaScript and CSS to the webpage,
+ if those files are blocked, the webpage will not function properly.
+ See https://github.com/AUTOMATIC1111/stable-diffusion-webui/issues/13292
+
+ This is a security feature was added to Gradio 3.32.0 and is removed in later versions,
+ this function replicates Gradio file access blocking logic.
+
+ This check should be removed when it's no longer applicable.
+ """
+ from packaging.version import parse
+ from pathlib import Path
+ import gradio
+
+ if parse('3.32.0') <= parse(gradio.__version__) < parse('4'):
+
+ def abspath(path):
+ """modified from Gradio 3.41.2 gradio.utils.abspath()"""
+ if path.is_absolute():
+ return path
+ is_symlink = path.is_symlink() or any(parent.is_symlink() for parent in path.parents)
+ return Path.cwd() / path if (is_symlink or path == path.resolve()) else path.resolve()
+
+ webui_root = Path(__file__).parent
+ if any(part.startswith(".") for part in abspath(webui_root).parts):
+ print(f'''{"!"*25} Warning {"!"*25}
+WebUI is installed in a directory that has a leading dot (.) in one of its parent directories.
+This will prevent WebUI from functioning properly.
+Please move the installation to a different directory.
+Current path: "{webui_root}"
+For more information see: https://github.com/AUTOMATIC1111/stable-diffusion-webui/issues/13292
+{"!"*25} Warning {"!"*25}''')
+
+
+def _webui():
from modules.shared_cmd_options import cmd_opts
launch_api = cmd_opts.api
- initialize.initialize()
from modules import shared, ui_tempdir, script_callbacks, ui, progress, ui_extra_networks
+ warning_if_invalid_install_dir()
+
while 1:
if shared.opts.clean_temp_dir_at_start:
ui_tempdir.cleanup_tmpdr()
@@ -137,6 +178,7 @@ def webui():
print("Stopping server...")
# If we catch a keyboard interrupt, we want to stop the server and exit.
shared.demo.close()
+ manager.task.stop()
break
# disable auto launch webui in browser for subsequent UI Reload
@@ -153,6 +195,13 @@ def webui():
initialize.initialize_rest(reload_script_modules=True)
+def api_only():
+ Thread(target=_api_only, daemon=True).start()
+
+
+def webui():
+ Thread(target=_webui, daemon=True).start()
+
if __name__ == "__main__":
from modules.shared_cmd_options import cmd_opts
@@ -160,3 +209,5 @@ def webui():
api_only()
else:
webui()
+
+ manager.task.main_loop()