diff --git a/.eslintrc.json b/.eslintrc.json
index 4ca2afbdb..05c467b50 100644
--- a/.eslintrc.json
+++ b/.eslintrc.json
@@ -103,6 +103,7 @@
// progressbar.js
"randomId": "readonly",
"requestProgress": "readonly",
+ "setRefreshInterval": "readonly",
// imageviewer.js
"modalPrevImage": "readonly",
"modalNextImage": "readonly",
diff --git a/.pylintrc b/.pylintrc
index bbe29b197..59585a8c5 100644
--- a/.pylintrc
+++ b/.pylintrc
@@ -40,6 +40,7 @@ ignore-paths=/usr/lib/.*$,
modules/xadapter,
modules/infiniteyou,
modules/flash_attn_triton_amd,
+ scripts/softfill.py,
repositories,
extensions-builtin/Lora,
extensions-builtin/sd-webui-agent-scheduler,
diff --git a/.ruff.toml b/.ruff.toml
index 6c77aa6f3..3057e3067 100644
--- a/.ruff.toml
+++ b/.ruff.toml
@@ -35,6 +35,7 @@ exclude = [
"modules/xadapter",
"modules/infiniteyou",
"modules/flash_attn_triton_amd",
+ "scripts/softfill.py",
"repositories",
"extensions-builtin/Lora",
"extensions-builtin/sd-extension-chainner/nodes",
diff --git a/CHANGELOG.md b/CHANGELOG.md
index ab70b2350..05df2f538 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,16 +1,67 @@
# Change Log for SD.Next
-## Update for 2025-04-04
+## Update for 2025-04-12
-- Video: add FasterCache and PAB support to WanDB and LTX models
-- ZLUDA: add more GPUs to recognized list
-- LoRA: obey configured device when performing calculations
-- Progress: add additional fields to progress API
-- Progress: use batch-count for progress
-- Grid: add of max-rows and max-columns in settings to control grid format
-- Gallery: add max-columns in settings for gradio gallery components
-- Styles: resize and bring quick-ui to forward on hover
-- Logging: fix debug logging
+### Highlights for 2025-04-12
+
+Last release was just over a week ago and here we are again with another update as a new high-end image model, [HiDream-I1](https://github.com/vladmandic/sdnext/wiki/HiDream) jumped out and generated a lot of buzz!
+There are quite a few other performance and quality-of-life improvements in this release and 40 commits, so please take a look at the full [ChangeLog](https://github.com/vladmandic/automatic/blob/master/CHANGELOG.md)
+
+[ReadMe](https://github.com/vladmandic/automatic/blob/master/README.md) | [ChangeLog](https://github.com/vladmandic/automatic/blob/master/CHANGELOG.md) | [Docs](https://vladmandic.github.io/sdnext-docs/) | [WiKi](https://github.com/vladmandic/automatic/wiki) | [Discord](https://discord.com/invite/sd-next-federal-batch-inspectors-1101998836328697867)
+
+### Details for 2025-04-12
+
+- **Models**
+ - [HiDream-I1](https://huggingface.co/HiDream-ai/HiDream-I1-Full) in fast, dev and full variants!
+ new absolutely massive image generative foundation model with **17B** parameters and 4 text-encoders with additional **8.3B** parameters
+ simply select from *networks -> models -> reference*
+ due to size (over 25B params in 58GB), offloading and on-the-fly quantization are pretty much a necessity
+ see [HiDream Wiki page](https://github.com/vladmandic/sdnext/wiki/HiDream) for details
+- **Features**
+ - Custom model loader
+ can be used to load any known diffusion model with default or custom model components
+ in models -> custom tab
+ see docs for details:
+ - Pipe: [SoftFill](https://github.com/zacheryvaughn/softfill-pipelines)
+- **Caching**
+ - add `TeaCache` support to *Flux, CogVideoX, Mochi, LTX*
+ - add `FasterCache` support to *WanAI, LTX* (other video models already supported)
+ - add `PyramidAttentionBroadcast` support to *WanAI, LTX* (other video models already supported)
+- **UI**
+ - client polling speeds up and slows down depending if client page is visible or not
+ client polling does not ask for live preview if page is not visible
+ significantly reduces server load if you hide or minimize the page
+ - progress: use batch-count for progress
+ - grid: add of max-rows and max-columns in settings to control grid format
+ - gallery: add max-columns in settings for gradio gallery components
+- **Other**
+ - ZLUDA: add more GPUs to recognized list
+ select in scripts, available for sdxl in inpaint model
+ - LoRA: add option to force-reload LoRA on every generate
+ - settings: add **Model options** sections as placeholder for per-model settings
+ - video: update *LTXVideo-0.9.5* pipeline
+ - te loader: allow free-form input in which case sdnext will attempt to load it as hf repo
+ - diag: add get-server-status to UI generate context menu
+ - diag: memory monitor detect gpu swapping
+ - use [hf-xet](https://huggingface.co/blog/xet-on-the-hub) for huggingface downloads where possible
+ - quant: update & fix `optimum-quanto` for transformers
+ - quant: update & fix `torchao`
+ - model load: new setting for model load initial device map
+ can be used to force gpu vs cpu when loading model to avoid oom before model offloading is even activated after load
+- **Changes**
+ - params: Reset default guidance-rescale from 0.7 to 0.0
+ - progress: add additional fields to progress API
+- **Fixes**
+ - styles: resize and bring quick-ui to forward on hover
+ - LoRA: obey configured device when performing calculations
+ - ZLUDA: startup issues
+ - offload: balanced offload remove non-blocking move op
+ - logging: debug causes invalid import
+ - logging: cleanup
+ - ROCm: flash attention repo with navi rotary fix
+ - prompt: prompt scheduling with te caching
+ - ui: progress allow for longer timeouts
+ - internal: cleanup defined pipelines
## Update for 2025-04-03
diff --git a/TODO.md b/TODO.md
index c6cef03da..fb1aec1cf 100644
--- a/TODO.md
+++ b/TODO.md
@@ -4,21 +4,20 @@ Main ToDo list can be found at [GitHub projects](https://github.com/users/vladma
## Current
+- ModernUI for custom model loader
+
### Issues/Limitations
N/A
## Future Candidates
-- Flux: NF4 loader:
- IPAdapter: negative guidance:
- Control: API enhance scripts compatibility
- Video: add generate context menu
- Video: API support
- Video: STG:
- Video: SmoothCache: https://github.com/huggingface/diffusers/issues/11135
-- SoftFill: https://github.com/zacheryvaughn/softfill-pipelines
-- SISO: https://github.com/yairshp/SISO
## Code TODO
@@ -31,6 +30,8 @@ N/A
- fc: autodetect distilled based on model
- processing: remove duplicate mask params
- model loader: implement model in-memory caching
+- custom: load receipe
+- custom: save receipe
- hypertile: vae breaks when using non-standard sizes
- model load: force-reloading entire model as loading transformers only leads to massive memory usage
- lora: add other quantization types
diff --git a/cli/zluda-python.py b/cli/zluda-python.py
index 894489b74..f9b7ea203 100644
--- a/cli/zluda-python.py
+++ b/cli/zluda-python.py
@@ -28,7 +28,6 @@ if __name__ == '__main__':
from modules import zluda_installer
zluda_installer.install()
- zluda_installer.make_copy()
zluda_installer.load()
import torch
diff --git a/html/reference.json b/html/reference.json
index b18818750..d7ac2cea5 100644
--- a/html/reference.json
+++ b/html/reference.json
@@ -333,7 +333,29 @@
"preview": "Alpha-VLLM--Lumina-Image-2.0.jpg",
"skip": true,
"extras": "sampler: Default"
- },
+ },
+
+ "HiDream-I1 Fast": {
+ "path": "HiDream-ai/HiDream-I1-Fast",
+ "desc": "HiDream-I1 is a new open-source image generative foundation model with 17B parameters that achieves state-of-the-art image generation quality within seconds.",
+ "preview": "HiDream-ai--HiDream-I1-Fast.jpg",
+ "skip": true,
+ "extras": "sampler: Default"
+ },
+ "HiDream-I1 Dev": {
+ "path": "HiDream-ai/HiDream-I1-Dev",
+ "desc": "HiDream-I1 is a new open-source image generative foundation model with 17B parameters that achieves state-of-the-art image generation quality within seconds.",
+ "preview": "HiDream-ai--HiDream-I1-Fast.jpg",
+ "skip": true,
+ "extras": "sampler: Default"
+ },
+ "HiDream-I1 Full": {
+ "path": "HiDream-ai/HiDream-I1-Full",
+ "desc": "HiDream-I1 is a new open-source image generative foundation model with 17B parameters that achieves state-of-the-art image generation quality within seconds.",
+ "preview": "HiDream-ai--HiDream-I1-Fast.jpg",
+ "skip": true,
+ "extras": "sampler: Default"
+ },
"Kwai Kolors": {
"path": "Kwai-Kolors/Kolors-diffusers",
diff --git a/installer.py b/installer.py
index 2e7ef987d..1750ba745 100644
--- a/installer.py
+++ b/installer.py
@@ -538,7 +538,7 @@ def check_diffusers():
t_start = time.time()
if args.skip_all or args.skip_git or args.experimental:
return
- sha = 'f10775b1b55cbebc58655b966b4ba3a6fc259ca3' # diffusers commit hash
+ sha = '0ef29355c9d65b78eabb6a4ac5bee73aa685e9a6' # diffusers commit hash
pkg = pkg_resources.working_set.by_key.get('diffusers', None)
minor = int(pkg.version.split('.')[1] if pkg is not None else 0)
cur = opts.get('diffusers_version', '') if minor > 0 else ''
@@ -649,7 +649,6 @@ def install_rocm_zluda():
if device is not None and zluda_installer.get_blaslt_enabled():
log.debug(f'ROCm hipBLASLt: arch={device.name} available={device.blaslt_supported}')
zluda_installer.set_blaslt_enabled(device.blaslt_supported)
- zluda_installer.make_copy()
zluda_installer.load()
torch_command = os.environ.get('TORCH_COMMAND', 'torch==2.6.0 torchvision --index-url https://download.pytorch.org/whl/cu118')
except Exception as e:
@@ -786,7 +785,11 @@ def install_torch_addons():
if opts.get('nncf_compress_weights', False) and not args.use_openvino:
install('nncf==2.7.0', 'nncf')
if opts.get('optimum_quanto_weights', False):
- install('optimum-quanto==0.2.6', 'optimum-quanto')
+ install('optimum-quanto==0.2.7', 'optimum-quanto')
+ if opts.get('torchao_quantization', False):
+ install('torchao==0.10.0', 'torchao')
+ if opts.get('samples_format', 'jpg') == 'jxl' or opts.get('grid_format', 'jpg') == 'jxl':
+ install('pillow-jxl-plugin==1.3.2', 'pillow-jxl-plugin')
if not args.experimental:
uninstall('wandb', quiet=True)
ts('addons', t_start)
@@ -1137,8 +1140,9 @@ def install_optional():
install('basicsr')
install('gfpgan')
install('clean-fid')
- install('pillow-jxl-plugin==1.3.1', ignore=True)
- install('optimum-quanto==0.2.6', ignore=True)
+ install('pillow-jxl-plugin==1.3.2', ignore=True)
+ install('optimum-quanto==0.2.7', ignore=True)
+ install('torchao==0.10.0', ignore=True)
install('bitsandbytes==0.45.1', ignore=True)
install('pynvml', ignore=True)
install('ultralytics==8.3.40', ignore=True)
@@ -1487,12 +1491,13 @@ def add_args(parser):
group_diag.add_argument('--test', default=os.environ.get("SD_TEST",False), action='store_true', help="Run test only and exit")
group_diag.add_argument('--version', default=False, action='store_true', help="Print version information")
group_diag.add_argument('--ignore', default=os.environ.get("SD_IGNORE",False), action='store_true', help="Ignore any errors and attempt to continue")
+ group_diag.add_argument("--monitor", default=os.environ.get("SD_MONITOR", 0), help="Run memory monitor, default: %(default)s")
+ group_diag.add_argument("--status", default=os.environ.get("SD_STATUS", 120), help="Run server is-alive status, default: %(default)s")
group_log = parser.add_argument_group('Logging')
group_log.add_argument("--log", type=str, default=os.environ.get("SD_LOG", None), help="Set log file, default: %(default)s")
group_log.add_argument('--debug', default=os.environ.get("SD_DEBUG",False), action='store_true', help="Run installer with debug logging, default: %(default)s")
group_log.add_argument("--profile", default=os.environ.get("SD_PROFILE", False), action='store_true', help="Run profiler, default: %(default)s")
- group_log.add_argument("--monitor", default=os.environ.get("SD_PROFILE", 0), help="Run memory monitor, default: %(default)s")
group_log.add_argument('--docs', default=os.environ.get("SD_DOCS", False), action='store_true', help="Mount API docs, default: %(default)s")
group_log.add_argument("--api-log", default=os.environ.get("SD_APILOG", True), action='store_true', help="Log all API requests")
diff --git a/javascript/contextMenus.js b/javascript/contextMenus.js
index b15615a2a..271f5a9e4 100644
--- a/javascript/contextMenus.js
+++ b/javascript/contextMenus.js
@@ -117,11 +117,35 @@ const reprocessClick = (tabId, state) => {
if (btn) btn.click();
};
+const getStatus = async () => {
+ const headers = new Headers();
+ const body = JSON.stringify({ id_task: -1, id_live_preview: false });
+ headers.set('Content-Type', 'application/json');
+ const tab = getUICurrentTabContent()?.id.replace('tab_', '') || '';
+ const el = gradioApp().querySelector(`#html_log_${tab} .performance p`);
+
+ let res;
+ let data;
+ res = await fetch('./internal/progress', { method: 'POST', headers, body });
+ if (res?.ok) {
+ data = await res.json();
+ log('progressInternal:', data);
+ if (el) el.innerText += '\nProgress internal:\n' + JSON.stringify(data, null, 2); // eslint-disable-line prefer-template
+ }
+ res = await fetch('./sdapi/v1/progress?skip_current_image=true', { method: 'GET', headers });
+ if (res?.ok) {
+ data = await res.json();
+ log('progressAPI:', data);
+ if (el) el.innerText += '\nProgress API:\n' + JSON.stringify(data, null, 2); // eslint-disable-line prefer-template
+ }
+};
+
async function initContextMenu() {
let id = '';
- for (const tab of ['txt2img', 'img2img', 'control']) {
+ for (const tab of ['txt2img', 'img2img', 'control', 'video']) {
id = `#${tab}_generate`;
- appendContextMenuOption(id, 'Copy to clipboard', () => navigator.clipboard.writeText(document.querySelector(`#${tab}_prompt > label > textarea`).value));
+ appendContextMenuOption(id, 'Get server status', getStatus);
+ appendContextMenuOption(id, 'Copy prompt to clipboard', () => navigator.clipboard.writeText(document.querySelector(`#${tab}_prompt > label > textarea`).value));
appendContextMenuOption(id, 'Generate forever', () => generateForever(`#${tab}_generate`));
appendContextMenuOption(id, 'Apply selected style', quickApplyStyle);
appendContextMenuOption(id, 'Quick save style', quickSaveStyle);
diff --git a/javascript/progressBar.js b/javascript/progressBar.js
index 1edcbb501..f0e19e725 100644
--- a/javascript/progressBar.js
+++ b/javascript/progressBar.js
@@ -1,4 +1,15 @@
let lastState = {};
+let refreshInterval = 10000;
+
+function setRefreshInterval() {
+ refreshInterval = opts.live_preview_refresh_period || 500;
+ log('refreshInterval', document.visibilityState, refreshInterval);
+ document.addEventListener('visibilitychange', () => {
+ if (document.hidden) refreshInterval = Math.max(2500, opts.live_preview_refresh_period || 1000);
+ else refreshInterval = opts.live_preview_refresh_period || 1000;
+ log('refreshInterval', document.visibilityState, refreshInterval);
+ });
+}
function pad2(x) {
return x < 10 ? `0${x}` : x;
@@ -122,14 +133,15 @@ function requestProgress(id_task, progressEl, galleryEl, atEnd = null, onProgres
const start = (id_task, id_live_preview) => { // eslint-disable-line no-shadow
if (!opts.live_previews_enable || opts.live_preview_refresh_period === 0 || opts.show_progress_every_n_steps === 0) return;
+ const request_id = document.hidden ? -1 : id_live_preview;
const onProgressHandler = (res) => {
- if (res?.debug) debug('livePreview:', dateStart, res);
+ if (res?.debug) debug('livePreview:', dateStart, request_id, res);
lastState = res;
const elapsedFromStart = (new Date() - dateStart) / 1000;
hasStarted |= res.active;
- if (res.completed || (!res.active && (hasStarted || once)) || (elapsedFromStart > 30 && !res.queued && res.progress === prevProgress)) {
- if (res?.debug) debug('livePreview end:', res);
+ if (res.completed || (!res.active && (hasStarted || once)) || (elapsedFromStart > 120 && !res.queued && res.progress === prevProgress)) {
+ debug('livePreview end:', res);
done();
return;
}
@@ -152,7 +164,7 @@ function requestProgress(id_task, progressEl, galleryEl, atEnd = null, onProgres
done();
};
- xhrPost('./internal/progress', { id_task, id_live_preview }, onProgressHandler, onProgressErrorHandler, false, 30000);
+ xhrPost('./internal/progress', { id_task, id_live_preview: request_id }, onProgressHandler, onProgressErrorHandler, false, 30000);
};
debug('livePreview start:', dateStart);
start(id_task, 0);
diff --git a/javascript/sdnext.css b/javascript/sdnext.css
index 8b52b1a47..5ef9aafd4 100644
--- a/javascript/sdnext.css
+++ b/javascript/sdnext.css
@@ -291,7 +291,11 @@ table.settings-value-table td { padding: 0.4em; border: 1px solid #ccc; max-widt
#pnginfo_html_info .gradio-html > div { margin: 0.5em; }
#models_image, #models_image > div { min-height: 0; }
#models_error { font-family: monospace; color: var(--body-text-color-subdued) }
-
+#model_loader_df button { display: none !important; }
+#model_loader_df table td:first-child { display: none; }
+#model_loader_df table th:first-child { display: none; }
+#model_loader_df table td:nth-child(2) { font-weight: bold; }
+#model_loader_df table td:nth-child(3) { color: pink; }
/* log monitor */
.log-monitor { display: none; justify-content: unset !important; overflow: hidden; padding: 0; margin-top: auto; font-family: monospace; font-size: var(--text-xxs); }
diff --git a/javascript/startup.js b/javascript/startup.js
index 328c167b6..842f925ba 100644
--- a/javascript/startup.js
+++ b/javascript/startup.js
@@ -35,6 +35,7 @@ async function initStartup() {
window.subpath = window.opts.subpath;
window.api = `${window.subpath}/sdapi/v1`;
}
+ setRefreshInterval();
executeCallbacks(uiReadyCallbacks);
initLogMonitor();
setupExtraNetworks();
diff --git a/launch.py b/launch.py
index d00b9ef22..440492586 100755
--- a/launch.py
+++ b/launch.py
@@ -274,7 +274,7 @@ def main():
alive = False
requests = 0
t_current = time.time()
- if t_current - t_server > 120:
+ if float(args.status) > 0 and t_current - t_server > float(args.status):
installer.log.trace(f'Server: alive={alive} requests={requests} memory={get_memory_stats()} {instance.state.status()}')
t_server = t_current
if float(args.monitor) > 0 and t_current - t_monitor > float(args.monitor):
diff --git a/models/Reference/HiDream-ai--HiDream-I1-Fast.jpg b/models/Reference/HiDream-ai--HiDream-I1-Fast.jpg
new file mode 100644
index 000000000..b2816b0b8
Binary files /dev/null and b/models/Reference/HiDream-ai--HiDream-I1-Fast.jpg differ
diff --git a/modules/cmd_args.py b/modules/cmd_args.py
index a8dec6748..d0d7a3217 100644
--- a/modules/cmd_args.py
+++ b/modules/cmd_args.py
@@ -37,7 +37,8 @@ def main_args():
group_diag.add_argument("--no-hashing", default=os.environ.get("SD_NOHASHING", False), action='store_true', help="Disable hashing of checkpoints, default: %(default)s")
group_diag.add_argument("--no-metadata", default=os.environ.get("SD_NOMETADATA", False), action='store_true', help="Disable reading of metadata from models, default: %(default)s")
group_diag.add_argument("--profile", default=os.environ.get("SD_PROFILE", False), action='store_true', help="Run profiler, default: %(default)s")
- group_diag.add_argument("--monitor", default=os.environ.get("SD_PROFILE", 0), help="Run memory monitor, default: %(default)s")
+ group_diag.add_argument("--monitor", default=os.environ.get("SD_MONITOR", 0), help="Run memory monitor, default: %(default)s")
+ group_diag.add_argument("--status", default=os.environ.get("SD_STATUS", 120), help="Run server is-alive status, default: %(default)s")
group_http = parser.add_argument_group('HTTP')
group_http.add_argument('--theme', type=str, default=os.environ.get("SD_THEME", None), help='Override UI theme')
diff --git a/modules/face/photomaker_model_v2.py b/modules/face/photomaker_model_v2.py
index 34704376f..304d708b4 100644
--- a/modules/face/photomaker_model_v2.py
+++ b/modules/face/photomaker_model_v2.py
@@ -323,7 +323,7 @@ class PhotoMakerIDEncoder_CLIPInsightfaceExtendtoken(CLIPVisionModelWithProjecti
self.num_tokens,
)
- def forward(self, id_pixel_values, prompt_embeds, class_tokens_mask, id_embeds): # pylint: disable=arguments-differ
+ def forward(self, id_pixel_values, prompt_embeds, class_tokens_mask, id_embeds): # pylint: disable=arguments-differ, arguments-renamed
b, num_inputs, c, h, w = id_pixel_values.shape
id_pixel_values = id_pixel_values.view(b * num_inputs, c, h, w)
diff --git a/modules/infiniteyou/pipeline_infu_flux.py b/modules/infiniteyou/pipeline_infu_flux.py
index 8ae9f6e95..2b4762d97 100644
--- a/modules/infiniteyou/pipeline_infu_flux.py
+++ b/modules/infiniteyou/pipeline_infu_flux.py
@@ -146,20 +146,19 @@ class InfUFluxPipeline:
self.infu_flux_version = infu_flux_version
self.model_version = model_version
-
- # Load pipeline
+ # Load controlnet
+ shared.log.debug(f'InfiniteYou: cls={shared.sd_model.__class__.__name__} loading')
local_path = snapshot_download(repo_id='ByteDance/InfiniteYou', cache_dir=shared.opts.hfcache_dir)
infiniteyou_path = os.path.join(local_path, f'infu_flux_{infu_flux_version}', model_version)
infusenet_path = os.path.join(infiniteyou_path, 'InfuseNetModel')
- quant_args = model_quant.create_config()
- # quant_args = {}
-
+ quant_args = model_quant.create_config(module='ControlNet')
+ shared.log.debug(f'InfiniteYou: fn="{infusenet_path}" load infusenet')
self.infusenet = FluxControlNetModel.from_pretrained(
infusenet_path,
torch_dtype=devices.dtype,
**quant_args,
)
-
+ # assemble pipeline
self.pipe = FluxInfuseNetPipeline(
vae=pipe.vae,
text_encoder=pipe.text_encoder,
@@ -170,11 +169,10 @@ class InfUFluxPipeline:
scheduler=pipe.scheduler,
controlnet=self.infusenet,
)
-
# Load image proj model
num_tokens = image_proj_num_tokens
image_emb_dim = 512
- image_proj_model = Resampler(
+ self.image_proj_model = Resampler(
dim=1280,
depth=4,
dim_head=64,
@@ -185,16 +183,15 @@ class InfUFluxPipeline:
ff_mult=4,
)
image_proj_model_path = os.path.join(infiniteyou_path, 'image_proj_model.bin')
+ shared.log.debug(f'InfiniteYou: fn="{image_proj_model_path}" load image projection')
ipm_state_dict = torch.load(image_proj_model_path, map_location="cpu")
- image_proj_model.load_state_dict(ipm_state_dict['image_proj'])
+ self.image_proj_model.load_state_dict(ipm_state_dict['image_proj'])
del ipm_state_dict
- image_proj_model.to(device=devices.device, dtype=devices.dtype)
- image_proj_model.eval()
-
- self.image_proj_model = image_proj_model
-
+ self.image_proj_model.to(device=devices.device, dtype=devices.dtype)
+ self.image_proj_model.eval()
# Load face encoder
insightface_root_path = os.path.join(local_path, 'supports', 'insightface')
+ shared.log.debug(f'InfiniteYou: fn="{insightface_root_path}" load face encoder')
self.app_640 = FaceAnalysis(name='antelopev2', root=insightface_root_path, providers=devices.onnx)
self.app_640.prepare(ctx_id=0, det_size=(640, 640))
self.app_320 = FaceAnalysis(name='antelopev2', root=insightface_root_path, providers=devices.onnx)
diff --git a/modules/intel/ipex/diffusers.py b/modules/intel/ipex/diffusers.py
index 7564c7226..8a10aa546 100644
--- a/modules/intel/ipex/diffusers.py
+++ b/modules/intel/ipex/diffusers.py
@@ -42,9 +42,33 @@ class FluxPosEmbed(torch.nn.Module):
return freqs_cos, freqs_sin
+def hidream_rope(pos: torch.Tensor, dim: int, theta: int) -> torch.Tensor:
+ assert dim % 2 == 0, "The dimension must be even."
+ return_device = pos.device
+ pos = pos.to("cpu")
+
+ scale = torch.arange(0, dim, 2, dtype=torch.float64, device=pos.device) / dim
+ omega = 1.0 / (theta**scale)
+
+ batch_size, seq_length = pos.shape
+ out = torch.einsum("...n,d->...nd", pos, omega)
+ cos_out = torch.cos(out)
+ sin_out = torch.sin(out)
+
+ stacked_out = torch.stack([cos_out, -sin_out, sin_out, cos_out], dim=-1)
+ out = stacked_out.view(batch_size, -1, dim // 2, 2, 2)
+ return out.to(return_device, dtype=torch.float32)
+
+
def ipex_diffusers(device_supports_fp64=False, can_allocate_plus_4gb=False):
+ # get around lazy imports
+ from diffusers.utils import torch_utils # pylint: disable=import-error, unused-import
diffusers.utils.torch_utils.fourier_filter = fourier_filter
if not device_supports_fp64:
+ # get around lazy imports
+ from diffusers.models import transformers as diffusers_transformers # pylint: disable=import-error, unused-import
+ from diffusers.models import controlnets as diffusers_controlnets # pylint: disable=import-error, unused-import
diffusers.models.embeddings.FluxPosEmbed = FluxPosEmbed
diffusers.models.transformers.transformer_flux.FluxPosEmbed = FluxPosEmbed
diffusers.models.controlnets.controlnet_flux.FluxPosEmbed = FluxPosEmbed
+ diffusers.models.transformers.transformer_hidream_image.rope = hidream_rope
diff --git a/modules/lora/extra_networks_lora.py b/modules/lora/extra_networks_lora.py
index 0f981d93a..a17914683 100644
--- a/modules/lora/extra_networks_lora.py
+++ b/modules/lora/extra_networks_lora.py
@@ -4,7 +4,7 @@ import re
import numpy as np
from modules.lora import networks, lora_overrides, lora_load
from modules.lora import lora_common as l
-from modules import extra_networks, shared
+from modules import extra_networks, shared, sd_models
debug = os.environ.get('SD_LORA_DEBUG', None) is not None
@@ -139,6 +139,8 @@ class ExtraNetworkLora(extra_networks.ExtraNetwork):
return [f'{name}:{te}:{unet}' for name, te, unet in zip(names, te_multipliers, unet_multipliers)]
def changed(self, requested: List[str], include: List[str], exclude: List[str]):
+ if shared.opts.lora_force_reload:
+ return True
sd_model = getattr(shared.sd_model, "pipe", shared.sd_model)
if not hasattr(sd_model, 'loaded_loras'):
sd_model.loaded_loras = {}
@@ -174,21 +176,24 @@ class ExtraNetworkLora(extra_networks.ExtraNetwork):
if force_diffusers:
has_changed = False # diffusers handle their own loading
if len(exclude) == 0:
- shared.state.begin('LoRA')
+ job = shared.state.job
+ shared.state.job = 'LoRA'
lora_load.network_load(names, te_multipliers, unet_multipliers, dyn_dims) # load only on first call
- shared.state.end()
+ sd_models.set_diffuser_offload(shared.sd_model, op="model")
+ shared.state.job = job
else:
lora_load.network_load(names, te_multipliers, unet_multipliers, dyn_dims) # load
has_changed = self.changed(requested, include, exclude)
if has_changed:
- shared.state.begin('LoRA')
+ job = shared.state.job
+ shared.state.job = 'LoRA'
if len(l.previously_loaded_networks) > 0:
shared.log.info(f'Network unload: type=LoRA apply={[n.name for n in l.previously_loaded_networks]} mode={"fuse" if shared.opts.lora_fuse_diffusers else "backup"}')
networks.network_deactivate(include, exclude)
networks.network_activate(include, exclude)
if len(exclude) > 0: # only update on last activation
l.previously_loaded_networks = l.loaded_networks.copy()
- shared.state.end()
+ shared.state.job = job
debug_log(f'Network load: type=LoRA previous={[n.name for n in l.previously_loaded_networks]} current={[n.name for n in l.loaded_networks]} changed')
if len(l.loaded_networks) > 0 and (len(networks.applied_layers) > 0 or force_diffusers) and step == 0:
diff --git a/modules/lora/lora_apply.py b/modules/lora/lora_apply.py
index 7164543d1..485581a4d 100644
--- a/modules/lora/lora_apply.py
+++ b/modules/lora/lora_apply.py
@@ -118,7 +118,6 @@ def network_add_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn.G
if model_weights is None: # weights are used if provided-from-backup else use self.weight
model_weights = self.weight
weight, new_weight = None, None
- device = device or devices.device
# TODO lora: add other quantization types
if self.__class__.__name__ == 'Linear4bit' and bnb is not None:
try:
@@ -134,7 +133,11 @@ def network_add_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn.G
new_weight = model_weights.to(devices.device) + lora_weights.to(devices.device)
except Exception as e:
shared.log.warning(f'Network load: {e}')
- new_weight = model_weights + lora_weights # try without device cast
+ if 'The size of tensor' in str(e):
+ shared.log.error(f'Network load: type=LoRA model={shared.sd_model.__class__.__name__} incompatible lora shape')
+ new_weight = model_weights
+ else:
+ new_weight = model_weights + lora_weights # try without device cast
weight = torch.nn.Parameter(new_weight.to(device), requires_grad=False)
if weight is not None:
if not bias:
@@ -147,7 +150,6 @@ def network_add_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn.G
def network_apply_direct(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn.GroupNorm, torch.nn.LayerNorm, diffusers.models.lora.LoRACompatibleLinear, diffusers.models.lora.LoRACompatibleConv], updown: torch.Tensor, ex_bias: torch.Tensor, deactivate: bool = False, device: torch.device = devices.device):
weights_backup = getattr(self, "network_weights_backup", False)
bias_backup = getattr(self, "network_bias_backup", False)
- device = device or devices.device
if not isinstance(weights_backup, bool): # remove previous backup if we switched settings
weights_backup = True
if not isinstance(bias_backup, bool):
diff --git a/modules/lora/lora_extract.py b/modules/lora/lora_extract.py
index 1217b952d..597b20d72 100644
--- a/modules/lora/lora_extract.py
+++ b/modules/lora/lora_extract.py
@@ -246,6 +246,8 @@ def create_ui():
return {"visible": visible, "__type__": "update"}
with gr.Tab(label="Extract LoRA"):
+ with gr.Row():
+ gr.HTML(' Extract currently loaded LoRA(s)
')
with gr.Row():
loaded = gr.Textbox(placeholder="Press refresh to query loaded LoRA", label="Loaded LoRA", interactive=False)
create_refresh_button(loaded, lambda: None, lambda: {'value': loaded_lora_str()}, "testid")
diff --git a/modules/lora/lora_load.py b/modules/lora/lora_load.py
index be7b127fc..f67e6b26b 100644
--- a/modules/lora/lora_load.py
+++ b/modules/lora/lora_load.py
@@ -2,7 +2,7 @@ from typing import Union
import os
import time
import concurrent
-from modules import shared, errors, devices, sd_models, sd_models_compile, files_cache
+from modules import shared, errors, sd_models, sd_models_compile, files_cache
from modules.lora import network, lora_overrides, lora_convert
from modules.lora import lora_common as l
@@ -72,7 +72,6 @@ def load_safetensors(name, network_on_disk) -> Union[network.Network, None]:
bundle_embeddings = {}
dtypes = []
convert = lora_convert.KeyConvert()
- device = devices.device if shared.opts.lora_apply_gpu else devices.cpu
for key_network, weight in state_dict.items():
parts = key_network.split('.')
if parts[0] == "bundle_emb":
@@ -116,7 +115,7 @@ def load_safetensors(name, network_on_disk) -> Union[network.Network, None]:
if l.debug:
shared.log.debug(f'Network load: type=LoRA name="{name}" unmatched={keys_failed_to_match}')
else:
- shared.log.debug(f'Network load: type=LoRA name="{name}" type={set(network_types)} keys={len(matched_networks)} device={device} dtypes={dtypes} direct={shared.opts.lora_fuse_diffusers}')
+ shared.log.debug(f'Network load: type=LoRA name="{name}" type={set(network_types)} keys={len(matched_networks)} dtypes={dtypes} direct={shared.opts.lora_fuse_diffusers}')
if len(matched_networks) == 0:
return None
lora_cache[name] = net
diff --git a/modules/lora/networks.py b/modules/lora/networks.py
index a36ba5631..107b6cff8 100644
--- a/modules/lora/networks.py
+++ b/modules/lora/networks.py
@@ -33,13 +33,12 @@ def network_activate(include=[], exclude=[]):
pbar = nullcontext()
applied_weight = 0
applied_bias = 0
- device = devices.device if shared.opts.lora_apply_gpu or shared.opts.diffusers_offload_mode == 'none' else devices.cpu
with devices.inference_context(), pbar:
wanted_names = tuple((x.name, x.te_multiplier, x.unet_multiplier, x.dyn_dim) for x in l.loaded_networks) if len(l.loaded_networks) > 0 else ()
applied_layers.clear()
backup_size = 0
for component in modules.keys():
- orig_device = getattr(sd_model, component, None).device
+ device = getattr(sd_model, component, None).device
for _, module in modules[component]:
network_layer_name = getattr(module, 'network_layer_name', None)
current_names = getattr(module, "network_current_names", ())
@@ -52,7 +51,7 @@ def network_activate(include=[], exclude=[]):
if shared.opts.lora_fuse_diffusers:
network_apply_direct(module, batch_updown, batch_ex_bias, device=device)
else:
- network_apply_weights(module, batch_updown, batch_ex_bias, device=orig_device)
+ network_apply_weights(module, batch_updown, batch_ex_bias, device=device)
if batch_updown is not None or batch_ex_bias is not None:
applied_layers.append(network_layer_name)
applied_weight += 1 if batch_updown is not None else 0
@@ -95,7 +94,6 @@ def network_deactivate(include=[], exclude=[]):
modules[name] = list(component.named_modules())
active_components.append(name)
total = sum(len(x) for x in modules.values())
- device = devices.device if shared.opts.lora_apply_gpu else devices.cpu
if len(l.previously_loaded_networks) > 0 and l.debug:
pbar = rp.Progress(rp.TextColumn('[cyan]Network: type=LoRA action=deactivate'), rp.BarColumn(), rp.TaskProgressColumn(), rp.TimeRemainingColumn(), rp.TimeElapsedColumn(), rp.TextColumn('[cyan]{task.description}'), console=shared.console)
task = pbar.add_task(description='', total=total)
@@ -105,7 +103,7 @@ def network_deactivate(include=[], exclude=[]):
with devices.inference_context(), pbar:
applied_layers.clear()
for component in modules.keys():
- orig_device = getattr(sd_model, component, None).device
+ device = getattr(sd_model, component, None).device
for _, module in modules[component]:
network_layer_name = getattr(module, 'network_layer_name', None)
if shared.state.interrupted or network_layer_name is None:
@@ -116,7 +114,7 @@ def network_deactivate(include=[], exclude=[]):
if shared.opts.lora_fuse_diffusers:
network_apply_direct(module, batch_updown, batch_ex_bias, device=device, deactivate=True)
else:
- network_apply_weights(module, batch_updown, batch_ex_bias, device=orig_device, deactivate=True)
+ network_apply_weights(module, batch_updown, batch_ex_bias, device=device, deactivate=True)
if batch_updown is not None or batch_ex_bias is not None:
applied_layers.append(network_layer_name)
del batch_updown, batch_ex_bias
diff --git a/modules/memstats.py b/modules/memstats.py
index 90512e870..f62645fca 100644
--- a/modules/memstats.py
+++ b/modules/memstats.py
@@ -69,6 +69,7 @@ def memory_stats():
'retries': stats.get('num_alloc_retries', 0),
'oom': stats.get('num_ooms', 0),
})
+ mem['swap'] = round(mem['active'] - mem['gpu']['used'], 2) if mem['active'] > mem['gpu']['used'] else 0
return mem
except Exception:
pass
diff --git a/modules/model_cogview.py b/modules/model_cogview.py
index a37dd88ca..c6b34fdc3 100644
--- a/modules/model_cogview.py
+++ b/modules/model_cogview.py
@@ -29,16 +29,16 @@ def load_cogview3(checkpoint_info, diffusers_load_config={}):
repo_id = sd_models.path_to_repo(checkpoint_info.name)
shared.log.debug(f'Load model: type=CogView3 model="{checkpoint_info.name}" repo="{repo_id}" offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype}')
- diffusers_load_config, quant_args = load_common(diffusers_load_config, module='Model')
+ load_args, quant_args = load_common(diffusers_load_config, module='Transformer')
transformer = diffusers.CogView3PlusTransformer2DModel.from_pretrained(
repo_id,
subfolder="transformer",
cache_dir=shared.opts.diffusers_dir,
- **diffusers_load_config,
+ **load_args,
**quant_args,
)
- diffusers_load_config, quant_args = load_common(diffusers_load_config, module='TE')
+ load_args, quant_args = load_common(diffusers_load_config, module='TE')
text_encoder = transformers.T5EncoderModel.from_pretrained(
repo_id,
subfolder="text_encoder",
@@ -47,12 +47,13 @@ def load_cogview3(checkpoint_info, diffusers_load_config={}):
**quant_args,
)
+ load_args, quant_args = load_common(diffusers_load_config, module='Transformer')
pipe = diffusers.CogView3PlusPipeline.from_pretrained(
repo_id,
text_encoder=text_encoder,
transformer=transformer,
cache_dir=shared.opts.diffusers_dir,
- **diffusers_load_config,
+ **load_args,
)
devices.torch_gc()
return pipe
@@ -62,7 +63,7 @@ def load_cogview4(checkpoint_info, diffusers_load_config={}):
repo_id = sd_models.path_to_repo(checkpoint_info.name)
shared.log.debug(f'Load model: type=CogView4 model="{checkpoint_info.name}" repo="{repo_id}" offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype}')
- diffusers_load_config, quant_args = load_common(diffusers_load_config, module='Model')
+ load_args, quant_args = load_common(diffusers_load_config, module='Transformer')
transformer = diffusers.CogView4Transformer2DModel.from_pretrained(
repo_id,
subfolder="transformer",
@@ -71,21 +72,22 @@ def load_cogview4(checkpoint_info, diffusers_load_config={}):
**quant_args,
)
- diffusers_load_config, quant_args = load_common(diffusers_load_config, module='TE')
+ load_args, quant_args = load_common(diffusers_load_config, module='TE')
text_encoder = transformers.AutoModelForCausalLM.from_pretrained(
repo_id,
subfolder="text_encoder",
cache_dir=shared.opts.diffusers_dir,
- **diffusers_load_config,
+ **load_args,
**quant_args,
)
+ load_args, quant_args = load_common(diffusers_load_config, module='Model')
pipe = diffusers.CogView4Pipeline.from_pretrained(
repo_id,
text_encoder=text_encoder,
transformer=transformer,
cache_dir=shared.opts.diffusers_dir,
- **diffusers_load_config,
+ **load_args,
)
if shared.opts.diffusers_eval:
pipe.text_encoder.eval()
diff --git a/modules/model_flux.py b/modules/model_flux.py
index 3dfe83ff2..7f1ae035d 100644
--- a/modules/model_flux.py
+++ b/modules/model_flux.py
@@ -223,6 +223,11 @@ def load_flux(checkpoint_info, diffusers_load_config): # triggered by opts.sd_ch
shared.sd_model = None
devices.torch_gc(force=True)
+ if shared.opts.teacache_enabled:
+ from modules import teacache
+ shared.log.debug(f'Transformers cache: type=teacache patch=forward cls={diffusers.FluxTransformer2DModel.__name__}')
+ diffusers.FluxTransformer2DModel.forward = teacache.teacache_flux_forward
+
# load overrides if any
if shared.opts.sd_unet != 'Default':
try:
diff --git a/modules/model_hidream.py b/modules/model_hidream.py
new file mode 100644
index 000000000..c8f4d842b
--- /dev/null
+++ b/modules/model_hidream.py
@@ -0,0 +1,82 @@
+import os
+import time
+import transformers
+import diffusers
+from modules import shared, devices, sd_models, timer, model_quant, modelloader
+
+
+def hijack_encode_prompt(*args, **kwargs):
+ t0 = time.time()
+ if 'max_sequence_length' in kwargs:
+ kwargs['max_sequence_length'] = os.environ.get('HIDREAM_MAX_SEQUENCE_LENGTH', 256)
+ res = shared.sd_model.orig_encode_prompt(*args, **kwargs)
+ t1 = time.time()
+ timer.process.add('te', t1-t0)
+ # shared.log.debug(f'Hijack: te={shared.sd_model.text_encoder.__class__.__name__} time={t1-t0:.2f}')
+ shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model)
+ return res
+
+
+def load_hidream(checkpoint_info, diffusers_load_config={}):
+ modelloader.hf_login()
+ repo_id = sd_models.path_to_repo(checkpoint_info.name)
+
+ load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='Transformer', device_map=True)
+ shared.log.debug(f'Load model: type=HiDream transformer="{repo_id}" quant="{model_quant.get_quant_type(quant_args)}" args={load_args}')
+ transformer = diffusers.HiDreamImageTransformer2DModel.from_pretrained(
+ repo_id,
+ subfolder="transformer",
+ cache_dir=shared.opts.hfcache_dir,
+ **load_args,
+ **quant_args,
+ )
+ if shared.opts.diffusers_offload_mode != 'none':
+ transformer = transformer.to(devices.cpu)
+
+ load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='TE', device_map=True)
+ shared.log.debug(f'Load model: type=HiDream te3="{repo_id}" quant="{model_quant.get_quant_type(quant_args)}" args={load_args}')
+ text_encoder_3 = transformers.T5EncoderModel.from_pretrained(
+ repo_id,
+ subfolder="text_encoder_3",
+ cache_dir=shared.opts.hfcache_dir,
+ **load_args,
+ **quant_args,
+ )
+ if shared.opts.diffusers_offload_mode != 'none':
+ text_encoder_3 = text_encoder_3.to(devices.cpu)
+
+ load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='LLM', device_map=True)
+ shared.log.debug(f'Load model: type=HiDream te4="{shared.opts.model_h1_llama_repo}" quant="{model_quant.get_quant_type(quant_args)}" args={load_args}')
+ tokenizer_4 = transformers.PreTrainedTokenizerFast.from_pretrained(
+ shared.opts.model_h1_llama_repo,
+ cache_dir=shared.opts.hfcache_dir,
+ **load_args,
+ )
+ text_encoder_4 = transformers.LlamaForCausalLM.from_pretrained(
+ shared.opts.model_h1_llama_repo,
+ output_hidden_states=True,
+ output_attentions=True,
+ cache_dir=shared.opts.hfcache_dir,
+ **load_args,
+ **quant_args,
+ )
+ if shared.opts.diffusers_offload_mode != 'none':
+ text_encoder_4 = text_encoder_4.to(devices.cpu)
+
+ load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config, module='Model')
+ shared.log.debug(f'Load model: type=HiDream model="{checkpoint_info.name}" repo="{repo_id}" offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}')
+ pipe = diffusers.HiDreamImagePipeline.from_pretrained(
+ repo_id,
+ text_encoder_3=text_encoder_3,
+ text_encoder_4=text_encoder_4,
+ tokenizer_4=tokenizer_4,
+ transformer=transformer,
+ cache_dir=shared.opts.diffusers_dir,
+ **load_args,
+ )
+
+ pipe.orig_encode_prompt = pipe.encode_prompt
+ pipe.encode_prompt = hijack_encode_prompt
+
+ devices.torch_gc()
+ return pipe
diff --git a/modules/model_lumina.py b/modules/model_lumina.py
index f19fcd7da..f9d3b9abd 100644
--- a/modules/model_lumina.py
+++ b/modules/model_lumina.py
@@ -3,23 +3,13 @@ import diffusers
def load_lumina(_checkpoint_info, diffusers_load_config={}):
- from modules import shared, devices, modelloader
+ from modules import shared, devices, modelloader, model_quant
modelloader.hf_login()
- # {'low_cpu_mem_usage': True, 'torch_dtype': torch.float16, 'load_connected_pipeline': True, 'safety_checker': None, 'requires_safety_checker': False}
- if 'torch_dtype' not in diffusers_load_config:
- diffusers_load_config['torch_dtype'] = 'torch.float16'
- if 'low_cpu_mem_usage' in diffusers_load_config:
- del diffusers_load_config['low_cpu_mem_usage']
- if 'load_connected_pipeline' in diffusers_load_config:
- del diffusers_load_config['load_connected_pipeline']
- if 'safety_checker' in diffusers_load_config:
- del diffusers_load_config['safety_checker']
- if 'requires_safety_checker' in diffusers_load_config:
- del diffusers_load_config['requires_safety_checker']
+ load_config, _quant_config = model_quant.get_dit_args(diffusers_load_config, allow_quant=False)
pipe = diffusers.LuminaText2ImgPipeline.from_pretrained(
'Alpha-VLLM/Lumina-Next-SFT-diffusers',
cache_dir = shared.opts.diffusers_dir,
- **diffusers_load_config,
+ **load_config,
)
devices.torch_gc(force=True)
return pipe
@@ -27,18 +17,35 @@ def load_lumina(_checkpoint_info, diffusers_load_config={}):
def load_lumina2(checkpoint_info, diffusers_load_config={}):
from modules import shared, devices, sd_models, model_quant
- quant_args = {}
- quant_args = model_quant.create_bnb_config(quant_args)
- if quant_args:
- model_quant.load_bnb(f'Load model: type=Lumina quant={quant_args}')
- if not quant_args:
- quant_args = model_quant.create_config()
- kwargs = {}
repo_id = sd_models.path_to_repo(checkpoint_info.name)
- if (('Model' in shared.opts.bnb_quantization or 'Model' in shared.opts.torchao_quantization or 'Model' in shared.opts.quanto_quantization) or ('Transformer' in shared.opts.bnb_quantization or 'Transformer' in shared.opts.torchao_quantization or 'Transformer' in shared.opts.quanto_quantization)):
- kwargs['transformer'] = diffusers.Lumina2Transformer2DModel.from_pretrained(repo_id, subfolder="transformer", cache_dir=shared.opts.diffusers_dir, torch_dtype=devices.dtype, **quant_args)
- if ('TE' in shared.opts.bnb_quantization or 'TE' in shared.opts.torchao_quantization or 'TE' in shared.opts.quanto_quantization):
- kwargs['text_encoder'] = transformers.AutoModel.from_pretrained(repo_id, subfolder="text_encoder", cache_dir=shared.opts.diffusers_dir, torch_dtype=devices.dtype, **quant_args)
- sd_model = diffusers.Lumina2Text2ImgPipeline.from_pretrained(checkpoint_info.path, cache_dir=shared.opts.diffusers_dir, **diffusers_load_config, **quant_args, **kwargs)
+
+ load_config, quant_config = model_quant.get_dit_args(diffusers_load_config, module='Transformer')
+ transformer = diffusers.Lumina2Transformer2DModel.from_pretrained(
+ repo_id,
+ subfolder="transformer",
+ cache_dir=shared.opts.hfcache_dir,
+ **load_config,
+ **quant_config,
+ )
+
+ load_config, quant_config = model_quant.get_dit_args(diffusers_load_config, module='TE', device_map=True)
+ text_encoder = transformers.AutoModel.from_pretrained(
+ repo_id,
+ subfolder="text_encoder",
+ cache_dir=shared.opts.hfcache_dir,
+ torch_dtype=devices.dtype,
+ **load_config,
+ **quant_config,
+ )
+
+ load_config, quant_config = model_quant.get_dit_args(diffusers_load_config, allow_quant=False)
+ pipe = diffusers.Lumina2Text2ImgPipeline.from_pretrained(
+ repo_id,
+ cache_dir=shared.opts.diffusers_dir,
+ text_encoder=text_encoder,
+ transformer=transformer,
+ **load_config,
+ )
+
devices.torch_gc(force=True)
- return sd_model
+ return pipe
diff --git a/modules/model_meissonic.py b/modules/model_meissonic.py
index d705a32d9..2be5c3bbe 100644
--- a/modules/model_meissonic.py
+++ b/modules/model_meissonic.py
@@ -3,12 +3,13 @@ import diffusers
def load_meissonic(checkpoint_info, diffusers_load_config={}):
- from modules import shared, devices, modelloader, sd_models
+ from modules import shared, devices, modelloader, sd_models, shared_items
from modules.meissonic.transformer import Transformer2DModel as TransformerMeissonic
from modules.meissonic.scheduler import Scheduler as MeissonicScheduler
from modules.meissonic.pipeline import Pipeline as PipelineMeissonic
from modules.meissonic.pipeline_img2img import Img2ImgPipeline as PipelineMeissonicImg2Img
from modules.meissonic.pipeline_inpaint import InpaintPipeline as PipelineMeissonicInpaint
+ shared_items.pipelines['Meissonic'] = PipelineMeissonic
modelloader.hf_login()
fn = sd_models.path_to_repo(checkpoint_info.path)
@@ -16,11 +17,29 @@ def load_meissonic(checkpoint_info, diffusers_load_config={}):
diffusers_load_config['variant'] = 'fp16'
diffusers_load_config['trust_remote_code'] = True
- model = TransformerMeissonic.from_pretrained(fn, subfolder="transformer", cache_dir=cache_dir, **diffusers_load_config)
- vqvae = diffusers.VQModel.from_pretrained(fn, subfolder="vqvae", cache_dir=cache_dir, **diffusers_load_config)
- text_encoder = transformers.CLIPTextModelWithProjection.from_pretrained(fn, subfolder="text_encoder", cache_dir=cache_dir)
- # text_encoder = transformers.CLIPTextModelWithProjection.from_pretrained("laion/CLIP-ViT-H-14-laion2B-s32B-b79K", cache_dir=cache_dir)
- tokenizer = transformers.CLIPTokenizer.from_pretrained(fn, subfolder="tokenizer", cache_dir=cache_dir)
+
+ model = TransformerMeissonic.from_pretrained(
+ fn,
+ subfolder="transformer",
+ cache_dir=cache_dir,
+ **diffusers_load_config,
+ )
+ vqvae = diffusers.VQModel.from_pretrained(
+ fn,
+ subfolder="vqvae",
+ cache_dir=cache_dir,
+ **diffusers_load_config,
+ )
+ text_encoder = transformers.CLIPTextModelWithProjection.from_pretrained(
+ fn,
+ subfolder="text_encoder",
+ cache_dir=cache_dir,
+ )
+ tokenizer = transformers.CLIPTokenizer.from_pretrained(
+ fn,
+ subfolder="tokenizer",
+ cache_dir=cache_dir,
+ )
scheduler = MeissonicScheduler.from_pretrained(fn, subfolder="scheduler", cache_dir=cache_dir)
pipe = PipelineMeissonic(
vqvae=vqvae.to(devices.dtype),
diff --git a/modules/model_omnigen.py b/modules/model_omnigen.py
index b7b6e3546..b7eb4684e 100644
--- a/modules/model_omnigen.py
+++ b/modules/model_omnigen.py
@@ -1,9 +1,10 @@
def load_omnigen(checkpoint_info, diffusers_load_config={}): # pylint: disable=unused-argument
- from modules import shared, devices, sd_models
+ from modules import shared, devices, sd_models, shared_items
repo_id = sd_models.path_to_repo(checkpoint_info.name)
# load
from modules.omnigen import OmniGenPipeline
+ shared_items.pipelines['OmniGen'] = OmniGenPipeline
pipe = OmniGenPipeline.from_pretrained(
model_name=repo_id,
vae_path='madebyollin/sdxl-vae-fp16-fix',
diff --git a/modules/model_pixart.py b/modules/model_pixart.py
index 0757a1216..58d3dbfa7 100644
--- a/modules/model_pixart.py
+++ b/modules/model_pixart.py
@@ -1,30 +1,36 @@
+import transformers
import diffusers
def load_pixart(checkpoint_info, diffusers_load_config={}):
- from modules import shared, devices, modelloader, model_te
+ from modules import shared, devices, modelloader, sd_models, model_quant
modelloader.hf_login()
- # shared.opts.data['cuda_dtype'] = 'FP32' # override
- # shared.opts.data['diffusers_offload_mode}'] = "model" # override
- # devices.set_cuda_params()
- fn = checkpoint_info.path.replace('huggingface/', '')
- t5 = model_te.load_t5(name=shared.opts.sd_text_encoder, cache_dir=shared.opts.diffusers_dir)
+ repo_id = sd_models.path_to_repo(checkpoint_info.name)
+
+ load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='Transformer')
transformer = diffusers.PixArtTransformer2DModel.from_pretrained(
- fn,
- subfolder = 'transformer',
- cache_dir = shared.opts.diffusers_dir,
- **diffusers_load_config,
+ repo_id,
+ subfolder='transformer',
+ cache_dir=shared.opts.hfcache_dir,
+ **load_args,
+ **quant_args,
)
- transformer.to(devices.device)
- kwargs = { 'transformer': transformer }
- if t5 is not None:
- kwargs['text_encoder'] = t5
- diffusers_load_config.pop('variant', None)
+ load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='TE', device_map=True)
+ text_encoder = transformers.T5EncoderModel.from_pretrained(
+ repo_id,
+ subfolder="text_encoder",
+ cache_dir=shared.opts.hfcache_dir,
+ **load_args,
+ **quant_args,
+ )
+
+ load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config, allow_quant=False)
pipe = diffusers.PixArtSigmaPipeline.from_pretrained(
'PixArt-alpha/PixArt-Sigma-XL-2-1024-MS',
- cache_dir = shared.opts.diffusers_dir,
- **kwargs,
- **diffusers_load_config,
+ cache_dir=shared.opts.diffusers_dir,
+ transformer=transformer,
+ text_encoder=text_encoder,
+ **load_args,
)
devices.torch_gc(force=True)
return pipe
diff --git a/modules/model_quant.py b/modules/model_quant.py
index 40a2ce3f8..a32a8e353 100644
--- a/modules/model_quant.py
+++ b/modules/model_quant.py
@@ -3,6 +3,7 @@ import sys
import copy
import time
import diffusers
+import transformers
from installer import installed, install, log, setup_logging
@@ -15,6 +16,12 @@ quant_last_model_device = None
debug = os.environ.get('SD_QUANT_DEBUG', None) is not None
+def get_quant_type(args):
+ if args is not None and "quantization_config" in args:
+ return args['quantization_config'].__class__.__name__
+ return None
+
+
def get_quant(name):
if "qint8" in name.lower():
return 'qint8'
@@ -34,7 +41,7 @@ def get_quant(name):
def create_bnb_config(kwargs = None, allow_bnb: bool = True, module: str = 'Model'):
from modules import shared, devices
if len(shared.opts.bnb_quantization) > 0 and allow_bnb:
- if 'Model' in shared.opts.bnb_quantization or (module is not None and module in shared.opts.bnb_quantization):
+ if 'Model' in shared.opts.bnb_quantization or (module is not None and module in shared.opts.bnb_quantization) or module == 'any':
load_bnb()
if bnb is None:
return kwargs
@@ -56,12 +63,15 @@ def create_bnb_config(kwargs = None, allow_bnb: bool = True, module: str = 'Mode
def create_ao_config(kwargs = None, allow_ao: bool = True, module: str = 'Model'):
from modules import shared
- if len(shared.opts.torchao_quantization) > 0 and shared.opts.torchao_quantization_mode == 'pre' and allow_ao:
- if 'Model' in shared.opts.torchao_quantization or (module is not None and module in shared.opts.torchao_quantization):
- load_torchao()
- if ao is None:
+ if len(shared.opts.torchao_quantization) > 0 and (shared.opts.torchao_quantization_mode == 'pre') and allow_ao:
+ if 'Model' in shared.opts.torchao_quantization or (module is not None and module in shared.opts.torchao_quantization) or module == 'any':
+ torchao = load_torchao()
+ if torchao is None:
return kwargs
- ao_config = diffusers.TorchAoConfig(shared.opts.torchao_quantization_type)
+ if module in {'TE', 'LLM'}:
+ ao_config = transformers.TorchAoConfig(quant_type=shared.opts.torchao_quantization_type)
+ else:
+ ao_config = diffusers.TorchAoConfig(shared.opts.torchao_quantization_type)
log.debug(f'Quantization: module="{module}" type=torchao dtype={shared.opts.torchao_quantization_type}')
if kwargs is None:
return ao_config
@@ -74,14 +84,13 @@ def create_ao_config(kwargs = None, allow_ao: bool = True, module: str = 'Model'
def create_quanto_config(kwargs = None, allow_quanto: bool = True, module: str = 'Model'):
from modules import shared
if len(shared.opts.quanto_quantization) > 0 and allow_quanto:
- if 'Model' in shared.opts.quanto_quantization or (module is not None and module in shared.opts.quanto_quantization):
+ if 'Model' in shared.opts.quanto_quantization or (module is not None and module in shared.opts.quanto_quantization) or module == 'any':
load_quanto(silent=True)
if optimum_quanto is None:
return kwargs
- quanto_config = diffusers.QuantoConfig(
- weights_dtype=shared.opts.quanto_quantization_type,
- )
+ quanto_config = diffusers.QuantoConfig(weights_dtype=shared.opts.quanto_quantization_type)
quanto_config.activations = None # patch so it works with transformers
+ quanto_config.weights = quanto_config.weights_dtype
log.debug(f'Quantization: module="{module}" type=quanto dtype={shared.opts.quanto_quantization_type}')
if kwargs is None:
return quanto_config
@@ -117,7 +126,7 @@ def load_torchao(msg='', silent=False):
if ao is not None:
return ao
if not installed('torchao'):
- install('torchao==0.8.0', quiet=True)
+ install('torchao==0.10.0', quiet=True)
log.warning('Quantization: torchao installed please restart')
try:
import torchao
@@ -174,6 +183,8 @@ def load_quanto(msg='', silent=False):
log.warning('Quantization: optimum-quanto installed please restart')
try:
from optimum import quanto # pylint: disable=no-name-in-module
+ # disable device specific tensors because the model can't be moved between cpu and gpu with them
+ quanto.tensor.weights.qbits.WeightQBitsTensor.create = lambda *args, **kwargs: quanto.tensor.weights.qbits.WeightQBitsTensor(*args, **kwargs)
optimum_quanto = quanto
fn = f'{sys._getframe(3).f_code.co_name}:{sys._getframe(2).f_code.co_name}:{sys._getframe(1).f_code.co_name}' # pylint: disable=protected-access
log.debug(f'Quantization: type=quanto version={quanto.__version__} fn={fn}') # pylint: disable=protected-access
@@ -372,7 +383,6 @@ def optimum_quanto_weights(sd_model):
log.info(f"Quantization: type=Optimum.quanto: modules={shared.opts.optimum_quanto_weights}")
global quant_last_model_name, quant_last_model_device # pylint: disable=global-statement
quanto = load_quanto()
- quanto.tensor.qbits.QBitsTensor.create = lambda *args, **kwargs: quanto.tensor.qbits.QBitsTensor(*args, **kwargs)
sd_model = sd_models.apply_function_to_model(sd_model, optimum_quanto_model, shared.opts.optimum_quanto_weights, op="optimum-quanto")
if quant_last_model_name is not None:
@@ -445,3 +455,32 @@ def torchao_quantization(sd_model):
log.error(f"Quantization: type=TorchAO {e}")
setup_logging() # torchao uses dynamo which messes with logging so reset is needed
return sd_model
+
+
+def get_dit_args(load_config:dict={}, module:str=None, device_map:bool=False, allow_quant:bool=True):
+ from modules import shared, devices
+ config = load_config.copy()
+ if 'torch_dtype' not in config:
+ config['torch_dtype'] = devices.dtype
+ if 'low_cpu_mem_usage' in config:
+ del config['low_cpu_mem_usage']
+ if 'load_connected_pipeline' in config:
+ del config['load_connected_pipeline']
+ if 'safety_checker' in config:
+ del config['safety_checker']
+ if 'requires_safety_checker' in config:
+ del config['requires_safety_checker']
+ if 'variant' in config:
+ del config['variant']
+ if device_map:
+ if shared.opts.device_map == 'cpu':
+ config['device_map'] = 'cpu'
+ if shared.opts.device_map == 'gpu':
+ config['device_map'] = devices.device
+ if devices.backend == "ipex" and os.environ.get('UR_L0_ENABLE_RELAXED_ALLOCATION_LIMITS', '0') != '1' and module in {'TE', 'LLM'}:
+ config['device_map'] = 'cpu' # alchemist gpus hits the 4GB allocation limit with transformers, UR_L0_ENABLE_RELAXED_ALLOCATION_LIMITS emulates above 4GB allocations
+ if allow_quant:
+ quant_args = create_config(module=module)
+ else:
+ quant_args = {}
+ return config, quant_args
diff --git a/modules/model_sana.py b/modules/model_sana.py
index 7f39f17e0..c2bc39119 100644
--- a/modules/model_sana.py
+++ b/modules/model_sana.py
@@ -20,9 +20,9 @@ def load_quants(kwargs, repo_id, cache_dir):
def load_sana(checkpoint_info, kwargs={}):
modelloader.hf_login()
-
fn = checkpoint_info if isinstance(checkpoint_info, str) else checkpoint_info.path
repo_id = sd_models.path_to_repo(fn)
+
kwargs.pop('load_connected_pipeline', None)
kwargs.pop('safety_checker', None)
kwargs.pop('requires_safety_checker', None)
diff --git a/modules/model_sd3.py b/modules/model_sd3.py
index e3774b291..c2ac92143 100644
--- a/modules/model_sd3.py
+++ b/modules/model_sd3.py
@@ -1,7 +1,7 @@
import os
import diffusers
import transformers
-from modules import shared, devices, sd_models, sd_unet, model_quant, model_tools
+from modules import shared, devices, errors, sd_models, sd_unet, model_quant, model_tools
def load_overrides(kwargs, cache_dir):
@@ -14,14 +14,15 @@ def load_overrides(kwargs, cache_dir):
shared.log.debug(f'Load model: type=SD3 unet="{shared.opts.sd_unet}" fmt=safetensors')
elif fn.endswith('.gguf'):
from modules import ggml
- # kwargs = load_gguf(kwargs, fn)
kwargs['transformer'] = ggml.load_gguf(fn, cls=diffusers.SD3Transformer2DModel, compute_dtype=devices.dtype)
sd_unet.loaded_unet = shared.opts.sd_unet
shared.log.debug(f'Load model: type=SD3 unet="{shared.opts.sd_unet}" fmt=gguf')
except Exception as e:
shared.log.error(f"Load model: type=SD3 failed to load UNet: {e}")
+ errors.display(e, 'UNet')
shared.opts.sd_unet = 'Default'
sd_unet.failed_unet.append(shared.opts.sd_unet)
+
if shared.opts.sd_text_encoder != 'Default':
try:
from modules.model_te import load_t5, load_vit_l, load_vit_g
@@ -36,7 +37,9 @@ def load_overrides(kwargs, cache_dir):
shared.log.debug(f'Load model: type=SD3 variant="t5" te="{shared.opts.sd_text_encoder}"')
except Exception as e:
shared.log.error(f"Load model: type=SD3 failed to load T5: {e}")
+ errors.display(e, 'TE')
shared.opts.sd_text_encoder = 'Default'
+
if shared.opts.sd_vae != 'Default' and shared.opts.sd_vae != 'Automatic':
try:
from modules import sd_vae
@@ -47,17 +50,17 @@ def load_overrides(kwargs, cache_dir):
shared.log.debug(f'Load model: type=SD3 vae="{shared.opts.sd_vae}"')
except Exception as e:
shared.log.error(f"Load model: type=SD3 failed to load VAE: {e}")
+ errors.display(e, 'VAE')
shared.opts.sd_vae = 'Default'
return kwargs
def load_quants(kwargs, repo_id, cache_dir):
- quant_args = model_quant.create_config()
- if not quant_args:
- return kwargs
- if 'transformer' not in kwargs and (('Model' in shared.opts.bnb_quantization or 'Model' in shared.opts.torchao_quantization or 'Model' in shared.opts.quanto_quantization) or ('Transformer' in shared.opts.bnb_quantization or 'Transformer' in shared.opts.torchao_quantization or 'Transformer' in shared.opts.quanto_quantization)):
+ quant_args = model_quant.create_config(module='Transformer')
+ if quant_args and 'quantization_config' in quant_args:
kwargs['transformer'] = diffusers.SD3Transformer2DModel.from_pretrained(repo_id, subfolder="transformer", cache_dir=cache_dir, torch_dtype=devices.dtype, **quant_args)
- if 'text_encoder_3' not in kwargs and ('TE' in shared.opts.bnb_quantization or 'TE' in shared.opts.torchao_quantization or 'TE' in shared.opts.quanto_quantization):
+ quant_args = model_quant.create_config(module='TE')
+ if quant_args and 'quantization_config' in quant_args:
kwargs['text_encoder_3'] = transformers.T5EncoderModel.from_pretrained(repo_id, subfolder="text_encoder_3", variant='fp16', cache_dir=cache_dir, torch_dtype=devices.dtype, **quant_args)
return kwargs
@@ -76,52 +79,19 @@ def load_missing(kwargs, fn, cache_dir):
kwargs['text_encoder_2'] = transformers.CLIPTextModelWithProjection.from_pretrained(repo_id, subfolder='text_encoder_2', cache_dir=cache_dir, torch_dtype=devices.dtype)
shared.log.debug(f'Load model: type=SD3 missing=te2 repo="{repo_id}"')
if 'text_encoder_3' not in kwargs and 'text_encoder_3' not in keys:
- kwargs['text_encoder_3'] = transformers.T5EncoderModel.from_pretrained(repo_id, subfolder="text_encoder_3", variant='fp16', cache_dir=cache_dir, torch_dtype=devices.dtype)
+ load_args, quant_args = model_quant.get_dit_args({}, module='TE', device_map=True)
+ kwargs['text_encoder_3'] = transformers.T5EncoderModel.from_pretrained(repo_id, subfolder="text_encoder_3", variant='fp16', cache_dir=cache_dir, **load_args, **quant_args)
shared.log.debug(f'Load model: type=SD3 missing=te3 repo="{repo_id}"')
if 'vae' not in kwargs and 'vae' not in keys:
kwargs['vae'] = diffusers.AutoencoderKL.from_pretrained(repo_id, subfolder='vae', cache_dir=cache_dir, torch_dtype=devices.dtype)
shared.log.debug(f'Load model: type=SD3 missing=vae repo="{repo_id}"')
- # if 'transformer' not in kwargs and 'transformer' not in keys:
- # kwargs['transformer'] = diffusers.SD3Transformer2DModel.from_pretrained(default_repo_id, subfolder="transformer", cache_dir=cache_dir, torch_dtype=devices.dtype)
return kwargs
-"""
-def load_gguf(kwargs, fn):
- ggml.install_gguf()
- from accelerate import init_empty_weights
- from diffusers.loaders.single_file_utils import convert_sd3_transformer_checkpoint_to_diffusers
- from modules import ggml, sd_hijack_accelerate
- with init_empty_weights():
- config = diffusers.SD3Transformer2DModel.load_config(os.path.join('configs', 'flux'), subfolder="transformer")
- transformer = diffusers.SD3Transformer2DModel.from_config(config).to(devices.dtype)
- expected_state_dict_keys = list(transformer.state_dict().keys())
- state_dict, stats = ggml.load_gguf_state_dict(fn, devices.dtype)
- state_dict = convert_sd3_transformer_checkpoint_to_diffusers(state_dict)
- applied, skipped = 0, 0
- for param_name, param in state_dict.items():
- if param_name not in expected_state_dict_keys:
- skipped += 1
- continue
- applied += 1
- sd_hijack_accelerate.hijack_set_module_tensor_simple(transformer, tensor_name=param_name, value=param, device=0)
- transformer.gguf = 'gguf'
- state_dict[param_name] = None
- shared.log.debug(f'Load model: type=Unet/Transformer applied={applied} skipped={skipped} stats={stats} compute={devices.dtype}')
- kwargs['transformer'] = transformer
- return kwargs
-"""
-
-
def load_sd3(checkpoint_info, cache_dir=None, config=None):
repo_id = sd_models.path_to_repo(checkpoint_info.name)
fn = checkpoint_info.path
- # unload current model
- sd_models.unload_model_weights()
- shared.sd_model = None
- devices.torch_gc(force=True)
-
kwargs = {}
kwargs = load_overrides(kwargs, cache_dir)
if (fn is None) or (not os.path.exists(fn) or os.path.isdir(fn)):
@@ -131,16 +101,10 @@ def load_sd3(checkpoint_info, cache_dir=None, config=None):
if fn is not None and os.path.exists(fn) and os.path.isfile(fn):
if fn.endswith('.safetensors'):
loader = diffusers.StableDiffusion3Pipeline.from_single_file
- # required_modules = model_tools.get_modules(diffusers.StableDiffusion3Pipeline)
- # have_modules = model_tools.get_safetensor_keys(fn)
- # loaded_modules = model_tools.load_modules('stabilityai/stable-diffusion-3.5-medium', required_modules)
- # kwargs = {**kwargs, **loaded_modules}
- # kwargs = load_missing(kwargs, fn, cache_dir)
repo_id = fn
elif fn.endswith('.gguf'):
from modules import ggml
kwargs['transformer'] = ggml.load_gguf(fn, cls=diffusers.SD3Transformer2DModel, compute_dtype=devices.dtype)
- # kwargs = load_gguf(kwargs, fn)
kwargs = load_missing(kwargs, fn, cache_dir)
kwargs['variant'] = 'fp16'
else:
@@ -148,7 +112,10 @@ def load_sd3(checkpoint_info, cache_dir=None, config=None):
shared.log.debug(f'Load model: type=SD3 kwargs={list(kwargs)} repo="{repo_id}"')
- kwargs = model_quant.create_config(kwargs)
+ if shared.opts.model_sd3_disable_te5:
+ shared.log.debug('Load model: type=SD3 option="disable-te5"')
+ kwargs['text_encoder_3'] = None
+
pipe = loader(
repo_id,
torch_dtype=devices.dtype,
diff --git a/modules/model_te.py b/modules/model_te.py
index 024dda47c..698d4506e 100644
--- a/modules/model_te.py
+++ b/modules/model_te.py
@@ -20,12 +20,14 @@ def load_t5(name=None, cache_dir=None):
modelloader.hf_login()
repo_id = 'stabilityai/stable-diffusion-3-medium-diffusers'
fn = te_dict.get(name) if name in te_dict else None
+
if fn is not None and name.lower().endswith('gguf'):
from modules import ggml
ggml.install_gguf()
with open(os.path.join('configs', 'flux', 'text_encoder_2', 'config.json'), encoding='utf8') as f:
t5_config = transformers.T5Config(**json.load(f))
t5 = transformers.T5EncoderModel.from_pretrained(None, gguf_file=fn, config=t5_config, device_map="auto", cache_dir=cache_dir, torch_dtype=devices.dtype)
+
elif fn is not None and 'fp8' in name.lower():
from accelerate.utils import set_module_tensor_to_device
with open(os.path.join('configs', 'flux', 'text_encoder_2', 'config.json'), encoding='utf8') as f:
@@ -45,28 +47,34 @@ def load_t5(name=None, cache_dir=None):
try:
t5 = t5.to(dtype=devices.dtype)
except Exception:
- shared.log.error(f"FLUX: Failed to cast text encoder to {devices.dtype}, set dtype to {t5.dtype}")
+ shared.log.error(f"T5: Failed to cast text encoder to {devices.dtype}, set dtype to {t5.dtype}")
raise
+
elif fn is not None:
with open(os.path.join('configs', 'flux', 'text_encoder_2', 'config.json'), encoding='utf8') as f:
t5_config = transformers.T5Config(**json.load(f))
state_dict = load_file(fn)
t5 = transformers.T5EncoderModel.from_pretrained(None, state_dict=state_dict, config=t5_config)
+
elif 'fp16' in name.lower():
t5 = transformers.T5EncoderModel.from_pretrained(repo_id, subfolder='text_encoder_3', cache_dir=cache_dir, torch_dtype=devices.dtype)
+
elif 'fp4' in name.lower():
model_quant.load_bnb('Load model: type=T5')
quantization_config = transformers.BitsAndBytesConfig(load_in_4bit=True)
t5 = transformers.T5EncoderModel.from_pretrained(repo_id, subfolder='text_encoder_3', quantization_config=quantization_config, cache_dir=cache_dir, torch_dtype=devices.dtype)
+
elif 'fp8' in name.lower():
model_quant.load_bnb('Load model: type=T5')
quantization_config = transformers.BitsAndBytesConfig(load_in_8bit=True)
t5 = transformers.T5EncoderModel.from_pretrained(repo_id, subfolder='text_encoder_3', quantization_config=quantization_config, cache_dir=cache_dir, torch_dtype=devices.dtype)
+
elif 'qint8' in name.lower():
model_quant.load_quanto('Load model: type=T5')
from modules.model_quant import optimum_quanto_model
t5 = transformers.T5EncoderModel.from_pretrained(repo_id, subfolder='text_encoder_3', cache_dir=cache_dir, torch_dtype=devices.dtype)
t5 = optimum_quanto_model(t5, weights="qint8", activations="none")
+
elif 'int8' in name.lower():
install('nncf==2.7.0', quiet=True)
from modules.model_quant import nncf_compress_model
@@ -78,8 +86,15 @@ def load_t5(name=None, cache_dir=None):
dtype=torch.float32 if devices.dtype != torch.bfloat16 else torch.bfloat16
)
t5 = nncf_compress_model(t5)
+
+ elif '/' in name:
+ shared.log.debug(f'Load model: type=T5 repo={name}')
+ quant_config = model_quant.create_config(module='TE')
+ t5 = transformers.T5EncoderModel.from_pretrained(name, cache_dir=cache_dir, torch_dtype=devices.dtype, **quant_config)
+
else:
t5 = None
+
if t5 is not None:
loaded_te = name
return t5
@@ -120,8 +135,8 @@ def load_vit_l():
config = transformers.PretrainedConfig.from_json_file('configs/sdxl/text_encoder/config.json')
state_dict = load_file(os.path.join(shared.opts.te_dir, f'{shared.opts.sd_text_encoder}.safetensors'))
te = transformers.CLIPTextModel.from_pretrained(pretrained_model_name_or_path=None, state_dict=state_dict, config=config)
- loaded_te = shared.opts.sd_text_encoder
te = te.to(dtype=devices.dtype)
+ loaded_te = shared.opts.sd_text_encoder
return te
@@ -130,8 +145,8 @@ def load_vit_g():
config = transformers.PretrainedConfig.from_json_file('configs/sdxl/text_encoder_2/config.json')
state_dict = load_file(os.path.join(shared.opts.te_dir, f'{shared.opts.sd_text_encoder}.safetensors'))
te = transformers.CLIPTextModelWithProjection.from_pretrained(pretrained_model_name_or_path=None, state_dict=state_dict, config=config)
- loaded_te = shared.opts.sd_text_encoder
te = te.to(dtype=devices.dtype)
+ loaded_te = shared.opts.sd_text_encoder
return te
diff --git a/modules/modeldata.py b/modules/modeldata.py
index f066eaa68..6f3314644 100644
--- a/modules/modeldata.py
+++ b/modules/modeldata.py
@@ -41,6 +41,8 @@ def get_model_type(pipe):
model_type = 'cogview4'
elif "Sana" in name:
model_type = 'sana'
+ elif "HiDream" in name:
+ model_type = 'h1'
# video models
elif "CogVideo" in name:
model_type = 'cogvideo'
diff --git a/modules/modelloader.py b/modules/modelloader.py
index a132b765a..edb26d84f 100644
--- a/modules/modelloader.py
+++ b/modules/modelloader.py
@@ -10,6 +10,7 @@ from urllib.parse import urlparse
from PIL import Image
import rich.progress as p
import huggingface_hub as hf
+from installer import install
from modules import shared, errors, files_cache
from modules.upscaler import Upscaler
from modules.paths import script_path, models_path
@@ -42,6 +43,7 @@ def hf_login(token=None):
line = [l for l in text.split('\n') if 'Token' in l]
shared.log.info(f'HF login: token="{hf.constants.HF_TOKEN_PATH}" {line[0] if len(line) > 0 else text}')
loggedin = token
+ install('hf_xet', quiet=True)
def download_civit_meta(model_path: str, model_id):
@@ -363,7 +365,7 @@ def find_diffuser(name: str, full=False):
if len(models) == 0:
models = list(hf_api.list_models(model_name=name, full=True, limit=20, sort="downloads", direction=-1)) # widen search
models = [m for m in models if m.id.startswith(name)] # filter exact
- shared.log.debug(f'Searching diffusers models: {name} {len(models) > 0}')
+ shared.log.debug(f'Search model: repo="{name}" {len(models) > 0}')
if len(models) > 0:
if not full:
return models[0].id
diff --git a/modules/options.py b/modules/options.py
index dc84b1a63..aeefc6db0 100644
--- a/modules/options.py
+++ b/modules/options.py
@@ -76,6 +76,8 @@ class OptionInfo:
value = [value]
for v in value:
if v not in choices:
+ if isinstance(choices, list) and ('All' in choices or 'all' in choices): # may be added dynamically
+ continue
log.debug(f'Setting validation: "{opt}"="{v}" default="{self.default}" choices={choices}')
# return False
minimum = args.get("minimum", None)
diff --git a/modules/pag/__init__.py b/modules/pag/__init__.py
index a72f7825d..ac17e7920 100644
--- a/modules/pag/__init__.py
+++ b/modules/pag/__init__.py
@@ -34,7 +34,7 @@ def apply(p: processing.StableDiffusionProcessing): # pylint: disable=arguments-
elif detect.is_f1(cls):
p.task_args['true_cfg_scale'] = p.pag_scale
else:
- shared.log.warning(f'PAG: pipeline={cls.__name__} required={StableDiffusionPipeline.__name__}')
+ # shared.log.warning(f'PAG: pipeline={cls.__name__} required={StableDiffusionPipeline.__name__}')
return None
p.task_args['pag_scale'] = p.pag_scale
diff --git a/modules/para_attention.py b/modules/para_attention.py
index 5ca53962c..f5c6e8635 100644
--- a/modules/para_attention.py
+++ b/modules/para_attention.py
@@ -4,17 +4,17 @@ from modules import shared
supported_models = ['Flux', 'HunyuanVideo', 'CogVideoX', 'Mochi']
-def apply_first_block_cache(p):
+def apply_first_block_cache():
if not shared.opts.para_cache_enabled or not shared.native:
return
- if not any(p.sd_model.__class__.__name__.startswith(x) for x in supported_models):
+ if not any(shared.sd_model.__class__.__name__.startswith(x) for x in supported_models):
return
from installer import install
install('para_attn')
try:
from para_attn.first_block_cache import diffusers_adapters
- diffusers_adapters.apply_cache_on_pipe(p.sd_model, residual_diff_threshold=shared.opts.para_diff_threshold)
- shared.log.info(f'Applying para-attn first-block-cache: diff-threshold={shared.opts.para_diff_threshold} cls={p.sd_model.__class__.__name__}')
+ diffusers_adapters.apply_cache_on_pipe(shared.sd_model, residual_diff_threshold=shared.opts.para_diff_threshold)
+ shared.log.info(f'Transformers cache: type=paraattn rdt={shared.opts.para_diff_threshold} cls={shared.sd_model.__class__.__name__}')
except Exception as e:
- shared.log.error(f'Applying para-attn first-block-cache: {e}')
+ shared.log.error(f'Transformers cache: type=paraattn {e}')
return
diff --git a/modules/processing.py b/modules/processing.py
index 619b544dc..f12fe1215 100644
--- a/modules/processing.py
+++ b/modules/processing.py
@@ -169,9 +169,10 @@ def process_images(p: StableDiffusionProcessing) -> Processed:
shared.prompt_styles.extract_comments(p)
if shared.opts.cuda_compile_backend == 'none':
token_merge.apply_token_merging(p.sd_model)
- from modules import sd_hijack_freeu, para_attention
+ from modules import sd_hijack_freeu, para_attention, teacache
sd_hijack_freeu.apply_freeu(p, not shared.native)
- para_attention.apply_first_block_cache(p)
+ para_attention.apply_first_block_cache()
+ teacache.apply_teacache(p)
if p.width is not None:
p.width = 8 * int(p.width / 8)
diff --git a/modules/processing_args.py b/modules/processing_args.py
index a839c9995..884f189ed 100644
--- a/modules/processing_args.py
+++ b/modules/processing_args.py
@@ -14,7 +14,7 @@ from modules.api import helpers
debug_enabled = os.environ.get('SD_DIFFUSERS_DEBUG', None)
-debug_log = shared.log.trace if os.environ.get('SD_DIFFUSERS_DEBUG', None) is not None else lambda *args, **kwargs: None
+debug_log = shared.log.trace if debug_enabled else lambda *args, **kwargs: None
disable_pbar = os.environ.get('SD_DISABLE_PBAR', None) is not None
diff --git a/modules/processing_class.py b/modules/processing_class.py
index e38a44fd9..6b10f07a7 100644
--- a/modules/processing_class.py
+++ b/modules/processing_class.py
@@ -41,7 +41,7 @@ class StableDiffusionProcessing:
# guidance
cfg_scale: float = 6.0,
cfg_end: float = 1,
- diffusers_guidance_rescale: float = 0.7,
+ diffusers_guidance_rescale: float = 0.0,
pag_scale: float = 0.0,
pag_adaptive: float = 0.5,
# styles
diff --git a/modules/processing_helpers.py b/modules/processing_helpers.py
index 8fffa8313..71e31f663 100644
--- a/modules/processing_helpers.py
+++ b/modules/processing_helpers.py
@@ -371,7 +371,7 @@ def validate_sample(tensor):
shared.log.error(f'Decode: sample={sample.shape} invalid={nans} dtype={dtype} vae={vae} upcast={upcast} failed to validate')
if upcast is not None and not upcast:
setattr(shared.sd_model.vae.config, 'force_upcast', True) # noqa: B010
- shared.log.warning('Decode: upcast=True set, retry operation')
+ shared.log.info('Decode: set upcast=True and attempt to retry operation')
t1 = time.time()
timer.process.add('validate', t1 - t0)
return cast
diff --git a/modules/progress.py b/modules/progress.py
index 6413e7188..83758cc7f 100644
--- a/modules/progress.py
+++ b/modules/progress.py
@@ -87,17 +87,21 @@ def api_progress(req: ProgressRequest):
id_live_preview = req.id_live_preview
live_preview = None
textinfo = shared.state.textinfo
- updated = shared.state.set_current_image()
if not active:
id_live_preview = -1
textinfo = "Queued..." if queued else "Waiting..."
- debug_log(f'Preview: job={shared.state.job} active={active} progress={step}/{steps}/{progress} image={shared.state.current_image_sampling_step} request={id_live_preview} last={shared.state.id_live_preview} enabled={shared.opts.live_previews_enable} job={shared.state.preview_job} updated={updated} image={shared.state.current_image} elapsed={elapsed:.3f}')
+ debug_log(f'Preview: job={shared.state.job} active={active} progress={step}/{steps}/{progress} image={shared.state.current_image_sampling_step} request={id_live_preview} last={shared.state.id_live_preview} enabled={shared.opts.live_previews_enable} job={shared.state.preview_job} elapsed={elapsed:.3f}')
- if shared.opts.live_previews_enable and active and (shared.state.id_live_preview != req.id_live_preview) and (shared.state.current_image is not None):
- buffered = io.BytesIO()
- shared.state.current_image.save(buffered, format='jpeg')
- live_preview = f'data:image/jpeg;base64,{base64.b64encode(buffered.getvalue()).decode("ascii")}'
+ if shared.opts.live_previews_enable and active and (req.id_live_preview != -1):
+ have_image = shared.state.set_current_image()
+ if have_image and shared.state.current_image is not None:
+ buffered = io.BytesIO()
+ shared.state.current_image.save(buffered, format='jpeg', quality=60)
+ b64 = base64.b64encode(buffered.getvalue())
+ live_preview = f'data:image/jpeg;base64,{b64.decode("ascii")}'
+ else:
+ live_preview = None
id_live_preview = shared.state.id_live_preview
diff --git a/modules/prompt_parser_diffusers.py b/modules/prompt_parser_diffusers.py
index 4aaf49a39..f6be32cdb 100644
--- a/modules/prompt_parser_diffusers.py
+++ b/modules/prompt_parser_diffusers.py
@@ -85,8 +85,12 @@ class PromptEmbedder:
def checkcache(self, p):
if shared.opts.sd_textencoder_cache_size == 0:
return False
+ if self.scheduled_prompt:
+ debug("Prompt cache: scheduled prompt")
+ cache.clear()
+ return False
if self.attention != shared.opts.prompt_attention:
- debug(f"Prompt change: parser={shared.opts.prompt_attention}")
+ debug(f"Prompt cache: parser={shared.opts.prompt_attention} changed")
cache.clear()
return False
@@ -284,7 +288,6 @@ class DiffusersTextualInversionManager(BaseTextualInversionManager):
def get_prompt_schedule(prompt, steps):
- t0 = time.time()
temp = []
schedule = prompt_parser.get_learned_conditioning_prompt_schedules([prompt], steps)[0]
if all(x == schedule[0] for x in schedule):
@@ -293,7 +296,6 @@ def get_prompt_schedule(prompt, steps):
for s in range(steps):
if len(temp) < s + 1 <= chunk[0]:
temp.append(chunk[1])
- debug(f'Prompt: schedule={temp} time={(time.time() - t0):.3f}')
return temp, len(schedule) > 1
diff --git a/modules/rocm.py b/modules/rocm.py
index a742556a2..f16809291 100644
--- a/modules/rocm.py
+++ b/modules/rocm.py
@@ -204,7 +204,9 @@ else:
def get_flash_attention_command(agent: Agent):
default = "git+https://github.com/ROCm/flash-attention"
if agent.gfx_version >= 0x1100 and os.environ.get("FLASH_ATTENTION_USE_TRITON_ROCM", "false").lower() != "true":
- default = "git+https://github.com/ROCm/flash-attention@howiejay/navi_support"
+ # use the navi_rotary_fix fork because the original doesn't support rotary_emb for transformers
+ # original: "git+https://github.com/ROCm/flash-attention@howiejay/navi_support"
+ default = "https://github.com/Disty0/flash-attention@navi_rotary_fix"
return os.environ.get("FLASH_ATTENTION_PACKAGE", default)
is_wsl: bool = os.environ.get('WSL_DISTRO_NAME', 'unknown' if spawn('wslpath -w /') else None) is not None
diff --git a/modules/sd_detect.py b/modules/sd_detect.py
index 642043f29..a49a40793 100644
--- a/modules/sd_detect.py
+++ b/modules/sd_detect.py
@@ -24,29 +24,29 @@ def detect_pipeline(f: str, op: str = 'model', warning=True, quiet=False):
elif (size >= 316 and size <= 324) or (size >= 156 and size <= 164): # 320 or 160
warn(f'Model detected as VAE model, but attempting to load as model: {op}={f} size={size} MB')
guess = 'VAE'
- elif (size >= 4970 and size <= 4976): # 4973
- guess = 'Stable Diffusion 2' # SD v2 but could be eps or v-prediction
- # elif size < 0: # unknown
- # guess = 'Stable Diffusion 2B'
- elif (size >= 5791 and size <= 5799): # 5795
- if op == 'model':
- warn(f'Model detected as SD-XL refiner model, but attempting to load a base model: {op}={f} size={size} MB')
- guess = 'Stable Diffusion XL Refiner'
- elif (size >= 6611 and size <= 7220): # 6617, HassakuXL is 6776, monkrenRealisticINT_v10 is 7217
+ elif (size >= 2002 and size <= 2038): # 2032
+ guess = 'Stable Diffusion 1.5'
+ elif (size >= 3138 and size <= 3142): #3140
guess = 'Stable Diffusion XL'
elif (size >= 3361 and size <= 3369): # 3368
guess = 'Stable Diffusion Upscale'
elif (size >= 4891 and size <= 4899): # 4897
guess = 'Stable Diffusion XL Inpaint'
- elif (size >= 9791 and size <= 9799): # 9794
- guess = 'Stable Diffusion XL Instruct'
- elif (size > 3138 and size < 3142): #3140
- guess = 'Stable Diffusion XL'
+ elif (size >= 4970 and size <= 4976): # 4973
+ guess = 'Stable Diffusion 2' # SD v2 but could be eps or v-prediction
+ elif (size >= 5791 and size <= 5799): # 5795
+ if op == 'model':
+ warn(f'Model detected as SD-XL refiner model, but attempting to load a base model: {op}={f} size={size} MB')
+ guess = 'Stable Diffusion XL Refiner'
elif (size > 5692 and size < 5698) or (size > 4134 and size < 4138) or (size > 10362 and size < 10366) or (size > 15028 and size < 15228):
guess = 'Stable Diffusion 3'
- elif (size > 18414 and size < 18420): # sd35-large aio
+ elif (size >= 6611 and size <= 7220): # 6617, HassakuXL is 6776, monkrenRealisticINT_v10 is 7217
+ guess = 'Stable Diffusion XL'
+ elif (size >= 9791 and size <= 9799): # 9794
+ guess = 'Stable Diffusion XL Instruct'
+ elif (size >= 18414 and size <= 18420): # sd35-large aio
guess = 'Stable Diffusion 3'
- elif (size > 20000 and size < 40000):
+ elif (size >= 20000 and size <= 40000):
guess = 'FLUX'
# guess by name
if 'instaflow' in f.lower():
@@ -56,7 +56,7 @@ def detect_pipeline(f: str, op: str = 'model', warning=True, quiet=False):
if 'hunyuandit' in f.lower():
guess = 'HunyuanDiT'
if 'pixart-xl' in f.lower():
- guess = 'PixArt-Alpha'
+ guess = 'PixArt Alpha'
if 'stable-diffusion-3' in f.lower():
guess = 'Stable Diffusion 3'
if 'stable-cascade' in f.lower() or 'stablecascade' in f.lower() or 'wuerstchen3' in f.lower() or ('sotediffusion' in f.lower() and "v2" in f.lower()):
@@ -64,7 +64,7 @@ def detect_pipeline(f: str, op: str = 'model', warning=True, quiet=False):
warn('Stable Cascade does not support Float16')
guess = 'Stable Cascade'
if 'pixart-sigma' in f.lower():
- guess = 'PixArt-Sigma'
+ guess = 'PixArt Sigma'
if 'sana' in f.lower():
guess = 'Sana'
if 'lumina-next' in f.lower():
@@ -76,9 +76,9 @@ def detect_pipeline(f: str, op: str = 'model', warning=True, quiet=False):
if 'auraflow' in f.lower():
guess = 'AuraFlow'
if 'cogview3' in f.lower():
- guess = 'CogView3'
+ guess = 'CogView 3'
if 'cogview4' in f.lower():
- guess = 'CogView4'
+ guess = 'CogView 4'
if 'meissonic' in f.lower():
guess = 'Meissonic'
pipeline = 'custom'
@@ -90,6 +90,8 @@ def detect_pipeline(f: str, op: str = 'model', warning=True, quiet=False):
pipeline = 'custom'
if 'sd3' in f.lower():
guess = 'Stable Diffusion 3'
+ if 'hidream' in f.lower():
+ guess = 'HiDream'
if 'flux' in f.lower() or 'flex.1' in f.lower():
guess = 'FLUX'
if size > 11000 and size < 16000:
diff --git a/modules/sd_models.py b/modules/sd_models.py
index d4fd53a0b..e724d2ad7 100644
--- a/modules/sd_models.py
+++ b/modules/sd_models.py
@@ -10,7 +10,7 @@ import diffusers.loaders.single_file_utils
import torch
from installer import log
-from modules import paths, shared, shared_state, modelloader, devices, script_callbacks, sd_vae, sd_unet, errors, sd_models_config, sd_models_compile, sd_hijack_accelerate, sd_detect, model_quant
+from modules import paths, shared, shared_state, shared_items, modelloader, devices, script_callbacks, sd_vae, sd_unet, errors, sd_models_config, sd_models_compile, sd_hijack_accelerate, sd_detect, model_quant
from modules.timer import Timer, process as process_timer
from modules.memstats import memory_stats
from modules.modeldata import model_data
@@ -31,6 +31,21 @@ debug_load = os.environ.get('SD_LOAD_DEBUG', None)
debug_process = log.trace if os.environ.get('SD_PROCESS_DEBUG', None) is not None else lambda *args, **kwargs: None
diffusers_version = int(diffusers.__version__.split('.')[1])
checkpoint_tiles = checkpoint_titles # legacy compatibility
+pipe_switch_task_exclude = [
+ 'StableDiffusionReferencePipeline',
+ 'StableDiffusionAdapterPipeline',
+ 'AnimateDiffPipeline',
+ 'AnimateDiffSDXLPipeline',
+ 'OmniGenPipeline',
+ 'StableDiffusion3ControlNetPipeline',
+ 'InstantIRPipeline',
+ 'FluxFillPipeline',
+ 'FluxControlPipeline',
+ 'PixelSmithXLPipeline',
+ 'PhotoMakerStableDiffusionXLPipeline',
+ 'StableDiffusionXLInstantIDPipeline',
+ 'LTXConditionPipeline',
+]
def change_backend():
@@ -262,18 +277,22 @@ def load_diffuser_initial(diffusers_load_config, op='model'):
def load_diffuser_force(model_type, checkpoint_info, diffusers_load_config, op='model'):
sd_model = None
+ unload_model_weights()
+ shared.sd_model = None
try:
if model_type in ['Stable Cascade']: # forced pipeline
from modules.model_stablecascade import load_cascade_combined
sd_model = load_cascade_combined(checkpoint_info, diffusers_load_config)
elif model_type in ['InstaFlow']: # forced pipeline
pipeline = diffusers.utils.get_class_from_dynamic_module('instaflow_one_step', module_file='pipeline.py')
+ shared_items.pipelines['InstaFlow'] = pipeline
sd_model = pipeline.from_pretrained(checkpoint_info.path, cache_dir=shared.opts.diffusers_dir, **diffusers_load_config)
elif model_type in ['SegMoE']: # forced pipeline
from modules.segmoe.segmoe_model import SegMoEPipeline
sd_model = SegMoEPipeline(checkpoint_info.path, cache_dir=shared.opts.diffusers_dir, **diffusers_load_config)
sd_model = sd_model.pipe # segmoe pipe does its stuff in __init__ and __call__ is the original pipeline
- elif model_type in ['PixArt-Sigma']: # forced pipeline
+ shared_items.pipelines['SegMoE'] = SegMoEPipeline
+ elif model_type in ['PixArt Sigma']: # forced pipeline
from modules.model_pixart import load_pixart
sd_model = load_pixart(checkpoint_info, diffusers_load_config)
elif model_type in ['Sana']: # forced pipeline
@@ -297,10 +316,10 @@ def load_diffuser_force(model_type, checkpoint_info, diffusers_load_config, op='
elif model_type in ['Stable Diffusion 3']:
from modules.model_sd3 import load_sd3
sd_model = load_sd3(checkpoint_info, cache_dir=shared.opts.diffusers_dir, config=diffusers_load_config.get('config', None))
- elif model_type in ['CogView3']: # forced pipeline
+ elif model_type in ['CogView 3']: # forced pipeline
from modules.model_cogview import load_cogview3
sd_model = load_cogview3(checkpoint_info, diffusers_load_config)
- elif model_type in ['CogView4']: # forced pipeline
+ elif model_type in ['CogView 4']: # forced pipeline
from modules.model_cogview import load_cogview4
sd_model = load_cogview4(checkpoint_info, diffusers_load_config)
elif model_type in ['Meissonic']: # forced pipeline
@@ -309,6 +328,9 @@ def load_diffuser_force(model_type, checkpoint_info, diffusers_load_config, op='
elif model_type in ['OmniGen']: # forced pipeline
from modules.model_omnigen import load_omnigen
sd_model = load_omnigen(checkpoint_info, diffusers_load_config)
+ elif model_type in ['HiDream']:
+ from modules.model_hidream import load_hidream
+ sd_model = load_hidream(checkpoint_info, diffusers_load_config)
except Exception as e:
shared.log.error(f'Load {op}: path="{checkpoint_info.path}" {e}')
if debug_load:
@@ -755,21 +777,6 @@ def clean_diffuser_pipe(pipe):
def set_diffuser_pipe(pipe, new_pipe_type):
- exclude = [
- 'StableDiffusionReferencePipeline',
- 'StableDiffusionAdapterPipeline',
- 'AnimateDiffPipeline',
- 'AnimateDiffSDXLPipeline',
- 'OmniGenPipeline',
- 'StableDiffusion3ControlNetPipeline',
- 'InstantIRPipeline',
- 'FluxFillPipeline',
- 'FluxControlPipeline',
- 'PixelSmithXLPipeline',
- 'PhotoMakerStableDiffusionXLPipeline',
- 'StableDiffusionXLInstantIDPipeline',
- ]
-
has_errors = False
if new_pipe_type == DiffusersTaskType.TEXT_2_IMAGE:
clean_diffuser_pipe(pipe)
@@ -779,7 +786,7 @@ def set_diffuser_pipe(pipe, new_pipe_type):
# skip specific pipelines
cls = pipe.__class__.__name__
- if cls in exclude:
+ if cls in pipe_switch_task_exclude:
return pipe
if 'Video' in cls:
return pipe
diff --git a/modules/sd_offload.py b/modules/sd_offload.py
index e27b70ac2..f98d2023a 100644
--- a/modules/sd_offload.py
+++ b/modules/sd_offload.py
@@ -136,10 +136,10 @@ class OffloadHook(accelerate.hooks.ModelHook):
if shared.opts.diffusers_offload_mode != 'balanced':
return
if shared.opts.diffusers_offload_min_gpu_memory < 0 or shared.opts.diffusers_offload_min_gpu_memory > 1:
- shared.opts.diffusers_offload_min_gpu_memory = 0.25
+ shared.opts.diffusers_offload_min_gpu_memory = 0.2
shared.log.warning(f'Offload: type=balanced op=validate: watermark low={shared.opts.diffusers_offload_min_gpu_memory} invalid value')
if shared.opts.diffusers_offload_max_gpu_memory < 0.1 or shared.opts.diffusers_offload_max_gpu_memory > 1:
- shared.opts.diffusers_offload_max_gpu_memory = 0.75
+ shared.opts.diffusers_offload_max_gpu_memory = 0.7
shared.log.warning(f'Offload: type=balanced op=validate: watermark high={shared.opts.diffusers_offload_max_gpu_memory} invalid value')
if shared.opts.diffusers_offload_min_gpu_memory > shared.opts.diffusers_offload_max_gpu_memory:
shared.opts.diffusers_offload_min_gpu_memory = shared.opts.diffusers_offload_max_gpu_memory
@@ -228,15 +228,15 @@ def apply_balanced_offload(sd_model=None, exclude=[]):
return modules
def apply_balanced_offload_to_module(pipe):
+ # shared.log.trace(f'Offload: type=balanced op=apply pipe={pipe.__class__.__name__}')
used_gpu, used_ram = devices.torch_gc(fast=True)
- if hasattr(pipe, "pipe"):
- apply_balanced_offload_to_module(pipe.pipe)
if hasattr(pipe, "_internal_dict"):
keys = pipe._internal_dict.keys() # pylint: disable=protected-access
else:
keys = get_signature(pipe).keys()
keys = [k for k in keys if k not in exclude and not k.startswith('_')]
for module_name, module_size in get_pipe_modules(pipe): # pylint: disable=protected-access
+ # shared.log.trace(f'Offload: type=balanced op=apply pipe={pipe.__class__.__name__} module={module_name} size={module_size:.3f}')
module = getattr(pipe, module_name, None)
if module is None:
continue
@@ -249,8 +249,7 @@ def apply_balanced_offload(sd_model=None, exclude=[]):
prev_gpu = used_gpu
do_offload = (perc_gpu > shared.opts.diffusers_offload_min_gpu_memory) and (module.device != devices.cpu)
if do_offload:
- non_blocking = devices.backend != "ipex" # non_blocking on ipex causes 2x slowdown
- module = module.to(devices.cpu, non_blocking=non_blocking)
+ module = module.to(devices.cpu)
used_gpu -= module_size
cls = module.__class__.__name__
quant = getattr(module, "quantization_method", None)
diff --git a/modules/sd_samplers.py b/modules/sd_samplers.py
index be375dbff..f73520027 100644
--- a/modules/sd_samplers.py
+++ b/modules/sd_samplers.py
@@ -12,7 +12,7 @@ samplers = all_samplers
samplers_for_img2img = all_samplers
samplers_map = {}
loaded_config = None
-flow_models = ['Flux', 'StableDiffusion3', 'Lumina', 'AuraFlow', 'Sana', 'CogView4']
+flow_models = ['Flux', 'StableDiffusion3', 'Lumina', 'AuraFlow', 'Sana', 'CogView4', 'HiDream']
flow_models += ['Hunyuan', 'LTX', 'Mochi']
diff --git a/modules/sd_vae_remote.py b/modules/sd_vae_remote.py
index 3c1846ea8..9153d5ebb 100644
--- a/modules/sd_vae_remote.py
+++ b/modules/sd_vae_remote.py
@@ -12,6 +12,7 @@ hf_decode_endpoints = {
'sd': 'https://q1bj3bpq6kzilnsu.us-east-1.aws.endpoints.huggingface.cloud',
'sdxl': 'https://x2dmsqunjd6k9prw.us-east-1.aws.endpoints.huggingface.cloud',
'f1': 'https://whhx50ex1aryqvw6.us-east-1.aws.endpoints.huggingface.cloud',
+ 'h1': 'https://whhx50ex1aryqvw6.us-east-1.aws.endpoints.huggingface.cloud',
'hunyuanvideo': 'https://o7ywnmrahorts457.us-east-1.aws.endpoints.huggingface.cloud',
}
hf_encode_endpoints = {
@@ -27,6 +28,13 @@ dtypes = {
}
+def h1_pack_latents(latents, _batch_size, _num_channels_latents, _height, _width): # TODO hidream: pack latents for remote vae
+ # latents = latents.view(batch_size, num_channels_latents, height // 2, 2, width // 2, 2)
+ # latents = latents.permute(0, 2, 4, 1, 3, 5)
+ # latents = latents.reshape(batch_size, (height // 2) * (width // 2) // (num_channels_latents * 4), num_channels_latents * 4)
+ return latents
+
+
def remote_decode(latents: torch.Tensor, width: int = 0, height: int = 0, model_type: str = None) -> Image.Image:
from modules import devices, shared, errors, modelloader
tensors = []
@@ -44,10 +52,14 @@ def remote_decode(latents: torch.Tensor, width: int = 0, height: int = 0, model_
latent_copy = latent_copy.unsqueeze(0)
for i in range(latent_copy.shape[0]):
+ params = {}
try:
latent = latent_copy[i]
if model_type != 'f1':
latent = latent.unsqueeze(0)
+ # if model_type == 'h1':
+ # num_channels_latents = shared.sd_model.transformer.config.in_channels
+ # latent = h1_pack_latents(latent, 1, num_channels_latents, height, width) # pylint: disable=protected-access
params = {
"input_tensor_type": "binary",
"shape": list(latent.shape),
@@ -72,7 +84,7 @@ def remote_decode(latents: torch.Tensor, width: int = 0, height: int = 0, model_
params["output_type"] = "pt"
params["output_tensor_type"] = "binary"
headers["Accept"] = "tensor/binary"
- if (model_type == 'f1') and (width > 0) and (height > 0):
+ if (model_type == 'f1' or model_type == 'h1') and (width > 0) and (height > 0):
params['width'] = width
params['height'] = height
if shared.sd_model.vae is not None and shared.sd_model.vae.config is not None:
diff --git a/modules/sd_vae_taesd.py b/modules/sd_vae_taesd.py
index a2a447a3e..a35c85a83 100644
--- a/modules/sd_vae_taesd.py
+++ b/modules/sd_vae_taesd.py
@@ -52,8 +52,10 @@ def get_model(model_type = 'decoder', variant = None):
global prev_cls, prev_type, prev_model # pylint: disable=global-statement
from modules import shared
cls = shared.sd_model_type
- if cls == 'ldm':
+ if cls == 'ldm': # original backend
cls = 'sd'
+ if cls == 'h1': # hidream uses flux vae
+ cls = 'f1'
variant = variant or shared.opts.taesd_variant
folder = os.path.join(paths.models_path, "TAESD")
os.makedirs(folder, exist_ok=True)
diff --git a/modules/shared.py b/modules/shared.py
index 8e84c71eb..9a2eb8cb4 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -341,7 +341,7 @@ def temp_disable_extensions():
def get_default_modes():
default_offload_mode = "none"
- default_diffusers_offload_min_gpu_memory = 0.25
+ default_diffusers_offload_min_gpu_memory = 0.2
if not (cmd_opts.lowvram or cmd_opts.medvram):
if "gpu" in mem_stat:
if gpu_memory <= 4:
@@ -356,7 +356,7 @@ def get_default_modes():
log.info(f"Device detect: memory={gpu_memory:.1f} default=balanced optimization=medvram")
else:
default_offload_mode = "balanced"
- default_diffusers_offload_min_gpu_memory = 0.25
+ default_diffusers_offload_min_gpu_memory = 0.2
log.info(f"Device detect: memory={gpu_memory:.1f} default=balanced")
elif cmd_opts.medvram:
default_offload_mode = "balanced"
@@ -403,16 +403,22 @@ options_templates.update(options_section(('sd', "Models & Loading"), {
"diffusers_offload_max_cpu_memory": OptionInfo(0.90, "Balanced offload CPU high watermark", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.01, "visible": False }),
"advanced_sep": OptionInfo("Advanced Options
", "", gr.HTML),
- "sd_checkpoint_autoload": OptionInfo(True, "Model autoload on start"),
+ "sd_checkpoint_autoload": OptionInfo(True, "Model auto-load on start"),
"sd_checkpoint_autodownload": OptionInfo(True, "Model auto-download on demand"),
"stream_load": OptionInfo(False, "Model load using streams", gr.Checkbox),
"diffusers_eval": OptionInfo(True, "Force model eval", gr.Checkbox, {"visible": False }),
- "diffusers_to_gpu": OptionInfo(False, "Model Load model direct to GPU"),
+ "diffusers_to_gpu": OptionInfo(False, "Model load model direct to GPU"),
+ "device_map": OptionInfo('default', "Model load device map", gr.Radio, {"choices": ['default', 'gpu', 'cpu'] }),
"disable_accelerate": OptionInfo(False, "Disable accelerate", gr.Checkbox, {"visible": False }),
"sd_model_dict": OptionInfo('None', "Use separate base dict", gr.Dropdown, lambda: {"choices": ['None'] + list_checkpoint_titles(), "visible": False}, refresh=refresh_checkpoints),
"sd_checkpoint_cache": OptionInfo(0, "Cached models", gr.Slider, {"minimum": 0, "maximum": 10, "step": 1, "visible": not native }),
}))
+options_templates.update(options_section(('model_options', "Models Options"), {
+ "model_sd3_disable_te5": OptionInfo(False, "StableDiffusion3: T5 disable encoder"),
+ "model_h1_llama_repo": OptionInfo("meta-llama/Meta-Llama-3.1-8B-Instruct", "HiDream: LLama repo", gr.Textbox),
+}))
+
options_templates.update(options_section(('vae_encoder', "Variable Auto Encoder"), {
"sd_vae": OptionInfo("Automatic", "VAE model", gr.Dropdown, lambda: {"choices": shared_items.sd_vae_items()}, refresh=shared_items.refresh_vae_list),
"diffusers_vae_upcast": OptionInfo("default", "VAE upcasting", gr.Radio, {"choices": ['default', 'true', 'false']}),
@@ -429,7 +435,7 @@ options_templates.update(options_section(('vae_encoder', "Variable Auto Encoder"
}))
options_templates.update(options_section(('text_encoder', "Text Encoder"), {
- "sd_text_encoder": OptionInfo('Default', "Text encoder model", gr.Dropdown, lambda: {"choices": shared_items.sd_te_items()}, refresh=shared_items.refresh_te_list),
+ "sd_text_encoder": OptionInfo('Default', "Text encoder model", DropdownEditable, lambda: {"choices": shared_items.sd_te_items()}, refresh=shared_items.refresh_te_list),
"prompt_attention": OptionInfo("native", "Prompt attention parser", gr.Radio, {"choices": ["native", "compel", "xhinker", "a1111", "fixed"] }),
"prompt_mean_norm": OptionInfo(False, "Prompt attention normalization", gr.Checkbox),
"sd_textencoder_cache": OptionInfo(True, "Cache text encoder results", gr.Checkbox, {"visible": False}),
@@ -509,12 +515,12 @@ options_templates.update(options_section(('backends', "Backend Settings"), {
options_templates.update(options_section(('quantization', "Quantization Settings"), {
"bnb_quantization_sep": OptionInfo("BitsAndBytes
", "", gr.HTML),
- "bnb_quantization": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "Transformer", "VAE", "TE", "Video", "LLM"], "visible": native}),
+ "bnb_quantization": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "Transformer", "VAE", "TE", "Video", "LLM", "ControlNet"], "visible": native}),
"bnb_quantization_type": OptionInfo("nf4", "Quantization type", gr.Dropdown, {"choices": ['nf4', 'fp8', 'fp4'], "visible": native}),
"bnb_quantization_storage": OptionInfo("uint8", "Backend storage", gr.Dropdown, {"choices": ["float16", "float32", "int8", "uint8", "float64", "bfloat16"], "visible": native}),
"quanto_quantization_sep": OptionInfo("Optimum Quanto
", "", gr.HTML),
- "quanto_quantization": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "Transformer", "VAE", "TE", "Video", "LLM"], "visible": native}),
+ "quanto_quantization": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "Transformer", "VAE", "TE", "Video", "LLM", "ControlNet"], "visible": native}),
"quanto_quantization_type": OptionInfo("int8", "Quantization weights type", gr.Dropdown, {"choices": ["float8", "int8", "int4", "int2"], "visible": native}),
"optimum_quanto_sep": OptionInfo("Optimum Quanto: post-load
", "", gr.HTML),
@@ -560,13 +566,13 @@ options_templates.update(options_section(('advanced', "Pipeline Modifiers"), {
"pag_apply_layers": OptionInfo("m0", "PAG layer names"),
"pab_sep": OptionInfo("PAB: Pyramid attention broadcast
", "", gr.HTML),
- "pab_enabled": OptionInfo(False, "Attention cache enabled"),
- "pab_spacial_skip_range": OptionInfo(2, "FC spacial skip range", gr.Slider, {"minimum": 1, "maximum": 4, "step": 1}),
- "pab_spacial_skip_start": OptionInfo(100, "FC spacial skip start", gr.Slider, {"minimum": 0, "maximum": 1000, "step": 1}),
- "pab_spacial_skip_end": OptionInfo(800, "FC spacial skip end", gr.Slider, {"minimum": 0, "maximum": 1000, "step": 1}),
+ "pab_enabled": OptionInfo(False, "PAB cache enabled"),
+ "pab_spacial_skip_range": OptionInfo(2, "PAB spacial skip range", gr.Slider, {"minimum": 1, "maximum": 4, "step": 1}),
+ "pab_spacial_skip_start": OptionInfo(100, "PAB spacial skip start", gr.Slider, {"minimum": 0, "maximum": 1000, "step": 1}),
+ "pab_spacial_skip_end": OptionInfo(800, "PAB spacial skip end", gr.Slider, {"minimum": 0, "maximum": 1000, "step": 1}),
"faster_cache__sep": OptionInfo("Faster Cache
", "", gr.HTML),
- "faster_cache_enabled": OptionInfo(False, "Faster cache enabled"),
+ "faster_cache_enabled": OptionInfo(False, "FC cache enabled"),
"fc_spacial_skip_range": OptionInfo(2, "FC spacial skip range", gr.Slider, {"minimum": 1, "maximum": 4, "step": 1}),
"fc_spacial_skip_start": OptionInfo(0, "FC spacial skip start", gr.Slider, {"minimum": 0, "maximum": 1000, "step": 1}),
"fc_spacial_skip_end": OptionInfo(681, "FC spacial skip end", gr.Slider, {"minimum": 0, "maximum": 1.0, "step": 0.01}),
@@ -581,6 +587,10 @@ options_templates.update(options_section(('advanced', "Pipeline Modifiers"), {
"para_cache_enabled": OptionInfo(False, "First-block cache enabled"),
"para_diff_threshold": OptionInfo(0.1, "Residual diff threshold", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01}),
+ "teacache_sep": OptionInfo("TeaCache
", "", gr.HTML),
+ "teacache_enabled": OptionInfo(False, "TC cache enabled"),
+ "teacache_thresh": OptionInfo(0.6, "TC L1 threshold", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01}),
+
"hypertile_sep": OptionInfo("HyperTile
", "", gr.HTML),
"hypertile_unet_enabled": OptionInfo(False, "UNet Enabled"),
"hypertile_hires_only": OptionInfo(False, "HiRes pass only"),
@@ -930,8 +940,8 @@ options_templates.update(options_section(('extra_networks', "Networks"), {
"lora_preferred_name": OptionInfo("filename", "LoRA preferred name", gr.Radio, {"choices": ["filename", "alias"], "visible": False}),
"lora_add_hashes_to_infotext": OptionInfo(False, "LoRA add hash info to metadata"),
"lora_fuse_diffusers": OptionInfo(True, "LoRA fuse directly to model"),
- "lora_apply_gpu": OptionInfo(False, "LoRA load directly on GPU"),
"lora_legacy": OptionInfo(not native, "LoRA load using legacy method"),
+ "lora_force_reload": OptionInfo(False, "LoRA force reload always"),
"lora_force_diffusers": OptionInfo(False if not cmd_opts.use_openvino else True, "LoRA load using Diffusers method"),
"lora_maybe_diffusers": OptionInfo(False, "LoRA load using Diffusers method for selected models"),
"lora_apply_tags": OptionInfo(0, "LoRA auto-apply tags", gr.Slider, {"minimum": -1, "maximum": 32, "step": 1}),
diff --git a/modules/shared_items.py b/modules/shared_items.py
index 91299850e..ac66e44d2 100644
--- a/modules/shared_items.py
+++ b/modules/shared_items.py
@@ -1,3 +1,53 @@
+import diffusers
+
+
+pipelines = {
+ # note: not all pipelines can be used manually as they require prior pipeline next to decoder pipeline
+ 'Autodetect': None,
+ 'Custom Diffusers Pipeline': getattr(diffusers, 'DiffusionPipeline', None),
+
+ # standard pipelines
+ 'Stable Diffusion 1.5': getattr(diffusers, 'StableDiffusionPipeline', None),
+ 'Stable Diffusion 2.x': getattr(diffusers, 'StableDiffusionPipeline', None),
+ 'Stable Diffusion Upscale': getattr(diffusers, 'StableDiffusionUpscalePipeline', None),
+ 'Stable Diffusion XL': getattr(diffusers, 'StableDiffusionXLPipeline', None),
+ 'Stable Cascade': getattr(diffusers, 'StableCascadeCombinedPipeline', None),
+ 'Stable Diffusion 3.x': getattr(diffusers, 'StableDiffusion3Pipeline', None),
+ 'Latent Consistency Model': getattr(diffusers, 'LatentConsistencyModelPipeline', None),
+ 'PixArt Alpha': getattr(diffusers, 'PixArtAlphaPipeline', None),
+ 'PixArt Sigma': getattr(diffusers, 'PixArtSigmaPipeline', None),
+ 'HunyuanDiT': getattr(diffusers, 'HunyuanDiTPipeline', None),
+ 'DeepFloyd IF': getattr(diffusers, 'IFPipeline', None),
+ 'FLUX': getattr(diffusers, 'FluxPipeline', None),
+ 'Sana': getattr(diffusers, 'SanaPipeline', None),
+ 'Lumina-Next': getattr(diffusers, 'LuminaText2ImgPipeline', None),
+ 'Lumina 2': getattr(diffusers, 'Lumina2Text2ImgPipeline', None),
+ 'AuraFlow': getattr(diffusers, 'AuraFlowPipeline', None),
+ 'Kandinsky 2.1': getattr(diffusers, 'KandinskyCombinedPipeline', None),
+ 'Kandinsky 2.2': getattr(diffusers, 'KandinskyV22CombinedPipeline', None),
+ 'Kandinsky 3.0': getattr(diffusers, 'Kandinsky3Pipeline', None),
+ 'Wuerstchen': getattr(diffusers, 'WuerstchenCombinedPipeline', None),
+ 'Kolors': getattr(diffusers, 'KolorsPipeline', None),
+ 'CogView 3': getattr(diffusers, 'CogView3PlusPipeline', None),
+ 'CogView 4': getattr(diffusers, 'CogView4Pipeline', None),
+ 'UniDiffuser': getattr(diffusers, 'UniDiffuserPipeline', None),
+ 'Amused': getattr(diffusers, 'AmusedPipeline', None),
+ 'HiDream': getattr(diffusers, 'HiDreamImagePipeline', None),
+
+ # dynamically imported and redefined later
+ 'Meissonic': getattr(diffusers, 'DiffusionPipeline', None), # dynamically redefined and loaded in sd_models.load_diffuser
+ 'OmniGenPipeline': getattr(diffusers, 'DiffusionPipeline', None), # dynamically redefined and loaded in sd_models.load_diffuser
+ 'InstaFlow': getattr(diffusers, 'DiffusionPipeline', None), # dynamically redefined and loaded in sd_models.load_diffuser
+ 'SegMoE': getattr(diffusers, 'DiffusionPipeline', None), # dynamically redefined and loaded in sd_models.load_diffuser
+}
+onnx_pipelines = {
+ 'ONNX Stable Diffusion': getattr(diffusers, 'OnnxStableDiffusionPipeline', None),
+ 'ONNX Stable Diffusion Img2Img': getattr(diffusers, 'OnnxStableDiffusionImg2ImgPipeline', None),
+ 'ONNX Stable Diffusion Inpaint': getattr(diffusers, 'OnnxStableDiffusionInpaintPipeline', None),
+ 'ONNX Stable Diffusion Upscale': getattr(diffusers, 'OnnxStableDiffusionUpscalePipeline', None),
+}
+
+
def postprocessing_scripts():
import modules.scripts
return modules.scripts.scripts_postproc.scripts
@@ -29,7 +79,7 @@ def refresh_unet_list():
def sd_te_items():
import modules.model_te
- predefined = ['None', 'T5 FP4', 'T5 FP8', 'T5 INT8', 'T5 QINT8', 'T5 FP16']
+ predefined = ['None']
return predefined + list(modules.model_te.te_dict)
@@ -38,8 +88,8 @@ def refresh_te_list():
modules.model_te.refresh_te_list()
-def list_crossattention(diffusers=False):
- if diffusers:
+def list_crossattention(native:bool=True):
+ if native:
return [
"Disabled",
"Scaled-Dot-Product",
@@ -60,68 +110,23 @@ def list_crossattention(diffusers=False):
]
def get_pipelines():
- import diffusers
from installer import log
-
- pipelines = { # note: not all pipelines can be used manually as they require prior pipeline next to decoder pipeline
- 'Autodetect': None,
- 'Stable Diffusion': getattr(diffusers, 'StableDiffusionPipeline', None),
- 'Stable Diffusion 2': getattr(diffusers, 'StableDiffusionPipeline', None),
- 'Stable Diffusion Inpaint': getattr(diffusers, 'StableDiffusionInpaintPipeline', None),
- 'Stable Diffusion Img2Img': getattr(diffusers, 'StableDiffusionImg2ImgPipeline', None),
- 'Stable Diffusion Instruct': getattr(diffusers, 'StableDiffusionInstructPix2PixPipeline', None),
- 'Stable Diffusion Upscale': getattr(diffusers, 'StableDiffusionUpscalePipeline', None),
- 'Stable Diffusion XL': getattr(diffusers, 'StableDiffusionXLPipeline', None),
- 'Stable Diffusion XL Refiner': getattr(diffusers, 'StableDiffusionXLImg2ImgPipeline', None),
- 'Stable Diffusion XL Img2Img': getattr(diffusers, 'StableDiffusionXLImg2ImgPipeline', None),
- 'Stable Diffusion XL Inpaint': getattr(diffusers, 'StableDiffusionXLInpaintPipeline', None),
- 'Stable Diffusion XL Instruct': getattr(diffusers, 'StableDiffusionXLInstructPix2PixPipeline', None),
- 'Latent Consistency Model': getattr(diffusers, 'LatentConsistencyModelPipeline', None),
- 'PixArt-Alpha': getattr(diffusers, 'PixArtAlphaPipeline', None),
- 'UniDiffuser': getattr(diffusers, 'UniDiffuserPipeline', None),
- 'Wuerstchen': getattr(diffusers, 'WuerstchenCombinedPipeline', None),
- 'Kandinsky 2.1': getattr(diffusers, 'KandinskyPipeline', None),
- 'Kandinsky 2.2': getattr(diffusers, 'KandinskyV22Pipeline', None),
- 'Kandinsky 3': getattr(diffusers, 'Kandinsky3Pipeline', None),
- 'DeepFloyd IF': getattr(diffusers, 'IFPipeline', None),
- 'Custom Diffusers Pipeline': getattr(diffusers, 'DiffusionPipeline', None),
- 'InstaFlow': getattr(diffusers, 'StableDiffusionPipeline', None), # dynamically redefined and loaded in sd_models.load_diffuser
- 'SegMoE': getattr(diffusers, 'StableDiffusionPipeline', None), # dynamically redefined and loaded in sd_models.load_diffuser
- 'Kolors': getattr(diffusers, 'KolorsPipeline', None),
- 'AuraFlow': getattr(diffusers, 'AuraFlowPipeline', None),
- 'CogView3': getattr(diffusers, 'CogView3PlusPipeline', None),
- 'CogView4': getattr(diffusers, 'CogView4Pipeline', None),
- 'Stable Cascade': getattr(diffusers, 'StableCascadeCombinedPipeline', None),
- 'PixArt-Sigma': getattr(diffusers, 'PixArtSigmaPipeline', None),
- 'HunyuanDiT': getattr(diffusers, 'HunyuanDiTPipeline', None),
- 'Stable Diffusion 3': getattr(diffusers, 'StableDiffusion3Pipeline', None),
- 'Stable Diffusion 3 Img2Img': getattr(diffusers, 'StableDiffusion3Img2ImgPipeline', None),
- 'Lumina-Next': getattr(diffusers, 'LuminaText2ImgPipeline', None),
- 'FLUX': getattr(diffusers, 'FluxPipeline', None),
- 'Sana': getattr(diffusers, 'SanaPAGPipeline', None),
- }
- if hasattr(diffusers, 'OnnxStableDiffusionPipeline'):
- onnx_pipelines = {
- 'ONNX Stable Diffusion': getattr(diffusers, 'OnnxStableDiffusionPipeline', None),
- 'ONNX Stable Diffusion Img2Img': getattr(diffusers, 'OnnxStableDiffusionImg2ImgPipeline', None),
- 'ONNX Stable Diffusion Inpaint': getattr(diffusers, 'OnnxStableDiffusionInpaintPipeline', None),
- 'ONNX Stable Diffusion Upscale': getattr(diffusers, 'OnnxStableDiffusionUpscalePipeline', None),
- }
+ if hasattr(diffusers, 'OnnxStableDiffusionPipeline') and 'ONNX Stable Diffusion' not in list(pipelines):
pipelines.update(onnx_pipelines)
- if hasattr(diffusers, 'OnnxStableDiffusionXLPipeline'):
- onnx_pipelines = {
- 'ONNX Stable Diffusion XL': getattr(diffusers, 'OnnxStableDiffusionXLPipeline', None),
- 'ONNX Stable Diffusion XL Img2Img': getattr(diffusers, 'OnnxStableDiffusionXLImg2ImgPipeline', None),
- }
- pipelines.update(onnx_pipelines)
-
- # items that may rely on diffusers dev version
- """
- if hasattr(diffusers, 'FluxPipeline'):
- pipelines['FLUX'] = getattr(diffusers, 'FluxPipeline', None)
- """
-
for k, v in pipelines.items():
if k != 'Autodetect' and v is None:
log.error(f'Not available: pipeline={k} diffusers={diffusers.__version__} path={diffusers.__file__}')
return pipelines
+
+
+def get_repo(model):
+ if model == 'StableDiffusionPipeline' or model == 'Stable Diffusion 1.5':
+ return 'stable-diffusion-v1-5/stable-diffusion-v1-5'
+ elif model == 'StableDiffusionXLPipeline' or model == 'Stable Diffusion XL':
+ return 'stabilityai/stable-diffusion-xl-base-1.0'
+ elif model == 'StableDiffusion3Pipeline' or model == 'Stable Diffusion 3.x':
+ return 'stabilityai/stable-diffusion-3.5-medium'
+ elif model == 'FluxPipeline' or model == 'FLUX':
+ return 'black-forest-labs/FLUX.1-dev'
+ else:
+ return None
diff --git a/modules/teacache/__init__.py b/modules/teacache/__init__.py
new file mode 100644
index 000000000..89e3cc0c8
--- /dev/null
+++ b/modules/teacache/__init__.py
@@ -0,0 +1,25 @@
+from .teacache_flux import teacache_flux_forward
+from .teacache_ltx import teacache_ltx_forward
+from .teacache_mochi import teacache_mochi_forward
+from .teacache_cogvideox import teacache_cog_forward
+
+
+supported_models = ['Flux', 'CogVideoX', 'Mochi', 'LTX']
+
+
+def apply_teacache(p):
+ from modules import shared
+ if not shared.opts.teacache_enabled:
+ return
+ if not any(shared.sd_model.__class__.__name__.startswith(x) for x in supported_models):
+ return
+ if not hasattr(shared.sd_model, 'transformer'):
+ return
+ shared.sd_model.transformer.__class__.enable_teacache = shared.opts.teacache_thresh > 0
+ shared.sd_model.transformer.__class__.cnt = 0
+ shared.sd_model.transformer.__class__.num_steps = p.steps
+ shared.sd_model.transformer.__class__.rel_l1_thresh = shared.opts.teacache_thresh # 0.25 for 1.5x speedup, 0.4 for 1.8x speedup, 0.6 for 2.0x speedup, 0.8 for 2.25x speedup
+ shared.sd_model.transformer.__class__.accumulated_rel_l1_distance = 0
+ shared.sd_model.transformer.__class__.previous_modulated_input = None
+ shared.sd_model.transformer.__class__.previous_residual = None
+ shared.log.info(f'Transformers cache: type=teacache cls={shared.sd_model.__class__.__name__} thresh={shared.opts.teacache_thresh}')
diff --git a/modules/teacache/teacache_cogvideox.py b/modules/teacache/teacache_cogvideox.py
new file mode 100644
index 000000000..436338b37
--- /dev/null
+++ b/modules/teacache/teacache_cogvideox.py
@@ -0,0 +1,182 @@
+from typing import Any, Dict, Optional, Union, Tuple
+import torch
+import numpy as np
+from diffusers.utils import USE_PEFT_BACKEND, is_torch_version, scale_lora_layers, unscale_lora_layers, logging
+from diffusers.models.modeling_outputs import Transformer2DModelOutput
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+def teacache_cog_forward(
+ self,
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: torch.Tensor,
+ timestep: Union[int, float, torch.LongTensor],
+ timestep_cond: Optional[torch.Tensor] = None,
+ ofs: Optional[Union[int, float, torch.LongTensor]] = None,
+ image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+ attention_kwargs: Optional[Dict[str, Any]] = None,
+ return_dict: bool = True,
+ ):
+ if attention_kwargs is not None:
+ attention_kwargs = attention_kwargs.copy()
+ lora_scale = attention_kwargs.pop("scale", 1.0)
+ else:
+ lora_scale = 1.0
+
+ if USE_PEFT_BACKEND:
+ # weight the lora layers by setting `lora_scale` for each PEFT layer
+ scale_lora_layers(self, lora_scale)
+ else:
+ if attention_kwargs is not None and attention_kwargs.get("scale", None) is not None:
+ logger.warning(
+ "Passing `scale` via `attention_kwargs` when not using the PEFT backend is ineffective."
+ )
+
+ batch_size, num_frames, channels, height, width = hidden_states.shape
+
+ # 1. Time embedding
+ timesteps = timestep
+ t_emb = self.time_proj(timesteps)
+
+ # timesteps does not contain any weights and will always return f32 tensors
+ # but time_embedding might actually be running in fp16. so we need to cast here.
+ # there might be better ways to encapsulate this.
+ t_emb = t_emb.to(dtype=hidden_states.dtype)
+ emb = self.time_embedding(t_emb, timestep_cond)
+
+ if self.ofs_embedding is not None:
+ ofs_emb = self.ofs_proj(ofs)
+ ofs_emb = ofs_emb.to(dtype=hidden_states.dtype)
+ ofs_emb = self.ofs_embedding(ofs_emb)
+ emb = emb + ofs_emb
+
+ # 2. Patch embedding
+ hidden_states = self.patch_embed(encoder_hidden_states, hidden_states)
+ hidden_states = self.embedding_dropout(hidden_states)
+
+ text_seq_length = encoder_hidden_states.shape[1]
+ encoder_hidden_states = hidden_states[:, :text_seq_length]
+ hidden_states = hidden_states[:, text_seq_length:]
+
+ if self.enable_teacache:
+ if self.cnt == 0 or self.cnt == self.num_steps-1:
+ should_calc = True
+ self.accumulated_rel_l1_distance = 0
+ else:
+ if not self.config.use_rotary_positional_embeddings:
+ # CogVideoX-2B
+ coefficients = [-3.10658903e+01, 2.54732368e+01, -5.92380459e+00, 1.75769064e+00, -3.61568434e-03]
+ else:
+ # CogVideoX-5B and CogvideoX1.5-5B
+ coefficients = [-1.53880483e+03, 8.43202495e+02, -1.34363087e+02, 7.97131516e+00, -5.23162339e-02]
+ rescale_func = np.poly1d(coefficients)
+ self.accumulated_rel_l1_distance += rescale_func(((emb-self.previous_modulated_input).abs().mean() / self.previous_modulated_input.abs().mean()).cpu().item())
+ if self.accumulated_rel_l1_distance < self.rel_l1_thresh:
+ should_calc = False
+ else:
+ should_calc = True
+ self.accumulated_rel_l1_distance = 0
+ self.previous_modulated_input = emb
+ self.cnt += 1
+ if self.cnt == self.num_steps:
+ self.cnt = 0
+
+ if self.enable_teacache:
+ if not should_calc:
+ hidden_states += self.previous_residual
+ encoder_hidden_states += self.previous_residual_encoder
+ else:
+ ori_hidden_states = hidden_states.clone()
+ ori_encoder_hidden_states = encoder_hidden_states.clone()
+ # 4. Transformer blocks
+ for i, block in enumerate(self.transformer_blocks):
+ if torch.is_grad_enabled() and self.gradient_checkpointing:
+
+ def create_custom_forward(module):
+ def custom_forward(*inputs):
+ return module(*inputs)
+
+ return custom_forward
+
+ ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
+ hidden_states, encoder_hidden_states = torch.utils.checkpoint.checkpoint(
+ create_custom_forward(block),
+ hidden_states,
+ encoder_hidden_states,
+ emb,
+ image_rotary_emb,
+ **ckpt_kwargs,
+ )
+ else:
+ hidden_states, encoder_hidden_states = block(
+ hidden_states=hidden_states,
+ encoder_hidden_states=encoder_hidden_states,
+ temb=emb,
+ image_rotary_emb=image_rotary_emb,
+ )
+
+ self.previous_residual = hidden_states - ori_hidden_states
+ self.previous_residual_encoder = encoder_hidden_states - ori_encoder_hidden_states
+ else:
+ # 4. Transformer blocks
+ for i, block in enumerate(self.transformer_blocks):
+ if torch.is_grad_enabled() and self.gradient_checkpointing:
+
+ def create_custom_forward(module):
+ def custom_forward(*inputs):
+ return module(*inputs)
+
+ return custom_forward
+
+ ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
+ hidden_states, encoder_hidden_states = torch.utils.checkpoint.checkpoint(
+ create_custom_forward(block),
+ hidden_states,
+ encoder_hidden_states,
+ emb,
+ image_rotary_emb,
+ **ckpt_kwargs,
+ )
+ else:
+ hidden_states, encoder_hidden_states = block(
+ hidden_states=hidden_states,
+ encoder_hidden_states=encoder_hidden_states,
+ temb=emb,
+ image_rotary_emb=image_rotary_emb,
+ )
+
+ if not self.config.use_rotary_positional_embeddings:
+ # CogVideoX-2B
+ hidden_states = self.norm_final(hidden_states)
+ else:
+ # CogVideoX-5B and CogvideoX1.5-5B
+ hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
+ hidden_states = self.norm_final(hidden_states)
+ hidden_states = hidden_states[:, text_seq_length:]
+
+ # 5. Final block
+ hidden_states = self.norm_out(hidden_states, temb=emb)
+ hidden_states = self.proj_out(hidden_states)
+
+ # 6. Unpatchify
+ p = self.config.patch_size
+ p_t = self.config.patch_size_t
+
+ if p_t is None:
+ output = hidden_states.reshape(batch_size, num_frames, height // p, width // p, -1, p, p)
+ output = output.permute(0, 1, 4, 2, 5, 3, 6).flatten(5, 6).flatten(3, 4)
+ else:
+ output = hidden_states.reshape(
+ batch_size, (num_frames + p_t - 1) // p_t, height // p, width // p, -1, p_t, p, p
+ )
+ output = output.permute(0, 1, 5, 4, 2, 6, 3, 7).flatten(6, 7).flatten(4, 5).flatten(1, 2)
+
+ if USE_PEFT_BACKEND:
+ # remove `lora_scale` from each PEFT layer
+ unscale_lora_layers(self, lora_scale)
+
+ if not return_dict:
+ return (output,)
+ return Transformer2DModelOutput(sample=output)
diff --git a/modules/teacache/teacache_flux.py b/modules/teacache/teacache_flux.py
new file mode 100644
index 000000000..aa750d3f0
--- /dev/null
+++ b/modules/teacache/teacache_flux.py
@@ -0,0 +1,308 @@
+from typing import Any, Dict, Optional, Union
+import torch
+import numpy as np
+from diffusers.models.modeling_outputs import Transformer2DModelOutput
+from diffusers.utils import USE_PEFT_BACKEND, is_torch_version, logging, scale_lora_layers, unscale_lora_layers
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+def teacache_flux_forward(
+ self,
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: torch.Tensor = None,
+ pooled_projections: torch.Tensor = None,
+ timestep: torch.LongTensor = None,
+ img_ids: torch.Tensor = None,
+ txt_ids: torch.Tensor = None,
+ guidance: torch.Tensor = None,
+ joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+ controlnet_block_samples=None,
+ controlnet_single_block_samples=None,
+ return_dict: bool = True,
+ controlnet_blocks_repeat: bool = False,
+ ) -> Union[torch.FloatTensor, Transformer2DModelOutput]:
+ """
+ The [`FluxTransformer2DModel`] forward method.
+
+ Args:
+ hidden_states (`torch.FloatTensor` of shape `(batch size, channel, height, width)`):
+ Input `hidden_states`.
+ encoder_hidden_states (`torch.FloatTensor` of shape `(batch size, sequence_len, embed_dims)`):
+ Conditional embeddings (embeddings computed from the input conditions such as prompts) to use.
+ pooled_projections (`torch.FloatTensor` of shape `(batch_size, projection_dim)`): Embeddings projected
+ from the embeddings of input conditions.
+ timestep ( `torch.LongTensor`):
+ Used to indicate denoising step.
+ block_controlnet_hidden_states: (`list` of `torch.Tensor`):
+ A list of tensors that if specified are added to the residuals of transformer blocks.
+ joint_attention_kwargs (`dict`, *optional*):
+ A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
+ `self.processor` in
+ [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
+ return_dict (`bool`, *optional*, defaults to `True`):
+ Whether or not to return a [`~models.transformer_2d.Transformer2DModelOutput`] instead of a plain
+ tuple.
+
+ Returns:
+ If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a
+ `tuple` where the first element is the sample tensor.
+ """
+ if joint_attention_kwargs is not None:
+ joint_attention_kwargs = joint_attention_kwargs.copy()
+ lora_scale = joint_attention_kwargs.pop("scale", 1.0)
+ else:
+ lora_scale = 1.0
+
+ if USE_PEFT_BACKEND:
+ # weight the lora layers by setting `lora_scale` for each PEFT layer
+ scale_lora_layers(self, lora_scale)
+ else:
+ if joint_attention_kwargs is not None and joint_attention_kwargs.get("scale", None) is not None:
+ logger.warning(
+ "Passing `scale` via `joint_attention_kwargs` when not using the PEFT backend is ineffective."
+ )
+
+ hidden_states = self.x_embedder(hidden_states)
+
+ timestep = timestep.to(hidden_states.dtype) * 1000
+ if guidance is not None:
+ guidance = guidance.to(hidden_states.dtype) * 1000
+ else:
+ guidance = None
+
+ temb = (
+ self.time_text_embed(timestep, pooled_projections)
+ if guidance is None
+ else self.time_text_embed(timestep, guidance, pooled_projections)
+ )
+ encoder_hidden_states = self.context_embedder(encoder_hidden_states)
+
+ if txt_ids.ndim == 3:
+ logger.warning(
+ "Passing `txt_ids` 3d torch.Tensor is deprecated."
+ "Please remove the batch dimension and pass it as a 2d torch Tensor"
+ )
+ txt_ids = txt_ids[0]
+ if img_ids.ndim == 3:
+ logger.warning(
+ "Passing `img_ids` 3d torch.Tensor is deprecated."
+ "Please remove the batch dimension and pass it as a 2d torch Tensor"
+ )
+ img_ids = img_ids[0]
+
+ ids = torch.cat((txt_ids, img_ids), dim=0)
+ image_rotary_emb = self.pos_embed(ids)
+
+ if joint_attention_kwargs is not None and "ip_adapter_image_embeds" in joint_attention_kwargs:
+ ip_adapter_image_embeds = joint_attention_kwargs.pop("ip_adapter_image_embeds")
+ ip_hidden_states = self.encoder_hid_proj(ip_adapter_image_embeds)
+ joint_attention_kwargs.update({"ip_hidden_states": ip_hidden_states})
+
+ if self.enable_teacache:
+ inp = hidden_states.clone()
+ temb_ = temb.clone()
+ modulated_inp, _gate_msa, _shift_mlp, _scale_mlp, _gate_mlp = self.transformer_blocks[0].norm1(inp, emb=temb_)
+ if self.cnt == 0 or self.cnt == self.num_steps-1:
+ should_calc = True
+ self.accumulated_rel_l1_distance = 0
+ else:
+ coefficients = [4.98651651e+02, -2.83781631e+02, 5.58554382e+01, -3.82021401e+00, 2.64230861e-01]
+ rescale_func = np.poly1d(coefficients)
+ self.accumulated_rel_l1_distance += rescale_func(((modulated_inp-self.previous_modulated_input).abs().mean() / self.previous_modulated_input.abs().mean()).cpu().item())
+ if self.accumulated_rel_l1_distance < self.rel_l1_thresh:
+ should_calc = False
+ else:
+ should_calc = True
+ self.accumulated_rel_l1_distance = 0
+ self.previous_modulated_input = modulated_inp
+ self.cnt += 1
+ if self.cnt == self.num_steps:
+ self.cnt = 0
+
+ if self.enable_teacache:
+ if not should_calc:
+ hidden_states += self.previous_residual
+ else:
+ ori_hidden_states = hidden_states.clone()
+ for index_block, block in enumerate(self.transformer_blocks):
+ if torch.is_grad_enabled() and self.gradient_checkpointing:
+
+ def create_custom_forward4(module, return_dict=None):
+ def custom_forward(*inputs):
+ if return_dict is not None:
+ return module(*inputs, return_dict=return_dict)
+ else:
+ return module(*inputs)
+
+ return custom_forward
+
+ ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
+ encoder_hidden_states, hidden_states = torch.utils.checkpoint.checkpoint(
+ create_custom_forward4(block),
+ hidden_states,
+ encoder_hidden_states,
+ temb,
+ image_rotary_emb,
+ **ckpt_kwargs,
+ )
+
+ else:
+ encoder_hidden_states, hidden_states = block(
+ hidden_states=hidden_states,
+ encoder_hidden_states=encoder_hidden_states,
+ temb=temb,
+ image_rotary_emb=image_rotary_emb,
+ joint_attention_kwargs=joint_attention_kwargs,
+ )
+
+ # controlnet residual
+ if controlnet_block_samples is not None:
+ interval_control = len(self.transformer_blocks) / len(controlnet_block_samples)
+ interval_control = int(np.ceil(interval_control))
+ # For Xlabs ControlNet.
+ if controlnet_blocks_repeat:
+ hidden_states = (
+ hidden_states + controlnet_block_samples[index_block % len(controlnet_block_samples)]
+ )
+ else:
+ hidden_states = hidden_states + controlnet_block_samples[index_block // interval_control]
+ hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
+
+ for index_block, block in enumerate(self.single_transformer_blocks):
+ if torch.is_grad_enabled() and self.gradient_checkpointing:
+
+ def create_custom_forward2(module, return_dict=None):
+ def custom_forward(*inputs):
+ if return_dict is not None:
+ return module(*inputs, return_dict=return_dict)
+ else:
+ return module(*inputs)
+
+ return custom_forward
+
+ ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
+ hidden_states = torch.utils.checkpoint.checkpoint(
+ create_custom_forward2(block),
+ hidden_states,
+ temb,
+ image_rotary_emb,
+ **ckpt_kwargs,
+ )
+
+ else:
+ hidden_states = block(
+ hidden_states=hidden_states,
+ temb=temb,
+ image_rotary_emb=image_rotary_emb,
+ joint_attention_kwargs=joint_attention_kwargs,
+ )
+
+ # controlnet residual
+ if controlnet_single_block_samples is not None:
+ interval_control = len(self.single_transformer_blocks) / len(controlnet_single_block_samples)
+ interval_control = int(np.ceil(interval_control))
+ hidden_states[:, encoder_hidden_states.shape[1] :, ...] = (
+ hidden_states[:, encoder_hidden_states.shape[1] :, ...]
+ + controlnet_single_block_samples[index_block // interval_control]
+ )
+
+ hidden_states = hidden_states[:, encoder_hidden_states.shape[1] :, ...]
+ self.previous_residual = hidden_states - ori_hidden_states
+ else:
+ for index_block, block in enumerate(self.transformer_blocks):
+ if torch.is_grad_enabled() and self.gradient_checkpointing:
+
+ def create_custom_forward1(module, return_dict=None):
+ def custom_forward(*inputs):
+ if return_dict is not None:
+ return module(*inputs, return_dict=return_dict)
+ else:
+ return module(*inputs)
+
+ return custom_forward
+
+ ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
+ encoder_hidden_states, hidden_states = torch.utils.checkpoint.checkpoint(
+ create_custom_forward1(block),
+ hidden_states,
+ encoder_hidden_states,
+ temb,
+ image_rotary_emb,
+ **ckpt_kwargs,
+ )
+
+ else:
+ encoder_hidden_states, hidden_states = block(
+ hidden_states=hidden_states,
+ encoder_hidden_states=encoder_hidden_states,
+ temb=temb,
+ image_rotary_emb=image_rotary_emb,
+ joint_attention_kwargs=joint_attention_kwargs,
+ )
+
+ # controlnet residual
+ if controlnet_block_samples is not None:
+ interval_control = len(self.transformer_blocks) / len(controlnet_block_samples)
+ interval_control = int(np.ceil(interval_control))
+ # For Xlabs ControlNet.
+ if controlnet_blocks_repeat:
+ hidden_states = (
+ hidden_states + controlnet_block_samples[index_block % len(controlnet_block_samples)]
+ )
+ else:
+ hidden_states = hidden_states + controlnet_block_samples[index_block // interval_control]
+ hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
+
+ for index_block, block in enumerate(self.single_transformer_blocks):
+ if torch.is_grad_enabled() and self.gradient_checkpointing:
+
+ def create_custom_forward3(module, return_dict=None):
+ def custom_forward(*inputs):
+ if return_dict is not None:
+ return module(*inputs, return_dict=return_dict)
+ else:
+ return module(*inputs)
+
+ return custom_forward
+
+ ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
+ hidden_states = torch.utils.checkpoint.checkpoint(
+ create_custom_forward3(block),
+ hidden_states,
+ temb,
+ image_rotary_emb,
+ **ckpt_kwargs,
+ )
+
+ else:
+ hidden_states = block(
+ hidden_states=hidden_states,
+ temb=temb,
+ image_rotary_emb=image_rotary_emb,
+ joint_attention_kwargs=joint_attention_kwargs,
+ )
+
+ # controlnet residual
+ if controlnet_single_block_samples is not None:
+ interval_control = len(self.single_transformer_blocks) / len(controlnet_single_block_samples)
+ interval_control = int(np.ceil(interval_control))
+ hidden_states[:, encoder_hidden_states.shape[1] :, ...] = (
+ hidden_states[:, encoder_hidden_states.shape[1] :, ...]
+ + controlnet_single_block_samples[index_block // interval_control]
+ )
+
+ hidden_states = hidden_states[:, encoder_hidden_states.shape[1] :, ...]
+
+ hidden_states = self.norm_out(hidden_states, temb)
+ output = self.proj_out(hidden_states)
+
+ if USE_PEFT_BACKEND:
+ # remove `lora_scale` from each PEFT layer
+ unscale_lora_layers(self, lora_scale)
+
+ if not return_dict:
+ return (output,)
+
+ return Transformer2DModelOutput(sample=output)
diff --git a/modules/teacache/teacache_ltx.py b/modules/teacache/teacache_ltx.py
index f7f9cd83d..8a4e1b392 100644
--- a/modules/teacache/teacache_ltx.py
+++ b/modules/teacache/teacache_ltx.py
@@ -1,15 +1,14 @@
-"""
-source: https://github.com/ali-vilab/TeaCache/blob/main/TeaCache4LTX-Video/teacache_ltx.py
-"""
-
from typing import Any, Dict, Optional, Tuple
-import numpy as np
import torch
+from diffusers.utils import USE_PEFT_BACKEND, is_torch_version, scale_lora_layers, unscale_lora_layers, logging
from diffusers.models.modeling_outputs import Transformer2DModelOutput
-from diffusers.utils import is_torch_version, scale_lora_layers, unscale_lora_layers
+import numpy as np
-def teacache_forward(
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+def teacache_ltx_forward(
self,
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor,
@@ -22,104 +21,73 @@ def teacache_forward(
attention_kwargs: Optional[Dict[str, Any]] = None,
return_dict: bool = True,
) -> torch.Tensor:
- if attention_kwargs is not None:
- attention_kwargs = attention_kwargs.copy()
- lora_scale = attention_kwargs.pop("scale", 1.0)
- else:
- lora_scale = 1.0
+ if attention_kwargs is not None:
+ attention_kwargs = attention_kwargs.copy()
+ lora_scale = attention_kwargs.pop("scale", 1.0)
+ else:
+ lora_scale = 1.0
+ if USE_PEFT_BACKEND:
+ # weight the lora layers by setting `lora_scale` for each PEFT layer
scale_lora_layers(self, lora_scale)
+ else:
+ if attention_kwargs is not None and attention_kwargs.get("scale", None) is not None:
+ logger.warning(
+ "Passing `scale` via `attention_kwargs` when not using the PEFT backend is ineffective."
+ )
- image_rotary_emb = self.rope(hidden_states, num_frames, height, width, rope_interpolation_scale)
+ image_rotary_emb = self.rope(hidden_states, num_frames, height, width, rope_interpolation_scale)
- # convert encoder_attention_mask to a bias the same way we do for attention_mask
- if encoder_attention_mask is not None and encoder_attention_mask.ndim == 2:
- encoder_attention_mask = (1 - encoder_attention_mask.to(hidden_states.dtype)) * -10000.0
- encoder_attention_mask = encoder_attention_mask.unsqueeze(1)
+ # convert encoder_attention_mask to a bias the same way we do for attention_mask
+ if encoder_attention_mask is not None and encoder_attention_mask.ndim == 2:
+ encoder_attention_mask = (1 - encoder_attention_mask.to(hidden_states.dtype)) * -10000.0
+ encoder_attention_mask = encoder_attention_mask.unsqueeze(1)
- batch_size = hidden_states.size(0)
- hidden_states = self.proj_in(hidden_states)
+ batch_size = hidden_states.size(0)
+ hidden_states = self.proj_in(hidden_states)
- temb, embedded_timestep = self.time_embed(
- timestep.flatten(),
- batch_size=batch_size,
- hidden_dtype=hidden_states.dtype,
- )
+ temb, embedded_timestep = self.time_embed(
+ timestep.flatten(),
+ batch_size=batch_size,
+ hidden_dtype=hidden_states.dtype,
+ )
- temb = temb.view(batch_size, -1, temb.size(-1))
- embedded_timestep = embedded_timestep.view(batch_size, -1, embedded_timestep.size(-1))
+ temb = temb.view(batch_size, -1, temb.size(-1))
+ embedded_timestep = embedded_timestep.view(batch_size, -1, embedded_timestep.size(-1))
- encoder_hidden_states = self.caption_projection(encoder_hidden_states)
- encoder_hidden_states = encoder_hidden_states.view(batch_size, -1, hidden_states.size(-1))
+ encoder_hidden_states = self.caption_projection(encoder_hidden_states)
+ encoder_hidden_states = encoder_hidden_states.view(batch_size, -1, hidden_states.size(-1))
- if self.enable_teacache:
- inp = hidden_states.clone()
- temb_ = temb.clone()
- inp = self.transformer_blocks[0].norm1(inp)
- num_ada_params = self.transformer_blocks[0].scale_shift_table.shape[0]
- ada_values = self.transformer_blocks[0].scale_shift_table[None, None] + temb_.reshape(batch_size, temb_.size(1), num_ada_params, -1)
- shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ada_values.unbind(dim=2)
- modulated_inp = inp * (1 + scale_msa) + shift_msa
- if self.cnt == 0 or self.cnt == self.num_steps-1:
+ if self.enable_teacache:
+ inp = hidden_states.clone()
+ temb_ = temb.clone()
+ inp = self.transformer_blocks[0].norm1(inp)
+ num_ada_params = self.transformer_blocks[0].scale_shift_table.shape[0]
+ ada_values = self.transformer_blocks[0].scale_shift_table[None, None] + temb_.reshape(batch_size, temb_.size(1), num_ada_params, -1)
+ shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ada_values.unbind(dim=2)
+ modulated_inp = inp * (1 + scale_msa) + shift_msa
+ if self.cnt == 0 or self.cnt == self.num_steps-1:
+ should_calc = True
+ self.accumulated_rel_l1_distance = 0
+ else:
+ coefficients = [2.14700694e+01, -1.28016453e+01, 2.31279151e+00, 7.92487521e-01, 9.69274326e-03]
+ rescale_func = np.poly1d(coefficients)
+ self.accumulated_rel_l1_distance += rescale_func(((modulated_inp-self.previous_modulated_input).abs().mean() / self.previous_modulated_input.abs().mean()).cpu().item())
+ if self.accumulated_rel_l1_distance < self.rel_l1_thresh:
+ should_calc = False
+ else:
should_calc = True
self.accumulated_rel_l1_distance = 0
- else:
- coefficients = [2.14700694e+01, -1.28016453e+01, 2.31279151e+00, 7.92487521e-01, 9.69274326e-03]
- rescale_func = np.poly1d(coefficients)
- self.accumulated_rel_l1_distance += rescale_func(((modulated_inp-self.previous_modulated_input).abs().mean() / self.previous_modulated_input.abs().mean()).cpu().item())
- if self.accumulated_rel_l1_distance < self.rel_l1_thresh:
- should_calc = False
- else:
- should_calc = True
- self.accumulated_rel_l1_distance = 0
- self.previous_modulated_input = modulated_inp
- self.cnt += 1
- if self.cnt == self.num_steps:
- self.cnt = 0
-
- if self.enable_teacache:
- if not should_calc:
- hidden_states += self.previous_residual
- else:
- ori_hidden_states = hidden_states.clone()
- for block in self.transformer_blocks:
- if torch.is_grad_enabled() and self.gradient_checkpointing:
-
- def create_custom_forward(module, return_dict=None):
- def custom_forward(*inputs):
- if return_dict is not None:
- return module(*inputs, return_dict=return_dict)
- else:
- return module(*inputs)
-
- return custom_forward
-
- ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
- hidden_states = torch.utils.checkpoint.checkpoint(
- create_custom_forward(block),
- hidden_states,
- encoder_hidden_states,
- temb,
- image_rotary_emb,
- encoder_attention_mask,
- **ckpt_kwargs,
- )
- else:
- hidden_states = block(
- hidden_states=hidden_states,
- encoder_hidden_states=encoder_hidden_states,
- temb=temb,
- image_rotary_emb=image_rotary_emb,
- encoder_attention_mask=encoder_attention_mask,
- )
-
- scale_shift_values = self.scale_shift_table[None, None] + embedded_timestep[:, :, None]
- shift, scale = scale_shift_values[:, :, 0], scale_shift_values[:, :, 1]
-
- hidden_states = self.norm_out(hidden_states)
- hidden_states = hidden_states * (1 + scale) + shift
- self.previous_residual = hidden_states - ori_hidden_states
+ self.previous_modulated_input = modulated_inp
+ self.cnt += 1
+ if self.cnt == self.num_steps:
+ self.cnt = 0
+
+ if self.enable_teacache:
+ if not should_calc:
+ hidden_states += self.previous_residual
else:
+ ori_hidden_states = hidden_states.clone()
for block in self.transformer_blocks:
if torch.is_grad_enabled() and self.gradient_checkpointing:
@@ -156,12 +124,52 @@ def teacache_forward(
hidden_states = self.norm_out(hidden_states)
hidden_states = hidden_states * (1 + scale) + shift
+ self.previous_residual = hidden_states - ori_hidden_states
+ else:
+ for block in self.transformer_blocks:
+ if torch.is_grad_enabled() and self.gradient_checkpointing:
+
+ def create_custom_forward(module, return_dict=None):
+ def custom_forward(*inputs):
+ if return_dict is not None:
+ return module(*inputs, return_dict=return_dict)
+ else:
+ return module(*inputs)
+
+ return custom_forward
+
+ ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
+ hidden_states = torch.utils.checkpoint.checkpoint(
+ create_custom_forward(block),
+ hidden_states,
+ encoder_hidden_states,
+ temb,
+ image_rotary_emb,
+ encoder_attention_mask,
+ **ckpt_kwargs,
+ )
+ else:
+ hidden_states = block(
+ hidden_states=hidden_states,
+ encoder_hidden_states=encoder_hidden_states,
+ temb=temb,
+ image_rotary_emb=image_rotary_emb,
+ encoder_attention_mask=encoder_attention_mask,
+ )
+
+ scale_shift_values = self.scale_shift_table[None, None] + embedded_timestep[:, :, None]
+ shift, scale = scale_shift_values[:, :, 0], scale_shift_values[:, :, 1]
+
+ hidden_states = self.norm_out(hidden_states)
+ hidden_states = hidden_states * (1 + scale) + shift
- output = self.proj_out(hidden_states)
+ output = self.proj_out(hidden_states)
+ if USE_PEFT_BACKEND:
+ # remove `lora_scale` from each PEFT layer
unscale_lora_layers(self, lora_scale)
- if not return_dict:
- return (output,)
- return Transformer2DModelOutput(sample=output)
+ if not return_dict:
+ return (output,)
+ return Transformer2DModelOutput(sample=output)
diff --git a/modules/teacache/teacache_mochi.py b/modules/teacache/teacache_mochi.py
new file mode 100644
index 000000000..e899fb164
--- /dev/null
+++ b/modules/teacache/teacache_mochi.py
@@ -0,0 +1,157 @@
+from typing import Any, Dict, Optional
+import torch
+import numpy as np
+from diffusers.utils import USE_PEFT_BACKEND, is_torch_version, scale_lora_layers, unscale_lora_layers, logging
+from diffusers.models.modeling_outputs import Transformer2DModelOutput
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+def teacache_mochi_forward(
+ self,
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: torch.Tensor,
+ timestep: torch.LongTensor,
+ encoder_attention_mask: torch.Tensor,
+ attention_kwargs: Optional[Dict[str, Any]] = None,
+ return_dict: bool = True,
+ ) -> torch.Tensor:
+ if attention_kwargs is not None:
+ attention_kwargs = attention_kwargs.copy()
+ lora_scale = attention_kwargs.pop("scale", 1.0)
+ else:
+ lora_scale = 1.0
+
+ if USE_PEFT_BACKEND:
+ # weight the lora layers by setting `lora_scale` for each PEFT layer
+ scale_lora_layers(self, lora_scale)
+ else:
+ if attention_kwargs is not None and attention_kwargs.get("scale", None) is not None:
+ logger.warning(
+ "Passing `scale` via `attention_kwargs` when not using the PEFT backend is ineffective."
+ )
+
+ batch_size, num_channels, num_frames, height, width = hidden_states.shape
+ p = self.config.patch_size
+
+ post_patch_height = height // p
+ post_patch_width = width // p
+
+ temb, encoder_hidden_states = self.time_embed(
+ timestep,
+ encoder_hidden_states,
+ encoder_attention_mask,
+ hidden_dtype=hidden_states.dtype,
+ )
+
+ hidden_states = hidden_states.permute(0, 2, 1, 3, 4).flatten(0, 1)
+ hidden_states = self.patch_embed(hidden_states)
+ hidden_states = hidden_states.unflatten(0, (batch_size, -1)).flatten(1, 2)
+
+ image_rotary_emb = self.rope(
+ self.pos_frequencies,
+ num_frames,
+ post_patch_height,
+ post_patch_width,
+ device=hidden_states.device,
+ dtype=torch.float32,
+ )
+
+ if self.enable_teacache:
+ inp = hidden_states.clone()
+ temb_ = temb.clone()
+ modulated_inp, gate_msa, scale_mlp, gate_mlp = self.transformer_blocks[0].norm1(inp, temb_)
+ if self.cnt == 0 or self.cnt == self.num_steps-1:
+ should_calc = True
+ self.accumulated_rel_l1_distance = 0
+ else:
+ coefficients = [-3.51241319e+03, 8.11675948e+02, -6.09400215e+01, 2.42429681e+00, 3.05291719e-03]
+ rescale_func = np.poly1d(coefficients)
+ self.accumulated_rel_l1_distance += rescale_func(((modulated_inp-self.previous_modulated_input).abs().mean() / self.previous_modulated_input.abs().mean()).cpu().item())
+ if self.accumulated_rel_l1_distance < self.rel_l1_thresh:
+ should_calc = False
+ else:
+ should_calc = True
+ self.accumulated_rel_l1_distance = 0
+ self.previous_modulated_input = modulated_inp
+ self.cnt += 1
+ if self.cnt == self.num_steps:
+ self.cnt = 0
+
+ if self.enable_teacache:
+ if not should_calc:
+ hidden_states += self.previous_residual
+ else:
+ ori_hidden_states = hidden_states.clone()
+ for i, block in enumerate(self.transformer_blocks):
+ if torch.is_grad_enabled() and self.gradient_checkpointing:
+
+ def create_custom_forward(module):
+ def custom_forward(*inputs):
+ return module(*inputs)
+
+ return custom_forward
+
+ ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
+ hidden_states, encoder_hidden_states = torch.utils.checkpoint.checkpoint(
+ create_custom_forward(block),
+ hidden_states,
+ encoder_hidden_states,
+ temb,
+ encoder_attention_mask,
+ image_rotary_emb,
+ **ckpt_kwargs,
+ )
+ else:
+ hidden_states, encoder_hidden_states = block(
+ hidden_states=hidden_states,
+ encoder_hidden_states=encoder_hidden_states,
+ temb=temb,
+ encoder_attention_mask=encoder_attention_mask,
+ image_rotary_emb=image_rotary_emb,
+ )
+ hidden_states = self.norm_out(hidden_states, temb)
+ self.previous_residual = hidden_states - ori_hidden_states
+ else:
+ for i, block in enumerate(self.transformer_blocks):
+ if torch.is_grad_enabled() and self.gradient_checkpointing:
+ def create_custom_forward(module):
+ def custom_forward(*inputs):
+ return module(*inputs)
+
+ return custom_forward
+
+ ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
+ hidden_states, encoder_hidden_states = torch.utils.checkpoint.checkpoint(
+ create_custom_forward(block),
+ hidden_states,
+ encoder_hidden_states,
+ temb,
+ encoder_attention_mask,
+ image_rotary_emb,
+ **ckpt_kwargs,
+ )
+ else:
+ hidden_states, encoder_hidden_states = block(
+ hidden_states=hidden_states,
+ encoder_hidden_states=encoder_hidden_states,
+ temb=temb,
+ encoder_attention_mask=encoder_attention_mask,
+ image_rotary_emb=image_rotary_emb,
+ )
+ hidden_states = self.norm_out(hidden_states, temb)
+
+ hidden_states = self.proj_out(hidden_states)
+
+ hidden_states = hidden_states.reshape(batch_size, num_frames, post_patch_height, post_patch_width, p, p, -1)
+ hidden_states = hidden_states.permute(0, 6, 1, 2, 4, 3, 5)
+ output = hidden_states.reshape(batch_size, -1, num_frames, height, width)
+
+ if USE_PEFT_BACKEND:
+ # remove `lora_scale` from each PEFT layer
+ unscale_lora_layers(self, lora_scale)
+
+ if not return_dict:
+ return (output,)
+ return Transformer2DModelOutput(sample=output)
diff --git a/modules/ui_models.py b/modules/ui_models.py
index 59d8a9196..de8d1e729 100644
--- a/modules/ui_models.py
+++ b/modules/ui_models.py
@@ -26,6 +26,7 @@ def create_ui():
gr.HTML(elem_id="models_progress", value="")
models_image = gr.Image(elem_id="models_image", show_label=False, interactive=False, type='pil')
models_outcome = gr.HTML(elem_id="models_error", value="")
+ models_file = gr.File(label='', type='file', help='', visible=False)
with gr.Column(elem_id='models_input_container', scale=3):
@@ -38,23 +39,31 @@ def create_ui():
components = [(m.name, m.cls, m.device, m.dtype, m.params, m.modules, str(m.config)) for m in model.modules]
return [desc, components, meta]
+ with gr.Row():
+ gr.HTML(' Analyze currently loaded model
')
with gr.Row():
model_analyze = gr.Button(value="Analyze", variant='primary')
with gr.Row():
model_desc = gr.HTML(value="", elem_id="model_desc")
with gr.Row():
module_headers = ['Module', 'Class', 'Device', 'DType', 'Params', 'Modules', 'Config']
- model_types = ['str', 'str', 'str', 'str', 'number', 'number', 'str']
- model_modules = gr.DataFrame(value=None, label=None, show_label=False, interactive=False, wrap=True, headers=module_headers, datatype=model_types, type='array')
+ module_types = ['str', 'str', 'str', 'str', 'number', 'number', 'str']
+ model_modules = gr.DataFrame(value=None, label=None, show_label=False, interactive=False, wrap=True, headers=module_headers, datatype=module_types, type='array')
with gr.Row():
model_meta = gr.JSON(label="Metadata", value={}, elem_id="model_meta")
model_analyze.click(fn=analyze, inputs=[], outputs=[model_desc, model_modules, model_meta])
+ with gr.Tab(label="Loader"):
+ from modules import ui_models_load
+ ui_models_load.create_ui(models_outcome, models_file)
+
with gr.Tab(label="Merge"):
def sd_model_choices():
return ['None'] + sd_models.checkpoint_titles()
+ with gr.Row():
+ gr.HTML(' Merge multiple models
')
with gr.Row(equal_height=False):
with gr.Column(variant='compact'):
with gr.Row():
@@ -290,6 +299,8 @@ def create_ui():
)
with gr.Tab(label="Modules"):
+ with gr.Row():
+ gr.HTML(' Replace model components
')
with gr.Row():
with gr.Column(scale=3):
model_type = gr.Dropdown(label="Model type", choices=['sd15', 'sdxl', 'sd21', 'sd35', 'flux.1'], value='sdxl', interactive=False)
@@ -363,6 +374,8 @@ def create_ui():
model_headers = ['name', 'type', 'filename', 'hash', 'added', 'size', 'metadata']
model_data = []
+ with gr.Row():
+ gr.HTML(' List all models
')
with gr.Row():
model_list_btn = gr.Button(value="List model details", variant='primary')
model_checkhash_btn = gr.Button(value="Calculate hash for all models", variant='primary')
@@ -434,7 +447,8 @@ def create_ui():
opts.save()
with gr.Column(scale=6):
- gr.HTML('Search for models
Select a model from the search results to download
')
+ with gr.Row():
+ gr.HTML(' Download model from huggingface
')
with gr.Row():
hf_search_text = gr.Textbox('', label='Search models', placeholder='search huggingface models')
hf_search_btn = ToolButton(value=ui_symbols.search, label="Search")
@@ -646,7 +660,8 @@ def create_ui():
opts.save()
with gr.Row():
- gr.HTML('Fetch information
Fetches preview and metadata information for all models with missing information
Models with existing previews and information are not updated
')
+ gr.HTML(' CivitAI fetch metadata
')
+ gr.HTML('Fetches preview and metadata information for all models with missing information
Models with existing previews and information are not updated
')
with gr.Row():
civit_previews_btn = gr.Button(value="Start", variant='primary')
with gr.Row():
@@ -665,7 +680,7 @@ def create_ui():
with gr.Row():
civit_search_res = gr.HTML('')
with gr.Row():
- gr.HTML('Download model
')
+ gr.HTML(' CivitAI download model
')
with gr.Row():
civit_download_model_btn = gr.Button(value="Download", variant='primary')
gr.HTML('Select a model, model version and and model variant from the search results to download or enter model URL manually
')
@@ -719,7 +734,7 @@ def create_ui():
with gr.Tab(label="Update"):
with gr.Row():
- gr.HTML('Fetch most recent information about all installed models
')
+ gr.HTML(' Scan CivitAI for information on latest available model versions
')
with gr.Row():
civit_update_btn = gr.Button(value="Update", variant='primary')
with gr.Row():
diff --git a/modules/ui_models_load.py b/modules/ui_models_load.py
new file mode 100644
index 000000000..902769365
--- /dev/null
+++ b/modules/ui_models_load.py
@@ -0,0 +1,318 @@
+import os
+import re
+import json # pylint: disable=unused-import
+import inspect
+import gradio as gr
+import torch
+import diffusers
+from huggingface_hub import hf_hub_download
+from modules import shared, errors, shared_items, sd_models, sd_checkpoint, devices, model_quant, modelloader
+
+
+debug_enabled = os.environ.get('SD_LOAD_DEBUG', None)
+debug_log = shared.log.trace if debug_enabled else lambda *args, **kwargs: None
+components = []
+
+
+def load_model(model: str, cls: str, repo: str, dataframes: list):
+ if cls is None:
+ shared.log.error('Model load: class is None')
+ return 'Model load: class is None'
+ if repo is None:
+ shared.log.error('Model load: repo is None')
+ return 'Model load: repo is None'
+ cls = getattr(diffusers, cls, None)
+ if cls is None:
+ cls = diffusers.AutoPipelineForText2Image
+ shared.log.info(f'Model load: name="{model}" cls={cls.__name__} repo="{repo}"')
+ kwargs = {}
+ for df in dataframes:
+ c = [x for x in components if x.id == df[0]]
+ if len(c) != 1:
+ debug_log(f'Model load component: id={df[0]} not found')
+ continue
+ c = c[0]
+ if not c.loadable: # not loadable
+ debug_log(f'Model load component: name={c.name} not loadable')
+ continue
+ if c.type != 'class':
+ debug_log(f'Model load component: name={c.name} not class')
+ continue
+ if len(c.local or '') == 0 and len(c.remote or '') == 0:
+ debug_log(f'Model load component: name={c.name} no local or remote')
+ continue
+ instance = c.load()
+ if instance is not None:
+ kwargs[c.name] = instance
+ shared.log.info(f'Model component: instance={instance.__class__.__name__}')
+ shared.log.info(f'Model load: name="{model}" cls={cls.__name__} repo="{repo}" preload={kwargs.keys()}')
+ pipe = None
+ if model == 'Current':
+ for k, v in kwargs.items():
+ debug_log(f'Model replace component={k}')
+ setattr(shared.sd_model, k, v)
+ sd_models.set_diffuser_options(shared.sd_model)
+ return f'Model load: name="{model}" cls={cls.__name__} repo="{repo}" preload={kwargs.keys()}'
+ else:
+ try:
+ pipe = cls.from_pretrained(
+ repo,
+ dtype=devices.dtype,
+ cache_dir=shared.opts.diffusers_dir,
+ **kwargs,
+ )
+ except Exception as e:
+ shared.log.error(f'Model load: name="{model}" {e}')
+ errors.display(e, 'Model load')
+ return f'Model load failed: {e}'
+ if pipe is not None:
+ shared.log.info(f'Model load: name="{model}" cls={cls.__name__} repo="{repo}" instance={pipe.__class__.__name__}')
+ shared.sd_model = pipe
+ shared.sd_model.sd_checkpoint_info = sd_checkpoint.CheckpointInfo(repo)
+ shared.sd_model.sd_model_hash = None
+ sd_models.set_diffuser_options(shared.sd_model)
+ return f'Model load: name="{model}" cls={cls.__name__} repo="{repo}" preload={kwargs.keys()}'
+ return 'Model load: no model'
+
+
+def unload_model():
+ sd_models.unload_model_weights(op='model')
+ return 'Model unloaded'
+
+
+def process_huggingface_url(url):
+ if url is None or len(url) == 0:
+ return None, None, None, False
+ url = url.replace('https://huggingface.co/', '').strip() # remove absolute url
+ url = re.sub(r'/blob/[^/]+/', '/', url) # remove /blob//
+ parts = url.split('/')
+ repo = f"{parts[0]}/{parts[1]}" if len(parts) >= 2 else url # get repo
+ subfolder = None
+ fn = None
+ if len(parts) == 3: # can be subfolder or filename
+ if '.' in parts[-1]:
+ fn = parts[-1]
+ else:
+ subfolder = parts[-1]
+ elif len(parts) > 3: # There's at least one subfolder
+ subfolder = '/'.join(parts[2:-1])
+ fn = parts[-1]
+ download = fn is not None
+ return repo, subfolder, fn, download
+
+
+class Component():
+ def __init__(self, signature, name=None, cls=None, val=None, local=None, remote=None, typ=None, dtype=None, quant=False, loadable=None):
+ self.id = len(components) + 1
+ self.name = signature.name if signature else name
+ self.cls = signature.annotation if signature else cls
+ self.str = str(signature.annotation) if signature else str(cls)
+ self.val = signature.default if signature and signature.default is not inspect.Parameter.empty else val
+ self.remote = remote
+ self.repo, self.subfolder, self.local, self.download = process_huggingface_url(self.remote)
+ self.local = local or self.local
+ self.dtype = str(dtype or devices.dtype).rsplit('.', maxsplit=1)[-1]
+ self.quant = quant
+ self.revision = None
+ self.enum = None
+ if typ is not None:
+ self.type = typ
+ else:
+ if self.cls in [str, int, float, bool]:
+ self.type = 'variable'
+ elif 'enum' in self.str:
+ self.type = 'enum'
+ self.enum = [v.name for v in self.cls]
+ elif inspect.isclass(signature.annotation):
+ self.type = 'class'
+ elif inspect.ismodule(signature.annotation):
+ self.type = 'module'
+ elif inspect.isfunction(signature.annotation):
+ self.type = 'function'
+ elif 'typing.Optional' in self.str:
+ self.type = 'optional'
+ self.cls = signature.annotation.__args__[0]
+ self.str = str(self.cls)
+ self.val = None
+ else:
+ self.type = 'unknown'
+ self.str = re.search(r"'(.*?)'", self.str).group(1) if re.search(r"'(.*?)'", self.str) else self.str
+ if '.' in self.str:
+ self.str = self.str.split('.')
+ self.str = self.str[0] + '.' + self.str[-1]
+ self.loadable = loadable if loadable is not None else (self.type == 'class' and hasattr(self.cls, 'from_pretrained'))
+ if not self.loadable:
+ self.dtype = None
+ self.quant = None
+
+ def __str__(self):
+ return f'id={self.id} name="{self.name}" cls={self.cls} type={self.type} loadable={self.loadable} val="{self.val}" str="{self.str}" enum="{self.enum}" local="{self.local}" remote="{self.remote}" repo="{self.repo}" subfolder="{self.subfolder}" dtype={self.dtype} quant={self.quant} revision={self.revision}'
+
+ def save(self):
+ return [self.name, self.local, self.remote, self.dtype, self.quant]
+
+ def dataframe(self):
+ return [self.id, self.name, self.loadable, self.val, self.str, self.local, self.remote, self.dtype, self.quant]
+
+ def load(self):
+ if not self.loadable:
+ return None
+ modelloader.hf_login()
+
+ load_args = {}
+ if self.subfolder is not None:
+ load_args['subfolder'] = self.subfolder
+ if self.revision is not None:
+ load_args['revision'] = self.revision
+ if self.dtype is not None:
+ load_args['torch_dtype'] = getattr(torch, self.dtype)
+ if not hasattr(self.cls, 'from_pretrained'):
+ debug_log(f'Model load component: name="{self.name}" cls={self.cls} not loadable')
+ return None
+ quant_args = model_quant.create_config(module='any', allow=self.quant)
+ quant_type = model_quant.get_quant_type(quant_args)
+
+ try:
+ if self.download:
+ debug_log(f'Model load component: url="{self.remote}" args={load_args} quant={quant_type}')
+ self.local = hf_hub_download(
+ repo_id=self.repo,
+ subfolder=self.subfolder,
+ filename=self.local,
+ revision=self.revision,
+ cache_dir=shared.opts.hfcache_dir,
+ )
+ if os.path.exists(self.local):
+ self.download = False
+ if self.local is not None and len(self.local) > 0:
+ if not os.path.exists(self.local):
+ debug_log(f'Model load component: local="{self.local}" file not found')
+ elif hasattr(self.cls, 'from_single_file') and os.path.isfile(self.local) and self.local.endswith('.safetensors'):
+ debug_log(f'Model load component: local="{self.local}" type=file args={load_args} quant={quant_type}')
+ return self.cls.from_single_file(self.local, **load_args, **quant_args, cache_dir=shared.opts.hfcache_dir)
+ elif os.path.isfile(self.local) and self.local.endswith('.gguf'):
+ debug_log(f'Model load component: local="{self.local}" type=gguf args={load_args} quant={quant_type}')
+ from modules import ggml
+ return ggml.load_gguf(self.local, cls=self.cls, compute_dtype=self.dtype)
+ else:
+ debug_log(f'Model load component: local="{self.local}" type=folder args={load_args} quant={quant_type}')
+ return self.cls.from_pretrained(self.local, **load_args, **quant_args, cache_dir=shared.opts.hfcache_dir)
+ elif self.repo is not None and len(self.repo) > 0:
+ debug_log(f'Model load component: repo="{self.repo}" args={load_args} quant={quant_type}')
+ return self.cls.from_pretrained(self.repo, **load_args, **quant_args, cache_dir=shared.opts.hfcache_dir)
+ elif self.val is not None and len(self.val) > 0:
+ debug_log(f'Model load component: default="{self.val}" args={load_args} quant={quant_type}')
+ return self.cls.from_pretrained(self.val, **load_args, **quant_args, cache_dir=shared.opts.hfcache_dir)
+ else:
+ debug_log(f'Model load component: name="{self.name}" cls={self.cls} no handler')
+ return None
+ except Exception as e:
+ shared.log.error(f'Model load component: name="{self.name}" {e}')
+ errors.display(e, 'Model load component')
+ return None
+
+
+def create_ui(gr_status, gr_file):
+ def get_components(cls):
+ if cls is None:
+ return []
+ signature = inspect.signature(cls.__init__, follow_wrapped=True)
+ components.clear()
+ for param in signature.parameters.values():
+ if param.name == 'self' or param.name == 'args' or param.name == 'kwargs':
+ continue
+ component = Component(param)
+ debug_log(f'Model component: {str(component)}')
+ components.append(component)
+ return components
+
+ def get_model(model):
+ if model == 'Current':
+ cls = shared.sd_model.__class__
+ else:
+ cls = shared_items.pipelines.get(model, None)
+ if cls is None:
+ cls = diffusers.AutoPipelineForText2Image
+ name = cls.__name__
+ repo = shared_items.get_repo(name) or shared_items.get_repo(model)
+ link = f'Link
{repo}' if repo else ''
+ get_components(cls)
+ dataframes = [c.dataframe() for c in components]
+ shared.log.debug(f'Model select: name="{model}" cls={name} repo="{repo}" link={link} components={len(components)}')
+ return [name, repo, link, dataframes]
+
+ def update_component(dataframes):
+ for df in dataframes:
+ c = [x for x in components if x.id == df[0]]
+ if len(c) != 1:
+ continue
+ c = c[0]
+ c.local = df[5].strip()
+ c.remote = df[6].strip()
+ c.dtype = df[7]
+ c.quant = df[8]
+ if c.remote and len(c.remote) > 0:
+ c.repo, c.subfolder, c.local, c.download = process_huggingface_url(c.remote)
+
+ # TODO loader: load receipe
+ def load_receipe(file_select):
+ if file_select is not None and 'name' in file_select:
+ fn = file_select['name']
+ shared.log.debug(f'Load receipe: fn={fn}')
+ return ['Load receipe not implemented yet', gr.update(label='Receipe .json file', file_types=['json'], visible=True)]
+
+ # TODO loader: save receipe
+ def save_receipe(model: str, repo: str):
+ receipe = {
+ 'model': model,
+ 'repo': repo,
+ 'components': []
+ }
+ for c in components:
+ if c.loadable:
+ receipe['components'].append(c.save())
+ # with open('/tmp/receipe.json', 'w', encoding='utf8') as f:
+ # json.dump(receipe, f, indent=2)
+ return 'Save receipe not implemented yet'
+
+ with gr.Row():
+ gr.HTML('')
+ with gr.Row():
+ choices = list(shared_items.pipelines)
+ choices = ['Current' if x.startswith('Custom') else x for x in choices]
+ model = gr.Dropdown(label="Model type", choices=choices, value='Autodetect')
+ cls = gr.Textbox(label="Model class", placeholder="Class name", interactive=False)
+ with gr.Row():
+ repo = gr.Textbox(label="Model repo", placeholder="Repo name", interactive=True)
+ link = gr.HTML(value="", interactive=False)
+ with gr.Row():
+ headers = ['ID', 'Name', 'Loadable', 'Default', 'Class', 'Local', 'Remote', 'Dtype', 'Quant']
+ datatype = ['number', 'str', 'bool', 'str', 'str', 'str', 'str', 'str', 'bool']
+ dataframes = gr.DataFrame(
+ value=None,
+ label=None,
+ show_label=False,
+ interactive=True,
+ wrap=True,
+ headers=headers,
+ datatype=datatype,
+ max_rows=None,
+ max_cols=None,
+ type='array',
+ elem_id="model_loader_df",
+ )
+ dataframes.change(fn=update_component, inputs=[dataframes], outputs=[])
+
+ model.change(get_model, inputs=[model], outputs=[cls, repo, link, dataframes])
+
+ with gr.Row():
+ btn_load_receipe = gr.Button(value="Load receipe")
+ btn_save_receipe = gr.Button(value="Save receipe")
+ with gr.Row():
+ btn_load_model = gr.Button(value="Load model")
+ btn_unload_model = gr.Button(value="Unload model")
+
+ btn_load_receipe.click(fn=load_receipe, inputs=[gr_file], outputs=[gr_status, gr_file])
+ btn_save_receipe.click(fn=save_receipe, inputs=[model, repo], outputs=[gr_status])
+ btn_load_model.click(fn=load_model, inputs=[model, cls, repo, dataframes], outputs=[gr_status])
+ btn_unload_model.click(fn=unload_model, inputs=[], outputs=[gr_status])
diff --git a/modules/ui_sections.py b/modules/ui_sections.py
index 2785c40b6..418b5679d 100644
--- a/modules/ui_sections.py
+++ b/modules/ui_sections.py
@@ -181,7 +181,7 @@ def create_advanced_inputs(tab, base=True):
cfg_scale, cfg_end = None, None
with gr.Row():
image_cfg_scale = gr.Slider(minimum=0.0, maximum=30.0, step=0.1, label='Refine guidance', value=6.0, elem_id=f"{tab}_image_cfg_scale")
- diffusers_guidance_rescale = gr.Slider(minimum=0.0, maximum=1.0, step=0.05, label='Rescale guidance', value=0.7, elem_id=f"{tab}_image_cfg_rescale", visible=shared.native)
+ diffusers_guidance_rescale = gr.Slider(minimum=0.0, maximum=1.0, step=0.05, label='Rescale guidance', value=0.0, elem_id=f"{tab}_image_cfg_rescale", visible=shared.native)
with gr.Row():
diffusers_pag_scale = gr.Slider(minimum=0.0, maximum=30.0, step=0.05, label='Attention guidance', value=0.0, elem_id=f"{tab}_pag_scale", visible=shared.native)
diffusers_pag_adaptive = gr.Slider(minimum=0.0, maximum=1.0, step=0.05, label='Adaptive scaling', value=0.5, elem_id=f"{tab}_pag_adaptive", visible=shared.native)
diff --git a/modules/video_models/models_def.py b/modules/video_models/models_def.py
index 69287bb65..31c6354f3 100644
--- a/modules/video_models/models_def.py
+++ b/modules/video_models/models_def.py
@@ -74,7 +74,7 @@ models = {
Model(name='LTXVideo 0.9.5 T2V', # https://github.com/huggingface/diffusers/pull/10968
url='https://huggingface.co/Lightricks/LTX-Video-0.9.5',
repo='Lightricks/LTX-Video-0.9.5',
- repo_cls=diffusers.LTXPipeline,
+ repo_cls=diffusers.LTXConditionPipeline,
te_cls=transformers.T5EncoderModel,
dit_cls=diffusers.LTXVideoTransformer3DModel),
Model(name='LTXVideo 0.9.5 I2V',
diff --git a/modules/video_models/video_cache.py b/modules/video_models/video_cache.py
new file mode 100644
index 000000000..01f541be3
--- /dev/null
+++ b/modules/video_models/video_cache.py
@@ -0,0 +1,16 @@
+import diffusers
+from modules import shared
+
+
+def apply_teacache_patch(cls):
+ if shared.opts.teacache_enabled:
+ from modules import teacache
+ shared.log.debug(f'Transformers cache: type=teacache patch=forward cls={cls.__name__}')
+ if cls.__name__ == 'LTXVideoTransformer3DModel':
+ cls.forward = teacache.teacache_ltx_forward
+ elif cls.__name__ == 'MochiTransformer3DModel':
+ cls.forward = teacache.teacache_mochi_forward
+ elif cls.__name__ == 'CogVideoXTransformer3DModel':
+ cls.forward = teacache.teacache_cog_forward
+
+ diffusers.FluxTransformer2DModel.forward = teacache.teacache_flux_forward
diff --git a/modules/video_models/video_load.py b/modules/video_models/video_load.py
index 13d0dc935..12f02caf1 100644
--- a/modules/video_models/video_load.py
+++ b/modules/video_models/video_load.py
@@ -1,7 +1,7 @@
import os
import time
from modules import shared, errors, sd_models, sd_checkpoint, model_quant, devices
-from modules.video_models import models_def, video_utils, video_vae, video_overrides
+from modules.video_models import models_def, video_utils, video_vae, video_overrides, video_cache
loaded_model = None
@@ -17,10 +17,12 @@ def load_model(selected: models_def.Model):
sd_models.unload_model_weights()
t0 = time.time()
+ video_cache.apply_teacache_patch(selected.dit_cls)
+
# text encoder
try:
quant_args = model_quant.create_config(module='TE')
- debug(f'Video load: module=te repo="{selected.te or selected.repo}" folder="{selected.te_folder}" cls={selected.te_cls.__name__} quant={video_utils.get_quant(quant_args)}')
+ debug(f'Video load: module=te repo="{selected.te or selected.repo}" folder="{selected.te_folder}" cls={selected.te_cls.__name__} quant={model_quant.get_quant_type(quant_args)}')
text_encoder = selected.te_cls.from_pretrained(
pretrained_model_name_or_path=selected.te or selected.repo,
subfolder=selected.te_folder,
@@ -36,7 +38,7 @@ def load_model(selected: models_def.Model):
# transformer
try:
quant_args = model_quant.create_config(module='Video')
- debug(f'Video load: module=transformer repo="{selected.dit or selected.repo}" folder="{selected.dit_folder}" cls={selected.dit_cls.__name__} quant={video_utils.get_quant(quant_args)}')
+ debug(f'Video load: module=transformer repo="{selected.dit or selected.repo}" folder="{selected.dit_folder}" cls={selected.dit_cls.__name__} quant={model_quant.get_quant_type(quant_args)}')
transformer = selected.dit_cls.from_pretrained(
pretrained_model_name_or_path=selected.dit or selected.repo,
subfolder=selected.dit_folder,
diff --git a/modules/video_models/video_run.py b/modules/video_models/video_run.py
index 2b4fbcbea..b0248e914 100644
--- a/modules/video_models/video_run.py
+++ b/modules/video_models/video_run.py
@@ -56,6 +56,10 @@ def generate(*args, **kwargs):
if init_image is None:
return video_utils.queue_err('init image not set')
p.task_args['image'] = images.resize_image(resize_mode=2, im=init_image, width=p.width, height=p.height, upscaler_name=None, output_type='pil')
+ shared.log.debug(f'Video: op=I2V init={init_image} resized={p.task_args["image"]}')
+ elif 'T2V' in model:
+ if init_image is not None:
+ shared.log.debug('Video: op=T2V init image not supported')
# cleanup memory
shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model)
diff --git a/modules/video_models/video_utils.py b/modules/video_models/video_utils.py
index d0a80d42a..16edb070d 100644
--- a/modules/video_models/video_utils.py
+++ b/modules/video_models/video_utils.py
@@ -11,12 +11,6 @@ def queue_err(msg):
return [], None, '', '', f'Error: {msg}'
-def get_quant(args):
- if args is not None and "quantization_config" in args:
- return args['quantization_config'].__class__.__name__
- return None
-
-
def get_url(url):
return f' {url}
' if url else '
'
diff --git a/modules/zluda.py b/modules/zluda.py
index 431ab2c8c..82b825936 100644
--- a/modules/zluda.py
+++ b/modules/zluda.py
@@ -33,8 +33,8 @@ def initialize_zluda():
from modules.zluda_hijacks import do_hijack
do_hijack()
- torch.backends.cudnn.enabled = zluda_installer.MIOpen_available
- if not zluda_installer.MIOpen_available:
+ torch.backends.cudnn.enabled = zluda_installer.MIOpen_enabled
+ if not zluda_installer.MIOpen_enabled:
torch.backends.cuda.enable_cudnn_sdp(False)
torch.backends.cuda.enable_cudnn_sdp = do_nothing
torch.backends.cuda.enable_flash_sdp(False)
diff --git a/modules/zluda_installer.py b/modules/zluda_installer.py
index 54a0f9234..2f097f707 100644
--- a/modules/zluda_installer.py
+++ b/modules/zluda_installer.py
@@ -19,8 +19,7 @@ DLL_MAPPING = {
}
HIPSDK_TARGETS = ['rocblas.dll', 'rocsolver.dll', 'hipfft.dll',]
-hipBLASLt_available = False
-MIOpen_available = False
+MIOpen_enabled = False
path = os.path.abspath(os.environ.get('ZLUDA', '.zluda'))
default_agent: Union[rocm.Agent, None] = None
@@ -65,36 +64,16 @@ core = None
ml = None
-def load_core_modules():
- global core, ml # pylint: disable=global-statement
- core = Core(ctypes.windll.LoadLibrary(os.path.join(path, 'nvcuda.dll')))
- ml = ZLUDALibrary(ctypes.windll.LoadLibrary(os.path.join(path, 'nvml.dll')))
-
-
def set_default_agent(agent: rocm.Agent):
global default_agent # pylint: disable=global-statement
default_agent = agent
- is_nightly = False
- try:
- load_core_modules()
- is_nightly = core.get_nightly_flag() == 1
- except Exception:
- pass
-
- global hipBLASLt_available, hipBLASLt_enabled # pylint: disable=global-statement
- hipBLASLt_available = is_nightly and os.path.exists(rocm.blaslt_tensile_libpath)
- hipBLASLt_enabled = hipBLASLt_available and os.path.exists(os.path.join(rocm.path, "bin", "hipblaslt.dll"))
-
- global MIOpen_available # pylint: disable=global-statement
- MIOpen_available = is_nightly and os.path.exists(os.path.join(rocm.path, "bin", "MIOpen.dll"))
-
def is_reinstall_needed() -> bool: # ZLUDA<3.8.7
return not os.path.exists(os.path.join(path, 'cufftw.dll'))
-def install() -> None:
+def install():
if os.path.exists(path):
return
@@ -115,7 +94,7 @@ def install() -> None:
os.remove('_zluda')
-def uninstall() -> None:
+def uninstall():
if os.path.exists(path):
shutil.rmtree(path)
@@ -139,7 +118,14 @@ def link_or_copy(src: os.PathLike, dst: os.PathLike):
shutil.copyfile(src, dst)
-def make_copy() -> None:
+def load():
+ global core, ml, hipBLASLt_enabled, MIOpen_enabled # pylint: disable=global-statement
+ core = Core(ctypes.windll.LoadLibrary(os.path.join(path, 'nvcuda.dll')))
+ ml = ZLUDALibrary(ctypes.windll.LoadLibrary(os.path.join(path, 'nvml.dll')))
+ is_nightly = core.get_nightly_flag() == 1
+ hipBLASLt_enabled = is_nightly and os.path.exists(rocm.blaslt_tensile_libpath) and os.path.exists(os.path.join(rocm.path, "bin", "hipblaslt.dll"))
+ MIOpen_enabled = is_nightly and os.path.exists(os.path.join(rocm.path, "bin", "MIOpen.dll"))
+
for k, v in DLL_MAPPING.items():
if not os.path.exists(os.path.join(path, v)):
link_or_copy(os.path.join(path, k), os.path.join(path, v))
@@ -147,17 +133,14 @@ def make_copy() -> None:
if hipBLASLt_enabled and not os.path.exists(os.path.join(path, 'cublasLt64_11.dll')):
link_or_copy(os.path.join(path, 'cublasLt.dll'), os.path.join(path, 'cublasLt64_11.dll'))
- if MIOpen_available and not os.path.exists(os.path.join(path, 'cudnn64_9.dll')):
+ if MIOpen_enabled and not os.path.exists(os.path.join(path, 'cudnn64_9.dll')):
link_or_copy(os.path.join(path, 'cudnn.dll'), os.path.join(path, 'cudnn64_9.dll'))
-
-def load() -> None:
log.info(f"ZLUDA load: path='{path}' nightly={bool(core.get_nightly_flag())}")
os.environ["ZLUDA_COMGR_LOG_LEVEL"] = "1"
os.environ["ZLUDA_NVRTC_LIB"] = os.path.join([v for v in site.getsitepackages() if v.endswith("site-packages")][0], "torch", "lib", "nvrtc64_112_0.dll")
- load_core_modules()
for v in HIPSDK_TARGETS:
ctypes.windll.LoadLibrary(os.path.join(rocm.path, 'bin', v))
for v in DLL_MAPPING.values():
@@ -170,7 +153,7 @@ def load() -> None:
else:
os.environ["DISABLE_ADDMM_CUDA_LT"] = "1"
- if MIOpen_available:
+ if MIOpen_enabled:
ctypes.windll.LoadLibrary(os.path.join(rocm.path, 'bin', 'MIOpen.dll'))
ctypes.windll.LoadLibrary(os.path.join(path, 'cudnn64_9.dll'))
diff --git a/noised_image.png b/noised_image.png
new file mode 100644
index 000000000..70c882cfe
Binary files /dev/null and b/noised_image.png differ
diff --git a/requirements.txt b/requirements.txt
index f1de3abd7..0bc763167 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -34,7 +34,7 @@ pi-heif
rich==13.9.4
safetensors==0.5.3
tensordict==0.1.2
-peft==0.14.0
+peft==0.15.1
httpx==0.24.1
compel==2.0.3
torchsde==0.2.6
@@ -52,7 +52,7 @@ numba==0.59.1
protobuf==4.25.3
pytorch_lightning==1.9.4
tokenizers==0.21.1
-transformers==4.50.3
+transformers==4.51.1
urllib3==1.26.19
Pillow==10.4.0
timm==0.9.16
diff --git a/scripts/allegrovideo.py b/scripts/allegrovideo.py
index cf35500fb..f1d1ab45f 100644
--- a/scripts/allegrovideo.py
+++ b/scripts/allegrovideo.py
@@ -21,7 +21,7 @@ def hijack_decode(*args, **kwargs):
def hijack_encode_prompt(*args, **kwargs):
t0 = time.time()
- res = shared.sd_model.vae.orig_encode_prompt(*args, **kwargs)
+ res = shared.sd_model.orig_encode_prompt(*args, **kwargs)
t1 = time.time()
timer.process.add('te', t1-t0)
shared.log.debug(f'Video: te={shared.sd_model.text_encoder.__class__.__name__} time={t1-t0:.2f}')
@@ -92,7 +92,7 @@ class Script(scripts.Script):
shared.sd_model.sd_checkpoint_info = sd_checkpoint.CheckpointInfo(repo_id)
shared.sd_model.sd_model_hash = None
shared.sd_model.vae.orig_decode = shared.sd_model.vae.decode
- shared.sd_model.vae.orig_encode_prompt = shared.sd_model.encode_prompt
+ shared.sd_model.orig_encode_prompt = shared.sd_model.encode_prompt
shared.sd_model.vae.decode = hijack_decode
shared.sd_model.encode_prompt = hijack_encode_prompt
shared.sd_model.vae.enable_tiling()
diff --git a/scripts/hunyuanvideo.py b/scripts/hunyuanvideo.py
index c39cec688..50cbf567f 100644
--- a/scripts/hunyuanvideo.py
+++ b/scripts/hunyuanvideo.py
@@ -50,7 +50,7 @@ def hijack_decode(*args, **kwargs):
def hijack_encode_prompt(*args, **kwargs):
t0 = time.time()
- res = shared.sd_model.vae.orig_encode_prompt(*args, **kwargs)
+ res = shared.sd_model.orig_encode_prompt(*args, **kwargs)
t1 = time.time()
timer.process.add('te', t1-t0)
shared.log.debug(f'Video: te={shared.sd_model.text_encoder.__class__.__name__} time={t1-t0:.2f}')
@@ -133,7 +133,7 @@ class Script(scripts.Script):
shared.sd_model.sd_checkpoint_info = sd_checkpoint.CheckpointInfo(models.get(model)['repo'])
shared.sd_model.sd_model_hash = None
shared.sd_model.vae.orig_decode = shared.sd_model.vae.decode
- shared.sd_model.vae.orig_encode_prompt = shared.sd_model.encode_prompt
+ shared.sd_model.orig_encode_prompt = shared.sd_model.encode_prompt
shared.sd_model.vae.decode = hijack_decode
shared.sd_model.encode_prompt = hijack_encode_prompt
shared.sd_model.vae.enable_slicing()
diff --git a/scripts/infiniteyou_ext.py b/scripts/infiniteyou_ext.py
index 59432cd40..a8ca9c9c3 100644
--- a/scripts/infiniteyou_ext.py
+++ b/scripts/infiniteyou_ext.py
@@ -80,6 +80,8 @@ class Script(scripts.Script):
if shared.sd_model_type != 'f1':
shared.log.error(f'{prefix}: invalid model type: {shared.sd_model_type}')
return None
+ if scale <= 0:
+ return None
global orig_pipeline, orig_prompt_attention # pylint: disable=global-statement
orig_pipeline = shared.sd_model
@@ -92,9 +94,9 @@ class Script(scripts.Script):
processing.fix_seed(p)
p.task_args['id_image'] = id_image
p.task_args['control_image'] = control_image
- p.task_args['infusenet_conditioning_scale'] = scale
- p.task_args['infusenet_guidance_start'] = start
- p.task_args['infusenet_guidance_end'] = end
+ p.task_args['infusenet_conditioning_scale'] = p.task_args.get('infusenet_conditioning_scale', scale)
+ p.task_args['infusenet_guidance_start'] = p.task_args.get('infusenet_guidance_start', start)
+ p.task_args['infusenet_guidance_end'] = p.task_args.get('infusenet_guidance_end', end)
p.task_args['seed'] = p.seed
p.task_args['negative_prompt'] = None
p.task_args['guidance_scale'] = id_guidance
diff --git a/scripts/legacy_allegrovideo.py b/scripts/legacy_allegrovideo.py
index cf35500fb..f1d1ab45f 100644
--- a/scripts/legacy_allegrovideo.py
+++ b/scripts/legacy_allegrovideo.py
@@ -21,7 +21,7 @@ def hijack_decode(*args, **kwargs):
def hijack_encode_prompt(*args, **kwargs):
t0 = time.time()
- res = shared.sd_model.vae.orig_encode_prompt(*args, **kwargs)
+ res = shared.sd_model.orig_encode_prompt(*args, **kwargs)
t1 = time.time()
timer.process.add('te', t1-t0)
shared.log.debug(f'Video: te={shared.sd_model.text_encoder.__class__.__name__} time={t1-t0:.2f}')
@@ -92,7 +92,7 @@ class Script(scripts.Script):
shared.sd_model.sd_checkpoint_info = sd_checkpoint.CheckpointInfo(repo_id)
shared.sd_model.sd_model_hash = None
shared.sd_model.vae.orig_decode = shared.sd_model.vae.decode
- shared.sd_model.vae.orig_encode_prompt = shared.sd_model.encode_prompt
+ shared.sd_model.orig_encode_prompt = shared.sd_model.encode_prompt
shared.sd_model.vae.decode = hijack_decode
shared.sd_model.encode_prompt = hijack_encode_prompt
shared.sd_model.vae.enable_tiling()
diff --git a/scripts/ltxvideo.py b/scripts/ltxvideo.py
index 697e64021..84b5953c4 100644
--- a/scripts/ltxvideo.py
+++ b/scripts/ltxvideo.py
@@ -5,7 +5,6 @@ import gradio as gr
import diffusers
import transformers
from modules import scripts, processing, shared, images, devices, sd_models, sd_checkpoint, model_quant, timer
-from modules.teacache.teacache_ltx import teacache_forward
repos = {
@@ -42,7 +41,7 @@ def hijack_decode(*args, **kwargs):
def hijack_encode_prompt(*args, **kwargs):
t0 = time.time()
- res = shared.sd_model.vae.orig_encode_prompt(*args, **kwargs)
+ res = shared.sd_model.orig_encode_prompt(*args, **kwargs)
t1 = time.time()
timer.process.add('te', t1-t0)
shared.log.debug(f'Video: te={shared.sd_model.text_encoder.__class__.__name__} time={t1-t0:.2f}')
@@ -113,7 +112,6 @@ class Script(scripts.Script):
if shared.sd_model.__class__ != cls:
sd_models.unload_model_weights()
kwargs = model_quant.create_config()
- diffusers.LTXVideoTransformer3DModel.forward = teacache_forward
if os.path.isfile(repo_id):
shared.sd_model = cls.from_single_file(
repo_id,
@@ -131,7 +129,7 @@ class Script(scripts.Script):
)
sd_models.set_diffuser_options(shared.sd_model)
shared.sd_model.vae.orig_decode = shared.sd_model.vae.decode
- shared.sd_model.vae.orig_encode_prompt = shared.sd_model.encode_prompt
+ shared.sd_model.orig_encode_prompt = shared.sd_model.encode_prompt
shared.sd_model.vae.decode = hijack_decode
shared.sd_model.encode_prompt = hijack_encode_prompt
shared.sd_model.sd_checkpoint_info = sd_checkpoint.CheckpointInfo(repo_id)
diff --git a/scripts/softfill.py b/scripts/softfill.py
new file mode 100644
index 000000000..7b90a845d
--- /dev/null
+++ b/scripts/softfill.py
@@ -0,0 +1,1671 @@
+# pylint: skip-file
+
+"""
+credits: https://github.com/zacheryvaughn/softfill-pipelines
+code from: https://github.com/zacheryvaughn/softfill-pipelines/blob/main/pipeline_stable_diffusion_xl_softfill.py
+sdnext implementation follows after pipeline-end
+"""
+
+pnoise2 = None # dynamically instlled and imported module
+
+### pipeline start
+
+import inspect
+import random
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+
+import cv2
+import numpy as np
+from PIL import Image, ImageFilter
+import torch
+import torchvision
+from torchvision import transforms
+from transformers import (
+ CLIPImageProcessor,
+ CLIPTextModel,
+ CLIPTextModelWithProjection,
+ CLIPTokenizer,
+ CLIPVisionModelWithProjection,
+)
+
+from diffusers.image_processor import PipelineImageInput, VaeImageProcessor
+from diffusers.loaders import (
+ FromSingleFileMixin,
+ IPAdapterMixin,
+ StableDiffusionXLLoraLoaderMixin,
+ TextualInversionLoaderMixin,
+)
+from diffusers.models import AutoencoderKL, ImageProjection, UNet2DConditionModel
+from diffusers.models.attention_processor import (
+ AttnProcessor2_0,
+ XFormersAttnProcessor,
+)
+from diffusers.models.lora import adjust_lora_scale_text_encoder
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline, StableDiffusionMixin
+from diffusers.pipelines.stable_diffusion_xl.pipeline_output import StableDiffusionXLPipelineOutput
+from diffusers.schedulers import KarrasDiffusionSchedulers
+from diffusers.utils import (
+ USE_PEFT_BACKEND,
+ deprecate,
+ is_torch_xla_available,
+ logging,
+ replace_example_docstring,
+ scale_lora_layers,
+ unscale_lora_layers,
+)
+from diffusers.utils.torch_utils import randn_tensor
+
+
+if is_torch_xla_available():
+ import torch_xla.core.xla_model as xm
+
+ XLA_AVAILABLE = True
+else:
+ XLA_AVAILABLE = False
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+EXAMPLE_DOC_STRING = """
+ Examples:
+ ```py
+ >>> import torch
+ >>> from diffusers import StableDiffusionXLImg2ImgPipeline
+ >>> from diffusers.utils import load_image
+
+ >>> pipe = StableDiffusionXLImg2ImgPipeline.from_pretrained(
+ ... "stabilityai/stable-diffusion-xl-refiner-1.0", torch_dtype=torch.float16
+ ... )
+ >>> pipe = pipe.to("cuda")
+ >>> url = "https://huggingface.co/datasets/patrickvonplaten/images/resolve/main/aa_xl/000000009.png"
+
+ >>> init_image = load_image(url).convert("RGB")
+ >>> prompt = "a photo of an astronaut riding a horse on mars"
+ >>> image = pipe(prompt, image=init_image).images[0]
+ ```
+"""
+
+
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.rescale_noise_cfg
+def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0):
+ """
+ Rescale `noise_cfg` according to `guidance_rescale`. Based on findings of [Common Diffusion Noise Schedules and
+ Sample Steps are Flawed](https://arxiv.org/pdf/2305.08891.pdf). See Section 3.4
+ """
+ std_text = noise_pred_text.std(dim=list(range(1, noise_pred_text.ndim)), keepdim=True)
+ std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True)
+ # rescale the results from guidance (fixes overexposure)
+ noise_pred_rescaled = noise_cfg * (std_text / std_cfg)
+ # mix with the original results from guidance by factor guidance_rescale to avoid "plain looking" images
+ noise_cfg = guidance_rescale * noise_pred_rescaled + (1 - guidance_rescale) * noise_cfg
+ return noise_cfg
+
+
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
+def retrieve_latents(
+ encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+):
+ if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
+ return encoder_output.latent_dist.sample(generator)
+ elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax":
+ return encoder_output.latent_dist.mode()
+ elif hasattr(encoder_output, "latents"):
+ return encoder_output.latents
+ else:
+ raise AttributeError("Could not access latents of provided encoder_output")
+
+
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
+def retrieve_timesteps(
+ scheduler,
+ num_inference_steps: Optional[int] = None,
+ device: Optional[Union[str, torch.device]] = None,
+ timesteps: Optional[List[int]] = None,
+ **kwargs,
+):
+ """
+ Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
+ custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
+
+ Args:
+ scheduler (`SchedulerMixin`):
+ The scheduler to get timesteps from.
+ num_inference_steps (`int`):
+ The number of diffusion steps used when generating samples with a pre-trained model. If used,
+ `timesteps` must be `None`.
+ device (`str` or `torch.device`, *optional*):
+ The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+ timesteps (`List[int]`, *optional*):
+ Custom timesteps used to support arbitrary spacing between timesteps. If `None`, then the default
+ timestep spacing strategy of the scheduler is used. If `timesteps` is passed, `num_inference_steps`
+ must be `None`.
+
+ Returns:
+ `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ second element is the number of inference steps.
+ """
+ if timesteps is not None:
+ accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+ if not accepts_timesteps:
+ raise ValueError(
+ f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+ f" timestep schedules. Please check whether you are using the correct scheduler."
+ )
+ scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
+ timesteps = scheduler.timesteps
+ num_inference_steps = len(timesteps)
+ else:
+ scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
+ timesteps = scheduler.timesteps
+ return timesteps, num_inference_steps
+
+
+class StableDiffusionXLSoftFillPipeline(
+ DiffusionPipeline,
+ StableDiffusionMixin,
+ TextualInversionLoaderMixin,
+ FromSingleFileMixin,
+ StableDiffusionXLLoraLoaderMixin,
+ IPAdapterMixin,
+):
+ r"""
+ Pipeline for text-to-image generation using Stable Diffusion XL.
+
+ This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the
+ library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)
+
+ In addition the pipeline inherits the following loading methods:
+ - *Textual-Inversion*: [`loaders.TextualInversionLoaderMixin.load_textual_inversion`]
+ - *LoRA*: [`loaders.StableDiffusionLoraLoaderMixin.load_lora_weights`]
+ - *Ckpt*: [`loaders.FromSingleFileMixin.from_single_file`]
+
+ as well as the following saving methods:
+ - *LoRA*: [`loaders.StableDiffusionLoraLoaderMixin.save_lora_weights`]
+
+ Args:
+ vae ([`AutoencoderKL`]):
+ Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.
+ text_encoder ([`CLIPTextModel`]):
+ Frozen text-encoder. Stable Diffusion XL uses the text portion of
+ [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModel), specifically
+ the [clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14) variant.
+ text_encoder_2 ([` CLIPTextModelWithProjection`]):
+ Second frozen text-encoder. Stable Diffusion XL uses the text and pool portion of
+ [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModelWithProjection),
+ specifically the
+ [laion/CLIP-ViT-bigG-14-laion2B-39B-b160k](https://huggingface.co/laion/CLIP-ViT-bigG-14-laion2B-39B-b160k)
+ variant.
+ tokenizer (`CLIPTokenizer`):
+ Tokenizer of class
+ [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer).
+ tokenizer_2 (`CLIPTokenizer`):
+ Second Tokenizer of class
+ [CLIPTokenizer](https://huggingface.co/docs/transformers/v4.21.0/en/model_doc/clip#transformers.CLIPTokenizer).
+ unet ([`UNet2DConditionModel`]): Conditional U-Net architecture to denoise the encoded image latents.
+ scheduler ([`SchedulerMixin`]):
+ A scheduler to be used in combination with `unet` to denoise the encoded image latents. Can be one of
+ [`DDIMScheduler`], [`LMSDiscreteScheduler`], or [`PNDMScheduler`].
+ """
+
+ model_cpu_offload_seq = "text_encoder->text_encoder_2->image_encoder->unet->vae"
+ _optional_components = [
+ "tokenizer",
+ "tokenizer_2",
+ "text_encoder",
+ "text_encoder_2",
+ "image_encoder",
+ "feature_extractor",
+ ]
+ _callback_tensor_inputs = [
+ "latents",
+ "prompt_embeds",
+ "negative_prompt_embeds",
+ "add_text_embeds",
+ "add_time_ids",
+ "negative_pooled_prompt_embeds",
+ "add_neg_time_ids",
+ ]
+
+ def __init__(
+ self,
+ vae: AutoencoderKL,
+ text_encoder: CLIPTextModel,
+ text_encoder_2: CLIPTextModelWithProjection,
+ tokenizer: CLIPTokenizer,
+ tokenizer_2: CLIPTokenizer,
+ unet: UNet2DConditionModel,
+ scheduler: KarrasDiffusionSchedulers,
+ image_encoder: CLIPVisionModelWithProjection = None,
+ feature_extractor: CLIPImageProcessor = None,
+ requires_aesthetics_score: bool = False,
+ force_zeros_for_empty_prompt: bool = True,
+ add_watermarker: Optional[bool] = None,
+ ):
+ super().__init__()
+
+ self.register_modules(
+ vae=vae,
+ text_encoder=text_encoder,
+ text_encoder_2=text_encoder_2,
+ tokenizer=tokenizer,
+ tokenizer_2=tokenizer_2,
+ unet=unet,
+ image_encoder=image_encoder,
+ feature_extractor=feature_extractor,
+ scheduler=scheduler,
+ )
+ self.register_to_config(force_zeros_for_empty_prompt=force_zeros_for_empty_prompt)
+ self.register_to_config(requires_aesthetics_score=requires_aesthetics_score)
+ self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8
+ self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
+ self.watermark = None
+
+ # Copied from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl.StableDiffusionXLPipeline.encode_prompt
+ def encode_prompt(
+ self,
+ prompt: str,
+ prompt_2: Optional[str] = None,
+ device: Optional[torch.device] = None,
+ num_images_per_prompt: int = 1,
+ do_classifier_free_guidance: bool = True,
+ negative_prompt: Optional[str] = None,
+ negative_prompt_2: Optional[str] = None,
+ prompt_embeds: Optional[torch.Tensor] = None,
+ negative_prompt_embeds: Optional[torch.Tensor] = None,
+ pooled_prompt_embeds: Optional[torch.Tensor] = None,
+ negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
+ lora_scale: Optional[float] = None,
+ clip_skip: Optional[int] = None,
+ ):
+ r"""
+ Encodes the prompt into text encoder hidden states.
+
+ Args:
+ prompt (`str` or `List[str]`, *optional*):
+ prompt to be encoded
+ prompt_2 (`str` or `List[str]`, *optional*):
+ The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
+ used in both text-encoders
+ device: (`torch.device`):
+ torch device
+ num_images_per_prompt (`int`):
+ number of images that should be generated per prompt
+ do_classifier_free_guidance (`bool`):
+ whether to use classifier free guidance or not
+ negative_prompt (`str` or `List[str]`, *optional*):
+ The prompt or prompts not to guide the image generation. If not defined, one has to pass
+ `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
+ less than `1`).
+ negative_prompt_2 (`str` or `List[str]`, *optional*):
+ The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
+ `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders
+ prompt_embeds (`torch.Tensor`, *optional*):
+ Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+ provided, text embeddings will be generated from `prompt` input argument.
+ negative_prompt_embeds (`torch.Tensor`, *optional*):
+ Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+ weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
+ argument.
+ pooled_prompt_embeds (`torch.Tensor`, *optional*):
+ Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting.
+ If not provided, pooled text embeddings will be generated from `prompt` input argument.
+ negative_pooled_prompt_embeds (`torch.Tensor`, *optional*):
+ Pre-generated negative pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+ weighting. If not provided, pooled negative_prompt_embeds will be generated from `negative_prompt`
+ input argument.
+ lora_scale (`float`, *optional*):
+ A lora scale that will be applied to all LoRA layers of the text encoder if LoRA layers are loaded.
+ clip_skip (`int`, *optional*):
+ Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that
+ the output of the pre-final layer will be used for computing the prompt embeddings.
+ """
+ device = device or self._execution_device
+
+ # set lora scale so that monkey patched LoRA
+ # function of text encoder can correctly access it
+ if lora_scale is not None and isinstance(self, StableDiffusionXLLoraLoaderMixin):
+ self._lora_scale = lora_scale
+
+ # dynamically adjust the LoRA scale
+ if self.text_encoder is not None:
+ if not USE_PEFT_BACKEND:
+ adjust_lora_scale_text_encoder(self.text_encoder, lora_scale)
+ else:
+ scale_lora_layers(self.text_encoder, lora_scale)
+
+ if self.text_encoder_2 is not None:
+ if not USE_PEFT_BACKEND:
+ adjust_lora_scale_text_encoder(self.text_encoder_2, lora_scale)
+ else:
+ scale_lora_layers(self.text_encoder_2, lora_scale)
+
+ prompt = [prompt] if isinstance(prompt, str) else prompt
+
+ if prompt is not None:
+ batch_size = len(prompt)
+ else:
+ batch_size = prompt_embeds.shape[0]
+
+ # Define tokenizers and text encoders
+ tokenizers = [self.tokenizer, self.tokenizer_2] if self.tokenizer is not None else [self.tokenizer_2]
+ text_encoders = (
+ [self.text_encoder, self.text_encoder_2] if self.text_encoder is not None else [self.text_encoder_2]
+ )
+
+ if prompt_embeds is None:
+ prompt_2 = prompt_2 or prompt
+ prompt_2 = [prompt_2] if isinstance(prompt_2, str) else prompt_2
+
+ # textual inversion: process multi-vector tokens if necessary
+ prompt_embeds_list = []
+ prompts = [prompt, prompt_2]
+ for prompt, tokenizer, text_encoder in zip(prompts, tokenizers, text_encoders):
+ if isinstance(self, TextualInversionLoaderMixin):
+ prompt = self.maybe_convert_prompt(prompt, tokenizer)
+
+ text_inputs = tokenizer(
+ prompt,
+ padding="max_length",
+ max_length=tokenizer.model_max_length,
+ truncation=True,
+ return_tensors="pt",
+ )
+
+ text_input_ids = text_inputs.input_ids
+ untruncated_ids = tokenizer(prompt, padding="longest", return_tensors="pt").input_ids
+
+ if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(
+ text_input_ids, untruncated_ids
+ ):
+ removed_text = tokenizer.batch_decode(untruncated_ids[:, tokenizer.model_max_length - 1 : -1])
+ logger.warning(
+ "The following part of your input was truncated because CLIP can only handle sequences up to"
+ f" {tokenizer.model_max_length} tokens: {removed_text}"
+ )
+
+ prompt_embeds = text_encoder(text_input_ids.to(device), output_hidden_states=True)
+
+ # We are only ALWAYS interested in the pooled output of the final text encoder
+ if pooled_prompt_embeds is None and prompt_embeds[0].ndim == 2:
+ pooled_prompt_embeds = prompt_embeds[0]
+
+ if clip_skip is None:
+ prompt_embeds = prompt_embeds.hidden_states[-2]
+ else:
+ # "2" because SDXL always indexes from the penultimate layer.
+ prompt_embeds = prompt_embeds.hidden_states[-(clip_skip + 2)]
+
+ prompt_embeds_list.append(prompt_embeds)
+
+ prompt_embeds = torch.concat(prompt_embeds_list, dim=-1)
+
+ # get unconditional embeddings for classifier free guidance
+ zero_out_negative_prompt = negative_prompt is None and self.config.force_zeros_for_empty_prompt
+ if do_classifier_free_guidance and negative_prompt_embeds is None and zero_out_negative_prompt:
+ negative_prompt_embeds = torch.zeros_like(prompt_embeds)
+ negative_pooled_prompt_embeds = torch.zeros_like(pooled_prompt_embeds)
+ elif do_classifier_free_guidance and negative_prompt_embeds is None:
+ negative_prompt = negative_prompt or ""
+ negative_prompt_2 = negative_prompt_2 or negative_prompt
+
+ # normalize str to list
+ negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt
+ negative_prompt_2 = (
+ batch_size * [negative_prompt_2] if isinstance(negative_prompt_2, str) else negative_prompt_2
+ )
+
+ uncond_tokens: List[str]
+ if prompt is not None and type(prompt) is not type(negative_prompt):
+ raise TypeError(
+ f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
+ f" {type(prompt)}."
+ )
+ elif batch_size != len(negative_prompt):
+ raise ValueError(
+ f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
+ f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
+ " the batch size of `prompt`."
+ )
+ else:
+ uncond_tokens = [negative_prompt, negative_prompt_2]
+
+ negative_prompt_embeds_list = []
+ for negative_prompt, tokenizer, text_encoder in zip(uncond_tokens, tokenizers, text_encoders):
+ if isinstance(self, TextualInversionLoaderMixin):
+ negative_prompt = self.maybe_convert_prompt(negative_prompt, tokenizer)
+
+ max_length = prompt_embeds.shape[1]
+ uncond_input = tokenizer(
+ negative_prompt,
+ padding="max_length",
+ max_length=max_length,
+ truncation=True,
+ return_tensors="pt",
+ )
+
+ negative_prompt_embeds = text_encoder(
+ uncond_input.input_ids.to(device),
+ output_hidden_states=True,
+ )
+ # We are only ALWAYS interested in the pooled output of the final text encoder
+ if negative_pooled_prompt_embeds is None and negative_prompt_embeds[0].ndim == 2:
+ negative_pooled_prompt_embeds = negative_prompt_embeds[0]
+ negative_prompt_embeds = negative_prompt_embeds.hidden_states[-2]
+
+ negative_prompt_embeds_list.append(negative_prompt_embeds)
+
+ negative_prompt_embeds = torch.concat(negative_prompt_embeds_list, dim=-1)
+
+ if self.text_encoder_2 is not None:
+ prompt_embeds = prompt_embeds.to(dtype=self.text_encoder_2.dtype, device=device)
+ else:
+ prompt_embeds = prompt_embeds.to(dtype=self.unet.dtype, device=device)
+
+ bs_embed, seq_len, _ = prompt_embeds.shape
+ # duplicate text embeddings for each generation per prompt, using mps friendly method
+ prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
+ prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1)
+
+ if do_classifier_free_guidance:
+ # duplicate unconditional embeddings for each generation per prompt, using mps friendly method
+ seq_len = negative_prompt_embeds.shape[1]
+
+ if self.text_encoder_2 is not None:
+ negative_prompt_embeds = negative_prompt_embeds.to(dtype=self.text_encoder_2.dtype, device=device)
+ else:
+ negative_prompt_embeds = negative_prompt_embeds.to(dtype=self.unet.dtype, device=device)
+
+ negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)
+ negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
+
+ pooled_prompt_embeds = pooled_prompt_embeds.repeat(1, num_images_per_prompt).view(
+ bs_embed * num_images_per_prompt, -1
+ )
+ if do_classifier_free_guidance:
+ negative_pooled_prompt_embeds = negative_pooled_prompt_embeds.repeat(1, num_images_per_prompt).view(
+ bs_embed * num_images_per_prompt, -1
+ )
+
+ if self.text_encoder is not None:
+ if isinstance(self, StableDiffusionXLLoraLoaderMixin) and USE_PEFT_BACKEND:
+ # Retrieve the original scale by scaling back the LoRA layers
+ unscale_lora_layers(self.text_encoder, lora_scale)
+
+ if self.text_encoder_2 is not None:
+ if isinstance(self, StableDiffusionXLLoraLoaderMixin) and USE_PEFT_BACKEND:
+ # Retrieve the original scale by scaling back the LoRA layers
+ unscale_lora_layers(self.text_encoder_2, lora_scale)
+
+ return prompt_embeds, negative_prompt_embeds, pooled_prompt_embeds, negative_pooled_prompt_embeds
+
+ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs
+ def prepare_extra_step_kwargs(self, generator, eta):
+ # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
+ # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
+ # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
+ # and should be between [0, 1]
+
+ accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())
+ extra_step_kwargs = {}
+ if accepts_eta:
+ extra_step_kwargs["eta"] = eta
+
+ # check if the scheduler accepts generator
+ accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())
+ if accepts_generator:
+ extra_step_kwargs["generator"] = generator
+ return extra_step_kwargs
+
+ def check_inputs(
+ self,
+ prompt,
+ prompt_2,
+ strength,
+ num_inference_steps,
+ callback_steps,
+ negative_prompt=None,
+ negative_prompt_2=None,
+ prompt_embeds=None,
+ negative_prompt_embeds=None,
+ ip_adapter_image=None,
+ ip_adapter_image_embeds=None,
+ callback_on_step_end_tensor_inputs=None,
+ ):
+ if strength < 0 or strength > 1:
+ raise ValueError(f"The value of strength should in [0.0, 1.0] but is {strength}")
+ if num_inference_steps is None:
+ raise ValueError("`num_inference_steps` cannot be None.")
+ elif not isinstance(num_inference_steps, int) or num_inference_steps <= 0:
+ raise ValueError(
+ f"`num_inference_steps` has to be a positive integer but is {num_inference_steps} of type"
+ f" {type(num_inference_steps)}."
+ )
+ if callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0):
+ raise ValueError(
+ f"`callback_steps` has to be a positive integer but is {callback_steps} of type"
+ f" {type(callback_steps)}."
+ )
+
+ if callback_on_step_end_tensor_inputs is not None and not all(
+ k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs
+ ):
+ raise ValueError(
+ f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}"
+ )
+
+ if prompt is not None and prompt_embeds is not None:
+ raise ValueError(
+ f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
+ " only forward one of the two."
+ )
+ elif prompt_2 is not None and prompt_embeds is not None:
+ raise ValueError(
+ f"Cannot forward both `prompt_2`: {prompt_2} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
+ " only forward one of the two."
+ )
+ elif prompt is None and prompt_embeds is None:
+ raise ValueError(
+ "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
+ )
+ elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
+ raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
+ elif prompt_2 is not None and (not isinstance(prompt_2, str) and not isinstance(prompt_2, list)):
+ raise ValueError(f"`prompt_2` has to be of type `str` or `list` but is {type(prompt_2)}")
+
+ if negative_prompt is not None and negative_prompt_embeds is not None:
+ raise ValueError(
+ f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"
+ f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
+ )
+ elif negative_prompt_2 is not None and negative_prompt_embeds is not None:
+ raise ValueError(
+ f"Cannot forward both `negative_prompt_2`: {negative_prompt_2} and `negative_prompt_embeds`:"
+ f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
+ )
+
+ if prompt_embeds is not None and negative_prompt_embeds is not None:
+ if prompt_embeds.shape != negative_prompt_embeds.shape:
+ raise ValueError(
+ "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"
+ f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"
+ f" {negative_prompt_embeds.shape}."
+ )
+
+ if ip_adapter_image is not None and ip_adapter_image_embeds is not None:
+ raise ValueError(
+ "Provide either `ip_adapter_image` or `ip_adapter_image_embeds`. Cannot leave both `ip_adapter_image` and `ip_adapter_image_embeds` defined."
+ )
+
+ if ip_adapter_image_embeds is not None:
+ if not isinstance(ip_adapter_image_embeds, list):
+ raise ValueError(
+ f"`ip_adapter_image_embeds` has to be of type `list` but is {type(ip_adapter_image_embeds)}"
+ )
+ elif ip_adapter_image_embeds[0].ndim not in [3, 4]:
+ raise ValueError(
+ f"`ip_adapter_image_embeds` has to be a list of 3D or 4D tensors but is {ip_adapter_image_embeds[0].ndim}D"
+ )
+
+ def get_timesteps(self, num_inference_steps, strength, device, denoising_start=None):
+ # get the original timestep using init_timestep
+ if denoising_start is None:
+ init_timestep = min(int(num_inference_steps * strength), num_inference_steps)
+ t_start = max(num_inference_steps - init_timestep, 0)
+ else:
+ t_start = 0
+
+ timesteps = self.scheduler.timesteps[t_start * self.scheduler.order :]
+
+ # Strength is irrelevant if we directly request a timestep to start at;
+ # that is, strength is determined by the denoising_start instead.
+ if denoising_start is not None:
+ discrete_timestep_cutoff = int(
+ round(
+ self.scheduler.config.num_train_timesteps
+ - (denoising_start * self.scheduler.config.num_train_timesteps)
+ )
+ )
+
+ num_inference_steps = (timesteps < discrete_timestep_cutoff).sum().item()
+ if self.scheduler.order == 2 and num_inference_steps % 2 == 0:
+ # if the scheduler is a 2nd order scheduler we might have to do +1
+ # because `num_inference_steps` might be even given that every timestep
+ # (except the highest one) is duplicated. If `num_inference_steps` is even it would
+ # mean that we cut the timesteps in the middle of the denoising step
+ # (between 1st and 2nd derivative) which leads to incorrect results. By adding 1
+ # we ensure that the denoising process always ends after the 2nd derivate step of the scheduler
+ num_inference_steps = num_inference_steps + 1
+
+ # because t_n+1 >= t_n, we slice the timesteps starting from the end
+ timesteps = timesteps[-num_inference_steps:]
+ return timesteps, num_inference_steps
+
+ return timesteps, num_inference_steps - t_start
+
+ def prepare_latents(
+ self, image, timestep, batch_size, num_images_per_prompt, dtype, device, generator=None, add_noise=True
+ ):
+ if not isinstance(image, (torch.Tensor, Image.Image, list)):
+ raise ValueError(
+ f"`image` has to be of type `torch.Tensor`, `PIL.Image.Image` or list but is {type(image)}"
+ )
+
+ # Offload text encoder if `enable_model_cpu_offload` was enabled
+ if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:
+ self.text_encoder_2.to("cpu")
+ torch.cuda.empty_cache()
+
+ image = image.to(device=device, dtype=dtype)
+
+ batch_size = batch_size * num_images_per_prompt
+
+ if image.shape[1] == 4:
+ init_latents = image
+
+ else:
+ # make sure the VAE is in float32 mode, as it overflows in float16
+ if self.vae.config.force_upcast:
+ image = image.float()
+ self.vae.to(dtype=torch.float32)
+
+ if isinstance(generator, list) and len(generator) != batch_size:
+ raise ValueError(
+ f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
+ f" size of {batch_size}. Make sure the batch size matches the length of the generators."
+ )
+
+ elif isinstance(generator, list):
+ init_latents = [
+ retrieve_latents(self.vae.encode(image[i : i + 1]), generator=generator[i])
+ for i in range(batch_size)
+ ]
+ init_latents = torch.cat(init_latents, dim=0)
+ else:
+ init_latents = retrieve_latents(self.vae.encode(image), generator=generator)
+
+ if self.vae.config.force_upcast:
+ self.vae.to(dtype)
+
+ init_latents = init_latents.to(dtype)
+ init_latents = self.vae.config.scaling_factor * init_latents
+
+ if batch_size > init_latents.shape[0] and batch_size % init_latents.shape[0] == 0:
+ # expand init_latents for batch_size
+ additional_image_per_prompt = batch_size // init_latents.shape[0]
+ init_latents = torch.cat([init_latents] * additional_image_per_prompt, dim=0)
+ elif batch_size > init_latents.shape[0] and batch_size % init_latents.shape[0] != 0:
+ raise ValueError(
+ f"Cannot duplicate `image` of batch size {init_latents.shape[0]} to {batch_size} text prompts."
+ )
+ else:
+ init_latents = torch.cat([init_latents], dim=0)
+
+ if add_noise:
+ shape = init_latents.shape
+ noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+ # get latents
+ init_latents = self.scheduler.add_noise(init_latents, noise, timestep)
+
+ latents = init_latents
+
+ return latents
+
+ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.encode_image
+ def encode_image(self, image, device, num_images_per_prompt, output_hidden_states=None):
+ dtype = next(self.image_encoder.parameters()).dtype
+
+ if not isinstance(image, torch.Tensor):
+ image = self.feature_extractor(image, return_tensors="pt").pixel_values
+
+ image = image.to(device=device, dtype=dtype)
+ if output_hidden_states:
+ image_enc_hidden_states = self.image_encoder(image, output_hidden_states=True).hidden_states[-2]
+ image_enc_hidden_states = image_enc_hidden_states.repeat_interleave(num_images_per_prompt, dim=0)
+ uncond_image_enc_hidden_states = self.image_encoder(
+ torch.zeros_like(image), output_hidden_states=True
+ ).hidden_states[-2]
+ uncond_image_enc_hidden_states = uncond_image_enc_hidden_states.repeat_interleave(
+ num_images_per_prompt, dim=0
+ )
+ return image_enc_hidden_states, uncond_image_enc_hidden_states
+ else:
+ image_embeds = self.image_encoder(image).image_embeds
+ image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0)
+ uncond_image_embeds = torch.zeros_like(image_embeds)
+
+ return image_embeds, uncond_image_embeds
+
+ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_ip_adapter_image_embeds
+ def prepare_ip_adapter_image_embeds(
+ self, ip_adapter_image, ip_adapter_image_embeds, device, num_images_per_prompt, do_classifier_free_guidance
+ ):
+ if ip_adapter_image_embeds is None:
+ if not isinstance(ip_adapter_image, list):
+ ip_adapter_image = [ip_adapter_image]
+
+ if len(ip_adapter_image) != len(self.unet.encoder_hid_proj.image_projection_layers):
+ raise ValueError(
+ f"`ip_adapter_image` must have same length as the number of IP Adapters. Got {len(ip_adapter_image)} images and {len(self.unet.encoder_hid_proj.image_projection_layers)} IP Adapters."
+ )
+
+ image_embeds = []
+ for single_ip_adapter_image, image_proj_layer in zip(
+ ip_adapter_image, self.unet.encoder_hid_proj.image_projection_layers
+ ):
+ output_hidden_state = not isinstance(image_proj_layer, ImageProjection)
+ single_image_embeds, single_negative_image_embeds = self.encode_image(
+ single_ip_adapter_image, device, 1, output_hidden_state
+ )
+ single_image_embeds = torch.stack([single_image_embeds] * num_images_per_prompt, dim=0)
+ single_negative_image_embeds = torch.stack(
+ [single_negative_image_embeds] * num_images_per_prompt, dim=0
+ )
+
+ if do_classifier_free_guidance:
+ single_image_embeds = torch.cat([single_negative_image_embeds, single_image_embeds])
+ single_image_embeds = single_image_embeds.to(device)
+
+ image_embeds.append(single_image_embeds)
+ else:
+ repeat_dims = [1]
+ image_embeds = []
+ for single_image_embeds in ip_adapter_image_embeds:
+ if do_classifier_free_guidance:
+ single_negative_image_embeds, single_image_embeds = single_image_embeds.chunk(2)
+ single_image_embeds = single_image_embeds.repeat(
+ num_images_per_prompt, *(repeat_dims * len(single_image_embeds.shape[1:]))
+ )
+ single_negative_image_embeds = single_negative_image_embeds.repeat(
+ num_images_per_prompt, *(repeat_dims * len(single_negative_image_embeds.shape[1:]))
+ )
+ single_image_embeds = torch.cat([single_negative_image_embeds, single_image_embeds])
+ else:
+ single_image_embeds = single_image_embeds.repeat(
+ num_images_per_prompt, *(repeat_dims * len(single_image_embeds.shape[1:]))
+ )
+ image_embeds.append(single_image_embeds)
+
+ return image_embeds
+
+ def _get_add_time_ids(
+ self,
+ original_size,
+ crops_coords_top_left,
+ target_size,
+ aesthetic_score,
+ negative_aesthetic_score,
+ negative_original_size,
+ negative_crops_coords_top_left,
+ negative_target_size,
+ dtype,
+ text_encoder_projection_dim=None,
+ ):
+ if self.config.requires_aesthetics_score:
+ add_time_ids = list(original_size + crops_coords_top_left + (aesthetic_score,))
+ add_neg_time_ids = list(
+ negative_original_size + negative_crops_coords_top_left + (negative_aesthetic_score,)
+ )
+ else:
+ add_time_ids = list(original_size + crops_coords_top_left + target_size)
+ add_neg_time_ids = list(negative_original_size + crops_coords_top_left + negative_target_size)
+
+ passed_add_embed_dim = (
+ self.unet.config.addition_time_embed_dim * len(add_time_ids) + text_encoder_projection_dim
+ )
+ expected_add_embed_dim = self.unet.add_embedding.linear_1.in_features
+
+ if (
+ expected_add_embed_dim > passed_add_embed_dim
+ and (expected_add_embed_dim - passed_add_embed_dim) == self.unet.config.addition_time_embed_dim
+ ):
+ raise ValueError(
+ f"Model expects an added time embedding vector of length {expected_add_embed_dim}, but a vector of {passed_add_embed_dim} was created. Please make sure to enable `requires_aesthetics_score` with `pipe.register_to_config(requires_aesthetics_score=True)` to make sure `aesthetic_score` {aesthetic_score} and `negative_aesthetic_score` {negative_aesthetic_score} is correctly used by the model."
+ )
+ elif (
+ expected_add_embed_dim < passed_add_embed_dim
+ and (passed_add_embed_dim - expected_add_embed_dim) == self.unet.config.addition_time_embed_dim
+ ):
+ raise ValueError(
+ f"Model expects an added time embedding vector of length {expected_add_embed_dim}, but a vector of {passed_add_embed_dim} was created. Please make sure to disable `requires_aesthetics_score` with `pipe.register_to_config(requires_aesthetics_score=False)` to make sure `target_size` {target_size} is correctly used by the model."
+ )
+ elif expected_add_embed_dim != passed_add_embed_dim:
+ raise ValueError(
+ f"Model expects an added time embedding vector of length {expected_add_embed_dim}, but a vector of {passed_add_embed_dim} was created. The model has an incorrect config. Please check `unet.config.time_embedding_type` and `text_encoder_2.config.projection_dim`."
+ )
+
+ add_time_ids = torch.tensor([add_time_ids], dtype=dtype)
+ add_neg_time_ids = torch.tensor([add_neg_time_ids], dtype=dtype)
+
+ return add_time_ids, add_neg_time_ids
+
+ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_upscale.StableDiffusionUpscalePipeline.upcast_vae
+ def upcast_vae(self):
+ dtype = self.vae.dtype
+ self.vae.to(dtype=torch.float32)
+ use_torch_2_0_or_xformers = isinstance(
+ self.vae.decoder.mid_block.attentions[0].processor,
+ (
+ AttnProcessor2_0,
+ XFormersAttnProcessor,
+ ),
+ )
+ # if xformers or torch_2_0 is used attention block does not need
+ # to be in float32 which can save lots of memory
+ if use_torch_2_0_or_xformers:
+ self.vae.post_quant_conv.to(dtype)
+ self.vae.decoder.conv_in.to(dtype)
+ self.vae.decoder.mid_block.to(dtype)
+
+ # Copied from diffusers.pipelines.latent_consistency_models.pipeline_latent_consistency_text2img.LatentConsistencyModelPipeline.get_guidance_scale_embedding
+ def get_guidance_scale_embedding(
+ self, w: torch.Tensor, embedding_dim: int = 512, dtype: torch.dtype = torch.float32
+ ) -> torch.Tensor:
+ """
+ See https://github.com/google-research/vdm/blob/dc27b98a554f65cdc654b800da5aa1846545d41b/model_vdm.py#L298
+
+ Args:
+ w (`torch.Tensor`):
+ Generate embedding vectors with a specified guidance scale to subsequently enrich timestep embeddings.
+ embedding_dim (`int`, *optional*, defaults to 512):
+ Dimension of the embeddings to generate.
+ dtype (`torch.dtype`, *optional*, defaults to `torch.float32`):
+ Data type of the generated embeddings.
+
+ Returns:
+ `torch.Tensor`: Embedding vectors with shape `(len(w), embedding_dim)`.
+ """
+ assert len(w.shape) == 1
+ w = w * 1000.0
+
+ half_dim = embedding_dim // 2
+ emb = torch.log(torch.tensor(10000.0)) / (half_dim - 1)
+ emb = torch.exp(torch.arange(half_dim, dtype=dtype) * -emb)
+ emb = w.to(dtype)[:, None] * emb[None, :]
+ emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1)
+ if embedding_dim % 2 == 1: # zero pad
+ emb = torch.nn.functional.pad(emb, (0, 1))
+ assert emb.shape == (w.shape[0], embedding_dim)
+ return emb
+
+ @property
+ def guidance_scale(self):
+ return self._guidance_scale
+
+ @property
+ def guidance_rescale(self):
+ return self._guidance_rescale
+
+ @property
+ def clip_skip(self):
+ return self._clip_skip
+
+ # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
+ # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
+ # corresponds to doing no classifier free guidance.
+ @property
+ def do_classifier_free_guidance(self):
+ return self._guidance_scale > 1 and self.unet.config.time_cond_proj_dim is None
+
+ @property
+ def cross_attention_kwargs(self):
+ return self._cross_attention_kwargs
+
+ @property
+ def denoising_end(self):
+ return self._denoising_end
+
+ @property
+ def denoising_start(self):
+ return self._denoising_start
+
+ @property
+ def num_timesteps(self):
+ return self._num_timesteps
+
+ @property
+ def interrupt(self):
+ return self._interrupt
+
+ @torch.no_grad()
+ @replace_example_docstring(EXAMPLE_DOC_STRING)
+ def __call__(
+ self,
+ prompt: Union[str, List[str]] = None,
+ prompt_2: Optional[Union[str, List[str]]] = None,
+ image: Image.Image = None,
+ mask: Image.Image = None,
+ noise_fill_image: bool = True, # Adds noise to the image at the masks >0.8 area.
+ strength: float = 0.3,
+ num_inference_steps: int = 50,
+ timesteps: List[int] = None,
+ denoising_start: Optional[float] = None,
+ denoising_end: Optional[float] = None,
+ guidance_scale: float = 5.0,
+ negative_prompt: Optional[Union[str, List[str]]] = None,
+ negative_prompt_2: Optional[Union[str, List[str]]] = None,
+ num_images_per_prompt: Optional[int] = 1,
+ eta: float = 0.0,
+ generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+ latents: Optional[torch.Tensor] = None,
+ prompt_embeds: Optional[torch.Tensor] = None,
+ negative_prompt_embeds: Optional[torch.Tensor] = None,
+ pooled_prompt_embeds: Optional[torch.Tensor] = None,
+ negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
+ ip_adapter_image: Optional[PipelineImageInput] = None,
+ ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
+ output_type: Optional[str] = "pil",
+ return_dict: bool = True,
+ cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+ guidance_rescale: float = 0.0,
+ original_size: Tuple[int, int] = None,
+ crops_coords_top_left: Tuple[int, int] = (0, 0),
+ target_size: Tuple[int, int] = None,
+ negative_original_size: Optional[Tuple[int, int]] = None,
+ negative_crops_coords_top_left: Tuple[int, int] = (0, 0),
+ negative_target_size: Optional[Tuple[int, int]] = None,
+ aesthetic_score: float = 6.0,
+ negative_aesthetic_score: float = 2.5,
+ clip_skip: Optional[int] = None,
+ callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
+ callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ **kwargs,
+ ):
+ r"""
+ Function invoked when calling the pipeline for generation.
+
+ Args:
+ prompt (`str` or `List[str]`, *optional*):
+ The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
+ instead.
+ prompt_2 (`str` or `List[str]`, *optional*):
+ The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
+ used in both text-encoders
+ image (`torch.Tensor` or `PIL.Image.Image` or `np.ndarray` or `List[torch.Tensor]` or `List[PIL.Image.Image]` or `List[np.ndarray]`):
+ The image(s) to modify with the pipeline.
+ strength (`float`, *optional*, defaults to 0.3):
+ Conceptually, indicates how much to transform the reference `image`. Must be between 0 and 1. `image`
+ will be used as a starting point, adding more noise to it the larger the `strength`. The number of
+ denoising steps depends on the amount of noise initially added. When `strength` is 1, added noise will
+ be maximum and the denoising process will run for the full number of iterations specified in
+ `num_inference_steps`. A value of 1, therefore, essentially ignores `image`. Note that in the case of
+ `denoising_start` being declared as an integer, the value of `strength` will be ignored.
+ num_inference_steps (`int`, *optional*, defaults to 50):
+ The number of denoising steps. More denoising steps usually lead to a higher quality image at the
+ expense of slower inference.
+ denoising_start (`float`, *optional*):
+ When specified, indicates the fraction (between 0.0 and 1.0) of the total denoising process to be
+ bypassed before it is initiated. Consequently, the initial part of the denoising process is skipped and
+ it is assumed that the passed `image` is a partly denoised image. Note that when this is specified,
+ strength will be ignored. The `denoising_start` parameter is particularly beneficial when this pipeline
+ is integrated into a "Mixture of Denoisers" multi-pipeline setup, as detailed in [**Refining the Image
+ Output**](https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/stable_diffusion_xl#refining-the-image-output).
+ denoising_end (`float`, *optional*):
+ When specified, determines the fraction (between 0.0 and 1.0) of the total denoising process to be
+ completed before it is intentionally prematurely terminated. As a result, the returned sample will
+ still retain a substantial amount of noise (ca. final 20% of timesteps still needed) and should be
+ denoised by a successor pipeline that has `denoising_start` set to 0.8 so that it only denoises the
+ final 20% of the scheduler. The denoising_end parameter should ideally be utilized when this pipeline
+ forms a part of a "Mixture of Denoisers" multi-pipeline setup, as elaborated in [**Refining the Image
+ Output**](https://huggingface.co/docs/diffusers/api/pipelines/stable_diffusion/stable_diffusion_xl#refining-the-image-output).
+ guidance_scale (`float`, *optional*, defaults to 7.5):
+ Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
+ `guidance_scale` is defined as `w` of equation 2. of [Imagen
+ Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
+ 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
+ usually at the expense of lower image quality.
+ negative_prompt (`str` or `List[str]`, *optional*):
+ The prompt or prompts not to guide the image generation. If not defined, one has to pass
+ `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
+ less than `1`).
+ negative_prompt_2 (`str` or `List[str]`, *optional*):
+ The prompt or prompts not to guide the image generation to be sent to `tokenizer_2` and
+ `text_encoder_2`. If not defined, `negative_prompt` is used in both text-encoders
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
+ The number of images to generate per prompt.
+ eta (`float`, *optional*, defaults to 0.0):
+ Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to
+ [`schedulers.DDIMScheduler`], will be ignored for others.
+ generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+ One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
+ to make generation deterministic.
+ latents (`torch.Tensor`, *optional*):
+ Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
+ generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
+ tensor will ge generated by sampling using the supplied random `generator`.
+ prompt_embeds (`torch.Tensor`, *optional*):
+ Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+ provided, text embeddings will be generated from `prompt` input argument.
+ negative_prompt_embeds (`torch.Tensor`, *optional*):
+ Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+ weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
+ argument.
+ pooled_prompt_embeds (`torch.Tensor`, *optional*):
+ Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting.
+ If not provided, pooled text embeddings will be generated from `prompt` input argument.
+ negative_pooled_prompt_embeds (`torch.Tensor`, *optional*):
+ Pre-generated negative pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+ weighting. If not provided, pooled negative_prompt_embeds will be generated from `negative_prompt`
+ input argument.
+ ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters.
+ ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):
+ Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of IP-adapters.
+ Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. It should contain the negative image embedding
+ if `do_classifier_free_guidance` is set to `True`.
+ If not provided, embeddings are computed from the `ip_adapter_image` input argument.
+ output_type (`str`, *optional*, defaults to `"pil"`):
+ The output format of the generate image. Choose between
+ [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
+ return_dict (`bool`, *optional*, defaults to `True`):
+ Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionXLPipelineOutput`] instead of a
+ plain tuple.
+ callback (`Callable`, *optional*):
+ A function that will be called every `callback_steps` steps during inference. The function will be
+ called with the following arguments: `callback(step: int, timestep: int, latents: torch.Tensor)`.
+ callback_steps (`int`, *optional*, defaults to 1):
+ The frequency at which the `callback` function will be called. If not specified, the callback will be
+ called at every step.
+ cross_attention_kwargs (`dict`, *optional*):
+ A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
+ `self.processor` in
+ [diffusers.cross_attention](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/cross_attention.py).
+ guidance_rescale (`float`, *optional*, defaults to 0.7):
+ Guidance rescale factor proposed by [Common Diffusion Noise Schedules and Sample Steps are
+ Flawed](https://arxiv.org/pdf/2305.08891.pdf) `guidance_scale` is defined as `φ` in equation 16. of
+ [Common Diffusion Noise Schedules and Sample Steps are Flawed](https://arxiv.org/pdf/2305.08891.pdf).
+ Guidance rescale factor should fix overexposure when using zero terminal SNR.
+ original_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ If `original_size` is not the same as `target_size` the image will appear to be down- or upsampled.
+ `original_size` defaults to `(width, height)` if not specified. Part of SDXL's micro-conditioning as
+ explained in section 2.2 of
+ [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
+ crops_coords_top_left (`Tuple[int]`, *optional*, defaults to (0, 0)):
+ `crops_coords_top_left` can be used to generate an image that appears to be "cropped" from the position
+ `crops_coords_top_left` downwards. Favorable, well-centered images are usually achieved by setting
+ `crops_coords_top_left` to (0, 0). Part of SDXL's micro-conditioning as explained in section 2.2 of
+ [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
+ target_size (`Tuple[int]`, *optional*, defaults to (1024, 1024)):
+ For most cases, `target_size` should be set to the desired height and width of the generated image. If
+ not specified it will default to `(width, height)`. Part of SDXL's micro-conditioning as explained in
+ section 2.2 of [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
+ aesthetic_score (`float`, *optional*, defaults to 6.0):
+ Used to simulate an aesthetic score of the generated image by influencing the positive text condition.
+ Part of SDXL's micro-conditioning as explained in section 2.2 of
+ [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952).
+ negative_aesthetic_score (`float`, *optional*, defaults to 2.5):
+ Part of SDXL's micro-conditioning as explained in section 2.2 of
+ [https://huggingface.co/papers/2307.01952](https://huggingface.co/papers/2307.01952). Can be used to
+ simulate an aesthetic score of the generated image by influencing the negative text condition.
+
+ Examples:
+
+ Returns:
+ [`~pipelines.stable_diffusion.StableDiffusionXLPipelineOutput`] or `tuple`:
+ [`~pipelines.stable_diffusion.StableDiffusionXLPipelineOutput`] if `return_dict` is True, otherwise a
+ `tuple. When returning a tuple, the first element is a list with the generated images.
+ """
+
+ callback = kwargs.pop("callback", None)
+ callback_steps = kwargs.pop("callback_steps", None)
+
+ if callback is not None:
+ deprecate(
+ "callback",
+ "1.0.0",
+ "Passing `callback` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`",
+ )
+ if callback_steps is not None:
+ deprecate(
+ "callback_steps",
+ "1.0.0",
+ "Passing `callback_steps` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`",
+ )
+
+ # 1. Check inputs. Raise error if not correct
+ self.check_inputs(
+ prompt,
+ prompt_2,
+ strength,
+ num_inference_steps,
+ callback_steps,
+ negative_prompt,
+ negative_prompt_2,
+ prompt_embeds,
+ negative_prompt_embeds,
+ ip_adapter_image,
+ ip_adapter_image_embeds,
+ callback_on_step_end_tensor_inputs,
+ )
+
+ self._guidance_scale = guidance_scale
+ self._guidance_rescale = guidance_rescale
+ self._clip_skip = clip_skip
+ self._cross_attention_kwargs = cross_attention_kwargs
+ self._denoising_end = denoising_end
+ self._denoising_start = denoising_start
+ self._interrupt = False
+
+ # 2. Define call parameters
+ if prompt is not None and isinstance(prompt, str):
+ batch_size = 1
+ elif prompt is not None and isinstance(prompt, list):
+ batch_size = len(prompt)
+ else:
+ batch_size = prompt_embeds.shape[0]
+
+ device = self._execution_device
+
+ # 3. Encode input prompt
+ text_encoder_lora_scale = (
+ cross_attention_kwargs.get("scale", None) if cross_attention_kwargs is not None else None
+ )
+ (
+ prompt_embeds,
+ negative_prompt_embeds,
+ pooled_prompt_embeds,
+ negative_pooled_prompt_embeds,
+ ) = self.encode_prompt(
+ prompt=prompt,
+ prompt_2=prompt_2,
+ device=device,
+ num_images_per_prompt=num_images_per_prompt,
+ do_classifier_free_guidance=self.do_classifier_free_guidance,
+ negative_prompt=negative_prompt,
+ negative_prompt_2=negative_prompt_2,
+ prompt_embeds=prompt_embeds,
+ negative_prompt_embeds=negative_prompt_embeds,
+ pooled_prompt_embeds=pooled_prompt_embeds,
+ negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
+ lora_scale=text_encoder_lora_scale,
+ )
+
+ # 4. PREPARE TIMESTEPS
+ def denoising_value_valid(dnv):
+ return isinstance(dnv, float) and 0 < dnv < 1
+
+ timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps)
+
+ timesteps, num_inference_steps = self.get_timesteps(
+ num_inference_steps,
+ strength,
+ device,
+ denoising_start=self.denoising_start if denoising_value_valid(self.denoising_start) else None,
+ )
+ latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt)
+
+ add_noise = True if denoising_start is None else False
+
+
+
+ # --------------------------------------------
+ # IMAGE PREPARATION UTILITIES
+ # --------------------------------------------
+
+ def fbm(x, y, scale, octaves, lacunarity, gain):
+ """
+ Fractal Brownian Motion (fbm) noise generator.
+ Combines multiple octaves of Perlin noise.
+ """
+ total = 0.0
+ amplitude = 1.0
+ frequency = 1.0
+ for _ in range(octaves):
+ total += amplitude * pnoise2(x * frequency / scale, y * frequency / scale)
+ amplitude *= gain
+ frequency *= lacunarity
+ return total
+
+ def pattern(x, y, scale, octaves, lacunarity, gain):
+ """
+ Domain-warped pattern using fbm.
+ Warps coordinates before applying final fbm call.
+ """
+ q0 = fbm(x, y, scale, octaves, lacunarity, gain)
+ q1 = fbm(x + 5.2, y + 1.3, scale, octaves, lacunarity, gain)
+ return fbm(x + 80.0 * q0, y + 80.0 * q1, scale, octaves, lacunarity, gain)
+
+ def generate_pattern_noise(size=(512, 512), scale=80, octaves=5, lacunarity=2.0, gain=0.5,
+ saturation=1.5, brightness=1, seed=None):
+ """
+ Generate colored noise image using domain-warped fractal noise and random offsets.
+ """
+ width, height = size
+ img = np.zeros((height, width, 3), dtype=np.uint8)
+ rng = random.Random(seed)
+ offset_x = rng.uniform(-1000, 1000)
+ offset_y = rng.uniform(-1000, 1000)
+
+ for i in range(height):
+ for j in range(width):
+ x = i + offset_x
+ y = j + offset_y
+
+ r_val = pattern(x, y, scale, octaves, lacunarity, gain)
+ g_val = pattern(x + 100, y + 100, scale, octaves, lacunarity, gain)
+ b_val = pattern(x + 200, y + 200, scale, octaves, lacunarity, gain)
+
+ r, g, b = [(val + 1) / 2 for val in (r_val, g_val, b_val)]
+ avg = (r + g + b) / 3
+
+ r = np.clip(avg + (r - avg) * saturation, 0, 1) * brightness
+ g = np.clip(avg + (g - avg) * saturation, 0, 1) * brightness
+ b = np.clip(avg + (b - avg) * saturation, 0, 1) * brightness
+
+ img[i, j] = [int(np.clip(r, 0, 1) * 255), int(np.clip(g, 0, 1) * 255), int(np.clip(b, 0, 1) * 255)]
+
+ image = Image.fromarray(img).filter(ImageFilter.GaussianBlur(radius=2))
+ return image
+
+ def measure_fade_pixels(mask_np):
+ """
+ Estimate edge fade width from a grayscale mask using gradient analysis.
+ Attempts measurement from top, right, bottom, and left.
+ Returns fallback value if no valid result is found.
+ """
+ h, w = mask_np.shape
+
+ def measure_line(line):
+ grad = np.gradient(line)
+ max_grad = np.max(grad)
+ if max_grad == 0:
+ return None
+ half_max = max_grad / 2.0
+ indices = np.where(grad >= half_max)[0]
+ if len(indices) == 0:
+ return None
+ return (indices[-1] - indices[0]) / 2.0
+
+ lines = [
+ mask_np[:, w // 2], # Top
+ mask_np[h // 2, ::-1], # Right
+ mask_np[::-1, w // 2], # Bottom
+ mask_np[h // 2, :] # Left
+ ]
+
+ for line in lines:
+ result = measure_line(line)
+ if result and result > 0:
+ return result
+
+ return 16.0 # Fallback
+
+ def compute_fade_mask(binary_mask, fade_pixels=16):
+ """
+ Compute a smooth fade-out mask from a binary mask using distance transform.
+ Pixels within `fade_pixels` of the edge get values between 0 and 1.
+ """
+ mask_uint8 = (binary_mask * 255).astype(np.uint8)
+ dist = cv2.distanceTransform(mask_uint8, distanceType=cv2.DIST_L2, maskSize=5)
+ return np.clip(dist / fade_pixels, 0, 1)
+
+ def preprocess_image(image, mask, noise_fill_image=True, seed=None):
+ """
+ Preprocesses image with optional noise-based fill on masked areas.
+ Includes smoothing transitions and standard cropping and normalization.
+ """
+ image = image.convert("RGB")
+
+ if noise_fill_image:
+ mask = mask.convert("L").resize(image.size, Image.Resampling.NEAREST)
+ mask_blur = np.array(mask, dtype=np.float32) / 255.0
+ fade_pixels = measure_fade_pixels(mask_blur)
+ binary_mask = (mask_blur > 0.5).astype(np.float32)
+
+ noise_img = generate_pattern_noise(size=image.size, seed=seed)
+ image_np = np.array(image)
+ noise_np = np.array(noise_img)
+
+ fade_mask = compute_fade_mask(binary_mask, fade_pixels=fade_pixels)
+ fade_mask = binary_mask * fade_mask
+ fade_mask_3c = np.repeat(fade_mask[:, :, None], 3, axis=2)
+
+ alpha = 0.75
+ blended = (1 - alpha * fade_mask_3c) * image_np + alpha * fade_mask_3c * noise_np
+ image = Image.fromarray(blended.astype(np.uint8))
+ image.save("noised_image.png")
+
+ image = transforms.CenterCrop((image.size[1] // 64 * 64, image.size[0] // 64 * 64))(image)
+ image = transforms.ToTensor()(image)
+ image = image * 2 - 1 # Normalize to [-1, 1]
+ return image.unsqueeze(0)
+
+ def preprocess_map(map):
+ """
+ Convert mask to normalized, inverted grayscale tensor.
+ Applies value remapping and center crop.
+ """
+ map = map.convert("L")
+ map = transforms.CenterCrop((map.size[1] // 64 * 64, map.size[0] // 64 * 64))(map)
+ map = transforms.ToTensor()(map)
+ map = (map - 0.05) / (0.95 - 0.05)
+ map = torch.clamp(map, 0.0, 1.0)
+ return 1.0 - map
+
+ # --------------------------------------------
+ # APPLY PREPROCESSING
+ # --------------------------------------------
+
+ # Prepare original image with optional noise fill
+ original_image_tensor = preprocess_image(image, mask, noise_fill_image=noise_fill_image).to(device)
+ image = original_image_tensor.clone().to(device)
+
+ # Prepare mask as rescaled tensor map
+ map = preprocess_map(mask).to(device)
+ map = torchvision.transforms.Resize(
+ tuple(s // self.vae_scale_factor for s in original_image_tensor.shape[2:]), antialias=None
+ )(map)
+
+ # Generate latent tensor with noise
+ original_with_noise = self.prepare_latents(
+ original_image_tensor, timesteps, batch_size, num_images_per_prompt, prompt_embeds.dtype, device, generator
+ )
+
+ # Create thresholded masks over timesteps
+ thresholds = torch.arange(num_inference_steps, dtype=map.dtype) / num_inference_steps
+ thresholds = thresholds.unsqueeze(1).unsqueeze(1).to(device)
+ masks = map > (thresholds + (denoising_start or 0))
+
+
+
+ # 6. Prepare latent variables.
+ latents = self.prepare_latents(
+ image,
+ latent_timestep,
+ batch_size,
+ num_images_per_prompt,
+ prompt_embeds.dtype,
+ device,
+ generator,
+ add_noise,
+ )
+
+ # 7. Prepare extra step kwargs.
+ extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
+
+ height, width = latents.shape[-2:]
+ height = height * self.vae_scale_factor
+ width = width * self.vae_scale_factor
+
+ original_size = original_size or (height, width)
+ target_size = target_size or (height, width)
+
+ # 8. Prepare added time ids & embeddings
+ if negative_original_size is None:
+ negative_original_size = original_size
+ if negative_target_size is None:
+ negative_target_size = target_size
+
+ add_text_embeds = pooled_prompt_embeds
+ if self.text_encoder_2 is None:
+ text_encoder_projection_dim = int(pooled_prompt_embeds.shape[-1])
+ else:
+ text_encoder_projection_dim = self.text_encoder_2.config.projection_dim
+
+ add_time_ids, add_neg_time_ids = self._get_add_time_ids(
+ original_size,
+ crops_coords_top_left,
+ target_size,
+ aesthetic_score,
+ negative_aesthetic_score,
+ negative_original_size,
+ negative_crops_coords_top_left,
+ negative_target_size,
+ dtype=prompt_embeds.dtype,
+ text_encoder_projection_dim=text_encoder_projection_dim,
+ )
+ add_time_ids = add_time_ids.repeat(batch_size * num_images_per_prompt, 1)
+
+ if self.do_classifier_free_guidance:
+ prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
+ add_text_embeds = torch.cat([negative_pooled_prompt_embeds, add_text_embeds], dim=0)
+ add_neg_time_ids = add_neg_time_ids.repeat(batch_size * num_images_per_prompt, 1)
+ add_time_ids = torch.cat([add_neg_time_ids, add_time_ids], dim=0)
+
+ prompt_embeds = prompt_embeds.to(device)
+ add_text_embeds = add_text_embeds.to(device)
+ add_time_ids = add_time_ids.to(device)
+
+ if ip_adapter_image is not None or ip_adapter_image_embeds is not None:
+ image_embeds = self.prepare_ip_adapter_image_embeds(
+ ip_adapter_image,
+ ip_adapter_image_embeds,
+ device,
+ batch_size * num_images_per_prompt,
+ self.do_classifier_free_guidance,
+ )
+
+ # 9. Denoising loop
+ num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
+
+ # 9.1 Apply denoising_end
+ if (
+ denoising_end is not None
+ and denoising_start is not None
+ and denoising_value_valid(denoising_end)
+ and denoising_value_valid(denoising_start)
+ and denoising_start >= denoising_end
+ ):
+ raise ValueError(f"`denoising_start`: {denoising_start} cannot be larger than or equal to `denoising_end`: {denoising_end} when using type float."
+ )
+ elif denoising_end is not None and denoising_value_valid(denoising_end):
+ discrete_timestep_cutoff = int(
+ round(
+ self.scheduler.config.num_train_timesteps
+ - (denoising_end * self.scheduler.config.num_train_timesteps)
+ )
+ )
+ num_inference_steps = len(list(filter(lambda ts: ts >= discrete_timestep_cutoff, timesteps)))
+ timesteps = timesteps[:num_inference_steps]
+
+ # 9.2 Optionally get Guidance Scale Embedding
+ timestep_cond = None
+ if self.unet.config.time_cond_proj_dim is not None:
+ guidance_scale_tensor = torch.tensor(self.guidance_scale - 1).repeat(batch_size * num_images_per_prompt)
+ timestep_cond = self.get_guidance_scale_embedding(
+ guidance_scale_tensor, embedding_dim=self.unet.config.time_cond_proj_dim
+ ).to(device=device, dtype=latents.dtype)
+
+ self._num_timesteps = len(timesteps)
+ with self.progress_bar(total=num_inference_steps) as progress_bar:
+ for i, t in enumerate(timesteps):
+ if self.interrupt:
+ continue
+
+ # diff diff
+ if i == 0 and denoising_start is None:
+ latents = original_with_noise[:1]
+ else:
+ mask = masks[i].unsqueeze(0)
+ # cast mask to the same type as latents etc
+ mask = mask.to(latents.dtype)
+ mask = mask.unsqueeze(1) # fit shape
+ latents = original_with_noise[i] * mask + latents * (1 - mask)
+ # end diff diff
+
+ # expand the latents if we are doing classifier free guidance
+ latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents
+
+ latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
+
+ # predict the noise residual
+ added_cond_kwargs = {"text_embeds": add_text_embeds, "time_ids": add_time_ids}
+ if ip_adapter_image is not None or ip_adapter_image_embeds is not None:
+ added_cond_kwargs["image_embeds"] = image_embeds
+ noise_pred = self.unet(
+ latent_model_input,
+ t,
+ encoder_hidden_states=prompt_embeds,
+ timestep_cond=timestep_cond,
+ cross_attention_kwargs=cross_attention_kwargs,
+ added_cond_kwargs=added_cond_kwargs,
+ return_dict=False,
+ )[0]
+
+ # perform guidance
+ if self.do_classifier_free_guidance:
+ noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+ noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
+
+ if self.do_classifier_free_guidance and guidance_rescale > 0.0:
+ # Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf
+ noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=guidance_rescale)
+
+ # compute the previous noisy sample x_t -> x_t-1
+ latents_dtype = latents.dtype
+ latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
+ if latents.dtype != latents_dtype:
+ if torch.backends.mps.is_available():
+ # some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272
+ latents = latents.to(latents_dtype)
+ else:
+ raise ValueError(
+ "For the given accelerator, there seems to be an unexpected problem in type-casting. Please file an issue on the PyTorch GitHub repository. See also: https://github.com/huggingface/diffusers/pull/7446/."
+ )
+
+ if callback_on_step_end is not None:
+ callback_kwargs = {}
+ for k in callback_on_step_end_tensor_inputs:
+ callback_kwargs[k] = locals()[k]
+ callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
+
+ latents = callback_outputs.pop("latents", latents)
+ prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
+ negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds)
+ add_text_embeds = callback_outputs.pop("add_text_embeds", add_text_embeds)
+ negative_pooled_prompt_embeds = callback_outputs.pop(
+ "negative_pooled_prompt_embeds", negative_pooled_prompt_embeds
+ )
+ add_time_ids = callback_outputs.pop("add_time_ids", add_time_ids)
+ add_neg_time_ids = callback_outputs.pop("add_neg_time_ids", add_neg_time_ids)
+
+ # call the callback, if provided
+ if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
+ progress_bar.update()
+ if callback is not None and i % callback_steps == 0:
+ step_idx = i // getattr(self.scheduler, "order", 1)
+ callback(step_idx, t, latents)
+
+ if XLA_AVAILABLE:
+ xm.mark_step()
+
+ if output_type != "latent":
+ # make sure the VAE is in float32 mode, as it overflows in float16
+ needs_upcasting = self.vae.dtype == torch.float16 and self.vae.config.force_upcast
+
+ if needs_upcasting:
+ self.upcast_vae()
+ latents = latents.to(next(iter(self.vae.post_quant_conv.parameters())).dtype)
+ elif latents.dtype != self.vae.dtype:
+ if torch.backends.mps.is_available():
+ # some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272
+ self.vae = self.vae.to(latents.dtype)
+ else:
+ raise ValueError(
+ "For the given accelerator, there seems to be an unexpected problem in type-casting. Please file an issue on the PyTorch GitHub repository. See also: https://github.com/huggingface/diffusers/pull/7446/."
+ )
+ # unscale/denormalize the latents
+ # denormalize with the mean and std if available and not None
+ has_latents_mean = hasattr(self.vae.config, "latents_mean") and self.vae.config.latents_mean is not None
+ has_latents_std = hasattr(self.vae.config, "latents_std") and self.vae.config.latents_std is not None
+ if has_latents_mean and has_latents_std:
+ latents_mean = (
+ torch.tensor(self.vae.config.latents_mean).view(1, 4, 1, 1).to(latents.device, latents.dtype)
+ )
+ latents_std = (
+ torch.tensor(self.vae.config.latents_std).view(1, 4, 1, 1).to(latents.device, latents.dtype)
+ )
+ latents = latents * latents_std / self.vae.config.scaling_factor + latents_mean
+ else:
+ latents = latents / self.vae.config.scaling_factor
+
+ image = self.vae.decode(latents, return_dict=False)[0]
+
+ # cast back to fp16 if needed
+ if needs_upcasting:
+ self.vae.to(dtype=torch.float16)
+ else:
+ image = latents
+
+ # apply watermark if available
+ if self.watermark is not None:
+ image = self.watermark.apply_watermark(image)
+
+ image = self.image_processor.postprocess(image, output_type=output_type)
+
+ # Offload all models
+ self.maybe_free_model_hooks()
+
+ if not return_dict:
+ return (image,)
+
+ return StableDiffusionXLPipelineOutput(images=image)
+
+### pipeline end
+
+### script start
+
+import gradio as gr
+from installer import install
+from modules import shared, scripts, processing, sd_models
+
+
+class Script(scripts.Script):
+ orig_pipeline = None
+
+ def title(self):
+ return 'SoftFill: Inpaint with Differential diffusion'
+
+ def show(self, is_img2img):
+ return is_img2img if shared.native else False
+
+ def ui(self, _is_img2img):
+ with gr.Row():
+ gr.HTML(' SoftFill: Inpaint with Differential diffusion
')
+ with gr.Row():
+ enabled = gr.Checkbox(label='Enabled', value=True)
+ with gr.Row():
+ noise = gr.Checkbox(label='Apply noise', value=True)
+ strength = gr.Slider(minimum=0.0, maximum=1.0, value=0.65, label='Fill strength')
+ return enabled, noise, strength
+
+ def run(self, p: processing.StableDiffusionProcessingImg2Img, enabled, noise, strength): # pylint: disable=arguments-differ
+ if not enabled:
+ return
+ if shared.sd_model_type not in ['sdxl']:
+ shared.log.error(f'SoftFill: incorrect base model: {shared.sd_model.__class__.__name__}')
+ return
+ if not hasattr(p, 'init_images') or len(p.init_images) == 0:
+ shared.log.error('SoftFill: no input image')
+ return
+ if not hasattr(p, 'mask') or p.mask is None:
+ shared.log.error('SoftFill: no input mask')
+ return
+
+ try:
+ global pnoise2 # pylint: disable=global-statement
+ install('noise')
+ import noise as noise_module
+ pnoise2 = noise_module.pnoise2
+ except Exception as e:
+ shared.log.error(f'SoftFill: {e}')
+ return
+
+ self.orig_pipeline = shared.sd_model
+ try:
+ shared.sd_model = sd_models.switch_pipe(StableDiffusionXLSoftFillPipeline, shared.sd_model)
+ if shared.sd_model.__class__.__name__ not in sd_models.pipe_switch_task_exclude:
+ sd_models.pipe_switch_task_exclude.append(shared.sd_model.__class__.__name__)
+ except Exception as e:
+ shared.log.error(f'SoftFill: {e}')
+ shared.sd_model = self.orig_pipeline
+ self.orig_pipeline = None
+ return
+
+ p.task_args['noise_fill_image'] = noise
+ p.task_args['strength'] = strength
+ p.task_args['image'] = p.init_images[0]
+ p.task_args['mask'] = p.mask
+ shared.log.info(f'SoftFill: cls={shared.sd_model.__class__.__name__} {p.task_args}')
+
+ def after(self, p: processing.StableDiffusionProcessingImg2Img, *args, **kwargs): # pylint: disable=unused-argument
+ if self.orig_pipeline is not None:
+ shared.sd_model = self.orig_pipeline
+ self.orig_pipeline = None
diff --git a/scripts/xyz_grid_classes.py b/scripts/xyz_grid_classes.py
index 7c7cd1b48..cae9daef5 100644
--- a/scripts/xyz_grid_classes.py
+++ b/scripts/xyz_grid_classes.py
@@ -1,4 +1,4 @@
-from scripts.xyz_grid_shared import apply_field, apply_task_args, apply_setting, apply_prompt_primary, apply_prompt_refine, apply_prompt_detailer, apply_prompt_all, apply_order, apply_sampler, apply_hr_sampler_name, confirm_samplers, apply_checkpoint, apply_refiner, apply_unet, apply_dict, apply_clip_skip, apply_vae, list_lora, apply_lora, apply_lora_strength, apply_te, apply_styles, apply_upscaler, apply_context, apply_detailer, apply_override, apply_processing, apply_options, apply_seed, format_value_add_label, format_value, format_value_join_list, do_nothing, format_nothing, str_permutations # pylint: disable=no-name-in-module, unused-import
+from scripts.xyz_grid_shared import apply_field, apply_task_arg, apply_task_args, apply_setting, apply_prompt_primary, apply_prompt_refine, apply_prompt_detailer, apply_prompt_all, apply_order, apply_sampler, apply_hr_sampler_name, confirm_samplers, apply_checkpoint, apply_refiner, apply_unet, apply_dict, apply_clip_skip, apply_vae, list_lora, apply_lora, apply_lora_strength, apply_te, apply_styles, apply_upscaler, apply_context, apply_detailer, apply_override, apply_processing, apply_options, apply_seed, format_value_add_label, format_value, format_value_join_list, do_nothing, format_nothing, str_permutations # pylint: disable=no-name-in-module, unused-import
from modules import shared, shared_items, sd_samplers, ipadapter, sd_models, sd_vae, sd_unet
@@ -209,4 +209,8 @@ axis_options = [
AxisOption("[PAG] Attention scale", float, apply_field('pag_scale')),
AxisOption("[PAG] Adaptive scaling", float, apply_field('pag_adaptive')),
AxisOption("[PAG] Applied layers", str, apply_setting('pag_apply_layers')),
+ AxisOption("[IY] Scale", float, apply_task_arg('infusenet_conditioning_scale')),
+ AxisOption("[IY] Start", float, apply_task_arg('infusenet_guidance_start')),
+ AxisOption("[IY] End", float, apply_task_arg('infusenet_guidance_end')),
+ AxisOption("[TeaCache] Threshold", float, apply_setting('teacache_thresh')),
]
diff --git a/scripts/xyz_grid_shared.py b/scripts/xyz_grid_shared.py
index f624458d6..6b1d814b3 100644
--- a/scripts/xyz_grid_shared.py
+++ b/scripts/xyz_grid_shared.py
@@ -17,6 +17,13 @@ def apply_field(field):
return fun
+def apply_task_arg(field):
+ def fun(p, x, xs):
+ shared.log.debug(f'XYZ grid apply task-arg: {field}={x}')
+ p.task_args[field] = x
+ return fun
+
+
def apply_task_args(p, x, xs):
for section in x.split(';'):
k, v = section.split('=')
diff --git a/webui.py b/webui.py
index 1996c41b5..c24dbe179 100644
--- a/webui.py
+++ b/webui.py
@@ -156,7 +156,7 @@ def initialize():
def load_model():
if not shared.opts.sd_checkpoint_autoload and shared.cmd_opts.ckpt is None:
- log.info('Model auto load disabled')
+ log.info('Model: autoload=False')
else:
shared.state.begin('Load')
thread_model = Thread(target=lambda: shared.sd_model)
@@ -333,8 +333,8 @@ def start_ui():
if public_ip is not None:
shared.log.info(f'Public URL: {proto}://{public_ip}:{shared.cmd_opts.port}')
if shared.cmd_opts.docs:
- shared.log.info(f'API Docs: {local_url[:-1]}/docs') # pylint: disable=unsubscriptable-object
- shared.log.info(f'API ReDocs: {local_url[:-1]}/redocs') # pylint: disable=unsubscriptable-object
+ shared.log.info(f'API docs: {local_url[:-1]}/docs') # pylint: disable=unsubscriptable-object
+ shared.log.info(f'API redocs: {local_url[:-1]}/redocs') # pylint: disable=unsubscriptable-object
if share_url is not None:
shared.log.info(f'Share URL: {share_url}')
# shared.log.debug(f'Gradio functions: registered={len(shared.demo.fns)}')
diff --git a/wiki b/wiki
index 7d2b46a48..b9cb79112 160000
--- a/wiki
+++ b/wiki
@@ -1 +1 @@
-Subproject commit 7d2b46a482d20febe8179b955ca160cc6936515f
+Subproject commit b9cb791121d539af705909a49b452a0be40728b4