diff --git a/.pylintrc b/.pylintrc index b472d26a0..bbe29b197 100644 --- a/.pylintrc +++ b/.pylintrc @@ -38,6 +38,8 @@ ignore-paths=/usr/lib/.*$, modules/todo, modules/unipc, modules/xadapter, + modules/infiniteyou, + modules/flash_attn_triton_amd, repositories, extensions-builtin/Lora, extensions-builtin/sd-webui-agent-scheduler, diff --git a/.ruff.toml b/.ruff.toml index 8e3d13e64..6c77aa6f3 100644 --- a/.ruff.toml +++ b/.ruff.toml @@ -33,6 +33,8 @@ exclude = [ "modules/todo", "modules/unipc", "modules/xadapter", + "modules/infiniteyou", + "modules/flash_attn_triton_amd", "repositories", "extensions-builtin/Lora", "extensions-builtin/sd-extension-chainner/nodes", @@ -70,6 +72,7 @@ select = [ ignore = [ "B006", # Do not use mutable data structures for argument defaults "B008", # Do not perform function call in argument defaults + "C420", # Unnecessary dict comprehension for iterable; use `dict.fromkeys` instead "C408", # Unnecessary `dict` call "I001", # Import block is un-sorted or un-formatted "E402", # Module level import not at top of file @@ -80,10 +83,13 @@ ignore = [ "F401", # Imported by unused "NPY002", # replace legacy random "RUF005", # Consider iterable unpacking + "RUF008", # Do not use mutable default values for dataclass "RUF010", # Use explicit conversion flag "RUF012", # Mutable class attributes "RUF013", # PEP 484 prohibits implicit `Optional` "RUF015", # Prefer `next(...)` over single element slice + "RUF046", # Value being cast to `int` is already an integer + "RUF051", # Prefer pop over del ] fixable = ["ALL"] unfixable = [] diff --git a/CHANGELOG.md b/CHANGELOG.md index 32d906a7d..0e73825b3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,20 +1,187 @@ # Change Log for SD.Next -## Update for 2025-03-14 +## Update for 2025-04-03 -- fix installer not starting when older version of rich is installed -- fix circular imports when debug flags are enabled -- fix cuda errors with directml -- fix memory stats not displaying the ram usage -- fix runpod memory limit reporting -- fix remote vae not being stored in metadata, thanks @iDeNoh -- add --upgrade to torch_command when using --use-nightly for ipex and rocm -- **ipex** +### Highlights for 2025-04-03 + +Time for another major release with ~120 commits and [ChangeLog](https://github.com/vladmandic/automatic/blob/master/CHANGELOG.md) that spans several pages! + +*Highlights?* +Video...Brand new Video processing module with support for all latest models: **WAN21, Hunyuan, LTX, Cog, Allegro, Mochi1, Latte1** in both *T2V* and *I2V* workflows +And combined with *on-the-fly quantization*, support for *Local/Tiny/Remote* VAE, acceleration modules such as *FasterCache or PAB*, and more! +Models...And support for new models: **CogView-4**, **SANA 1.5**, + +*Plus...* +- New **Prompt Enhance** using LLM, +- New pipelines such as **InfiniteYou** +- New **CLiP** models, improvements to **remote VAE**, additional wiki/docs/guides +- More quantization options and granular control +- Pretty big performance updates to a) Any model using DiT based architecture due to new caching methods, b) ZLUDA with new attention methods, c) LoRA with much lower memory usage + +[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-03 + +- **Video tab** + - see [Video Wiki](https://github.com/vladmandic/sdnext/wiki/Video) for details! + - new top-level tab, replaces previous *video* script in text/image tabs + old scripts are still present, but will be removed in the future + - support for all latest models: + - [Hunyuan](https://huggingface.co/Tencent/HunyuanVideo): *HunyuanVideo, FastHunyuan, SkyReels* | *T2V, I2V* + - [WAN21](https://huggingface.co/Wan-AI/Wan2.1-T2V-1.3B-Diffusers): *1.3B, 14B* | *T2V, I2V* + - [LTXVideo](https://huggingface.co/Lightricks/LTX-Video): *0.9.0, 0.9.1, 0.9.5* | *T2V, I2V* + - [CogVideoX](https://huggingface.co/THUDM/CogVideoX-5b): *2B, 5B* | *T2V, I2V* + - [Allegro](https://huggingface.co/rhymes-ai/Allegro): *T2V* + - [Mochi1](https://huggingface.co/genmo/mochi-1-preview): *T2V* + - [Latte1](https://huggingface.co/maxin-cn/Latte-1): *T2V + - decoding: + - **Default**: use vae from model + - **Tiny VAE**: support for *Hunyuan, WAN, Mochi* + - **Remote VAE**: support for *Hunyuan* + - **LoRA** + - support for *Hunyuan, LTX, WAN, Mochi, Cog* + - add option to apply LoRA directly on GPU or use CPU first in low-memory scenarios + - improve metadata and preview parallel fetch + - support for mp4 so first frame is extracted as used as lora preview + - additional key points: + - all models are auto-downloaded upon first use + uses *system paths -> huggingface* folder + - support for many video types + - optional video interpolation while creating video files + - optional video preview in ui + present if video output is selected + - support for balanced offloading and model offloading + uses system settings + - on-the-fly quantization: *BnB, Quanto, TorchAO* + uses system settings, granular for *transformer* and *text-encoder* separately + - different video models support different video resolutions, frame counts, etc. + and may require specific settings - see model links for details + - see *ToDo/Limitations* section for additional notes +- **Models & Pipelines** + - [THUDM CogView 4](https://huggingface.co/THUDM/CogView4-6B) **6B** variant + new foundation model for image generation based o GLM-4 text encoder and a flow-based diffusion transformer + fully supports offloading and on-the-fly quantization + simply select from *networks -> models -> reference* + *note* cogview4 is compatible with flowmatching samplers + - [NVLabs SANA 1.5](https://huggingface.co/Efficient-Large-Model/SANA1.5_4.8B_1024px_diffusers) in **1.6B**, **4.8B** and [Sprint](https://huggingface.co/Efficient-Large-Model/Sana_Sprint_1.6B_1024px_diffusers) variations + big update to previous SANA model + fully supports offloading and on-the-fly quantization + simply select from *networks -> models -> reference* + - [ByteDance InfiniteYou](https://github.com/bytedance/InfiniteYou/): Flexible Photo Recrafting While Preserving Your Identity + face-transfer model for FLUX.1 + select from *Scripts -> InfiniteYou* + its large, ~12GB on top of FLUX.1 base model so make sure you have offloading and quantization setup + *note* model will be auto-downloaded on first use + - New [zer0int CLiP-L](https://huggingface.co/zer0int/CLIP-Registers-Gated_MLP-ViT-L-14) models: + download text encoders into folder set in settings -> system paths -> text encoders (default is *models/Text-encoder*) + load using *settings -> text encoder* + *tip*: add *sd_text_encoder* to your *settings -> user interface -> quicksettings* list to have it appear at the top of the ui +- **Prompt Enhance** + - see [Prompt Enhance Wiki](https://github.com/vladmandic/sdnext/wiki/Prompt-Enhance) for details! + - new built-in extension available in text/image/control tabs + - can be used to manually or automatically enhance prompts using LLM + - built-in presets for **Gemma-3, Qwen-2.5, Phi-4, Llama-3.2, SmolLM2, Dolphin-3** + - support for custom models + load any models hosted on huggingface + load either model in huggingface format or `gguf` format + *note*: any hf model in `transformers.AutoModelForCausalLM` standard should work + *note*: not all model architecture are supported for `gguf` format + - models are auto-downloaded on first use + - support quantization and offloading + - auto-detect censored output + - debug using `SD_LLM_DEBUG=true` env variable +- **Acceleration** + - Support for most DiT-based models, for example: *FLUX.1, SD35, Hunyuan, Mochi, Latte, Allegro, Cog* + - Enable and configure in *Settings -> Pipeline modifiers* + - [FasterCache](https://huggingface.co/papers/2410.19355) + - [PyramidAttentionBroadcast](https://huggingface.co/papers/2408.12588) +- **Remote VAE** + - add support for remote vae encode in addition to remote vae decode + - used by *img2img, inpaint, hires, detailer* + - remote vae encode is disabled by default, you can enable it in *settings -> variable auto-encoder* + - add remote vae info to metadata, thanks @iDeNoh + - remote vae use `scaling_factor` and `shift_factor` +- **Caption/VLM** + - [Google Gemma 3](https://huggingface.co/google/gemma-3-4b-it) 4B + simply select from list of available models in caption tab + - [ByteDance/Sa2VA](https://huggingface.co/ByteDance/Sa2VA-1B) 1B, 4B + simply select from list of available models in caption tab + - add option to set system prompt for vlm models that support it: *Gemma, Smol, Qwen* +- [NudeNet](https://github.com/vladmandic/sd-extension-nudenet/) extension updates + - add detection of prompt language and alphabet and filter based on those values + - add image policy checks using `LlavaGuard` VLM to detect policy violations (and reasons) + against top-10 standard harmful content categories + - add banned words/expressions check against prompt variations +- **LoRA** + - enable memory cache by default + - significantly reduce memory usage + - improve performance + - improve detection of lora changes + - unload lora only when changes are detected + - refactor code for modularity +- **IPEX** + - add `--upgrade` to torch_command when using `--use-nightly` - add xpu to profiler - fix untyped_storage, torch.eye and torch.cuda.device ops - fix torch 2.7 compatibility - fix performance with balanced offload - fix triton and torch.compile +- **ROCm** + - add `--upgrade` to torch_command when using `--use-nightly` + - disable fp16 for gfx1102 (rx 7600 and rx 7500 series) gpus +- **ZLUDA** + - [triton for ZLUDA v3.9.2](https://github.com/vladmandic/sdnext/wiki/ZLUDA#how-to-enable-triton) + - `torch.compile` is now available + - Flash Attention 2 is now available +- **Other** + - new command line option `--monitor PERIOD` to monitor CPU and GPU memory ever n seconds + - **upscale**: new [asymmetric vae v2](https://huggingface.co/Heasterian/AsymmetricAutoencoderKLUpscaler_v2) upscaling method + - **upscale**: new experimental support for `libvips` upscaling + - **quantization**: add support for `optimum-quanto` on-the-fly quantization during load for all models + note: previous method for quanto is still valid and is noted in settings as post-load quantization + - add quantization support to **CogView-3Plus** + - update `diffusers` and other requirements + - rename vae, unet and text-encoder settings *None* to *Default* to avoid confusion + - **CLI**: add `cli/api-grid.py` which can generate grids using params-from-file for x/y axis + - **Samplers** add ability to set sigma adjustment for each sampler + - **ModernUI** updates + - **CSS** updates + - settings vertiocal/dirty indicator restores to default setting instead to previous value + - video interpolate do not skip duplicate frames + - **settings UI** full refactor +- **Wiki/Docs** + - updated [Models](https://github.com/vladmandic/sdnext/wiki/Models) info + - new [Video](https://github.com/vladmandic/sdnext/wiki/Video) guide + - new [Caption](https://github.com/vladmandic/sdnext/wiki/Caption) guide + - new [VAE](https://github.com/vladmandic/sdnext/wiki/VAE) guide + - updated [SD3](https://github.com/vladmandic/sdnext/wiki/SD3) guide + - updated [ZLUDA](https://github.com/vladmandic/sdnext/wiki/ZLUDA) guide + - updated [OpenVINO](https://github.com/vladmandic/sdnext/wiki/OpenVINO) guide + - updated [AMD-ROCm](https://github.com/vladmandic/sdnext/wiki/AMD-ROCm) guide + - updated [Intel-ARC](https://github.com/vladmandic/sdnext/wiki/Intel-ARC) guide +- **Fixes** + - fix installer not starting when older version of `rich` is installed + - fix circular imports when debug flags are enabled + - fix cuda errors with *directml* + - fix memory stats not displaying the ram usage + - fix **RunPod** memory limit reporting + - fix flux ipadapter with start/stop values + - fix progress api `eta_relative` + - fix `insightface` loader + - fix remove vae for flux.1 + - guard against git returining invalid timestamp + - fix hires with latent upscale + - fix legacy diffusion latent upscalers + - fix upscaler selection in postprocessing + - fix sd35 with batch processing + - fix extra networks cover and inline views + - fix token counter error style with modernui + - fix sampler metadata when using default sampler + - fix paste incorrect float to int cast + - fix server restart from ui + - fix style apply params + - do not allow edit of built-in styles + - improve lora compatibility with balanced offload ## Update for 2025-02-28 @@ -238,7 +405,7 @@ Just one week after latest release and what a week it was with over 50 commits! with detailed defaults for each model type also configurable - select between 150+ *OpenCLiP* supported models, 20+ built-in *VLMs*, *DeepDanbooru* - **VLM**: now that we can use VLMs freely, we've also added support for few more out-of-the-box - [Alibaba Qwen VL2](https://huggingface.co/Qwen/Qwen2-VL-2B), [Huggingface Smol VL2](HuggingFaceTB/SmolVLM-Instruct), [ToriiGate 0.4](Minthy/ToriiGate-v0.4-2B) + [Alibaba Qwen VL2](https://huggingface.co/Qwen/Qwen2-VL-2B), [Huggingface Smol VL2](https://huggingface.co/HuggingFaceTB/SmolVLM-Instruct), [ToriiGate 0.4](https://huggingface.co/Minthy/ToriiGate-v0.4-2B) - **Postprocess** - new sota remove background model: [BEN2](https://huggingface.co/PramaLLC/BEN2) select in *process -> remove background* or enable postprocessing for txt2img/img2img operations @@ -346,7 +513,7 @@ Two weeks since last release, time for update! - piecewise rectified flow as model acceleration - use `perflow` scheduler combined with one of the available pre-trained [models](https://huggingface.co/hansyan) - **Other**: - - **upscale**: new [asymmetric vae](Heasterian/AsymmetricAutoencoderKLUpscaler) upscaling method + - **upscale**: new [asymmetric vae](https://huggingface.co/Heasterian/AsymmetricAutoencoderKLUpscaler) upscaling method - **gallery**: add http fallback for slow/unreliable links - **splash**: add legacy mode indicator on splash screen - **network**: extract thumbnail from model metadata if present diff --git a/TODO.md b/TODO.md index 6acb369c5..c6cef03da 100644 --- a/TODO.md +++ b/TODO.md @@ -4,28 +4,37 @@ Main ToDo list can be found at [GitHub projects](https://github.com/users/vladma ## Current -- +### Issues/Limitations + +N/A ## Future Candidates -- Redesign postprocessing -- Flux NF4 loader: -- IPAdapter negative: -- Control API enhance scripts compatibility -- CogView4 +- 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 -- flux: loader for civitai nf4 models (fixme) -- hypertile: vae breaks when using non-standard sizes (fixme) -- install: enable ROCm for windows when available (fixme) -- lora make support quantized flux (fixme) -- lora: add other quantization types (fixme) -- model load: force-reloading entire model as loading transformers only leads to massive memory usage (fixme) -- model loader: implement model in-memory caching (fixme) -- modernui: monkey-patch for missing tabs.select event (fixme) -- processing: remove duplicate mask params (fixme) -- resize image: enable full VAE mode for resize-latent (fixme) -- sana: fails when quantized (fixme) -- support scripts via api (fixme) -- transformer from-single-file with quant (fixme) +> pnpm lint | grep W0511 | awk -F'TODO ' '{print "- "$NF}' | sed 's/ (fixme)//g' + +- install: enable ROCm for windows when available +- resize image: enable full VAE mode for resize-latent +- infotext: handle using regex instead +- fc: autodetect tensor format based on model +- fc: autodetect distilled based on model +- processing: remove duplicate mask params +- model loader: implement model in-memory caching +- 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 +- lora: maybe force imediate quantization +- modules/lora/lora_extract.py:185:9: W0511: TODO: lora support pre-quantized flux +- control: support scripts via api +- modernui: monkey-patch for missing tabs.select event diff --git a/cli/api-grid.py b/cli/api-grid.py new file mode 100755 index 000000000..e41276bb9 --- /dev/null +++ b/cli/api-grid.py @@ -0,0 +1,198 @@ +#!/usr/bin/env python +from dataclasses import dataclass +import io +import os +import time +import math +import base64 +import logging +import argparse +import requests +import urllib3 +from PIL import Image, ImageDraw, ImageFont + + +@dataclass +class Options: # set default parameters here + prompt: str = '' + negative_prompt: str = '' + seed: int = -1 + steps: int = 20 + cfg_scale: float = 6.0 + sampler_name: str = "Default" + width: int = 1024 + height: int = 1024 + save_images: bool = False + send_images: bool = True + + +@dataclass +class Server: # set server and save options here or use command line arguments + url: str = 'http://127.0.0.1:7860' + api: str = '/sdapi/v1/txt2img' + user: str = None + password: str = None + folder: str = '/tmp' + name: str = str(round(time.time())) + images: bool = False + grids: bool = False + labels: bool = False + + +logging.basicConfig(level = logging.INFO, format = '%(asctime)s %(levelname)s: %(message)s') +log = logging.getLogger(__name__) +urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning) +server = Server() +options = Options() + + +def post(): + try: + req = requests.post(f'{server.url}{server.api}', + json=vars(options), + timeout=300, + verify=False, + auth=requests.auth.HTTPBasicAuth(server.user, server.password) if (server.user is not None) and (server.password is not None) else None) + return { 'error': req.status_code, 'reason': req.reason, 'url': req.url } if req.status_code != 200 else req.json() + except Exception as e: + return { 'error': 0, 'reason': str(e), 'url': server.url } + + +def generate(x: int, y: int): # pylint: disable=redefined-outer-name + t0 = time.time() + log.info(f'x={x} y={y} {options}') + data = post() + t1 = time.time() + images = [] + if 'images' in data: + for i in range(len(data['images'])): + b64 = data['images'][i].split(',',1)[0] + image = Image.open(io.BytesIO(base64.b64decode(b64))) + images.append(image) + info = data['info'] + fn = os.path.join(server.folder, f'{server.name}-{x}-{y}.jpg') if server.images else None + log.info(f'image: time={t1-t0:.2f} size={image.size} fn="{fn}" info="{info}"') + if fn is not None: + image.save(fn) + else: + log.warning(data) + return images + + +def merge(images: list[Image.Image], horizontal: bool, labels: list[str] = None): + rows = 1 if horizontal else len(images) + cols = math.ceil(len(images) / rows) + w = max([i.size[0] for i in images]) + h = max([i.size[1] for i in images]) + image = Image.new('RGB', size = (cols * w, rows * h), color = 'black') + font = ImageFont.truetype('DejaVuSansMono', 1024 // 32) + for i, img in enumerate(images): + x = i % cols * w + y = i // cols * h + img.thumbnail((w, h), Image.Resampling.LANCZOS) + image.paste(img, box=(x, y)) + if labels is not None and len(images) == len(labels): + ctx = ImageDraw.Draw(image) + ctx.text((x + 1, y + 1), labels[i], font = font, fill = (0, 0, 0)) + ctx.text((x, y), labels[i], font = font, fill = (255, 255, 255)) + # log.info({ 'grid': { 'images': len(images), 'rows': rows, 'cols': cols, 'cell': [w, h] } }) + return image + + +def grid(x_file: str, y_file: str): + def set_param(line): + param = line.split(':', maxsplit=1) + k = param[0].strip() + v = param[1].strip() if len(param) > 1 else '' + if k == 'prompt': + options.prompt += f'{v} ' # prompt is appended so its not overwritten + elif k == 'lora': + options.prompt += f' ' # lora is appended to prompt + else: + setattr(options, k, v) + + log.info(server) + os.makedirs(server.folder, exist_ok=True) + try: + x = open(x_file, encoding='utf8').read().splitlines() if x_file is not None else [] + y = open(y_file, encoding='utf8').read().splitlines() if y_file is not None else [] + except Exception as e: + log.error(f'read file: x={x_file} y={y_file} {e}') + return + x = [line for line in x if ':' in line] + y = [line for line in y if ':' in line] + t0 = time.time() + log.info(f'grid: x={len(x)} y={len(y)} prefix={server.name}') + vertical = [] + Image.MAX_IMAGE_PIXELS = None + for j in range(max(1, len(y))): + horizontal = [] + labels = [] + for i in range(max(1, len(x))): + if len(x) > i: + set_param(x[i]) + if len(y) > j: + set_param(y[j]) + images = generate(i, j) + if images is not None and len(images) > 0: + horizontal.extend(images) + labels.append(f'{x[i] if len(x) > i else ""}\n{y[j] if len(y) > j else ""}') + options.prompt = '' # reset prompt + if server.grids: + if len(horizontal) == 0: + log.warning(f'grid: empty row={j}') + continue + merged = merge(horizontal, horizontal=True, labels=labels if server.labels else None) + vertical.append(merged) + if server.grids: + if len(vertical) == 0: + log.warning('grid: empty grid') + return + merged = merge(vertical, horizontal=False) + fn = os.path.join(server.folder, f'{server.name}.jpg') + merged.save(fn) + log.info(f'grid: size={merged.size} fn="{fn}"') + t1 = time.time() + log.info(f'done: time={t1-t0:.2f}') + + +if __name__ == "__main__": + log.info(__file__) + parser = argparse.ArgumentParser(description = 'api-grid') + parser.add_argument('--x', type=str, required=False, default=None, help='file to use for x-axis values') + parser.add_argument('--y', type=str, required=False, default=None, help='file to use for y-axis values') + parser.add_argument('--folder', type=str, required=False, default='/tmp', help='folder to use for saving images') + parser.add_argument('--name', type=str, required=False, default=str(round(time.time())), help='name prefix to use for saving images and grids') + parser.add_argument('--image', type=bool, required=False, default=False, help='save individual images') + parser.add_argument('--grid', type=bool, required=False, default=True, help='save image grids') + parser.add_argument('--labels', type=bool, required=False, default=True, help='draw image labels') + parser.add_argument('--url', type=str, required=False, default='http://127.0.0.1:7860', help='server url') + parser.add_argument('--user', type=str, required=False, default=None, help='server user') + parser.add_argument('--password', type=str, required=False, default=None, help='server password') + parser.add_argument('--prompt', type=str, required=False, default='', help='generate prompt') + parser.add_argument('--negative', type=str, required=False, default='', help='generate negative prompt') + parser.add_argument('--sampler', type=str, required=False, default='Default', help='generate sampler') + parser.add_argument('--width', type=int, required=False, default=1024, help='generate width') + parser.add_argument('--height', type=int, required=False, default=1024, help='generate height') + parser.add_argument('--steps', type=int, required=False, default=20, help='generate steps') + parser.add_argument('--cfg', type=float, required=False, default=6.0, help='generate guidance scale') + parser.add_argument('--seed', type=int, required=False, default=-1, help='generate seed') + args = parser.parse_args() + log.info(args) + server.folder = args.folder + server.name = args.name + server.images = bool(args.image) + server.grids = bool(args.grid) + server.labels = bool(args.labels) + server.url = args.url + server.user = args.user + server.password = args.password + options.prompt = args.prompt + options.negative_prompt = args.negative + options.width = int(args.width) + options.height = int(args.height) + options.sampler_name = args.sampler + options.seed = int(args.seed) + options.steps = int(args.steps) + options.cfg_scale = float(args.cfg) + grid(args.x, args.y) diff --git a/cli/api-progress.py b/cli/api-progress.py index 00ed618d2..cb81293da 100755 --- a/cli/api-progress.py +++ b/cli/api-progress.py @@ -18,7 +18,7 @@ class Dot(dict): opts = Dot({ "timeout": 3600, - "frequency": 60, + "frequency": 1, "action": "sudo shutdown now", "url": "http://127.0.0.1:7860", "user": "", @@ -46,15 +46,22 @@ log.info(f'sdnext monitor started: {opts}') while True: try: status = progress() + # {'progress': 0.0, 'eta_relative': 0.0, 'state': {'skipped': False, 'interrupted': False, 'job': '', 'job_count': 0, 'job_timestamp': '20250316110822', 'job_no': 0, 'sampling_step': 20, 'sampling_steps': 20}, 'current_image': None, 'textinfo': None} state = status.get('state', {}) - last_job = state.get('job_timestamp', None) - if last_job is None: + job_timestamp = state.get('job_timestamp', None) + job_progress = status.get('progress', 0) + eta_relative = status.get('eta_relative', 0) + job = state.get('job', '') + job_timestamp = state.get('job_timestamp', None) + sampling_step = state.get('sampling_step', 0) + sampling_steps = state.get('sampling_steps', 0) + if job_timestamp is None: log.warning(f'sdnext montoring cannot get last job info: {status}') else: - last_job = datetime.datetime.strptime(last_job, "%Y%m%d%H%M%S") - elapsed = datetime.datetime.now() - last_job + job_timestamp = datetime.datetime.strptime(job_timestamp, "%Y%m%d%H%M%S") if job_timestamp != '0' else datetime.datetime.now() + elapsed = datetime.datetime.now() - job_timestamp timeout = round(opts.timeout - elapsed.total_seconds()) - log.info(f'sdnext: last_job={last_job} elapsed={elapsed} timeout={timeout}') + log.info(f'sdnext: last="{job_timestamp}" elapsed={elapsed} timeout={timeout} progress={job_progress} eta={eta_relative} step={sampling_step}/{sampling_steps} job="{job}"') if timeout < 0: log.warning(f'sdnext reached: timeout={opts.timeout} action={opts.action}') os.system(opts.action) diff --git a/cli/api-txt2img.py b/cli/api-txt2img.py index 868b13eee..02bb876d6 100755 --- a/cli/api-txt2img.py +++ b/cli/api-txt2img.py @@ -54,10 +54,12 @@ def generate(args): # pylint: disable=redefined-outer-name options['hr_sampler_name'] = args.sampler data = post('/sdapi/v1/txt2img', options) t1 = time.time() + images = [] if 'images' in data: for i in range(len(data['images'])): b64 = data['images'][i].split(',',1)[0] image = Image.open(io.BytesIO(base64.b64decode(b64))) + images.append(image) info = data['info'] log.info(f'image received: size={image.size} time={t1-t0:.2f} info="{info}"') if args.output: @@ -65,6 +67,7 @@ def generate(args): # pylint: disable=redefined-outer-name log.info(f'image saved: size={image.size} filename={args.output}') else: log.warning(f'no images received: {data}') + return images if __name__ == "__main__": diff --git a/extensions-builtin/Lora/extra_networks_lora.py b/extensions-builtin/Lora/extra_networks_lora.py index 76d490eda..2cbdaea60 100644 --- a/extensions-builtin/Lora/extra_networks_lora.py +++ b/extensions-builtin/Lora/extra_networks_lora.py @@ -55,7 +55,7 @@ class ExtraNetworkLora(extra_networks.ExtraNetwork): loaded.tags = loaded.tags[:shared.opts.lora_apply_tags] all_tags.extend(loaded.tags) if len(all_tags) > 0: - shared.log.debug(f"Load network: type=LoRA tags={all_tags} max={shared.opts.lora_apply_tags} apply") + shared.log.debug(f"Network load: type=LoRA tags={all_tags} max={shared.opts.lora_apply_tags} apply") all_tags = ', '.join(all_tags) p.extra_generation_params["LoRA tags"] = all_tags if '_tags_' in p.prompt: @@ -129,7 +129,7 @@ class ExtraNetworkLora(extra_networks.ExtraNetwork): if len(networks.loaded_networks) > 0 and step == 0: self.infotext(p) self.prompt(p) - shared.log.info(f'Load network: type=LoRA apply={[n.name for n in networks.loaded_networks]} method=legacy te={te_multipliers} unet={unet_multipliers} dims={dyn_dims} load={t1-t0:.2f}') + shared.log.info(f'Network load: type=LoRA apply={[n.name for n in networks.loaded_networks]} method=legacy te={te_multipliers} unet={unet_multipliers} dims={dyn_dims} load={t1-t0:.2f}') def deactivate(self, p): t0 = time.time() diff --git a/extensions-builtin/Lora/networks.py b/extensions-builtin/Lora/networks.py index 1f02f3846..e59555993 100644 --- a/extensions-builtin/Lora/networks.py +++ b/extensions-builtin/Lora/networks.py @@ -95,13 +95,13 @@ def load_diffusers(name, network_on_disk, lora_scale=shared.opts.extra_networks_ t0 = time.time() name = name.replace(".", "_") #cached = lora_cache.get(name, None) - shared.log.debug(f'Load network: type=LoRA name="{name}" file="{network_on_disk.filename}" detected={network_on_disk.sd_version} method=diffusers scale={lora_scale} fuse={shared.opts.lora_fuse_diffusers}') + shared.log.debug(f'Network load: type=LoRA name="{name}" file="{network_on_disk.filename}" detected={network_on_disk.sd_version} method=diffusers scale={lora_scale} fuse={shared.opts.lora_fuse_diffusers}') # if cached is not None: # return cached if not shared.native: return None if not hasattr(shared.sd_model, 'load_lora_weights'): - shared.log.error(f'Load network: type=LoRA class={shared.sd_model.__class__} does not implement load lora') + shared.log.error(f'Network load: type=LoRA class={shared.sd_model.__class__} does not implement load lora') return None try: shared.sd_model.load_lora_weights(network_on_disk.filename, adapter_name=name) @@ -110,9 +110,9 @@ def load_diffusers(name, network_on_disk, lora_scale=shared.opts.extra_networks_ pass else: if 'The following keys have not been correctly renamed' in str(e): - shared.log.error(f'Load network: type=LoRA name="{name}" diffusers unsupported format') + shared.log.error(f'Network load: type=LoRA name="{name}" diffusers unsupported format') else: - shared.log.error(f'Load network: type=LoRA name="{name}" {e}') + shared.log.error(f'Network load: type=LoRA name="{name}" {e}') if debug: errors.display(e, "LoRA") return None @@ -133,7 +133,7 @@ def load_network(name, network_on_disk) -> network.Network: t0 = time.time() cached = lora_cache.get(name, None) if debug: - shared.log.debug(f'Load network: type=LoRA name="{name}" file="{network_on_disk.filename}" type=lora {"cached" if cached else ""}') + shared.log.debug(f'Network load: type=LoRA name="{name}" file="{network_on_disk.filename}" type=lora {"cached" if cached else ""}') if cached is not None: return cached net = network.Network(name, network_on_disk) @@ -182,11 +182,11 @@ def load_network(name, network_on_disk) -> network.Network: else: net.modules[key] = net_module if len(keys_failed_to_match) > 0: - shared.log.warning(f'Load network: type=LoRA name="{name}" type={set(network_types)} unmatched={len(keys_failed_to_match)} matched={len(matched_networks)}') + shared.log.warning(f'Network load: type=LoRA name="{name}" type={set(network_types)} unmatched={len(keys_failed_to_match)} matched={len(matched_networks)}') if debug: - shared.log.debug(f'Load network: type=LoRA name="{name}" unmatched={keys_failed_to_match}') + shared.log.debug(f'Network load: type=LoRA name="{name}" unmatched={keys_failed_to_match}') else: - shared.log.debug(f'Load network: type=LoRA name="{name}" type={set(network_types)} keys={len(matched_networks)}') + shared.log.debug(f'Network load: type=LoRA name="{name}" type={set(network_types)} keys={len(matched_networks)}') if len(matched_networks) == 0: return None lora_cache[name] = net @@ -233,7 +233,7 @@ def load_networks(names, te_multipliers=None, unet_multipliers=None, dyn_dims=No if network_on_disk is not None: shorthash = getattr(network_on_disk, 'shorthash', '').lower() if debug: - shared.log.debug(f'Load network: type=LoRA name="{name}" file="{network_on_disk.filename}" hash="{shorthash}"') + shared.log.debug(f'Network load: type=LoRA name="{name}" file="{network_on_disk.filename}" hash="{shorthash}"') try: if recompile_model: shared.compiled_model_state.lora_model.append(f"{name}:{te_multipliers[i] if te_multipliers else shared.opts.extra_networks_default_multiplier}") @@ -245,13 +245,13 @@ def load_networks(names, te_multipliers=None, unet_multipliers=None, dyn_dims=No net.mentioned_name = name network_on_disk.read_hash() except Exception as e: - shared.log.error(f'Load network: type=LoRA file="{network_on_disk.filename}" {e}') + shared.log.error(f'Network load: type=LoRA file="{network_on_disk.filename}" {e}') if debug: errors.display(e, 'LoRA') continue if net is None: failed_to_load_networks.append(name) - shared.log.error(f'Load network: type=LoRA name="{name}" detected={network_on_disk.sd_version if network_on_disk is not None else None} failed') + shared.log.error(f'Network load: type=LoRA name="{name}" detected={network_on_disk.sd_version if network_on_disk is not None else None} failed') continue if shared.native: shared.sd_model.embedding_db.load_diffusers_embedding(None, net.bundle_embeddings) @@ -265,24 +265,24 @@ def load_networks(names, te_multipliers=None, unet_multipliers=None, dyn_dims=No lora_cache.pop(name, None) if len(diffuser_loaded) > 0: - shared.log.debug(f'Load network: type=LoRA loaded={diffuser_loaded} available={shared.sd_model.get_list_adapters()} active={shared.sd_model.get_active_adapters()} scales={diffuser_scales}') + shared.log.debug(f'Network load: type=LoRA loaded={diffuser_loaded} available={shared.sd_model.get_list_adapters()} active={shared.sd_model.get_active_adapters()} scales={diffuser_scales}') try: shared.sd_model.set_adapters(adapter_names=diffuser_loaded, adapter_weights=diffuser_scales) if shared.opts.lora_fuse_diffusers: shared.sd_model.fuse_lora(adapter_names=diffuser_loaded, lora_scale=1.0, fuse_unet=True, fuse_text_encoder=True) # fuse uses fixed scale since later apply does the scaling shared.sd_model.unload_lora_weights() except Exception as e: - shared.log.error(f'Load network: type=LoRA {e}') + shared.log.error(f'Network load: type=LoRA {e}') if debug: errors.display(e, 'LoRA') if len(loaded_networks) > 0 and debug: - shared.log.debug(f'Load network: type=LoRA loaded={len(loaded_networks)} cache={list(lora_cache)}') + shared.log.debug(f'Network load: type=LoRA loaded={len(loaded_networks)} cache={list(lora_cache)}') devices.torch_gc() if recompile_model: - shared.log.info("Load network: type=LoRA recompiling model") + shared.log.info("Network load: type=LoRA recompiling model") backup_lora_model = shared.compiled_model_state.lora_model if 'Model' in shared.opts.cuda_compile: shared.sd_model = sd_models_compile.compile_diffusers(shared.sd_model) @@ -310,7 +310,7 @@ def network_restore_weights_from_backup(self: Union[torch.nn.Conv2d, torch.nn.Li self.weight = torch.nn.Parameter(weights_backup.to(self.weight.device, copy=True)) self.freeze() elif getattr(self, "quant_type", None) in ['nf4', 'fp4']: - bnb = model_quant.load_bnb('Load network: type=LoRA', silent=True) + bnb = model_quant.load_bnb('Network load: type=LoRA', silent=True) if bnb is not None: device = self.weight.device self.weight = bnb.nn.Params4bit(weights_backup, quant_state=self.quant_state, quant_type=self.quant_type, blocksize=self.blocksize) @@ -339,7 +339,7 @@ def maybe_backup_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn. if isinstance(self, torch.nn.MultiheadAttention): weights_backup = (self.in_proj_weight.clone().to(devices.cpu), self.out_proj.weight.clone().to(devices.cpu)) elif getattr(self.weight, "quant_type", None) in ['nf4', 'fp4']: - bnb = model_quant.load_bnb('Load network: type=LoRA', silent=True) + bnb = model_quant.load_bnb('Network load: type=LoRA', silent=True) if bnb is not None: with devices.inference_context(): weights_backup = bnb.functional.dequantize_4bit(self.weight, quant_state=self.weight.quant_state, quant_type=self.weight.quant_type, blocksize=self.weight.blocksize,).to(devices.cpu) @@ -390,7 +390,7 @@ def network_apply_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn # inpainting model. zero pad updown to make channel[1] 4 to 9 updown = torch.nn.functional.pad(updown, (0, 0, 0, 0, 0, 5)) # pylint: disable=not-callable if getattr(self.weight, "quant_type", None) in ['nf4', 'fp4']: # or self.weight.numel() != updown.numel(): - bnb = model_quant.load_bnb('Load network: type=LoRA', silent=True) + bnb = model_quant.load_bnb('Network load: type=LoRA', silent=True) if bnb is not None: device = self.weight.device weight = bnb.functional.dequantize_4bit(self.weight, quant_state=self.weight.quant_state, quant_type=self.weight.quant_type, blocksize=self.weight.blocksize) diff --git a/extensions-builtin/sd-extension-system-info b/extensions-builtin/sd-extension-system-info index 8c7edb3be..ce373b9c2 160000 --- a/extensions-builtin/sd-extension-system-info +++ b/extensions-builtin/sd-extension-system-info @@ -1 +1 @@ -Subproject commit 8c7edb3be11b8b8c2d2dcd0421e93345bd20fcae +Subproject commit ce373b9c27544f56ad73a1f7fe2c5530a89c1c32 diff --git a/extensions-builtin/sdnext-modernui b/extensions-builtin/sdnext-modernui index 7fc52bb97..770db0076 160000 --- a/extensions-builtin/sdnext-modernui +++ b/extensions-builtin/sdnext-modernui @@ -1 +1 @@ -Subproject commit 7fc52bb976783322bdf381e046ceb689261dbe2c +Subproject commit 770db007688d5be9df0def02af64a1fe6449c04e diff --git a/html/licenses.html b/html/licenses.html index 6597fa3ae..dc0e1fdbe 100644 --- a/html/licenses.html +++ b/html/licenses.html @@ -637,6 +637,40 @@ SOFTWARE. limitations under the License. +

Flash Attention

+Fast and memory-efficient exact attention +
+BSD 3-Clause License
+
+Copyright (c) 2022, the respective contributors, as shown by the AUTHORS file.
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+* Redistributions of source code must retain the above copyright notice, this
+   list of conditions and the following disclaimer.
+
+* Redistributions in binary form must reproduce the above copyright notice,
+   this list of conditions and the following disclaimer in the documentation
+   and/or other materials provided with the distribution.
+
+* Neither the name of the copyright holder nor the names of its
+   contributors may be used to endorse or promote products derived from
+   this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+

Curated transformers

The MPS workaround for nn.Linear on macOS 13.2.X is based on the MPS workaround for nn.Linear created by danieldk for Curated transformers 
diff --git a/html/reference.json b/html/reference.json
index 4cc2edb28..b18818750 100644
--- a/html/reference.json
+++ b/html/reference.json
@@ -187,25 +187,43 @@
     "extras": "sampler: Default, cfg_scale: 3.5"
   },
 
-  "NVLabs Sana 1.6B 4k": {
+  "NVLabs Sana 1.5 1.6B 1k": {
+    "path": "Efficient-Large-Model/SANA1.5_1.6B_1024px_diffusers",
+    "desc": "Sana is an efficient model with scaling of training-time and inference time techniques. SANA-1.5 delivers: efficient model growth from 1.6B Sana-1.0 model to 4.8B, achieving similar or better performance than training from scratch and saving 60% training cost; efficient model depth pruning, slimming any model size as you want; powerful VLM selection based inference scaling, smaller model+inference scaling > larger model.",
+    "preview": "Efficient-Large-Model--Sana_1600M_1024px_diffusers.jpg",
+    "skip": true
+  }, 
+  "NVLabs Sana 1.5 4.8B 1k": {
+    "path": "Efficient-Large-Model/SANA1.5_4.8B_1024px_diffusers",
+    "desc": "Sana is an efficient model with scaling of training-time and inference time techniques. SANA-1.5 delivers: efficient model growth from 1.6B Sana-1.0 model to 4.8B, achieving similar or better performance than training from scratch and saving 60% training cost; efficient model depth pruning, slimming any model size as you want; powerful VLM selection based inference scaling, smaller model+inference scaling > larger model.",
+    "preview": "Efficient-Large-Model--Sana_1600M_1024px_diffusers.jpg",
+    "skip": true
+  }, 
+  "NVLabs Sana 1.5 1.6B 1k Sprint": {
+    "path": "Efficient-Large-Model/Sana_Sprint_1.6B_1024px_diffusers",
+    "desc": "SANA-Sprint is an ultra-efficient diffusion model for text-to-image (T2I) generation, reducing inference steps from 20 to 1-4 while achieving state-of-the-art performance.",
+    "preview": "Efficient-Large-Model--Sana_1600M_1024px_diffusers.jpg",
+    "skip": true
+  }, 
+  "NVLabs Sana 1.0 1.6B 4k": {
     "path": "Efficient-Large-Model/Sana_1600M_4Kpx_BF16_diffusers",
     "desc": "Sana is a text-to-image framework that can efficiently generate images up to 4096 × 4096 resolution. Sana can synthesize high-resolution, high-quality images with strong text-image alignment at a remarkably fast speed, deployable on laptop GPU.",
     "preview": "Efficient-Large-Model--Sana_1600M_1024px_diffusers.jpg",
     "skip": true
   }, 
-  "NVLabs Sana 1.6B 2k": {
+  "NVLabs Sana 1.0 1.6B 2k": {
     "path": "Efficient-Large-Model/Sana_1600M_2Kpx_BF16_diffusers",
     "desc": "Sana is a text-to-image framework that can efficiently generate images up to 4096 × 4096 resolution. Sana can synthesize high-resolution, high-quality images with strong text-image alignment at a remarkably fast speed, deployable on laptop GPU.",
     "preview": "Efficient-Large-Model--Sana_1600M_1024px_diffusers.jpg",
     "skip": true
   }, 
-  "NVLabs Sana 1.6B 1k": {
+  "NVLabs Sana 1.0 1.6B 1k": {
     "path": "Efficient-Large-Model/Sana_1600M_1024px_diffusers",
     "desc": "Sana is a text-to-image framework that can efficiently generate images up to 4096 × 4096 resolution. Sana can synthesize high-resolution, high-quality images with strong text-image alignment at a remarkably fast speed, deployable on laptop GPU.",
     "preview": "Efficient-Large-Model--Sana_1600M_1024px_diffusers.jpg",
     "skip": true
   }, 
-  "NVLabs Sana 0.6B 0.5k": {
+  "NVLabs Sana 1.0 0.6B 0.5k": {
     "path": "Efficient-Large-Model/Sana_600M_512px_diffusers",
     "desc": "Sana is a text-to-image framework that can efficiently generate images up to 4096 × 4096 resolution. Sana can synthesize high-resolution, high-quality images with strong text-image alignment at a remarkably fast speed, deployable on laptop GPU.",
     "preview": "Efficient-Large-Model--Sana_1600M_1024px_diffusers.jpg",
@@ -376,9 +394,15 @@
     "extras": "sampler: DPM++ 2M EDM"
   },
 
+  "CogView 4": {
+    "path": "THUDM/CogView4-6B",
+    "desc": "An innovative cascaded framework that enhances the performance of text-to-image diffusion. CogView is the first model implementing relay diffusion in the realm of text-to-image generation, executing the task by first creating low-resolution images and subsequently applying relay-based super-resolution.",
+    "preview": "THUDM--CogView4-6B.jpg",
+    "skip": true
+  },
   "CogView 3 Plus": {
     "path": "THUDM/CogView3-Plus-3B",
-    "desc": "This model is the DiT version of CogView3, a text-to-image generation model, supporting image generation from 512 to 2048px. Resolution: Width and height must meet the range from 512px to 2048px and must be divisible by 32.",
+    "desc": "An innovative cascaded framework that enhances the performance of text-to-image diffusion. CogView is the first model implementing relay diffusion in the realm of text-to-image generation, executing the task by first creating low-resolution images and subsequently applying relay-based super-resolution.",
     "preview": "THUDM--CogView3-Plus-3B.jpg",
     "skip": true
   },
diff --git a/installer.py b/installer.py
index 153dae193..36c84e549 100644
--- a/installer.py
+++ b/installer.py
@@ -517,7 +517,7 @@ def check_python(supported_minors=[9, 10, 11, 12], reason=None):
         log.error(f"Python version incompatible: {sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro} required 3.{supported_minors}")
         if reason is not None:
             log.error(reason)
-        if not args.ignore:
+        if not args.ignore and not args.experimental:
             sys.exit(1)
     if int(sys.version_info.minor) == 12:
         os.environ.setdefault('SETUPTOOLS_USE_DISTUTILS', 'local') # hack for python 3.11 setuptools
@@ -536,9 +536,9 @@ def check_python(supported_minors=[9, 10, 11, 12], reason=None):
 # check diffusers version
 def check_diffusers():
     t_start = time.time()
-    if args.skip_all or args.skip_git:
+    if args.skip_all or args.skip_git or args.experimental:
         return
-    sha = 'b75b204a584e29ebf4e80a61be11458e9ed56e3e' # diffusers commit hash
+    sha = 'e5c6027ef89ec1a2800c0421599da89d4820f2e4' # 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 ''
@@ -652,7 +652,6 @@ def install_rocm_zluda():
                 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')
-                log.info(f'Using ZLUDA in {zluda_installer.path}')
             except Exception as e:
                 error = e
                 log.warning(f'Failed to load ZLUDA: {e}')
@@ -1302,7 +1301,7 @@ def check_ui(ver):
     def same(ver):
         core = ver['branch'] if ver is not None and 'branch' in ver else 'unknown'
         ui = ver['ui'] if ver is not None and 'ui' in ver else 'unknown'
-        return core == ui or (core == 'master' and ui == 'main')
+        return (core == ui) or (core == 'master' and ui == 'main') or (core == 'dev' and ui == 'dev')
 
     t_start = time.time()
     if not same(ver):
@@ -1433,14 +1432,16 @@ def check_timestamp():
             if 'Setup complete without errors' in line:
                 setup_time = int(line.split(' ')[-1])
     try:
-        version_time = int(git('log -1 --pretty=format:"%at"'))
+        version_time = git('log -1 --pretty=format:"%at"')
+        version_time = ''.join(filter(str.isdigit, version_time))
+        version_time = int(version_time) if len(version_time) > 0 else -1
+        log.debug(f'Timestamp repository update time: {time.ctime(version_time)}')
     except Exception as e:
         log.error(f'Timestamp local repository version: {e}')
-    log.debug(f'Timestamp repository update time: {time.ctime(int(version_time))}')
     if setup_time == -1:
         return False
     log.debug(f'Timestamp previous setup time: {time.ctime(setup_time)}')
-    if setup_time < version_time:
+    if setup_time < version_time or version_time == -1:
         ok = False
     extension_time = check_extensions()
     log.debug(f'Timestamp latest extensions time: {time.ctime(extension_time)}')
@@ -1491,6 +1492,7 @@ def add_args(parser):
     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/amethyst-nightfall.css b/javascript/amethyst-nightfall.css
index ef582848d..929bd9e77 100644
--- a/javascript/amethyst-nightfall.css
+++ b/javascript/amethyst-nightfall.css
@@ -88,8 +88,6 @@ svg.feather.feather-image, .feather .feather-image { display: none }
 /* gradio elements overrides */
 #div.gradio-container { overflow-x: hidden; }
 #img2img_label_copy_to_img2img { font-weight: normal; }
-#txt2img_prompt, #txt2img_neg_prompt, #img2img_prompt, #img2img_neg_prompt { background-color: var(--background-color); box-shadow: 4px 4px 4px 0px #333333  !important; }
-#txt2img_prompt > label > textarea, #txt2img_neg_prompt > label > textarea, #img2img_prompt > label > textarea, #img2img_neg_prompt > label > textarea { font-size: 1.1rem; }
 #img2img_settings { min-width: calc(2 * var(--left-column)); max-width: calc(2 * var(--left-column)); background-color: #111111; padding-top: 16px; }
 #interrogate, #deepbooru { margin: 0 0px 10px 0px; max-width: 80px; max-height: 80px; font-weight: normal; font-size: 0.95em; }
 #quicksettings .gr-button-tool { font-size: 1.6rem; box-shadow: none; margin-left: -20px; margin-top: -2px; height: 2.4em; }
diff --git a/javascript/base.css b/javascript/base.css
index 8f89685c2..fe06e6732 100644
--- a/javascript/base.css
+++ b/javascript/base.css
@@ -4,11 +4,11 @@
 .gradio-button.tool { max-width: min-content; min-width: min-content !important; align-self: end; font-size: 1.4em; color: var(--body-text-color) !important; }
 
 /* token counters */
-.block.token-counter { position: absolute; display: inline-block; right: 0; min-width: 0 !important; width: auto; z-index: 100; top: 0; }
-.block.token-counter span { background: var(--input-background-fill) !important; box-shadow: 0 0 0.0 0.3em rgba(192,192,192,0.15), inset 0 0 0.6em rgba(192,192,192,0.075); border: 2px solid rgba(192,192,192,0.4) !important; }
-.block.token-counter.error span{ box-shadow: 0 0 0.0 0.3em rgba(255,0,0,0.15), inset 0 0 0.6em rgba(255,0,0,0.075); border: 2px solid rgba(255,0,0,0.4) !important; }
-.block.token-counter div { display: inline; }
-.block.token-counter span { padding: 0.1em 0.75em; }
+.token-counter { position: absolute; display: inline-block; right: 0; min-width: 0 !important; width: auto; z-index: 100; top: 0; }
+.token-counter span { background: var(--input-background-fill) !important; box-shadow: 0 0 0.0 0.3em rgba(192,192,192,0.15), inset 0 0 0.6em rgba(192,192,192,0.075); border: 2px solid rgba(192,192,192,0.4) !important; }
+.token-counter.error span{ box-shadow: 0 0 0.0 0.3em rgba(255,0,0,0.15), inset 0 0 0.6em rgba(255,0,0,0.075); border: 2px solid rgba(255,0,0,0.4) !important; }
+.token-counter div { display: inline; }
+.token-counter span { padding: 0.1em 0.75em; }
 
 /* tooltips and statuses */ 
 .infotext { overflow-wrap: break-word; }
diff --git a/javascript/black-gray.css b/javascript/black-gray.css
index c262a3bf4..784626d60 100644
--- a/javascript/black-gray.css
+++ b/javascript/black-gray.css
@@ -103,11 +103,6 @@ svg.feather.feather-image, .feather .feather-image { display: none }
 /* gradio elements overrides */
 #div.gradio-container { overflow-x: hidden; }
 #img2img_label_copy_to_img2img { font-weight: normal; }
-
-#txt2img_prompt, #txt2img_neg_prompt, #img2img_prompt, #img2img_neg_prompt, #control_prompt, #control_neg_prompt { background-color: var(--background-color); box-shadow: none !important; }
-#txt2img_prompt > label > textarea, #txt2img_neg_prompt > label > textarea, #img2img_prompt > label > textarea, #img2img_neg_prompt > label > textarea, #control_prompt > label > textarea, #control_neg_prompt > label > textarea { font-size: 1.0em; line-height: 1.4em; }
-#txt2img_styles, #img2img_styles, #control_styles { padding: 0; }
-#txt2img_styles_refresh, #img2img_styles_refresh, #control_styles_refresh { padding: 0; margin-top: 1em; }
 #img2img_settings { min-width: calc(2 * var(--left-column)); max-width: calc(2 * var(--left-column)); background-color: var(--primary-950); padding-top: 16px; }
 #interrogate, #deepbooru { margin: 0 0px 10px 0px; max-width: 80px; max-height: 80px; font-weight: normal; font-size: 0.95em; }
 #quicksettings .gr-button-tool { font-size: 1.6rem; box-shadow: none; margin-left: -20px; margin-top: -2px; height: 2.4em; }
diff --git a/javascript/black-orange.css b/javascript/black-orange.css
index 54b98b1df..467c53d1d 100644
--- a/javascript/black-orange.css
+++ b/javascript/black-orange.css
@@ -105,8 +105,6 @@ svg.feather.feather-image, .feather .feather-image { display: none }
 /* gradio elements overrides */
 #div.gradio-container { overflow-x: hidden; }
 #img2img_label_copy_to_img2img { font-weight: normal; }
-#txt2img_prompt, #txt2img_neg_prompt, #img2img_prompt, #img2img_neg_prompt { background-color: var(--background-color); box-shadow: 4px 4px 4px 0px #333333  !important; }
-#txt2img_prompt > label > textarea, #txt2img_neg_prompt > label > textarea, #img2img_prompt > label > textarea, #img2img_neg_prompt > label > textarea { font-size: 1.1rem; }
 #img2img_settings { min-width: calc(2 * var(--left-column)); max-width: calc(2 * var(--left-column)); background-color: #111111; padding-top: 16px; }
 #interrogate, #deepbooru { margin: 0 0px 10px 0px; max-width: 80px; max-height: 80px; font-weight: normal; font-size: 0.95em; }
 #quicksettings .gr-button-tool { font-size: 1.6rem; box-shadow: none; margin-left: -20px; margin-top: -2px; height: 2.4em; }
diff --git a/javascript/black-teal-reimagined.css b/javascript/black-teal-reimagined.css
index be6176ac4..ba8957c3e 100644
--- a/javascript/black-teal-reimagined.css
+++ b/javascript/black-teal-reimagined.css
@@ -816,6 +816,8 @@ svg.feather.feather-image,
 #txt2img_extra_search,
 #img2img_description,
 #img2img_extra_search,
+#video_description,
+#video_extra_search,
 #control_description,
 #control_extra_search {
   margin-top: 50px;
diff --git a/javascript/black-teal.css b/javascript/black-teal.css
index 0a6db4fa2..8a99cb91c 100644
--- a/javascript/black-teal.css
+++ b/javascript/black-teal.css
@@ -54,15 +54,21 @@
   --line-md: 1.4em;
   --line-lg: 1.5em;
   --range-shadow:
-     -20em 0 0 0 hsl(180, 54%,  2%), -19em 0 0 0 hsl(180, 54%,  5%), -18em 0 0 0 hsl(180, 54%,  0%), -17em 0 0 0 hsl(180, 54%, 11%),
-     -16em 0 0 0 hsl(180, 54%, 14%), -15em 0 0 0 hsl(180, 54%, 17%), -14em 0 0 0 hsl(180, 54%, 20%), -13em 0 0 0 hsl(180, 54%, 23%),
+     -32em 0 0 0 hsl(180, 54%, 6%), -31em 0 0 0 hsl(180, 54%, 7%), -30em 0 0 0 hsl(180, 54%, 8%), -29em 0 0 0 hsl(180, 54%, 9%),
+     -28em 0 0 0 hsl(180, 54%, 10%), -27em 0 0 0 hsl(180, 54%, 11%), -26em 0 0 0 hsl(180, 54%, 12%), -25em 0 0 0 hsl(180, 54%, 13%),
+     -24em 0 0 0 hsl(180, 54%, 14%), -23em 0 0 0 hsl(180, 54%, 15%), -22em 0 0 0 hsl(180, 54%, 16%), -21em 0 0 0 hsl(180, 54%, 17%),
+     -20em 0 0 0 hsl(180, 54%, 18%), -19em 0 0 0 hsl(180, 54%, 19%), -18em 0 0 0 hsl(180, 54%, 20%), -17em 0 0 0 hsl(180, 54%, 21%),
+     -16em 0 0 0 hsl(180, 54%, 22%), -15em 0 0 0 hsl(180, 54%, 23%), -14em 0 0 0 hsl(180, 54%, 24%), -13em 0 0 0 hsl(180, 54%, 25%),
      -12em 0 0 0 hsl(180, 54%, 26%), -11em 0 0 0 hsl(180, 54%, 29%), -10em 0 0 0 hsl(180, 54%, 32%),  -9em 0 0 0 hsl(180, 54%, 35%),
       -8em 0 0 0 hsl(180, 54%, 38%),  -7em 0 0 0 hsl(180, 54%, 41%), -6em 0 0 0 hsl(180, 54%, 44%),  -5em 0 0 0 hsl(180, 54%, 47%),
       -4em 0 0 0 hsl(180, 54%, 50%),  -3em 0 0 0 hsl(180, 54%, 53%), -2em 0 0 0 hsl(180, 54%, 56%),  -1em 0 0 0 hsl(180, 54%, 59%),
       1em 0 0 0 var(--neutral-800),  2em 0 0 0 var(--neutral-800), 3em 0 0 0 var(--neutral-800),  4em 0 0 0 var(--neutral-800),
       5em 0 0 0 var(--neutral-800),  6em 0 0 0 var(--neutral-800), 7em 0 0 0 var(--neutral-800),  8em 0 0 0 var(--neutral-800),
       9em 0 0 0 var(--neutral-800), 10em 0 0 0 var(--neutral-800), 11em 0 0 0 var(--neutral-800), 12em 0 0 0 var(--neutral-800),
-     13em 0 0 0 var(--neutral-800), 14em 0 0 0 var(--neutral-800), 15em 0 0 0 var(--neutral-800), 16em 0 0 0 var(--neutral-800);
+     13em 0 0 0 var(--neutral-800), 14em 0 0 0 var(--neutral-800), 15em 0 0 0 var(--neutral-800), 16em 0 0 0 var(--neutral-800),
+     17em 0 0 0 var(--neutral-800), 18em 0 0 0 var(--neutral-800), 19em 0 0 0 var(--neutral-800), 20em 0 0 0 var(--neutral-800),
+     21em 0 0 0 var(--neutral-800), 22em 0 0 0 var(--neutral-800), 23em 0 0 0 var(--neutral-800), 24em 0 0 0 var(--neutral-800),
+     25em 0 0 0 var(--neutral-800), 26em 0 0 0 var(--neutral-800), 27em 0 0 0 var(--neutral-800), 28em 0 0 0 var(--neutral-800);
 }
 
 html { font-size: var(--font-size); font-family: var(--font); }
@@ -70,13 +76,6 @@ body, button, input, select, textarea { font-family: var(--font); }
 button { max-width: 400px; white-space: nowrap; }
 img { background-color: var(--background-color); }
 
-/*
-input[type=range] { height: var(--line-xs) !important; appearance: none !important; margin-top: 0 !important; min-width: max(4em, 100%) !important; background-color: var(--background-color) !important; width: 100% !important; background: transparent !important; }
-input[type=range]::-webkit-slider-runnable-track { width: 100% !important; height: 6px !important; cursor: pointer !important; background: var(--input-background-fill) !important; border-radius: var(--radius-lg) !important; border: 0px solid var(--neutral-900) !important; }
-input[type=range]::-moz-range-track              { width: 100% !important; height: 6px !important; cursor: pointer !important; background: var(--input-background-fill) !important; border-radius: var(--radius-lg) !important; border: 0px solid var(--neutral-900) !important; }
-input[type=range]::-webkit-slider-thumb { border: 0px solid #000000 !important; height: var(--line-xs) !important; width: var(--line-md) !important; border-radius: var(--radius-lg) !important; background: var(--highlight-color) !important; cursor: pointer !important; appearance: none !important; margin-top: -4px !important; }
-input[type=range]::-moz-range-thumb     { border: 0px solid #000000 !important; height: var(--line-xs) !important; width: var(--line-md) !important; border-radius: var(--radius-lg) !important; background: var(--highlight-color) !important; cursor: pointer !important; appearance: none !important; margin-top: -4px !important; }
-*/
 input[type='range'] { display: block; margin: 0; padding: 0; height: 1em; background-color: transparent; overflow: hidden; cursor: pointer; box-shadow: 0 0 0 0 transparent; -webkit-appearance: none; appearance: none; }
 input[type='range']::-webkit-slider-thumb { height: .9em; width: .9em; background-color: hsl(180, 54%, 61%); box-shadow: var(--range-shadow); border-radius: var(--radius-xs); }
 input[type='range']::-webkit-slider-runnable-track, input[type='range']::-webkit-slider-thumb { -webkit-appearance: none; }
@@ -142,11 +141,6 @@ svg.feather.feather-image, .feather .feather-image { display: none }
 /* gradio elements overrides */
 #div.gradio-container { overflow-x: hidden; }
 #img2img_label_copy_to_img2img { font-weight: normal; }
-
-#txt2img_prompt, #txt2img_neg_prompt, #img2img_prompt, #img2img_neg_prompt, #control_prompt, #control_neg_prompt { background-color: var(--background-color); box-shadow: none !important; }
-#txt2img_prompt > label > textarea, #txt2img_neg_prompt > label > textarea, #img2img_prompt > label > textarea, #img2img_neg_prompt > label > textarea, #control_prompt > label > textarea, #control_neg_prompt > label > textarea { font-size: 1.0em; line-height: 1.4em; }
-#txt2img_styles, #img2img_styles, #control_styles { padding: 0; margin-top: 2px; }
-#txt2img_styles_refresh, #img2img_styles_refresh, #control_styles_refresh { padding: 0; margin-top: 1em; }
 #img2img_settings { min-width: calc(2 * var(--left-column)); max-width: calc(2 * var(--left-column)); background-color: var(--neutral-950); padding-top: 16px; }
 #interrogate, #deepbooru { margin: 0 0px 10px 0px; max-width: 80px; max-height: 80px; font-weight: normal; font-size: 0.95em; }
 #quicksettings .gr-button-tool { font-size: 1.6rem; box-shadow: none; margin-left: -20px; margin-top: -2px; height: 2.4em; }
diff --git a/javascript/emerald-paradise.css b/javascript/emerald-paradise.css
index f951356cc..411b37774 100644
--- a/javascript/emerald-paradise.css
+++ b/javascript/emerald-paradise.css
@@ -101,8 +101,6 @@ button.selected {background: var(--button-primary-background-fill);}
 /* gradio elements overrides */
 #div.gradio-container { overflow-x: hidden; }
 #img2img_label_copy_to_img2img { font-weight: normal; }
-#txt2img_prompt, #txt2img_neg_prompt, #img2img_prompt, #img2img_neg_prompt { background-color: var(--background-color); box-shadow: 4px 4px 4px 0px #333333  !important; }
-#txt2img_prompt > label > textarea, #txt2img_neg_prompt > label > textarea, #img2img_prompt > label > textarea, #img2img_neg_prompt > label > textarea { font-size: 1.1rem; }
 #img2img_settings { min-width: calc(2 * var(--left-column)); max-width: calc(2 * var(--left-column)); background-color: #111111; padding-top: 16px; }
 #interrogate, #deepbooru { margin: 0 0px 10px 0px; max-width: 80px; max-height: 80px; font-weight: normal; font-size: 0.95em; }
 #quicksettings .gr-button-tool { font-size: 1.6rem; box-shadow: none; margin-top: -2px; height: 2.4em; }
diff --git a/javascript/extraNetworks.js b/javascript/extraNetworks.js
index a9f055a32..90701cdff 100644
--- a/javascript/extraNetworks.js
+++ b/javascript/extraNetworks.js
@@ -9,6 +9,7 @@ const getENActiveTab = () => {
   if (gradioApp().getElementById('tab_txt2img').style.display === 'block') tabName = 'txt2img';
   else if (gradioApp().getElementById('tab_img2img').style.display === 'block') tabName = 'img2img';
   else if (gradioApp().getElementById('tab_control').style.display === 'block') tabName = 'control';
+  else if (gradioApp().getElementById('tab_video').style.display === 'block') tabName = 'video';
   // log('getENActiveTab', tabName);
   return tabName;
 };
@@ -417,6 +418,9 @@ function setupExtraNetworksForTab(tabname) {
         if (h <= 0) return;
         const vh = opts.logmonitor_show ? '55vh' : '68vh';
         if (window.opts.extra_networks_card_cover === 'sidebar' && window.opts.theme_type === 'Standard') el.style.height = `max(${vh}, ${h - 90}px)`;
+        else if (window.opts.extra_networks_card_cover === 'inline' && window.opts.theme_type === 'Standard') el.style.height = '25vh';
+        else if (window.opts.extra_networks_card_cover === 'cover' && window.opts.theme_type === 'Standard') el.style.height = '50vh';
+        else el.style.height = 'unset';
         // log(`${tabname} height: ${entry.target.id}=${h} ${el.id}=${el.clientHeight}`);
       }
     }
@@ -456,6 +460,8 @@ function setupExtraNetworksForTab(tabname) {
           en.style.height = 'unset';
           en.style.width = 'unset';
           en.style.right = 'unset';
+          en.style.maxWidth = 'unset';
+          en.style.maxHeight = '58vh';
           en.style.top = '13em';
           en.style.transition = '';
           en.style.zIndex = 100;
@@ -465,6 +471,7 @@ function setupExtraNetworksForTab(tabname) {
           en.style.height = 'auto';
           en.style.width = `${window.opts.extra_networks_sidebar_width}vw`;
           en.style.maxWidth = '50vw';
+          en.style.maxHeight = 'unset';
           en.style.right = '0';
           en.style.top = '13em';
           en.style.transition = 'width 0.3s ease';
@@ -476,6 +483,8 @@ function setupExtraNetworksForTab(tabname) {
           en.style.height = 'unset';
           en.style.width = 'unset';
           en.style.right = 'unset';
+          en.style.maxWidth = 'unset';
+          en.style.maxHeight = '33vh';
           en.style.top = 0;
           en.style.transition = '';
           en.style.zIndex = 0;
@@ -491,7 +500,7 @@ function setupExtraNetworksForTab(tabname) {
 }
 
 async function showNetworks() {
-  for (const tabname of ['txt2img', 'img2img', 'control']) {
+  for (const tabname of ['txt2img', 'img2img', 'control', 'video']) {
     if (window.opts.extra_networks_show) gradioApp().getElementById(`${tabname}_extra_networks_btn`).click();
   }
   log('showNetworks');
@@ -501,6 +510,7 @@ async function setupExtraNetworks() {
   setupExtraNetworksForTab('txt2img');
   setupExtraNetworksForTab('img2img');
   setupExtraNetworksForTab('control');
+  setupExtraNetworksForTab('video');
 
   function registerPrompt(tabname, id) {
     const textarea = gradioApp().querySelector(`#${id} > label > textarea`);
@@ -515,6 +525,8 @@ async function setupExtraNetworks() {
   registerPrompt('img2img', 'img2img_neg_prompt');
   registerPrompt('control', 'control_prompt');
   registerPrompt('control', 'control_neg_prompt');
+  registerPrompt('video', 'video_prompt');
+  registerPrompt('video', 'video_neg_prompt');
   log('initNetworks', window.opts.extra_networks_card_size);
   document.documentElement.style.setProperty('--card-size', `${window.opts.extra_networks_card_size}px`);
 }
diff --git a/javascript/imageParams.js b/javascript/imageParams.js
index 9ac42ace2..69057ba5e 100644
--- a/javascript/imageParams.js
+++ b/javascript/imageParams.js
@@ -9,6 +9,7 @@ async function initDragDrop() {
     if (tab === 0) promptTarget = 'txt2img_prompt_image';
     else if (tab === 1) promptTarget = 'img2img_prompt_image';
     else if (tab === 2) promptTarget = 'control_prompt_image';
+    else if (tab === 3) promptTarget = 'video_prompt_image';
     else return;
     const imgParent = gradioApp().getElementById(promptTarget);
     const fileInput = imgParent.querySelector('input[type="file"]');
diff --git a/javascript/invoked.css b/javascript/invoked.css
index 72e78d31a..a5954adbf 100644
--- a/javascript/invoked.css
+++ b/javascript/invoked.css
@@ -101,8 +101,6 @@ button.selected {background: var(--button-primary-background-fill);}
 /* gradio elements overrides */
 #div.gradio-container { overflow-x: hidden; }
 #img2img_label_copy_to_img2img { font-weight: normal; }
-#txt2img_prompt, #txt2img_neg_prompt, #img2img_prompt, #img2img_neg_prompt { background-color: var(--background-color); box-shadow: 4px 4px 4px 0px #333333  !important; }
-#txt2img_prompt > label > textarea, #txt2img_neg_prompt > label > textarea, #img2img_prompt > label > textarea, #img2img_neg_prompt > label > textarea { font-size: 1.1rem; }
 #img2img_settings { min-width: calc(2 * var(--left-column)); max-width: calc(2 * var(--left-column)); background-color: #111111; padding-top: 16px; }
 #interrogate, #deepbooru { margin: 0 0px 10px 0px; max-width: 80px; max-height: 80px; font-weight: normal; font-size: 0.95em; }
 #quicksettings .gr-button-tool { font-size: 1.6rem; box-shadow: none; margin-top: -2px; height: 2.4em; }
diff --git a/javascript/light-teal.css b/javascript/light-teal.css
index 7dc9e5950..df8a3ab51 100644
--- a/javascript/light-teal.css
+++ b/javascript/light-teal.css
@@ -101,10 +101,6 @@ svg.feather.feather-image, .feather .feather-image { display: none }
 /* gradio elements overrides */
 #div.gradio-container { overflow-x: hidden; }
 #img2img_label_copy_to_img2img { font-weight: normal; }
-#txt2img_prompt, #txt2img_neg_prompt, #img2img_prompt, #img2img_neg_prompt, #control_prompt, #control_neg_prompt { background-color: var(--background-color); box-shadow: 4px 4px 4px 0px #333333  !important; }
-#txt2img_prompt > label > textarea, #txt2img_neg_prompt > label > textarea, #img2img_prompt > label > textarea, #img2img_neg_prompt > label > textarea, #control_prompt > label > textarea, #control_neg_prompt > label > textarea { font-size: 1.0em; line-height: 1.4em; }
-#txt2img_styles, #img2img_styles, #control_styles { padding: 0; }
-#txt2img_styles_refresh, #img2img_styles_refresh, #control_styles_refresh { padding: 0; margin-top: 1em; }
 #img2img_settings { min-width: calc(2 * var(--left-column)); max-width: calc(2 * var(--left-column)); background-color: #111111; padding-top: 16px; }
 #interrogate, #deepbooru { margin: 0 0px 10px 0px; max-width: 80px; max-height: 80px; font-weight: normal; font-size: 0.95em; }
 #quicksettings .gr-button-tool { font-size: 1.6rem; box-shadow: none; margin-left: -20px; margin-top: -2px; height: 2.4em; }
diff --git a/javascript/midnight-barbie.css b/javascript/midnight-barbie.css
index ea78e5cab..9facd698d 100644
--- a/javascript/midnight-barbie.css
+++ b/javascript/midnight-barbie.css
@@ -94,8 +94,6 @@ svg.feather.feather-image, .feather .feather-image { display: none }
 /* gradio elements overrides */
 #div.gradio-container { overflow-x: hidden; }
 #img2img_label_copy_to_img2img { font-weight: normal; }
-#txt2img_prompt, #txt2img_neg_prompt, #img2img_prompt, #img2img_neg_prompt { background-color: var(--background-color); box-shadow: 4px 4px 4px 0px #333333  !important; }
-#txt2img_prompt > label > textarea, #txt2img_neg_prompt > label > textarea, #img2img_prompt > label > textarea, #img2img_neg_prompt > label > textarea { font-size: 1.1rem; }
 #img2img_settings { min-width: calc(2 * var(--left-column)); max-width: calc(2 * var(--left-column)); background-color: #111111; padding-top: 16px; }
 #interrogate, #deepbooru { margin: 0 0px 10px 0px; max-width: 80px; max-height: 80px; font-weight: normal; font-size: 0.95em; }
 #quicksettings .gr-button-tool { font-size: 1.6rem; box-shadow: none; margin-left: -20px; margin-top: -2px; height: 2.4em; }
diff --git a/javascript/orchid-dreams.css b/javascript/orchid-dreams.css
index 4b121c761..915823bb3 100644
--- a/javascript/orchid-dreams.css
+++ b/javascript/orchid-dreams.css
@@ -101,8 +101,6 @@ button.selected {background: var(--button-primary-background-fill);}
 /* gradio elements overrides */
 #div.gradio-container { overflow-x: hidden; }
 #img2img_label_copy_to_img2img { font-weight: normal; }
-#txt2img_prompt, #txt2img_neg_prompt, #img2img_prompt, #img2img_neg_prompt { background-color: var(--background-color); box-shadow: 4px 4px 4px 0px #333333  !important; }
-#txt2img_prompt > label > textarea, #txt2img_neg_prompt > label > textarea, #img2img_prompt > label > textarea, #img2img_neg_prompt > label > textarea { font-size: 1.1rem; }
 #img2img_settings { min-width: calc(2 * var(--left-column)); max-width: calc(2 * var(--left-column)); background-color: #111111; padding-top: 16px; }
 #interrogate, #deepbooru { margin: 0 0px 10px 0px; max-width: 80px; max-height: 80px; font-weight: normal; font-size: 0.95em; }
 #quicksettings .gr-button-tool { font-size: 1.6rem; box-shadow: none; margin-top: -2px; height: 2.4em; }
diff --git a/javascript/progressBar.js b/javascript/progressBar.js
index 0bac99d6f..1edcbb501 100644
--- a/javascript/progressBar.js
+++ b/javascript/progressBar.js
@@ -14,12 +14,16 @@ function checkPaused(state) {
   lastState.paused = state ? !state : !lastState.paused;
   const t_el = document.getElementById('txt2img_pause');
   const i_el = document.getElementById('img2img_pause');
+  const c_el = document.getElementById('control_pause');
+  const v_el = document.getElementById('video_pause');
   if (t_el) t_el.innerText = lastState.paused ? 'Resume' : 'Pause';
   if (i_el) i_el.innerText = lastState.paused ? 'Resume' : 'Pause';
+  if (c_el) c_el.innerText = lastState.paused ? 'Resume' : 'Pause';
+  if (v_el) v_el.innerText = lastState.paused ? 'Resume' : 'Pause';
 }
 
 function setProgress(res) {
-  const elements = ['txt2img_generate', 'img2img_generate', 'extras_generate', 'control_generate'];
+  const elements = ['txt2img_generate', 'img2img_generate', 'extras_generate', 'control_generate', 'video_generate'];
   const progress = res?.progress || 0;
   const job = res?.job || '';
   let perc = '';
diff --git a/javascript/promptChecker.js b/javascript/promptChecker.js
index a02119e97..e6a269d78 100644
--- a/javascript/promptChecker.js
+++ b/javascript/promptChecker.js
@@ -34,4 +34,6 @@ async function initPromptChecker() {
   setupBracketChecking('img2img_neg_prompt', 'img2img_negative_token_counter');
   setupBracketChecking('control_prompt', 'control_token_counter');
   setupBracketChecking('control_neg_prompt', 'control_negative_token_counter');
+  setupBracketChecking('video_prompt', 'video_token_counter');
+  setupBracketChecking('video_neg_prompt', 'video_negative_token_counter');
 }
diff --git a/javascript/script.js b/javascript/script.js
index f943f4626..8c270bb63 100644
--- a/javascript/script.js
+++ b/javascript/script.js
@@ -65,6 +65,7 @@ function onOptionsChanged(callback) {
 function executeCallbacks(queue, arg) {
   // if (!uiLoaded) return
   for (const callback of queue) {
+    if (!callback) continue;
     try {
       callback(arg);
     } catch (e) {
diff --git a/javascript/sdnext.css b/javascript/sdnext.css
index 3a564679e..5ffc03d36 100644
--- a/javascript/sdnext.css
+++ b/javascript/sdnext.css
@@ -84,11 +84,11 @@ button.custom-button { border-radius: var(--button-large-radius); padding: var(-
 .theme-preview { display: none; position: fixed; border: var(--spacing-sm) solid var(--neutral-600); box-shadow: 2px 2px 2px 2px var(--neutral-700); top: 0; bottom: 0; left: 0; right: 0; margin: auto; max-width: 75vw; z-index: 999; }
 
 /* txt2img/img2img specific */
-.block.token-counter{ position: absolute; right: 1em; min-width: 0 !important; width: auto; z-index: 100; top: -0.5em; }
-.block.token-counter span{ background: var(--input-background-fill) !important; box-shadow: 0 0 0.0 0.3em rgba(192,192,192,0.15), inset 0 0 0.6em rgba(192,192,192,0.075); border: 2px solid rgba(192,192,192,0.4) !important; }
-.block.token-counter.error span{ box-shadow: 0 0 0.0 0.3em rgba(255,0,0,0.15), inset 0 0 0.6em rgba(255,0,0,0.075); border: 2px solid rgba(255,0,0,0.4) !important; }
-.block.token-counter div{ display: inline; }
-.block.token-counter span{ padding: 0.1em 0.75em; }
+.token-counter { position: absolute; right: 1em; min-width: 0 !important; width: auto; z-index: 100; top: -0.5em; }
+.token-counter span { background: var(--input-background-fill) !important; box-shadow: 0 0 0.0 0.3em rgba(192,192,192,0.15), inset 0 0 0.6em rgba(192,192,192,0.075); border: 2px solid rgba(192,192,192,0.4) !important; }
+.token-counter.error span { box-shadow: 0 0 0.0 0.3em rgba(255,0,0,0.15), inset 0 0 0.6em rgba(255,0,0,0.075); border: 2px solid rgba(255,0,0,0.4) !important; }
+.token-counter div { display: inline; }
+.token-counter span { padding: 0.1em 0.75em; }
 .performance { font-size: var(--text-xs); color: #444; }
 .performance p { display: inline-block; color: var(--primary-500) !important }
 .performance .time { margin-right: 0; }
@@ -99,20 +99,25 @@ button.custom-button { border-radius: var(--button-large-radius); padding: var(-
 #txt2img_gallery, #img2img_gallery { height: 50vh; }
 #control-result { background: var(--button-secondary-background-fill); padding: 0.2em; }
 #control-inputs { margin-top: 1em; }
-#txt2img_prompt_container, #img2img_prompt_container, #control_prompt_container { margin-right: var(--layout-gap) }
+#txt2img_prompt_container, #img2img_prompt_container, #control_prompt_container, #video_prompt_container { margin-right: var(--layout-gap) }
 #txt2img_footer, #img2img_footer, #control_footer { height: fit-content; display: none; }
 #txt2img_generate_box, #img2img_generate_box { gap: 0.5em; flex-wrap: unset; min-width: unset; width: 66.6%; }
-#control_generate_box { gap: 0.5em; flex-wrap: unset; min-width: unset; width: 100%; }
-#control_generate_box button:nth-child(1) { flex-grow: 2; }
-#control_generate_box button:nth-child(2) { flex-grow: 1; }
-#txt2img_actions_column, #img2img_actions_column, #control_actions_column { gap: 0.3em; height: fit-content; }
-#txt2img_generate_box>button, #img2img_generate_box>button, #control_generate_box>button, #txt2img_enqueue, #img2img_enqueue, #txt2img_enqueue>button, #img2img_enqueue>button { min-height: 44px !important; max-height: 44px !important; line-height: 1em; white-space: break-spaces; min-width: unset; }
+#control_generate_box, #video_generate_box { gap: 0.5em; flex-wrap: unset; min-width: unset; width: 100%; }
+#control_generate_box button:nth-child(1), #video_generate_box button:nth-child(1) { flex-grow: 2; }
+#control_generate_box button:nth-child(2), #video_generate_box button:nth-child(2) { flex-grow: 1; }
+#txt2img_actions_column, #img2img_actions_column, #control_actions_column, #video_actions_column { gap: 0.3em; height: fit-content; }
+#txt2img_generate_box>button, #img2img_generate_box>button, #control_generate_box>button, #video_generate_box>button, #txt2img_enqueue, #img2img_enqueue, #txt2img_enqueue>button, #img2img_enqueue>button { min-height: 44px !important; max-height: 44px !important; line-height: 1em; white-space: break-spaces; min-width: unset; }
 #txt2img_enqueue_wrapper, #img2img_enqueue_wrapper, #control_enqueue_wrapper { min-width: unset !important; width: 31%; }
-#txt2img_generate_line2, #img2img_generate_line2, #txt2img_tools, #img2img_tools, #control_generate_line2, #control_tools { display: flex; }
+#txt2img_generate_line2, #img2img_generate_line2, #txt2img_tools, #img2img_tools, #control_generate_line2, #control_tools, #video_generate_line2, #video_tools { display: flex; }
 #txt2img_generate_line2>button, #img2img_generate_line2>button, #extras_generate_box>button, #control_generate_line2>button, #txt2img_tools>button, #img2img_tools>button, #control_tools>button { height: 2em; line-height: 0; font-size: var(--text-md);
   min-width: unset; display: block !important; }
-#txt2img_prompt, #txt2img_neg_prompt, #img2img_prompt, #img2img_neg_prompt, #control_prompt, #control_neg_prompt { display: contents; }
-#txt2img_actions_column, #img2img_actions_column, #control_actions { flex-flow: wrap; justify-content: space-between; }
+#txt2img_prompt, #txt2img_neg_prompt, #img2img_prompt, #img2img_neg_prompt, #control_prompt, #control_neg_prompt, #video_prompt, #video_neg_prompt { display: contents; }
+#txt2img_actions_column, #img2img_actions_column, #control_actions, #video_actions { flex-flow: wrap; justify-content: space-between; }
+#txt2img_seed, #img2img_seed, #control_seed, #video_seed { min-width: 90px !important }
+#video_generate_box>button { max-width: unset; }
+#interrogate_output_prompt>textarea { resize: vertical; }
+#prompt_enhance_apply, #prompt_enhance_model, #prompt_enhance_custom_load { max-width: unset; min-width: 100% !important; }
+#prompt_enhance_system textarea { color: var(--body-text-color-subdued) !important }
 
 .interrogate { position: absolute; right: 2.8em; top: 0.2em; max-width: fit-content; background: none !important; z-index: 50; font-size: 1.5em !important; }
 .interrogate:hover { background: var(--button-primary-background-fill-hover) !important; }
@@ -140,32 +145,35 @@ div#extras_scale_to_tab div.form { flex-direction: row; }
 #txt2img_advanced_options, #img2img_advanced_options, #control_advanced_options { min-width: 100%; }
 #txt2img_advanced_options .gradio-checkbox, #img2img_advanced_options .gradio-checkbox, #control_advanced_options .gradio-checkbox { min-width: unset !important; max-width: fit-content; }
 
+#txt2img_prompt, #txt2img_neg_prompt, #img2img_prompt, #img2img_neg_prompt, #control_prompt, #control_neg_prompt, #video_prompt, #video_neg_prompt { background-color: var(--background-color); box-shadow: none !important; }
+#txt2img_prompt > label > textarea, #txt2img_neg_prompt > label > textarea, #img2img_prompt > label > textarea, #img2img_neg_prompt > label > textarea, #control_prompt > label > textarea, #control_neg_prompt > label > textarea, #video_prompt > label > textarea, #video_neg_prompt > label > textarea { font-size: 1.0em; line-height: 1.4em; }
+#txt2img_styles, #img2img_styles, #control_styles, #video_styles { padding: 0; margin-top: 2px; }
+#txt2img_styles_refresh, #img2img_styles_refresh, #control_styles_refresh, #video_styles_refresh { padding: 0; margin-top: 1em; }
+
 /* settings */
-#si-sparkline-memo, #si-sparkline-load { background-color: #111; }
 #quicksettings { width: fit-content; }
 #quicksettings>button { padding: 0 1em 0 0; align-self: end; margin-bottom: 6px; }
-#settings { display: flex; gap: var(--layout-gap); }
-#settings div { border: none; gap: 0; margin: 0 0 var(--layout-gap) 0px; padding: 0; }
-#settings>div.tab-content { flex: 10 0 75%; display: grid; }
-#settings>div.tab-content>div { border: none; padding: 0; }
+#settings { display: flex; margin-left: 0.5em; }
+#settings>div.tab-content { margin-top: 1em; }
+#settings>div.tab-content>div>div { gap: 0; }
 #settings>div.tab-content>div>div>div>div>div { flex-direction: unset; }
-#settings>div.tab-nav { display: grid; grid-template-columns: repeat(auto-fill, .5em minmax(10em, 1fr)); flex: 1 0 auto; width: 12em; align-self: flex-start; gap: 8px; }
-#settings>div.tab-nav button { display: block; border: none; text-align: left; white-space: initial; padding: 0; }
-#settings>div.tab-nav>#settings_show_all_pages { padding: var(--size-2) var(--size-4); }
+#settings>div.tab-nav { width: 14em; display: block; background: var(--neutral-900); border-radius: var(--block-radius); margin-right: 1em;}
+#settings>div.tab-nav button { width: 100%; height: 2em; text-align: left; border: none; border-radius: var(--block-radius); }
+#settings .dirtyable.hidden { visibility: hidden; }
+#settings .modification-indicator { background: none; border-radius: var(--radius-lg); padding: 0; width: 4px !important; height: 2em !important; position: absolute; float: left; left: -6px; }
+#settings .modification-indicator:disabled { background: none; }
+#settings .modification-indicator.saved { background: var(--color-accent-soft); }
+#settings .modification-indicator.changed { background: var(--color-accent);  }
+#settings .modification-indicator.changed.unsaved { background: var(--color-warning); }
 #settings .block.gradio-checkbox { margin: 0; width: auto; }
-#settings .dirtyable { gap: .5em; }
-#settings .dirtyable.hidden { display: none; }
-#settings .modification-indicator { height: 1.2em; border-radius: 1em !important; padding: 0; width: 0; margin-right: 0.5em; border-left: inset; }
-#settings .modification-indicator:disabled { visibility: hidden; }
-#settings .modification-indicator.saved { background: var(--color-accent-soft); width: var(--spacing-sm); }
-#settings .modification-indicator.changed { background: var(--color-accent); width: var(--spacing-sm); }
-#settings .modification-indicator.changed.unsaved { background-image: linear-gradient(var(--color-accent) 25%, var(--color-accent-soft) 75%); width: var(--spacing-sm); }
-#settings_result { margin: 0 1.2em; }
-#tab_settings .gradio-slider, #tab_settings .gradio-dropdown { width: 300px !important; max-width: 300px; }
-#tab_settings textarea { max-width: 500px; }
+#settings .block.gradio-number { min-width: 500px !important; }
+#settings .gradio-slider, #tab_settings .gradio-dropdown { width: 500px !important; max-width: 500px !important; }
+#settings .gradio-radio { padding: var(--block-padding) !important; }
+#settings textarea { width: 500px !important; max-width: 500px !important; }
 .licenses { display: block !important; }
 
 /* live preview */
+#si-sparkline-memo, #si-sparkline-load { background-color: #111; }
 .progressDiv { position: relative; height: 20px; background: #b4c0cc; margin-bottom: -3px; }
 .dark .progressDiv { background: #424c5b; }
 .progressDiv .progress { width: 0%; height: 20px; background: #0060df; color: white; font-weight: bold; line-height: 20px; padding: 0 8px 0 0; text-align: right; overflow: visible; white-space: nowrap; padding: 0 0.5em; }
@@ -388,13 +396,14 @@ div:has(>#tab-gallery-folders) { flex-grow: 0 !important; background-color: var(
       #txt2img_results, #extras_results, #txt2im g_footer p { text-wrap: wrap; max-width: 100% !important; } /* maintain side by side split on larger mobile displays for from text */
     }
     #scripts_alwayson_txt2img div, #scripts_alwayson_img2img div { max-width: 100%; }
-    #txt2img_prompt_container, #img2img_prompt_container, #control_prompt_container { resize: vertical !important; }
+    #txt2img_prompt_container, #img2img_prompt_container, #control_prompt_container, #video_prompt_container { resize: vertical !important; }
     #txt2img_generate_box, #txt2img_enqueue_wrapper { min-width: 100% !important;} /* make generate and enqueue buttons take up the entire width of their rows. */
     #img2img_toprow>div.gradio-column { flex-grow: 1 !important;} /*make interrogate buttons take up appropriate space. */
     #img2img_actions_column { display: flex; min-width: fit-content !important; flex-direction: row;justify-content: space-evenly; align-items: center;}
     #txt2img_generate_box, #img2img_generate_box, #txt2img_enqueue_wrapper,#img2img_enqueue_wrapper {display: flex;flex-direction: column;height: 4em !important;align-items: stretch;justify-content: space-evenly;}
     #img2img_interface, #img2img_results, #img2img_footer p { text-wrap: wrap; min-width: 100% !important; max-width: 100% !important;} /* maintain single column for from image operations on larger mobile devices */
-    #txt2img_sampler, #txt2img_batch, #txt2img_seed_group, #txt2img_advanced, #txt2img_second_pass, #img2img_sampling_group, #img2img_resize_group, #img2img_batch_group, #img2img_seed_group, #img2img_denoise_group, #img2img_advanced_group { width: 100% !important; } /* fix from text/image UI elements to prevent them from moving around within the UI */
+    #txt2img_sampler, #txt2img_batch, #txt2img_seed_group, #txt2img_advanced, #txt2img_second_pass, #img2img_sampling_group, #img2img_resize_group, #img2img_batch_group, #img2img_seed_group, #img2img_denoise_group, #img2img_advanced_group { width: 100% !important; } /* fix from text/image UI 
+    elements to prevent them from moving around within the UI */
     #img2img_resize_group .gradio-radio>div { display: flex; flex-direction: column; width: unset !important; }
     #inpaint_controls div { display:flex;flex-direction: row;}
     #inpaint_controls .gradio-radio>div { display: flex; flex-direction: column !important; }
diff --git a/javascript/settings.js b/javascript/settings.js
index 1a891c844..99902fb02 100644
--- a/javascript/settings.js
+++ b/javascript/settings.js
@@ -26,7 +26,7 @@ async function updateOpts(json_string) {
     const key = Object.keys(op)[0];
     const callback = op[key];
     if (opts[key] && opts[key] !== settings_data.values[key]) {
-      log('updateOpts', key, opts[key], settings_data.values[key]);
+      log('updateOpt', key, opts[key], settings_data.values[key]);
       if (callback) callback(new_opts[key], opts[key]);
     }
   }
@@ -37,7 +37,8 @@ async function updateOpts(json_string) {
     if (callback) callback(new_opts[key], opts[key]);
   }
 
-  opts = new_opts;
+  window.opts = new_opts;
+  log('updateOpts', Object.keys(new_opts).length);
   Object.entries(opts_metadata).forEach(([opt, meta]) => {
     if (!opts_tabs[meta.tab_name]) opts_tabs[meta.tab_name] = {};
     if (!opts_tabs[meta.tab_name].unsaved_keys) opts_tabs[meta.tab_name].unsaved_keys = new Set();
@@ -49,8 +50,8 @@ async function updateOpts(json_string) {
 
 function showAllSettings() {
   // Try to ensure that the show all settings tab is opened by clicking on its tab button
-  const tab_dirty_indicator = gradioApp().getElementById('modification_indicator_show_all_pages');
-  if (tab_dirty_indicator && tab_dirty_indicator.nextSibling) tab_dirty_indicator.nextSibling.click();
+  // const tab_dirty_indicator = gradioApp().getElementById('modification_indicator_show_all_pages');
+  // if (tab_dirty_indicator && tab_dirty_indicator.nextSibling) tab_dirty_indicator.nextSibling.click();
   getSettingsTabs().forEach((elem) => {
     if (elem.id === 'settings_tab_licenses' || elem.id === 'settings_show_all_pages') return;
     elem.style.display = 'block';
@@ -192,9 +193,12 @@ async function initSettings() {
   tabContentWrapper.className = 'tab-content';
   tabNavElements.parentElement.insertBefore(tabContentWrapper, tabNavElements.nextSibling);
   tabElements.forEach((elem, index) => {
-    const tabName = elem.id.replace('settings_', '');
+    const tabName = elem.id.replace('settings_section_tab_', '');
     const indicator = gradioApp().getElementById(`modification_indicator_${tabName}`);
-    tabNavElements.insertBefore(indicator, tabNavButtons[index]);
+    if (indicator) {
+      tabNavElements.insertBefore(document.createElement('br'), tabNavButtons[index]);
+      tabNavElements.insertBefore(indicator, tabNavButtons[index]);
+    }
     tabContentWrapper.appendChild(elem);
     observer.observe(elem, { attributes: true, attributeFilter: ['style'] });
   });
diff --git a/javascript/startup.js b/javascript/startup.js
index d8407391d..328c167b6 100644
--- a/javascript/startup.js
+++ b/javascript/startup.js
@@ -3,6 +3,7 @@ window.api = '/sdapi/v1';
 window.subpath = '';
 
 async function initStartup() {
+  const t0 = performance.now();
   log('initStartup');
   if (window.setupLogger) await setupLogger();
 
@@ -24,7 +25,11 @@ async function initStartup() {
   await reconnectUI();
 
   // make sure all of the ui is ready and options are loaded
-  while (Object.keys(window.opts).length === 0) await sleep(50);
+  let t1 = performance.now();
+  while ((Object.keys(window.opts).length === 0) && (t1 - t0 < 10000)) {
+    t1 = performance.now();
+    await sleep(50);
+  }
   log('mountURL', window.opts.subpath);
   if (window.opts.subpath?.length > 0) {
     window.subpath = window.opts.subpath;
@@ -43,6 +48,8 @@ async function initStartup() {
   setHints();
   applyStyles();
   initIndexDB();
+  t1 = performance.now();
+  log('initStartup', Math.round(1000 * (t1 - t0) / 1000000));
 }
 
 onUiLoaded(initStartup);
diff --git a/javascript/timeless-beige.css b/javascript/timeless-beige.css
index 4b0f7d9e4..a8a9c1536 100644
--- a/javascript/timeless-beige.css
+++ b/javascript/timeless-beige.css
@@ -101,8 +101,6 @@ button.selected {background: var(--button-primary-background-fill);}
 /* gradio elements overrides */
 #div.gradio-container { overflow-x: hidden; }
 #img2img_label_copy_to_img2img { font-weight: normal; }
-#txt2img_prompt, #txt2img_neg_prompt, #img2img_prompt, #img2img_neg_prompt { background-color: var(--background-color); box-shadow: 4px 4px 4px 0px #333333  !important; }
-#txt2img_prompt > label > textarea, #txt2img_neg_prompt > label > textarea, #img2img_prompt > label > textarea, #img2img_neg_prompt > label > textarea { font-size: 1.1rem; }
 #img2img_settings { min-width: calc(2 * var(--left-column)); max-width: calc(2 * var(--left-column)); background-color: #111111; padding-top: 16px; }
 #interrogate, #deepbooru { margin: 0 0px 10px 0px; max-width: 80px; max-height: 80px; font-weight: normal; font-size: 0.95em; }
 #quicksettings .gr-button-tool { font-size: 1.6rem; box-shadow: none; margin-top: -2px; height: 2.4em; }
diff --git a/javascript/ui.js b/javascript/ui.js
index 9df1a26e4..9a5f2b794 100644
--- a/javascript/ui.js
+++ b/javascript/ui.js
@@ -155,6 +155,11 @@ function switch_to_control(...args) {
   return Array.from(arguments);
 }
 
+function switch_to_video(...args) {
+  switchToTab('Video');
+  return Array.from(arguments);
+}
+
 function switch_to_caption(...args) {
   switchToTab('Caption');
   return Array.from(arguments);
@@ -240,6 +245,18 @@ function submit_control(...args) {
   return res;
 }
 
+function submit_video(...args) {
+  log('submitVideo');
+  clearGallery('video');
+  const id = randomId();
+  requestProgress(id, null, gradioApp().getElementById('video_gallery'));
+  const res = create_submit_args(args);
+  res[0] = id;
+  res[1] = window.submit_state;
+  window.submit_state = '';
+  return res;
+}
+
 function submit_postprocessing(...args) {
   log('SubmitExtras');
   clearGallery('extras');
diff --git a/launch.py b/launch.py
index c80840036..d00b9ef22 100755
--- a/launch.py
+++ b/launch.py
@@ -150,10 +150,14 @@ def run_extension_installer(ext_dir): # compatbility function
     installer.run_extension_installer(ext_dir)
 
 
-def get_memory_stats():
-    from modules.memstats import ram_stats
-    res = ram_stats()
-    return f'{res["used"]}/{res["total"]}'
+def get_memory_stats(detailed:bool=False):
+    from modules.memstats import ram_stats, memory_stats
+    if not detailed:
+        res = ram_stats()
+        return f'{res["used"]}/{res["total"]}'
+    else:
+        res = memory_stats()
+        return res
 
 
 def start_server(immediate=True, server=None):
@@ -260,6 +264,8 @@ def main():
     get_custom_args()
 
     uv, instance = start_server(immediate=True, server=None)
+    t_server = time.time()
+    t_monitor = time.time()
     while True:
         try:
             alive = uv.thread.is_alive()
@@ -267,8 +273,13 @@ def main():
         except Exception:
             alive = False
             requests = 0
-        if round(time.time()) % 120 == 0:
-            installer.log.debug(f'Server: alive={alive} requests={requests} memory={get_memory_stats()} {instance.state.status()}')
+        t_current = time.time()
+        if t_current - t_server > 120:
+            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):
+            installer.log.trace(f'Monitor: {get_memory_stats(detailed=True)}')
+            t_monitor = t_current
         if not alive:
             if uv is not None and uv.wants_restart:
                 installer.log.info('Server restarting...')
diff --git a/models/Reference/Alpha-VLLM--Lumina-Image-2.0.jpg b/models/Reference/Alpha-VLLM--Lumina-Image-2.0.jpg
old mode 100755
new mode 100644
diff --git a/models/Reference/THUDM--CogView4-6B.jpg b/models/Reference/THUDM--CogView4-6B.jpg
new file mode 100644
index 000000000..5876b935b
Binary files /dev/null and b/models/Reference/THUDM--CogView4-6B.jpg differ
diff --git a/modules/api/endpoints.py b/modules/api/endpoints.py
index 7bdc31c6a..87362b4c9 100644
--- a/modules/api/endpoints.py
+++ b/modules/api/endpoints.py
@@ -41,10 +41,10 @@ def get_embeddings():
     return {"loaded": convert_embeddings(db.word_embeddings), "skipped": convert_embeddings(db.skipped_embeddings)}
 
 def get_loras():
-    from modules.lora import network, networks
+    from modules.lora import network, lora_load
     def create_lora_json(obj: network.NetworkOnDisk):
         return { "name": obj.name, "alias": obj.alias, "path": obj.filename, "metadata": obj.metadata }
-    return [create_lora_json(obj) for obj in networks.available_networks.values()]
+    return [create_lora_json(obj) for obj in lora_load.available_networks.values()]
 
 def get_extra_networks(page: Optional[str] = None, name: Optional[str] = None, filename: Optional[str] = None, title: Optional[str] = None, fullname: Optional[str] = None, hash: Optional[str] = None): # pylint: disable=redefined-builtin
     res = []
@@ -113,7 +113,7 @@ def post_vqa(req: models.ReqVQA):
     image = helpers.decode_base64_to_image(req.image)
     image = image.convert('RGB')
     from modules.interrogate import vqa
-    answer = vqa.interrogate(req.question, '', image, req.model)
+    answer = vqa.interrogate(req.question, req.system, '', image, req.model)
     return models.ResVQA(answer=answer)
 
 def post_unload_checkpoint():
@@ -134,8 +134,8 @@ def post_refresh_vae():
     return shared.refresh_vaes()
 
 def post_refresh_loras():
-    from modules.lora import networks
-    return networks.list_available_networks()
+    from modules.lora import lora_load
+    return lora_load.list_available_networks()
 
 def get_extensions_list():
     from modules import extensions
diff --git a/modules/api/gallery.py b/modules/api/gallery.py
index e56bb00ac..e1add81af 100644
--- a/modules/api/gallery.py
+++ b/modules/api/gallery.py
@@ -74,7 +74,7 @@ def register_api(app: FastAPI): # register api
     manager = ConnectionManager()
 
     def get_video_thumbnail(filepath):
-        from modules.ui_control_helpers import get_video_params
+        from modules.video import get_video_params
         try:
             stat = os.stat(filepath)
             frames, fps, duration, width, height, codec, frame = get_video_params(filepath, capture=True)
diff --git a/modules/api/server.py b/modules/api/server.py
index 828d3fd95..28c429e29 100644
--- a/modules/api/server.py
+++ b/modules/api/server.py
@@ -91,10 +91,10 @@ def get_progress(req: models.ReqProgress = Depends()):
     step_y = max(shared.state.sampling_steps, 1)
     current = step_y * batch_x + step_x
     total = step_y * batch_y
-    progress = current / total if current > 0 and total > 0 else 0
+    progress = min((current / total) if current > 0 and total > 0 else 0, 1)
     time_since_start = time.time() - shared.state.time_start
     eta_relative = (time_since_start / progress) - time_since_start if progress > 0 else 0
-    res = models.ResProgress(progress=progress, eta_relative=eta_relative, state=shared.state.dict(), current_image=current_image, textinfo=shared.state.textinfo)
+    res = models.ResProgress(progress=round(progress, 2), eta_relative=round(eta_relative, 2), current_image=current_image, textinfo=shared.state.textinfo, state=shared.state.dict(), )
     return res
 
 def get_status():
diff --git a/modules/call_queue.py b/modules/call_queue.py
index 8368e5759..b33a9b6b2 100644
--- a/modules/call_queue.py
+++ b/modules/call_queue.py
@@ -81,8 +81,12 @@ def wrap_gradio_call(func, extra_outputs=None, add_stats=False, name=None):
             ooms = mem_mon_read.pop("oom")
             retries = mem_mon_read.pop("retries")
             vram = {k: v//1048576 for k, v in mem_mon_read.items()}
-            peak = max(vram['active_peak'], vram['reserved_peak'], vram['used'])
-            used = round(100.0 * peak / vram['total']) if vram['total'] > 0 else 0
+            if 'active_peak' in vram:
+                peak = max(vram['active_peak'], vram['reserved_peak'], vram['used'])
+                used = round(100.0 * peak / vram['total']) if vram['total'] > 0 else 0
+            else:
+                peak = 0
+                used = 0
             if peak > 0:
                 gpu += f"| GPU {peak} MB"
                 gpu += f" {used}%" if used > 0 else ''
diff --git a/modules/cmd_args.py b/modules/cmd_args.py
index 5e5e21054..a8dec6748 100644
--- a/modules/cmd_args.py
+++ b/modules/cmd_args.py
@@ -37,6 +37,7 @@ 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_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/devices.py b/modules/devices.py
index caa544652..c35c8b909 100644
--- a/modules/devices.py
+++ b/modules/devices.py
@@ -52,7 +52,8 @@ def has_zluda() -> bool:
         return False
     try:
         dev = torch.device("cuda")
-        return torch.cuda.get_device_name(dev).endswith("[ZLUDA]")
+        cc = torch.cuda.get_device_capability(dev)
+        return cc == (8, 8)
     except Exception:
         return False
 
@@ -319,9 +320,18 @@ def test_fp16():
     global fp16_ok # pylint: disable=global-statement
     if fp16_ok is not None:
         return fp16_ok
-    if sys.platform == "darwin" or backend == 'openvino': # override
-        fp16_ok = False
-        return fp16_ok
+    if opts.cuda_dtype != 'FP16': # don't override if the user sets it
+        if sys.platform == "darwin" or backend == 'openvino': # override
+            fp16_ok = False
+            return fp16_ok
+        elif backend == 'rocm':
+            # gfx1102 (RX 7600, 7500, 7650 and 7700S) causes segfaults with fp16
+            # agent can be overriden to gfx1100 to get gfx1102 working with ROCm so check the gpu name as well
+            agent = getattr(torch.cuda.get_device_properties(device), "gcnArchName", "gfx0000")
+            agent_name = getattr(torch.cuda.get_device_properties(device), "name", "AMD Radeon RX 0000")
+            if agent == "gfx1102" or (agent == "gfx1100" and any(i in agent_name for i in ("7600", "7500", "7650", "7700S"))):
+                fp16_ok = False
+                return fp16_ok
     try:
         x = torch.tensor([[1.5,.0,.0,.0]]).to(device=device, dtype=torch.float16)
         layerNorm = torch.nn.LayerNorm(4, eps=0.00001, elementwise_affine=True, dtype=torch.float16, device=device)
@@ -377,8 +387,6 @@ def set_cudnn_params():
         torch.backends.cuda.matmul.allow_tf32 = True
         torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = True
         torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = True
-        if hasattr(torch.backends.cuda, "allow_fp16_bf16_reduction_math_sdp"): # only valid for torch >= 2.5
-            torch.backends.cuda.allow_fp16_bf16_reduction_math_sdp(True)
     except Exception as e:
         log.warning(f'Torch matmul: {e}')
     if torch.backends.cudnn.is_available():
@@ -401,6 +409,7 @@ def override_ipex_math():
         try:
             if hasattr(torch.xpu, "set_fp32_math_mode"): # not available with pure torch+xpu, requires ipex
                 torch.xpu.set_fp32_math_mode(mode=torch.xpu.FP32MathMode.TF32)
+            torch.backends.mkldnn.allow_tf32 = True
         except Exception as e:
             log.warning(f'Torch ipex: {e}')
 
@@ -422,6 +431,8 @@ def set_sdpa_params():
             torch.backends.cuda.enable_flash_sdp('Flash attention' in opts.sdp_options)
             torch.backends.cuda.enable_mem_efficient_sdp('Memory attention' in opts.sdp_options)
             torch.backends.cuda.enable_math_sdp('Math attention' in opts.sdp_options)
+            if hasattr(torch.backends.cuda, "allow_fp16_bf16_reduction_math_sdp"): # only valid for torch >= 2.5
+                torch.backends.cuda.allow_fp16_bf16_reduction_math_sdp(True)
             log.debug(f'Torch attention: type="sdpa" flash={"Flash attention" in opts.sdp_options} memory={"Memory attention" in opts.sdp_options} math={"Math attention" in opts.sdp_options}')
         except Exception as err:
             log.warning(f'Torch attention: type="sdpa" {err}')
@@ -452,7 +463,21 @@ def set_sdpa_params():
                 @wraps(sdpa_pre_flash_atten)
                 def sdpa_flash_atten(query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False, scale=None):
                     if query.shape[-1] <= 128 and attn_mask is None and query.dtype != torch.float32:
-                        return flash_attn_func(q=query.transpose(1, 2), k=key.transpose(1, 2), v=value.transpose(1, 2), dropout_p=dropout_p, causal=is_causal, softmax_scale=scale).transpose(1, 2)
+                        is_unsqueezed = False
+                        if query.dim() == 3:
+                            query = query.unsqueeze(0)
+                            is_unsqueezed = True
+                            if key.dim() == 3:
+                                key = key.unsqueeze(0)
+                            if value.dim() == 3:
+                                value = value.unsqueeze(0)
+                        query = query.transpose(1, 2)
+                        key = key.transpose(1, 2)
+                        value = value.transpose(1, 2)
+                        attn_output = flash_attn_func(q=query, k=key, v=value, dropout_p=dropout_p, causal=is_causal, softmax_scale=scale).transpose(1, 2)
+                        if is_unsqueezed:
+                            attn_output = attn_output.squeeze(0)
+                        return attn_output
                     else:
                         return sdpa_pre_flash_atten(query=query, key=key, value=value, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal, scale=scale)
                 torch.nn.functional.scaled_dot_product_attention = sdpa_flash_atten
diff --git a/modules/extra_networks.py b/modules/extra_networks.py
index 420f3beda..d8b638b85 100644
--- a/modules/extra_networks.py
+++ b/modules/extra_networks.py
@@ -19,8 +19,9 @@ def register_default_extra_networks():
     from modules.ui_extra_networks_styles import ExtraNetworkStyles
     register_extra_network(ExtraNetworkStyles())
     if not shared.opts.lora_legacy:
-        from modules.lora.networks import extra_network_lora
-        register_extra_network(extra_network_lora)
+        from modules.lora import lora_common, extra_networks_lora
+        lora_common.extra_network_lora = extra_networks_lora.ExtraNetworkLora()
+        register_extra_network(lora_common.extra_network_lora)
     if shared.opts.hypernetwork_enabled:
         from modules.ui_extra_networks_hypernet import ExtraNetworkHypernet
         register_extra_network(ExtraNetworkHypernet())
@@ -87,7 +88,7 @@ def activate(p, extra_network_data=None, step=0, include=[], exclude=[]):
         stepwise = stepwise or is_stepwise(extra_network_args)
     functional = shared.opts.lora_functional
     if shared.opts.lora_force_diffusers and stepwise:
-        shared.log.warning("Load network: type=LoRA method=composable loader=diffusers not compatible")
+        shared.log.warning("Network load: type=LoRA method=composable loader=diffusers not compatible")
         stepwise = False
     shared.opts.data['lora_functional'] = stepwise or functional
 
diff --git a/modules/face/faceswap.py b/modules/face/faceswap.py
index d7d0a32a5..df3765fb2 100644
--- a/modules/face/faceswap.py
+++ b/modules/face/faceswap.py
@@ -13,13 +13,19 @@ swapper = None
 
 
 def face_swap(p: processing.StableDiffusionProcessing, app, input_images: List[Image.Image], source_image: Image.Image, cache: bool):
-    import insightface.model_zoo
     global swapper # pylint: disable=global-statement
     if swapper is None:
-        model_path = hf.hf_hub_download(repo_id='ezioruan/inswapper_128.onnx', filename='inswapper_128.onnx', cache_dir=shared.opts.hfcache_dir)
+        import insightface.model_zoo
+        repo_id = 'ezioruan/inswapper_128.onnx'
+        model_path = hf.hf_hub_download(repo_id=repo_id, filename='inswapper_128.onnx', cache_dir=shared.opts.hfcache_dir)
+        shared.log.debug(f'FaceSwap load: repo="{repo_id}" path="{model_path}"')
         # model_path = hf.hf_hub_download(repo_id='somanchiu/reswapper', filename='reswapper_256-1567500_originalInswapperClassCompatible.onnx', cache_dir=shared.opts.hfcache_dir)
-        router: insightface.model_zoo.model_zoo.INSwapper = insightface.model_zoo.model_zoo.ModelRouter(model_path)
-        swapper = router.get_model()
+        try:
+            router: insightface.model_zoo.model_zoo.INSwapper = insightface.model_zoo.model_zoo.ModelRouter(model_path)
+            swapper = router.get_model()
+        except Exception as e:
+            shared.log.error(f'FaceSwap load: {e}')
+            return None
 
     np_image = cv2.cvtColor(np.array(source_image), cv2.COLOR_RGB2BGR)
     faces = app.get(np_image)
diff --git a/modules/face/insightface.py b/modules/face/insightface.py
index 529e4be32..7c9ec5af9 100644
--- a/modules/face/insightface.py
+++ b/modules/face/insightface.py
@@ -20,13 +20,14 @@ def get_app(mp_name, threshold=0.5, resolution=640):
         install('git+https://github.com/tencent-ailab/IP-Adapter.git', 'ip_adapter', ignore=False)
 
     if insightface_app is None or mp_name != instightface_mp:
+        import insightface
         from insightface.model_zoo import model_zoo
         from insightface.app import face_analysis
         model_zoo.print = lambda *args, **kwargs: None
         face_analysis.print = lambda *args, **kwargs: None
         import huggingface_hub as hf
         import zipfile
-        log.debug(f"InsightFace: mp={mp_name} provider={devices.onnx}")
+        log.debug(f"InsightFace: version={insightface.__version__} mp={mp_name} provider={devices.onnx}")
         root_dir = os.path.join(opts.diffusers_dir, 'models--vladmandic--insightface-faceanalysis')
         local_dir = os.path.join(root_dir, 'models')
         extract_dir = os.path.join(local_dir, mp_name)
diff --git a/modules/flash_attn_triton_amd/__init__.py b/modules/flash_attn_triton_amd/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/modules/flash_attn_triton_amd/bwd_prefill.py b/modules/flash_attn_triton_amd/bwd_prefill.py
new file mode 100644
index 000000000..7f5be379b
--- /dev/null
+++ b/modules/flash_attn_triton_amd/bwd_prefill.py
@@ -0,0 +1,606 @@
+import torch
+import triton
+import triton.language as tl
+from modules.flash_attn_triton_amd.utils import get_shape_from_layout, get_strides_from_layout
+
+
+@triton.jit
+def _bwd_preprocess_use_o(
+    Out,
+    DO,
+    Delta,
+    stride_oz, stride_oh, stride_om, stride_ok,
+    stride_doz, stride_doh, stride_dom, stride_dok, # pylint: disable=unused-argument
+    stride_deltaz, stride_deltah, stride_deltam,
+    cu_seqlens_q,
+    cu_seqlens_k,
+    max_seqlen_q,
+    max_seqlen_k,
+    BLOCK_M: tl.constexpr,
+    BLOCK_DMODEL: tl.constexpr,
+    ACTUAL_BLOCK_DMODEL: tl.constexpr,
+    N_CTX_Q: tl.constexpr,
+    Z: tl.constexpr, # pylint: disable=unused-argument
+    H: tl.constexpr,
+    IS_VARLEN: tl.constexpr
+):
+    pid_m = tl.program_id(0)
+    pid_bh = tl.program_id(1)
+
+    # Compute batch and head indices
+    off_z = pid_bh // H
+    off_h = pid_bh % H
+
+    if IS_VARLEN:
+        # Compute sequence lengths for the current batch
+        q_start = tl.load(cu_seqlens_q + off_z)
+        q_end = tl.load(cu_seqlens_q + off_z + 1)
+        k_start = tl.load(cu_seqlens_k + off_z)
+        k_end = tl.load(cu_seqlens_k + off_z + 1)
+
+        # Compute actual sequence lengths
+        N_CTX_Q = q_end - q_start
+        N_CTX_K = k_end - k_start # pylint: disable=unused-variable
+    else:
+        q_start = 0
+        k_start = 0
+        N_CTX_Q = max_seqlen_q
+        N_CTX_K = max_seqlen_k # pylint: disable=unused-variable
+
+    off_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    off_d = tl.arange(0, BLOCK_DMODEL)
+
+    # create masks
+    mask_m = off_m < N_CTX_Q
+    mask_d = off_d < ACTUAL_BLOCK_DMODEL
+
+    # compute offsets
+    o_offset = Out + off_z * stride_oz + off_h * stride_oh + q_start * stride_om
+    do_offset = DO + off_z * stride_oz + off_h * stride_oh + q_start * stride_om
+
+    # compute pointers
+    out_ptrs = o_offset + off_m[:, None] * stride_om + off_d[None, :] * stride_ok
+    do_ptrs = do_offset + off_m[:, None] * stride_dom + off_d[None, :] * stride_dok
+
+    # load
+    o = tl.load(out_ptrs, mask=mask_m[:, None] & mask_d[None, :], other=0.0).to(tl.float32)
+    do = tl.load(do_ptrs, mask=mask_m[:, None] & mask_d[None, :], other=0.0).to(tl.float32)
+
+    # compute delta
+    delta = tl.sum(o * do, axis=1)
+
+    # write-back delta
+    delta_offset = Delta + off_z * stride_deltaz + off_h * stride_deltah + q_start * stride_deltam
+    delta_ptrs = delta_offset + off_m * stride_deltam
+    tl.store(delta_ptrs, delta, mask=mask_m)
+
+
+@triton.jit
+def _bwd_kernel_one_col_block(
+    Q,
+    K,
+    V,
+    sm_scale,
+    Out, DO, DQ, DK, DV, L, D, # pylint: disable=unused-argument
+    q_offset,
+    k_offset,
+    v_offset,
+    do_offset,
+    dq_offset,
+    dk_offset,
+    dv_offset,
+    d_offset,
+    l_offset,
+    stride_dq_all, stride_qz, stride_qh, # pylint: disable=unused-argument
+    stride_qm,
+    stride_qk,
+    stride_kz, stride_kh, # pylint: disable=unused-argument
+    stride_kn,
+    stride_kk,
+    stride_vz, stride_vh, # pylint: disable=unused-argument
+    stride_vn,
+    stride_vk,
+    stride_deltaz,  stride_deltah, # pylint: disable=unused-argument
+    stride_deltam,
+    Z, H, # pylint: disable=unused-argument
+    N_CTX_Q,
+    N_CTX_K,
+    off_h, off_z, off_hz, # pylint: disable=unused-argument
+    start_n,
+    num_block_m,
+    num_block_n, # pylint: disable=unused-argument
+    BLOCK_M: tl.constexpr,
+    BLOCK_DMODEL: tl.constexpr,
+    ACTUAL_BLOCK_DMODEL: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    SEQUENCE_PARALLEL: tl.constexpr,
+    CAUSAL: tl.constexpr,
+    USE_EXP2: tl.constexpr,
+):
+    if CAUSAL:
+        # TODO: Causal can skip more blocks with something like lo = start_m * BLOCK_M
+        lo = 0
+    else:
+        lo = 0
+
+    # initialize col and head offsets
+    offs_n = start_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    offs_d = tl.arange(0, BLOCK_DMODEL)
+
+    # masks
+    mask_n = offs_n < N_CTX_K
+    mask_d = offs_d < ACTUAL_BLOCK_DMODEL
+    kv_mask = mask_n[:, None] & mask_d[None, :]
+
+    # initialize grad accumulators
+    dv = tl.zeros([BLOCK_N, BLOCK_DMODEL], dtype=tl.float32)
+    dk = tl.zeros([BLOCK_N, BLOCK_DMODEL], dtype=tl.float32)
+
+    # load k and v once per column block
+    k_ptrs = k_offset + offs_n[:, None] * stride_kn + offs_d[None, :] * stride_kk
+    v_ptrs = v_offset + offs_n[:, None] * stride_vn + offs_d[None, :] * stride_vk
+    k = tl.load(k_ptrs, mask=kv_mask, other=0.0)
+    v = tl.load(v_ptrs, mask=kv_mask, other=0.0)
+
+    # loop over rows
+    for start_m in range(lo, num_block_m * BLOCK_M, BLOCK_M):
+        offs_m = start_m + tl.arange(0, BLOCK_M)
+        q_ptrs = q_offset + offs_m[:, None] * stride_qm + offs_d[None, :] * stride_qk
+        dq_ptrs = dq_offset + offs_m[:, None] * stride_qm + offs_d[None, :] * stride_qk
+        do_ptrs = do_offset + offs_m[:, None] * stride_qm + offs_d[None, :] * stride_qk
+
+        # update mask as row block changes
+        mask_m = offs_m < N_CTX_Q
+        q_mask = mask_m[:, None] & mask_d[None, :]
+
+        # load q, k, v, do on-chip
+        q = tl.load(q_ptrs, mask=q_mask, other=0.0)
+        do = tl.load(do_ptrs, mask=q_mask, other=0.0)
+
+        # recompute p = softmax(qk, dim=-1).T
+        qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
+        qk += tl.dot(q, tl.trans(k))
+
+        if CAUSAL:
+            col_offset = N_CTX_Q - N_CTX_K
+            causal_mask = offs_m[:, None] >= (col_offset + offs_n[None, :])
+            qk = tl.where(causal_mask, qk, float("-inf"))
+
+        l_ptrs = l_offset + offs_m * stride_deltam
+        l_i = tl.load(l_ptrs, mask=mask_m)
+
+        # compute p
+        if USE_EXP2:
+            RCP_LN2: tl.constexpr = 1.4426950408889634
+            qk *= sm_scale * RCP_LN2
+            l_i *= RCP_LN2
+            p = tl.math.exp2(qk - l_i[:, None])
+        else:
+            qk *= sm_scale
+            p = tl.math.exp(qk - l_i[:, None])
+
+        # mask block in the cases where the data is smaller the block size
+        p_mask = mask_m[:, None] & mask_n[None, :]
+        p = tl.where(p_mask, p, 0.0)
+
+        # compute dv
+        dv += tl.dot(tl.trans(p.to(Q.dtype.element_ty)), do)
+
+        # compute dp
+        dp = tl.dot(do, tl.trans(v))
+
+        # compute ds , ds = p * (dp - delta[:, None])
+        d_ptrs = d_offset + offs_m * stride_deltam
+        Di = tl.load(d_ptrs, mask=mask_m)
+        ds = (p * (dp - Di[:, None])) * sm_scale
+        ds = tl.where(p_mask, ds, 0.0).to(Q.dtype.element_ty)
+
+        # compute dk = dot(ds.T, q)
+        dk += tl.dot(tl.trans(ds), q)
+
+        # compute dq
+        if SEQUENCE_PARALLEL:
+            dq = tl.dot(ds, k)
+        else:
+            dq = tl.load(dq_ptrs, mask=q_mask, other=0.0)
+            dq += tl.dot(ds, k)
+        tl.store(dq_ptrs, dq.to(Q.dtype.element_ty), mask=q_mask)
+
+    # write-back dv and dk
+    dk_ptrs = dk_offset + offs_n[:, None] * stride_kn + offs_d[None, :] * stride_kk
+    dv_ptrs = dv_offset + offs_n[:, None] * stride_vn + offs_d[None, :] * stride_vk
+
+    # write-back
+    tl.store(dk_ptrs, dk.to(K.dtype.element_ty), mask=kv_mask)
+    tl.store(dv_ptrs, dv.to(V.dtype.element_ty), mask=kv_mask)
+
+@triton.jit
+def _bwd_kernel(
+    Q,
+    K,
+    V,
+    sm_scale,
+    Out,
+    DO,
+    DQ,
+    DK,
+    DV,
+    L,
+    D,
+    stride_dq_all,
+    stride_qz,
+    stride_qh,
+    stride_qm,
+    stride_qk,
+    stride_kz,
+    stride_kh,
+    stride_kn,
+    stride_kk,
+    stride_vz,
+    stride_vh,
+    stride_vn,
+    stride_vk,
+    stride_deltaz,
+    stride_deltah,
+    stride_deltam,
+    Z,
+    H,
+    num_block_m,
+    num_block_n,
+    cu_seqlens_q,
+    cu_seqlens_k,
+    max_seqlen_q,
+    max_seqlen_k,
+    BLOCK_M: tl.constexpr,
+    BLOCK_DMODEL: tl.constexpr,
+    ACTUAL_BLOCK_DMODEL: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    SEQUENCE_PARALLEL: tl.constexpr,
+    CAUSAL: tl.constexpr,
+    USE_EXP2: tl.constexpr,
+    IS_VARLEN: tl.constexpr,
+):
+    # program ids
+    off_hz = tl.program_id(0)
+    if SEQUENCE_PARALLEL:
+        start_n = tl.program_id(1)
+    off_z = off_hz // H
+    off_h = off_hz % H
+
+    if IS_VARLEN:
+        # Compute sequence lengths for the current batch
+        q_start = tl.load(cu_seqlens_q + off_z)
+        q_end = tl.load(cu_seqlens_q + off_z + 1)
+        k_start = tl.load(cu_seqlens_k + off_z)
+        k_end = tl.load(cu_seqlens_k + off_z + 1)
+
+        # Compute actual sequence lengths
+        N_CTX_Q = q_end - q_start
+        N_CTX_K = k_end - k_start
+    else:
+        q_start = 0
+        k_start = 0
+        N_CTX_Q = max_seqlen_q
+        N_CTX_K = max_seqlen_k
+
+    # input tensor offsets
+    q_offset = Q + off_z * stride_qz + off_h * stride_qh + q_start * stride_qm
+    k_offset = K + off_z * stride_kz + off_h * stride_kh + k_start * stride_kn
+    v_offset = V + off_z * stride_vz + off_h * stride_vh + k_start * stride_vn
+    do_offset = DO + off_z * stride_qz + off_h * stride_qh + q_start * stride_qm
+    l_offset = L + off_z * stride_deltaz + off_h * stride_deltah + q_start * stride_deltam
+    d_offset = D + off_z * stride_deltaz + off_h * stride_deltah + q_start * stride_deltam
+
+    # output tensor offsets
+    dk_offset = DK + off_z * stride_kz + off_h * stride_kh + k_start * stride_kn
+    dv_offset = DV + off_z * stride_vz + off_h * stride_vh + k_start * stride_vn
+    if SEQUENCE_PARALLEL:
+        dq_offset = DQ + start_n * stride_dq_all + off_z * stride_qz + off_h * stride_qh + q_start * stride_qm
+    else:
+        dq_offset = DQ + off_z * stride_qz + off_h * stride_qh + q_start * stride_qm
+
+    # inner loop
+    if SEQUENCE_PARALLEL:
+        _bwd_kernel_one_col_block(
+            Q,
+            K,
+            V,
+            sm_scale,
+            Out,
+            DO,
+            DQ,
+            DK,
+            DV,
+            L,
+            D,
+            q_offset,
+            k_offset,
+            v_offset,
+            do_offset,
+            dq_offset,
+            dk_offset,
+            dv_offset,
+            d_offset,
+            l_offset,
+            stride_dq_all,
+            stride_qz,
+            stride_qh,
+            stride_qm,
+            stride_qk,
+            stride_kz,
+            stride_kh,
+            stride_kn,
+            stride_kk,
+            stride_vz,
+            stride_vh,
+            stride_vn,
+            stride_vk,
+            stride_deltaz,
+            stride_deltah,
+            stride_deltam,
+            Z,
+            H,
+            N_CTX_Q,
+            N_CTX_K,
+            off_h,
+            off_z,
+            off_hz,
+            start_n,
+            num_block_m,
+            num_block_n,
+            BLOCK_M=BLOCK_M,
+            BLOCK_DMODEL=BLOCK_DMODEL,
+            ACTUAL_BLOCK_DMODEL=ACTUAL_BLOCK_DMODEL,
+            BLOCK_N=BLOCK_N,
+            SEQUENCE_PARALLEL=SEQUENCE_PARALLEL,
+            CAUSAL=CAUSAL,
+            USE_EXP2=USE_EXP2,
+        )
+    else:
+        for start_n in range(0, num_block_n):
+            _bwd_kernel_one_col_block(
+                Q,
+                K,
+                V,
+                sm_scale,
+                Out,
+                DO,
+                DQ,
+                DK,
+                DV,
+                L,
+                D,
+                q_offset,
+                k_offset,
+                v_offset,
+                do_offset,
+                dq_offset,
+                dk_offset,
+                dv_offset,
+                d_offset,
+                l_offset,
+                stride_dq_all,
+                stride_qz,
+                stride_qh,
+                stride_qm,
+                stride_qk,
+                stride_kz,
+                stride_kh,
+                stride_kn,
+                stride_kk,
+                stride_vz,
+                stride_vh,
+                stride_vn,
+                stride_vk,
+                stride_deltaz,
+                stride_deltah,
+                stride_deltam,
+                Z,
+                H,
+                N_CTX_Q,
+                N_CTX_K,
+                off_h,
+                off_z,
+                off_hz,
+                start_n,
+                num_block_m,
+                num_block_n,
+                BLOCK_M=BLOCK_M,
+                BLOCK_DMODEL=BLOCK_DMODEL,
+                ACTUAL_BLOCK_DMODEL=ACTUAL_BLOCK_DMODEL,
+                BLOCK_N=BLOCK_N,
+                SEQUENCE_PARALLEL=SEQUENCE_PARALLEL,
+                CAUSAL=CAUSAL,
+                USE_EXP2=USE_EXP2,
+            )
+
+
+# NOTE: smaller blocks have lower accuracy. more accumlation error probably 128 * 128 seems good but leads to oom. 64 * 64 has accumlation errors but no oom.
+def attention_prefill_backward_triton_impl(
+    do,
+    q,
+    k,
+    v,
+    o,
+    softmax_lse,
+    dq,
+    dk,
+    dv,
+    sm_scale: float,
+    alibi_slopes, # pylint: disable=unused-argument
+    causal,
+    layout: str,
+    cu_seqlens_q,
+    cu_seqlens_k,
+    max_seqlen_q: int,
+    max_seqlen_k: int,
+    use_exp2: bool,
+    sequence_parallel = True,
+):
+    # make contigious
+    q = q.contiguous()
+    k = k.contiguous()
+    v = v.contiguous()
+    softmax_lse = softmax_lse.contiguous()
+
+    # get strides and shape
+    batch, nheads_q, nheads_k, head_size, max_seqlen_q, max_seqlen_k = get_shape_from_layout(q, k, layout, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k) # pylint: disable=unused-variable
+    q_strides, k_strides, v_strides, o_strides = get_strides_from_layout(q, k, v, o, layout)
+    stride_qz, stride_qh, stride_qm, stride_qk =  q_strides
+    stride_kz, stride_kh, stride_kn, stride_kk = k_strides
+    stride_vz, stride_vh, stride_vn, stride_vk = v_strides
+    stride_oz, stride_oh, stride_om, stride_ok = o_strides
+    batch_headsize = batch * nheads_q
+    is_varlen = layout == "thd"
+
+    # FIXME: some configs lead to oom for some reason when using 64 x 64 blocks
+    if max_seqlen_q <= 32 or max_seqlen_k <= 32:
+        BLOCK_M = 32
+        BLOCK_N = 32
+    else:
+        BLOCK_M = 64
+        BLOCK_N = 64
+    num_warps = 4 # NOTE: originial is 8. changing it to 1 caused issues be careful
+    num_stages = 1
+    waves_per_eu = 1
+
+    # divide up the problem
+    num_blocks_m = triton.cdiv(max_seqlen_q, BLOCK_M)
+    num_blocks_n = triton.cdiv(max_seqlen_k, BLOCK_N)
+
+    # get closest power of 2 over or equal to 32.
+    padded_d_model = 1 << (head_size - 1).bit_length()
+    padded_d_model = max(padded_d_model, 16)
+    BLOCK_DMODEL = padded_d_model
+    ACTUAL_BLOCK_DMODEL = head_size
+
+    do = do.contiguous()
+    # NOTE: we might need to copy the output tensor if they are not continuous or have other issues
+    copy_back = {"dq": False, "dk": False, "dv": False}
+
+    dq_og = None
+    # deal with dq
+    if dq is None:
+        if sequence_parallel:
+            dq = torch.zeros((num_blocks_n,) + q.shape, device=q.device, dtype=q.dtype)
+        else:
+            dq = torch.zeros(q.shape, device=q.device, dtype=q.dtype)
+    else:
+        dq_og = dq
+        if not dq.is_contiguous():
+            dq = dq.contiguous()
+            copy_back["dq"] = True
+
+        if sequence_parallel:
+            dq = torch.zeros((num_blocks_n,) + q.shape, device=q.device, dtype=q.dtype)
+            copy_back["dq"] = True
+        else:
+            # NOTE: the kernel does inplace accumlation so dq has to be zeros. This avoids the case where we are passed empty dq and it is not all zeros
+            dq.zero_()
+    stride_dq_all = dq.stride()[0]
+
+    dk_og = None
+    dv_og = None
+    # deal with dk, dv
+    if (dk is None) or (dv is None):
+        dk = torch.empty_like(k)
+        dv = torch.empty_like(v)
+    else:
+        if not dk.is_contiguous():
+            dk_og = dk
+            dk = dk.contiguous()
+            copy_back["dk"] = True
+
+        if not dv.is_contiguous():
+            dv_og = dv
+            dv = dv.contiguous()
+            copy_back["dv"] = True
+
+    # assert contigious
+    assert do.is_contiguous()
+    assert q.is_contiguous()
+    assert k.is_contiguous()
+    assert v.is_contiguous()
+    assert o.is_contiguous()
+    assert softmax_lse.is_contiguous()
+
+    # init delta
+    delta = torch.empty_like(softmax_lse)
+    if is_varlen:
+        stride_deltam, stride_deltah = delta.stride()
+        stride_deltaz = 0
+    else:
+        stride_deltaz, stride_deltah, stride_deltam = delta.stride()
+
+    _bwd_preprocess_use_o[(num_blocks_m, batch_headsize)](
+        o,
+        do,
+        delta,
+        stride_oz, stride_oh, stride_om, stride_ok,
+        stride_oz, stride_oh, stride_om, stride_ok,
+        stride_deltaz, stride_deltah, stride_deltam,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        max_seqlen_q,
+        max_seqlen_k,
+        BLOCK_M=BLOCK_M,
+        BLOCK_DMODEL=BLOCK_DMODEL,
+        ACTUAL_BLOCK_DMODEL=ACTUAL_BLOCK_DMODEL,
+        N_CTX_Q=max_seqlen_q,
+        Z=batch,
+        H=nheads_q,
+        IS_VARLEN=is_varlen
+    )
+
+    _bwd_kernel[(batch_headsize, num_blocks_n if sequence_parallel else 1)](
+        q,
+        k,
+        v,
+        sm_scale,
+        o,
+        do,
+        dq,
+        dk,
+        dv,
+        softmax_lse,
+        delta,
+        stride_dq_all,
+        stride_qz, stride_qh, stride_qm, stride_qk,
+        stride_kz, stride_kh, stride_kn, stride_kk,
+        stride_vz, stride_vh, stride_vn, stride_vk,
+        stride_deltaz, stride_deltah, stride_deltam,
+        batch,
+        nheads_q,
+        num_blocks_m,
+        num_blocks_n,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        max_seqlen_q,
+        max_seqlen_k,
+        BLOCK_M=BLOCK_M,
+        BLOCK_N=BLOCK_N,
+        BLOCK_DMODEL=BLOCK_DMODEL,
+        ACTUAL_BLOCK_DMODEL=ACTUAL_BLOCK_DMODEL,
+        SEQUENCE_PARALLEL=sequence_parallel,
+        CAUSAL=causal,
+        USE_EXP2=use_exp2,
+        num_warps=num_warps,
+        num_stages=num_stages,
+        waves_per_eu = waves_per_eu,
+        IS_VARLEN=is_varlen
+    )
+
+    if sequence_parallel:
+        dq = dq.sum(dim=0)
+
+    if copy_back["dq"]:
+        dq_og.copy_(dq)
+        dq = dq_og
+    if copy_back["dk"]:
+        dk_og.copy_(dk)
+        dk = dk_og
+    if copy_back["dv"]:
+        dv_og.copy_(dv)
+        dv = dv_og
+
+    return dq, dk, dv, delta, None, None
diff --git a/modules/flash_attn_triton_amd/bwd_ref.py b/modules/flash_attn_triton_amd/bwd_ref.py
new file mode 100644
index 000000000..2b1befd88
--- /dev/null
+++ b/modules/flash_attn_triton_amd/bwd_ref.py
@@ -0,0 +1,271 @@
+import math
+import torch
+
+
+def attention_backward_core_ref_impl(
+    do, q, k, v, o, softmax_lse, sm_scale, causal, use_exp2
+):
+    # cast to float32
+    do = do.to(torch.float32)
+    q = q.to(torch.float32)
+    k = k.to(torch.float32)
+    v = v.to(torch.float32)
+    o = o.to(torch.float32)
+    softmax_lse = softmax_lse.to(torch.float32)
+
+    # recompute attention_scores. Make sure it matches the forward impl. i.e. It use float32
+    attention_scores = torch.matmul(q.to(torch.float32), k.transpose(-2, -1).to(torch.float32))
+
+    # scale scores
+    attention_scaled_scores = sm_scale * attention_scores
+
+    # Apply causal mask if necessary
+    if causal:
+        L_q, L_k = q.shape[1], k.shape[1]
+        row_idx = torch.arange(L_q, device=q.device).unsqueeze(1)
+        col_idx = torch.arange(L_k, device=q.device).unsqueeze(0)
+        col_offset = L_q-L_k
+        causal_mask = row_idx >= (col_offset + col_idx)
+        # set -inf to places the causal mask is false
+        attention_scaled_scores = attention_scaled_scores.masked_fill(
+             torch.logical_not(causal_mask.unsqueeze(0)), float('-inf')
+        )
+
+    # compute probabilities using softmax_lse
+    if use_exp2:
+        RCP_LN = 1 / math.log(2)
+        attention_scaled_scores_base2 = attention_scaled_scores * RCP_LN
+        softmax_lse_base2 = softmax_lse * RCP_LN
+        softmax_lse_3d =  softmax_lse_base2.unsqueeze(-1)
+        p = torch.exp2(attention_scaled_scores_base2 - softmax_lse_3d)
+    else:
+        softmax_lse_3d =  softmax_lse.unsqueeze(-1)
+        p = torch.exp(attention_scaled_scores - softmax_lse_3d)
+
+    # compute gradient wrt v
+    dv = torch.matmul(p.transpose(-2, -1), do.to(torch.float32))
+
+    # compute dp
+    dp = torch.matmul(do, v.transpose(-2, -1))
+
+    # calculate ds using dp
+    delta = torch.sum(o * do, axis=-1).to(torch.float32)  # what OAI kernel uses
+    delta_3d = delta.unsqueeze(-1)
+    ds = (p * (dp - delta_3d)) * sm_scale
+
+    # compute gradient wrt k
+    dk = torch.matmul(ds.transpose(-2, -1), q.to(torch.float32))
+
+    # compute gradient wrt q
+    dq = torch.matmul(ds, k.to(torch.float32))
+
+    # cast back to original dtype
+    dq = dq.to(torch.float16)
+    dk = dk.to(torch.float16)
+    dv = dv.to(torch.float16)
+
+    # remove d dim with size 1
+    delta = delta_3d.squeeze(-1)
+
+    return dq, dk, dv, delta
+
+def attention_varlen_backward_pytorch_ref_impl(
+    do,
+    q,
+    k,
+    v,
+    o,
+    softmax_lse,
+    sm_scale,
+    causal,
+    layout,
+    cu_seqlens_q,
+    cu_seqlens_k,
+    max_seqlen_q, max_seqlen_k, # pylint: disable=unused-argument
+    use_exp2,
+):
+    # Ensure the layout is 'thd'
+    if layout != 'thd':
+        raise ValueError(f"Unsupported layout {layout}. Expected 'thd'.")
+
+    batch_size = cu_seqlens_q.shape[0] - 1
+    num_heads = q.shape[1]
+    head_dim = q.shape[2] # pylint: disable=unused-variable
+
+    # Pre-allocate outputs
+    total_L_q = q.shape[0]
+    total_L_k = k.shape[0] # pylint: disable=unused-variable
+
+    dq = torch.zeros_like(q)
+    dk = torch.zeros_like(k)
+    dv = torch.zeros_like(v)
+    # delta has the same shape as softmax_lse: [total_L_q, num_heads]
+    delta = torch.zeros((total_L_q, num_heads), dtype=torch.float32, device=o.device)
+
+    for i in range(batch_size):
+        # Get the start and end indices for the current sequence
+        start_q = cu_seqlens_q[i].item()
+        end_q = cu_seqlens_q[i + 1].item()
+        start_k = cu_seqlens_k[i].item()
+        end_k = cu_seqlens_k[i + 1].item()
+
+        # Extract q_i, k_i, v_i, do_i, o_i, softmax_lse_i
+        q_i = q[start_q:end_q, :, :]      # [L_q_i, num_heads, head_dim]
+        k_i = k[start_k:end_k, :, :]      # [L_k_i, num_heads, head_dim]
+        v_i = v[start_k:end_k, :, :]      # [L_k_i, num_heads, head_dim]
+        do_i = do[start_q:end_q, :, :]    # [L_q_i, num_heads, head_dim]
+        o_i = o[start_q:end_q, :, :]      # [L_q_i, num_heads, head_dim]
+        # softmax_lse has shape [total_L_q, num_heads]
+        softmax_lse_i = softmax_lse[start_q:end_q, :]  # [L_q_i, num_heads]
+        softmax_lse_i = softmax_lse_i.transpose(0, 1)  # [num_heads, L_q_i]
+
+        # Permute to [num_heads, L_q_i, head_dim]
+        q_i = q_i.permute(1, 0, 2)
+        k_i = k_i.permute(1, 0, 2)
+        v_i = v_i.permute(1, 0, 2)
+        do_i = do_i.permute(1, 0, 2)
+        o_i = o_i.permute(1, 0, 2)
+        # softmax_lse_i is already in [num_heads, L_q_i]
+
+        # Call the core backward function for this sequence
+        dq_i, dk_i, dv_i, delta_i = attention_backward_core_ref_impl(
+            do_i,
+            q_i,
+            k_i,
+            v_i,
+            o_i,
+            softmax_lse_i,
+            sm_scale,
+            causal,
+            use_exp2
+        )
+
+        # Convert back to 'thd' layout
+        dq_i = dq_i.permute(1, 0, 2)  # [L_q_i, num_heads, head_dim]
+        dk_i = dk_i.permute(1, 0, 2)  # [L_k_i, num_heads, head_dim]
+        dv_i = dv_i.permute(1, 0, 2)  # [L_k_i, num_heads, head_dim]
+
+        # Place outputs in pre-allocated tensors
+        dq[start_q:end_q, :, :] = dq_i
+        dk[start_k:end_k, :, :] += dk_i  # Accumulate gradients for shared keys
+        dv[start_k:end_k, :, :] += dv_i  # Accumulate gradients for shared values
+        # delta_i has shape [num_heads, L_q_i]
+        delta_i = delta_i.transpose(1, 0)  # [L_q_i, num_heads]
+        delta[start_q:end_q, :] = delta_i
+
+    return dq, dk, dv, delta
+
+def attention_vanilla_backward_pytorch_ref_impl(
+    do,
+    q,
+    k,
+    v,
+    o,
+    softmax_lse,
+    sm_scale,
+    causal,
+    layout,
+    use_exp2,
+):
+    if layout == "bshd":
+        do = do.transpose(1, 2).contiguous()
+        q = q.transpose(1, 2).contiguous()
+        k = k.transpose(1, 2).contiguous()
+        v = v.transpose(1, 2).contiguous()
+        o = o.transpose(1, 2).contiguous()
+    elif layout == "bhsd":
+        pass
+    else:
+        raise ValueError(f"Unknown layout {layout}")
+
+    # Prepare tensors in [batch_size * num_heads, seq_len, head_dim] format
+    batch_size, num_heads, seq_len_q, head_dim = q.shape
+    seq_len_k = k.shape[2]
+
+    # Merge batch and heads dimensions
+    do = do.reshape(batch_size * num_heads, seq_len_q, head_dim)
+    q = q.reshape(batch_size * num_heads, seq_len_q, head_dim)
+    k = k.reshape(batch_size * num_heads, seq_len_k, head_dim)
+    v = v.reshape(batch_size * num_heads, seq_len_k, head_dim)
+    softmax_lse = softmax_lse.reshape(batch_size * num_heads, seq_len_q)
+    o = o.reshape(batch_size * num_heads, seq_len_q, head_dim)
+
+    dq, dk, dv, delta = attention_backward_core_ref_impl(
+        do,
+        q,
+        k,
+        v,
+        o,
+        softmax_lse,
+        sm_scale,
+        causal,
+        use_exp2
+    )
+
+    # Reshape outputs back to [batch_size, num_heads, seq_len, head_dim]
+    dq = dq.reshape(batch_size, num_heads, seq_len_q, head_dim)
+    dk = dk.reshape(batch_size, num_heads, seq_len_k, head_dim)
+    dv = dv.reshape(batch_size, num_heads, seq_len_k, head_dim)
+    delta = delta.reshape(batch_size, num_heads, seq_len_q)
+
+    # Go back to original layout
+    if layout == "bshd":
+        dq = dq.transpose(1, 2)
+        dk = dk.transpose(1, 2)
+        dv = dv.transpose(1, 2)
+    elif layout == "bhsd":
+        pass
+    else:
+        raise ValueError(f"Unknown layout {layout}")
+
+    return dq, dk, dv, delta
+
+
+def attention_backward_pytorch_ref_impl(
+    do,
+    q,
+    k,
+    v,
+    o,
+    softmax_lse,
+    sm_scale,
+    causal,
+    layout,
+    cu_seqlens_q,
+    cu_seqlens_k,
+    max_seqlen_q,
+    max_seqlen_k,
+    use_exp2
+):
+    if layout == "thd":
+        dq, dk, dv, delta = attention_varlen_backward_pytorch_ref_impl(
+            do,
+            q,
+            k,
+            v,
+            o,
+            softmax_lse,
+            sm_scale,
+            causal,
+            layout,
+            cu_seqlens_q,
+            cu_seqlens_k,
+            max_seqlen_q,
+            max_seqlen_k,
+            use_exp2,
+        )
+    else:
+        dq, dk, dv, delta = attention_vanilla_backward_pytorch_ref_impl(
+            do,
+            q,
+            k,
+            v,
+            o,
+            softmax_lse,
+            sm_scale,
+            causal,
+            layout,
+            use_exp2,
+        )
+
+    return dq, dk, dv, delta
diff --git a/modules/flash_attn_triton_amd/fwd_decode.py b/modules/flash_attn_triton_amd/fwd_decode.py
new file mode 100644
index 000000000..7a2a234d6
--- /dev/null
+++ b/modules/flash_attn_triton_amd/fwd_decode.py
@@ -0,0 +1,700 @@
+import torch
+import triton
+import triton.language as tl
+from modules.flash_attn_triton_amd.utils import _strides, get_padded_headsize
+
+
+@triton.jit
+def _fwd_kernel_splitK(
+    Q,
+    K,
+    V,
+    sm_scale,
+    Out_splitK,  # [B, H, split_k, Mq, K]
+    Metadata,  # [B, H, 2, split_k, M_ceil] contains [mi, li]
+    K_new,
+    V_new,
+    Cache_seqlens,
+    Cache_batch_idx,
+    Alibi_slopes,
+    stride_qz,
+    stride_qm,
+    stride_qg,
+    stride_qh,
+    stride_qd,
+    stride_kz,
+    stride_kn,
+    stride_kg,
+    stride_kh,
+    stride_kd,
+    stride_vz,
+    stride_vn,
+    stride_vg,
+    stride_vh,
+    stride_vd,
+    stride_osk_zhg,
+    stride_osk_s,
+    stride_osk_m,
+    stride_osk_d, # pylint: disable=unused-argument
+    stride_mzhg,
+    stride_m2,
+    stride_ms,
+    stride_mm, # pylint: disable=unused-argument
+    stride_kn_z,
+    stride_kn_n,
+    stride_kn_g,
+    stride_kn_h,
+    stride_kn_d,
+    stride_vn_z,
+    stride_vn_n,
+    stride_vn_g,
+    stride_vn_h,
+    stride_vn_d,
+    stride_az,
+    stride_ah,
+    Z, # pylint: disable=unused-argument
+    N_CTX_Q,
+    N_CTX_K,
+    N_CTX_NEW,
+    BLOCK_N_PER_SPLIT,
+    H_q: tl.constexpr,
+    H_kv: tl.constexpr,
+    G_q: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_DMODEL: tl.constexpr,
+    ACTUAL_BLOCK_DMODEL: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BOUNDS_CHECKS_N: tl.constexpr,
+    USE_CACHE_SEQLENs: tl.constexpr,
+    USE_CACHE_BATCH_IDX: tl.constexpr,
+    NEW_KV: tl.constexpr,
+    IS_GQA: tl.constexpr,
+    IS_CAUSAL: tl.constexpr,
+    USE_ALIBI: tl.constexpr,
+):
+    # Padding
+    PADDED_HEAD: tl.constexpr = ACTUAL_BLOCK_DMODEL != BLOCK_DMODEL
+    if PADDED_HEAD:
+        d_mask = tl.arange(0, BLOCK_DMODEL) < ACTUAL_BLOCK_DMODEL
+
+    start_m = tl.program_id(0)
+    off_zhg = tl.program_id(1)
+    off_z = off_zhg // (H_q * G_q)
+    off_h_q = (off_zhg // G_q) % H_q
+    off_g_q = off_zhg % G_q
+    splitk_idx = tl.program_id(2)
+
+    # pick batch index
+    if USE_CACHE_BATCH_IDX:
+        cache_batch_idx = tl.load(Cache_batch_idx + off_z)
+    else:
+        cache_batch_idx = off_z
+
+    # Load ALiBi slope if enabled
+    if USE_ALIBI:
+        a_offset = off_z * stride_az + off_h_q * stride_ah
+        alibi_slope = tl.load(Alibi_slopes + a_offset)
+    else:
+        alibi_slope = None
+
+    lo = splitk_idx * BLOCK_N_PER_SPLIT
+    if USE_CACHE_SEQLENs:
+        cache_seqlen_last_idx = tl.load(Cache_seqlens + off_z)
+        if NEW_KV:
+            kv_len = cache_seqlen_last_idx + N_CTX_NEW
+        else:
+            kv_len = cache_seqlen_last_idx
+    else:
+        kv_len = N_CTX_K
+    hi = tl.minimum((splitk_idx + 1) * BLOCK_N_PER_SPLIT, kv_len)
+
+    HEAD_RATIO: tl.constexpr = H_q // H_kv
+    if IS_GQA:
+        k_head_idx = off_h_q // HEAD_RATIO
+        v_head_idx = k_head_idx
+    else:
+        k_head_idx = off_h_q
+        v_head_idx = off_h_q
+
+    # calculate base offset
+    k_base = K + k_head_idx * stride_kh + cache_batch_idx * stride_kz + off_g_q * stride_kg
+    v_base = V + v_head_idx * stride_vh + cache_batch_idx * stride_vz + off_g_q * stride_vg
+
+    # Copy new Keys and Values into Cache
+    if NEW_KV:
+        knew_base = K_new + k_head_idx * stride_kn_h + off_z * stride_kn_z + off_g_q * stride_kn_g
+
+        # Determine the starting position for new data in the cache
+        if USE_CACHE_SEQLENs:
+            start_idx = tl.load(Cache_seqlens + off_z)
+        else:
+            start_idx = N_CTX_K - N_CTX_NEW
+
+        # Copy new Keys
+        for i in range(0, N_CTX_NEW, BLOCK_N):
+            # Load from K_new
+            k_new_block = tl.load(
+                knew_base +
+                tl.arange(0, BLOCK_DMODEL)[:, None] * stride_kn_d +
+                (tl.arange(0, BLOCK_N) + i)[None, :] * stride_kn_n,
+                 mask=(tl.arange(0, BLOCK_N)[None, :] + i < N_CTX_NEW) &
+                     (tl.arange(0, BLOCK_DMODEL)[:, None] < ACTUAL_BLOCK_DMODEL),
+                other=0
+            )
+
+            # Store to K
+            tl.store(
+                k_base +
+                tl.arange(0, BLOCK_DMODEL)[:, None] * stride_kd +
+                (tl.arange(0, BLOCK_N) + i + start_idx)[None, :] * stride_kn,
+                k_new_block,
+                 mask=(tl.arange(0, BLOCK_N)[None, :] + i < N_CTX_NEW) &
+                     (tl.arange(0, BLOCK_DMODEL)[:, None] < ACTUAL_BLOCK_DMODEL),
+            )
+
+        # Copy new Values
+        vnew_base = V_new + v_head_idx * stride_vn_h + off_z * stride_vn_z + off_g_q * stride_vn_g
+        for i in range(0, N_CTX_NEW, BLOCK_N):
+            # Load from V_new
+            v_new_block = tl.load(
+                vnew_base +
+                (tl.arange(0, BLOCK_N) + i)[:, None] * stride_vn_n +
+                tl.arange(0, BLOCK_DMODEL)[None, :] * stride_vn_d,
+                mask=(tl.arange(0, BLOCK_N)[:, None] + i < N_CTX_NEW) &
+                     (tl.arange(0, BLOCK_DMODEL)[None, :] < ACTUAL_BLOCK_DMODEL),
+                other=0
+            )
+
+            # Store to V
+            tl.store(
+                v_base +
+                (tl.arange(0, BLOCK_N) + i + start_idx)[:, None] * stride_vn +
+                tl.arange(0, BLOCK_DMODEL)[None, :] * stride_vd,
+                v_new_block,
+                 mask=(tl.arange(0, BLOCK_N)[:, None] + i < N_CTX_NEW) &
+                     (tl.arange(0, BLOCK_DMODEL)[None, :] < ACTUAL_BLOCK_DMODEL),
+            )
+
+    Q_block_ptr = tl.make_block_ptr(
+        base=Q + off_h_q * stride_qh + off_z * stride_qz + off_g_q * stride_qg,
+        shape=(N_CTX_Q, ACTUAL_BLOCK_DMODEL),
+        strides=(stride_qm, stride_qd),
+        offsets=(start_m * BLOCK_M, 0),
+        block_shape=(BLOCK_M, BLOCK_DMODEL),
+        order=(1, 0),
+    )
+
+    K_block_ptr = tl.make_block_ptr(
+        base=k_base,
+        shape=(ACTUAL_BLOCK_DMODEL, hi),
+        strides=(stride_kd, stride_kn),
+        offsets=(0, lo),
+        block_shape=(BLOCK_DMODEL, BLOCK_N),
+        order=(0, 1),
+    )
+    V_block_ptr = tl.make_block_ptr(
+        base=v_base,
+        shape=(hi, ACTUAL_BLOCK_DMODEL),
+        strides=(stride_vn, stride_vd),
+        offsets=(lo, 0),
+        block_shape=(BLOCK_N, BLOCK_DMODEL),
+        order=(1, 0),
+    )
+
+    K_scale_shift_block_ptr = None
+    V_scale_shift_block_ptr = None
+
+    # initialize pointer to m and l
+    m_i = tl.full([BLOCK_M], float("-inf"), dtype=tl.float32)
+    l_i = tl.zeros([BLOCK_M], dtype=tl.float32)
+
+    acc = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32)  # noqa: F821
+
+    # scale sm_scale by log_2(e) and use
+    # 2^x instead of exp in the loop because CSE and LICM
+    # don't work as expected with `exp` in the loop
+    qk_scale = sm_scale * 1.44269504
+    # load q: it will stay in SRAM throughout
+    q = tl.load(  # noqa: F821
+        tl.advance(Q_block_ptr, (0, 0)), boundary_check=(0, ))
+    q = (q * qk_scale).to(q.dtype)
+    if PADDED_HEAD:
+        q = tl.where(d_mask[None, :], q, 0.0)
+
+    # loop over k, v and update accumulator
+    for start_n in range(lo, hi, BLOCK_N):
+        k, v = load_k_v_group(
+            K_block_ptr,
+            V_block_ptr,
+            K_scale_shift_block_ptr,
+            V_scale_shift_block_ptr,
+            BOUNDS_CHECKS_N,
+            1,
+            BLOCK_DMODEL,
+            ACTUAL_BLOCK_DMODEL,
+            Q.dtype.element_ty,
+            0,
+        )
+        if PADDED_HEAD:
+            k = tl.where(d_mask[:, None], k, 0.0)
+            v = tl.where(d_mask[None, :], v, 0.0)
+
+        # -- compute qk ---
+        qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
+        qk += tl.dot(q, k)  # noqa: F821
+
+        if USE_ALIBI:
+            row_idx = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
+            col_idx = start_n + tl.arange(0, BLOCK_N)
+
+            # Compute relative positions
+            relative_pos = row_idx[:, None] + kv_len - (N_CTX_Q + col_idx[None, :])
+            relative_pos = tl.abs(relative_pos)
+
+            # Compute ALiBi bias
+            alibi_bias = -1 * alibi_slope * relative_pos
+            qk += (alibi_bias * 1.44269504)
+
+        # Apply causal mask if IS_CAUSAL is True
+        if IS_CAUSAL:
+            row_idx = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
+            col_idx = start_n + tl.arange(0, BLOCK_N)
+
+            # create a N_CTX_Q x kv_len causal mask
+            col_offset = N_CTX_Q - kv_len
+            causal_mask = row_idx[:, None] >= (col_offset + col_idx[None, :])
+
+            # Apply the mask
+            qk = tl.where(causal_mask, qk, float("-inf"))
+
+        # TODO: This is slow, and only needed at the last iteration.
+        # Maybe we can unroll the last iteration instead?
+        if BOUNDS_CHECKS_N:
+            qk = tl.where(tl.arange(0, BLOCK_N) < hi - start_n, qk, float("-inf"))
+
+        # -- compute scaling constant ---
+        m_i_new = tl.maximum(m_i, tl.max(qk, 1))
+        if IS_CAUSAL:
+            alpha = tl.math.exp2(tl.where(m_i > float("-inf"), m_i - m_i_new, float("-inf")))
+        else:
+            alpha = tl.math.exp2(m_i - m_i_new)
+        # cause of nan because subtracting infs
+        if IS_CAUSAL:
+            qk = tl.where(qk > float("-inf"), qk - m_i_new[:, None], float("-inf"))
+        else:
+            qk = qk - m_i_new[:, None]
+
+        p = tl.math.exp2(qk)
+
+        # -- update m_i and l_i --
+        l_i = l_i * alpha + tl.sum(p, 1)
+        m_i = m_i_new
+        p = p.to(Q.dtype.element_ty)
+
+        # -- scale and update acc --
+        acc *= alpha[:, None]
+        acc += tl.dot(p.to(v.dtype), v)
+
+        # update pointers
+        K_block_ptr = tl.advance(K_block_ptr, (0, BLOCK_N))
+        V_block_ptr = tl.advance(V_block_ptr, (BLOCK_N, 0))
+
+    # write back O
+    O_block_ptr = tl.make_block_ptr(
+        base=Out_splitK + off_zhg * stride_osk_zhg + splitk_idx * stride_osk_s,
+        shape=(N_CTX_Q, BLOCK_DMODEL),
+        strides=(stride_osk_m, 1),
+        offsets=(start_m * BLOCK_M, 0),
+        block_shape=(BLOCK_M, BLOCK_DMODEL),
+        order=(1, 0),
+    )
+    tl.store(
+        tl.advance(O_block_ptr, (0, 0)),
+        acc,
+        boundary_check=(0, ),
+    )
+    # Write metadata for split-K reduction
+    Metadata_ptr = (Metadata + off_zhg * stride_mzhg + splitk_idx * stride_ms + start_m * BLOCK_M +
+                    tl.arange(0, BLOCK_M))
+    tl.store(Metadata_ptr, m_i)
+    tl.store(Metadata_ptr + stride_m2, l_i)
+
+
+@triton.jit
+def load_k_v_group(
+    K_block_ptr,
+    V_block_ptr,
+    K_scale_shift_block_ptr, V_scale_shift_block_ptr, # pylint: disable=unused-argument
+    BOUNDS_CHECKS_N: tl.constexpr,
+    PACKED_PER_VAL: tl.constexpr, BLOCK_DMODEL: tl.constexpr, # pylint: disable=unused-argument
+    ACTUAL_BLOCK_DMODEL: tl.constexpr,
+    dtype: tl.constexpr, # pylint: disable=unused-argument
+    group_id: tl.constexpr,
+):
+    # Load K/V for a given block
+    # Advance to the current quantization group
+    K_block_ptr = tl.advance(K_block_ptr, (ACTUAL_BLOCK_DMODEL * group_id, 0))
+    V_block_ptr = tl.advance(V_block_ptr, (0, ACTUAL_BLOCK_DMODEL * group_id))
+
+    # -- load k, v --
+    k = tl.load(K_block_ptr, boundary_check=(1, ) if BOUNDS_CHECKS_N else ())
+    v = tl.load(V_block_ptr, boundary_check=(0, ) if BOUNDS_CHECKS_N else ())
+
+    return k, v
+
+
+@triton.jit
+def cast_uint32_to_half2(scale_shift):
+    # Extract two float16 packed into one int32
+    scale = scale_shift & 0xFFFF
+    shift = scale_shift >> 16
+    scale = scale.to(tl.uint16).to(tl.float16, bitcast=True)
+    shift = shift.to(tl.uint16).to(tl.float16, bitcast=True)
+    return scale, shift
+
+
+@triton.jit
+def dequantize(
+    x_,
+    scale,
+    shift,
+    PACKED_PER_VAL: tl.constexpr = 8,
+):
+    # PACKED_PER_VAL is the number of values packed into
+    # each element x_. For example, for int4 quantization
+    #and x_ of type int32, PACKED_PER_VAL is 8.
+
+    BLOCK_N: tl.constexpr = x_.shape[0]
+    BLOCK_DMODEL_PACKED: tl.constexpr = x_.shape[1]
+    offsets = tl.arange(0, PACKED_PER_VAL) * 4
+    quant_offset = (x_[:, None, :] >> offsets[None, :, None])  # (BLOCK_N, PACKED_PER_VAL, D // PACKED_PER_VAL)
+
+    quant_offset = tl.view(quant_offset, (BLOCK_N, BLOCK_DMODEL_PACKED * PACKED_PER_VAL))
+    # Trick - instead of converting int4 to float16 we view it as float16
+    # and then multiply by 32768 * 512 == 2**24
+    quant_offset = (quant_offset & 0xF).to(tl.uint16).to(tl.float16, bitcast=True)
+    quant_offset = (quant_offset * 32768.0).to(tl.float16)
+    scale_512 = scale * 512
+
+    dequant = quant_offset * scale_512 + shift
+    return dequant
+
+
+@triton.jit
+def _splitK_reduce(
+    Out_splitK,  # [B, H, split_k, Mq, K]
+    Metadata,  # [B, H, 2, split_k, M_ceil] contains [mi, li]
+    Out,  # [B, H, M, K]
+    LSE,  # [B, H, M]
+    stride_osk_zhg,
+    stride_osk_s,
+    stride_osk_m,
+    stride_osk_k,
+    stride_mzhg,
+    stride_m2,
+    stride_ms,
+    stride_mm,
+    stride_oz,
+    stride_oh,
+    stride_og,
+    stride_om,
+    stride_ok, # pylint: disable=unused-argument
+    stride_lse_zhg,
+    stride_lse_m, M_ceil: tl.constexpr, # pylint: disable=unused-argument
+    BLOCK_SIZE: tl.constexpr,
+    H: tl.constexpr,
+    G: tl.constexpr,
+    split_k: tl.constexpr,
+    splitK_pow2: tl.constexpr,
+    use_mask: tl.constexpr,
+    IS_CAUSAL: tl.constexpr,
+):
+    off_zhg = tl.program_id(0)
+    off_z = off_zhg // (H * G)
+    off_h = (off_zhg // G) % H
+    off_g = off_zhg % G
+    off_m = tl.program_id(1)
+    off_k = tl.program_id(2)
+
+    # read  chunk
+    spk_idx = tl.arange(0, splitK_pow2)
+    kidx = tl.arange(0, BLOCK_SIZE)
+
+    Metadata_ptr = Metadata + stride_mzhg * off_zhg + spk_idx * stride_ms + off_m * stride_mm
+
+    o_ptr = (Out_splitK + off_zhg * stride_osk_zhg + stride_osk_m * off_m + off_k * BLOCK_SIZE +
+             stride_osk_s * spk_idx[:, None] + kidx[None, :] * stride_osk_k)
+
+    # read max values of each splitK
+    if use_mask:
+        spk_mask = spk_idx < split_k
+        l_m = tl.load(Metadata_ptr, mask=spk_mask, other=float("-inf"))
+        l_sum = tl.load(Metadata_ptr + stride_m2, mask=spk_mask, other=0.0)
+        acc = tl.load(o_ptr, mask=spk_mask[:, None], other=0.0)
+    else:
+        l_m = tl.load(Metadata_ptr)
+        l_sum = tl.load(Metadata_ptr + stride_m2)
+        acc = tl.load(o_ptr)
+
+    g_m = tl.max(l_m, axis=0)
+
+    if IS_CAUSAL:
+        l_m_offset = l_m - g_m
+        alpha = tl.where(l_m_offset > float("-inf"), tl.math.exp2(l_m_offset), 0.0)
+    else:
+        alpha = tl.math.exp2(l_m - g_m)
+
+    # read sum
+    l_sum *= alpha
+    g_sum = tl.sum(l_sum, axis=0)
+    acc = acc * alpha[:, None]
+
+    if IS_CAUSAL:
+        # Avoid division by zero
+        g_sum_safe = tl.where(g_sum > 0, g_sum, 1.0)
+        acc_out = tl.sum(acc, axis=0) / g_sum_safe
+    else:
+        acc_out = tl.sum(acc, axis=0) / g_sum
+
+    # Store output
+    Out_ptr = (Out + stride_oz * off_z + stride_oh * off_h + stride_og * off_g + stride_om * off_m +
+               off_k * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE))
+    tl.store(Out_ptr, acc_out)
+
+    # Store lse
+    l_ptrs = LSE + off_zhg * stride_lse_zhg + off_m
+    if IS_CAUSAL:
+        lse = tl.where(g_sum > 0, (g_m + tl.math.log2(g_sum)) / 1.44269504, g_m)
+        tl.store(l_ptrs, lse)
+    else:
+        tl.store(l_ptrs, (g_m + tl.math.log2(g_sum)) / 1.44269504)
+
+
+def quantize_kv_int4(k: torch.Tensor, num_groups: int = 1) -> torch.Tensor:
+    # Scale and shift are such that quantization linearly maps
+    # int4 values range [0..15] to input values range min(k)..max(k)
+    # individually for every row
+    k = k.reshape(*k.shape[:-1], num_groups, k.shape[-1] // num_groups)
+    max_vals = torch.max(k, dim=-1, keepdim=True).values
+    min_vals = torch.min(k, dim=-1, keepdim=True).values
+    scale_k: torch.Tensor = (max_vals - min_vals) / 15
+
+    shift_k = torch.min(k, dim=-1, keepdim=True).values
+    scale_k = scale_k.to(torch.float16)
+    shift_k = shift_k.to(torch.float16)
+
+    in_bytes = ((k - shift_k.expand(k.shape)) / scale_k.expand(k.shape)) + 0.5
+    in_bytes = in_bytes.to(torch.uint8)
+    in_int4 = in_bytes & 0xF
+    in_int4_packed = in_int4[..., ::2] + (in_int4[..., 1::2] << 4)
+    scale_shift = torch.concat([scale_k.view(torch.uint8), shift_k.view(torch.uint8)], dim=-1)
+    k_quant = torch.concat(
+        [
+            scale_shift.flatten(start_dim=-2),
+            in_int4_packed.flatten(start_dim=-2),
+        ],
+        dim=-1,
+    ).view(torch.int16)
+    return k_quant
+
+
+def dequantize_kv_fp16(quant_k: torch.Tensor, num_groups: int = 1) -> torch.Tensor:
+    k_i16 = quant_k.view(torch.int16)
+    k_ui8 = k_i16.view(torch.uint8)
+
+    ss_size = num_groups * 4
+    scale_shift_ui8 = k_ui8[..., 0:ss_size]
+    scale_shift_ui8 = scale_shift_ui8.reshape(*scale_shift_ui8.shape[:-1], num_groups, 4)
+    scale = scale_shift_ui8[..., 0:2].view(torch.float16)
+    shift = scale_shift_ui8[..., 2:4].view(torch.float16)
+
+    kv_ui8 = k_ui8[..., ss_size:]
+    k_ui8 = kv_ui8.reshape(*kv_ui8.shape[:-1], num_groups, -1)
+    k1_i4 = k_ui8 & 0xF
+    k2_i4 = (k_ui8 & 0xF0) >> 4
+    k_shape = k1_i4.shape
+    k1_f16 = k1_i4.to(torch.float16) * scale.expand(k_shape) + shift.expand(k_shape)
+    k2_f16 = k2_i4.to(torch.float16) * scale.expand(k_shape) + shift.expand(k_shape)
+
+    out = torch.empty((*k1_f16.shape[:-1], k1_f16.shape[-1] * 2), dtype=torch.float16, device=quant_k.device)
+    out[..., ::2] = k1_f16
+    out[..., 1::2] = k2_f16
+    out = out.reshape(*k_shape[:-2], -1)
+
+    return out
+
+
+def get_split_k(B: int, G: int, H: int, Mk: int) -> int:
+    """Heuristic for the number of splits"""
+    bh = max(B * H, 1)  # NOTE: Handle B*h=0 case
+    split_k = max(Mk, 1024) // bh
+    max_chunk_size = 64
+    while split_k > 0 and Mk / split_k < max_chunk_size:
+        split_k = split_k // 2
+    while B * H * G * split_k >= 1024:
+        split_k = split_k // 2
+    split_k = min(split_k, 512)
+    split_k = max(split_k, 1)
+    return split_k
+
+def attention_decode_forward_triton_impl(q, k, v, sm_scale, causal, alibi_slopes, layout, cache_seqlens, cache_batch_idx, new_kv, k_new, v_new):
+    # kernel config
+    BLOCK_M = 16
+    BLOCK_N = 64
+    SPLIT_K = None
+    NUM_QUANT_GROUPS = 1 # pylint: disable=unused-variable
+
+    # kernels expects "bsghd"
+    original_layout = layout
+    if layout == "bshd":
+        q = q.unsqueeze(2)
+        k = k.unsqueeze(2)
+        v = v.unsqueeze(2)
+        if new_kv:
+            k_new = k_new.unsqueeze(2)
+            v_new = v_new.unsqueeze(2)
+        layout = "bsghd"
+    elif layout == "bhsd":
+        q = q.permute(0, 2, 1, 3).unsqueeze(2)
+        k = k.permute(0, 2, 1, 3).unsqueeze(2)
+        v = v.permute(0, 2, 1, 3).unsqueeze(2)
+        if new_kv:
+            k_new = k_new.permute(0, 2, 1, 3).unsqueeze(2)
+            v_new = v_new.permute(0, 2, 1, 3).unsqueeze(2)
+        layout = "bsghd"
+    elif layout == "bsghd":
+        pass
+    elif layout is None:
+        raise ValueError("Layout not given")
+    assert layout == "bsghd"
+
+    # get dims
+    batch_size, seqlen_q, n_group_q, heads_per_group_q, dim_q = q.shape
+    _, seqlen_k, n_group_k, heads_per_group_k, dim_k = k.shape # pylint: disable=unused-variable
+    _, seqlen_v, n_group_v, heads_per_group_v, dim_v = v.shape # pylint: disable=unused-variable
+
+    assert dim_q == dim_k == dim_v, f"Dimensions must match: {dim_q}, {dim_k}, {dim_v}"
+
+    # get padded size
+    dim_padded  = get_padded_headsize(dim_k)
+
+    # Handle MQA/GQA case
+    if heads_per_group_q > heads_per_group_k:
+        is_gqa = True
+    elif heads_per_group_q < heads_per_group_k:
+        raise ValueError("heads_per_group_q < heads_per_group_k")
+    else:
+        is_gqa = False
+
+    assert dim_k == dim_q, f"Keys have head dim {dim_k} but queries have head dim {dim_q}"
+
+    if SPLIT_K is not None:
+        split_k = SPLIT_K
+    else:
+        # Use heuristics
+        split_k = get_split_k(batch_size, n_group_q, heads_per_group_q, seqlen_k) # NOTE: should the split think about seqlens?
+
+    seqlen_q_ceil = (seqlen_q + BLOCK_M - 1) // BLOCK_M * BLOCK_M
+    out_splitk = torch.empty([batch_size * n_group_q * heads_per_group_q, split_k, seqlen_q_ceil, dim_padded], dtype=torch.float32, device=q.device)
+    metadata = torch.empty([batch_size * n_group_q * heads_per_group_q, 2, split_k, seqlen_q_ceil], dtype=torch.float32, device=q.device)
+    lse = torch.empty((batch_size * n_group_q * heads_per_group_q, seqlen_q), device=q.device, dtype=torch.float32)
+    grid = (triton.cdiv(seqlen_q, BLOCK_M), batch_size * n_group_q * heads_per_group_q, split_k)
+
+    num_warps = 1
+    split_size = (seqlen_k + split_k - 1) // split_k
+    use_cache_seqlens = cache_seqlens is not None
+
+    # TODO: enable quantization
+    _fwd_kernel_splitK[grid](
+        Q=q,
+        K=k,
+        V=v,
+        sm_scale=sm_scale,
+        Out_splitK=out_splitk,
+        Metadata=metadata,
+        K_new = k_new,
+        V_new = v_new,
+        Cache_seqlens=cache_seqlens,
+        Cache_batch_idx=cache_batch_idx,
+        Alibi_slopes=alibi_slopes,
+        **_strides(q, "qz", "qm", "qg", "qh", "qd"),
+        **_strides(k, "kz", "kn", "kg", "kh", "kd"),
+        **_strides(v, "vz", "vn", "vg", "vh", "vd"),
+        **_strides(out_splitk, "osk_zhg", "osk_s", "osk_m", "osk_d"),
+        **_strides(metadata, "mzhg", "m2", "ms", "mm"),
+        **_strides(k_new, "kn_z", "kn_n", "kn_g", "kn_h", "kn_d"),
+        **_strides(v_new, "vn_z", "vn_n", "vn_g", "vn_h", "vn_d"),
+        **_strides(alibi_slopes, "az", "ah"),
+        Z=batch_size,
+        H_q=heads_per_group_q,
+        H_kv=heads_per_group_k,
+        G_q=n_group_q,
+        N_CTX_Q=seqlen_q,
+        N_CTX_K=seqlen_k,
+        N_CTX_NEW=k_new.shape[1] if new_kv else None,
+        BLOCK_N_PER_SPLIT=split_size,
+        BLOCK_M=BLOCK_M,
+        BLOCK_N=BLOCK_N,
+        BLOCK_DMODEL=dim_padded,
+        ACTUAL_BLOCK_DMODEL=dim_k,
+        BOUNDS_CHECKS_N=(split_size % BLOCK_N) > 0 or use_cache_seqlens,
+        USE_CACHE_SEQLENs=use_cache_seqlens,
+        USE_CACHE_BATCH_IDX=cache_batch_idx is not None,
+        NEW_KV=new_kv,
+        IS_GQA=is_gqa,
+        IS_CAUSAL=causal,
+        USE_ALIBI=False if alibi_slopes is None else True,
+        num_warps=num_warps,
+        num_stages=1,
+    )
+
+    out = torch.empty((batch_size, seqlen_q, n_group_q, heads_per_group_q, dim_padded), device=q.device, dtype=q.dtype)
+
+    # Merge together
+    splitK_pow2 = triton.next_power_of_2(split_k)
+    use_mask = splitK_pow2 > split_k
+    if batch_size * n_group_q * heads_per_group_q * seqlen_q >= 512:
+        k_block_num = 1
+    else:
+        k_block_num = 2
+    assert dim_padded % k_block_num == 0
+    k_block_size = dim_padded // k_block_num
+    grid = (batch_size * n_group_q * heads_per_group_q, seqlen_q, k_block_num)
+
+    _splitK_reduce[grid](
+        out_splitk,
+        metadata,
+        out,
+        lse,
+        **_strides(out_splitk, "osk_zhg", "osk_s", "osk_m", "osk_k"),
+        **_strides(metadata, "mzhg", "m2", "ms", "mm"),
+        **_strides(out, "oz", "om", "og", "oh", "ok"),
+        **_strides(lse, "lse_zhg", "lse_m"),
+        M_ceil=seqlen_q_ceil,
+        BLOCK_SIZE=k_block_size,
+        G=n_group_q,
+        H=heads_per_group_q,
+        # TODO: Tune num_warps
+        split_k=split_k,
+        splitK_pow2=splitK_pow2,
+        use_mask=use_mask,
+        IS_CAUSAL=causal,
+        num_warps=4)
+
+    lse = lse.reshape([batch_size, n_group_q, heads_per_group_q, seqlen_q])
+    if q.ndim == 4:
+        # BMGHK -> BMHK
+        assert n_group_q == 1
+        out = out[:, :, 0]
+        lse = lse[:, 0]
+    if seqlen_k == 0:
+        out.zero_()
+    out = out.reshape(batch_size, heads_per_group_q * n_group_q, -1, dim_padded).contiguous()
+
+    # output is batch_size, heads_per_group_q * group_q, seqlen_q, dim_q
+    if original_layout == "bshd":
+        # out=out.transpose(1, 2).contiguous() # this screws up heads and data.
+        # the data is laid out properly. Just need to reshape dims
+        out = out.reshape(batch_size, seqlen_q, -1, dim_padded)
+
+    return out.narrow(-1, 0, dim_k), lse
diff --git a/modules/flash_attn_triton_amd/fwd_prefill.py b/modules/flash_attn_triton_amd/fwd_prefill.py
new file mode 100644
index 000000000..3e2cd32af
--- /dev/null
+++ b/modules/flash_attn_triton_amd/fwd_prefill.py
@@ -0,0 +1,634 @@
+import torch
+import triton
+import triton.language as tl
+from modules.flash_attn_triton_amd.utils import get_shape_from_layout, get_strides_from_layout, is_cdna, is_rdna, AUTOTUNE
+
+
+@triton.jit
+def cdiv_fn(x, y):
+    return (x + y - 1) // y
+
+
+@triton.jit
+def dropout_offsets(philox_seed, philox_offset, dropout_p, m, n, stride): # pylint: disable=unused-argument
+    ms = tl.arange(0, m)
+    ns = tl.arange(0, n)
+    return philox_offset + ms[:, None] * stride + ns[None, :]
+
+
+@triton.jit
+def dropout_rng(philox_seed, philox_offset, dropout_p, m, n, stride):
+    rng_offsets = dropout_offsets(philox_seed, philox_offset, dropout_p, m, n, stride).to(tl.uint32)
+    # TODO: use tl.randint for better performance
+    return tl.rand(philox_seed, rng_offsets)
+
+
+@triton.jit
+def dropout_mask(philox_seed, philox_offset, dropout_p, m, n, stride):
+    rng_output = dropout_rng(philox_seed, philox_offset, dropout_p, m, n, stride)
+    rng_keep = rng_output > dropout_p
+    return rng_keep
+
+
+# Convenience function to load with optional boundary checks.
+# "First" is the major dim, "second" is the minor dim.
+@triton.jit
+def load_fn(ptrs, offset_first, offset_second, boundary_first, boundary_second):
+    if offset_first is not None and offset_second is not None:
+        mask = (offset_first[:, None] < boundary_first) & \
+               (offset_second[None, :] < boundary_second)
+        tensor = tl.load(ptrs, mask=mask, other=0.0)
+    elif offset_first is not None:
+        mask = offset_first[:, None] < boundary_first
+        tensor = tl.load(ptrs, mask=mask, other=0.0)
+    elif offset_second is not None:
+        mask = offset_second[None, :] < boundary_second
+        tensor = tl.load(ptrs, mask=mask, other=0.0)
+    else:
+        tensor = tl.load(ptrs)
+    return tensor
+
+
+@triton.jit
+def compute_alibi_block(alibi_slope, seqlen_q, seqlen_k, offs_m, offs_n, transpose=False):
+    # when seqlen_k and seqlen_q are different we want the diagonal to stick to the bottom right of the attention matrix
+    # for casual mask we want something like this where (1 is kept and 0 is masked)
+    # seqlen_q = 2 and seqlen_k = 5
+    #   1 1 1 1 0
+    #   1 1 1 1 1
+    # seqlen_q = 5 and seqlen_k = 2
+    #        0 0
+    #        0 0
+    #        0 0
+    #        1 0
+    #        1 1
+    # for alibi the diagonal is 0 indicating no penalty for attending to that spot and increasing penalty for attending further from the diagonal
+    # e.g. alibi_slope = 1, seqlen_q = 2, seqlen_k = 5, offs_m = [0, 1, 2, 3], offs_n = [0, 1, 2, 3, 4], transpose = False
+    # 1. offs_m[:,None] = [[0],
+    #                       [1],
+    # 2. offs_m[:,None] + seqlen_k = [[5],
+    #                                  [6],
+    # 3. offs_m[:,None] + seqlen_k - seqlen_q = [[3],
+    #                                             [4],
+    # 4. offs_m[:,None] + seqlen_k - seqlen_q - offs_n[None,:] = [[3], - [[0, 1, 2, 3, 4]] =  [[ 3, 2, 1, 0,-1],
+    #                                                            [4],                           [ 4, 3, 2, 1, 0]]
+    # 5. -1 * alibi_slope * tl.abs(relative_pos_block) = [[ -3, -2, -1, 0,-1],
+    #                                                     [ -4, -3, -2, -1, 0]],
+    relative_pos_block = offs_m[:, None] + seqlen_k - seqlen_q - offs_n[None, :]
+    alibi_block = -1 * alibi_slope * tl.abs(relative_pos_block)
+    if transpose:
+        return alibi_block.T
+    else:
+        return alibi_block
+
+
+@triton.jit
+def _attn_fwd_inner(acc, l_i, m_i, q, k_ptrs, v_ptrs, bias_ptrs, stride_kn, stride_vk, stride_bn, start_m,
+                    actual_seqlen_k, actual_seqlen_q, dropout_p, philox_seed, batch_philox_offset, exp_scores_ptrs,
+                    block_min, block_max, offs_n_causal, masked_blocks, n_extra_tokens, alibi_slope, score_ptrs, scores_scaled_shifted_ptrs, # pylint: disable=unused-argument
+                    IS_CAUSAL: tl.constexpr, BLOCK_M: tl.constexpr, BLOCK_DMODEL: tl.constexpr, BLOCK_N: tl.constexpr,
+                    OFFS_M: tl.constexpr, OFFS_N: tl.constexpr, PRE_LOAD_V: tl.constexpr, MASK_STEPS: tl.constexpr,
+                    ENABLE_DROPOUT: tl.constexpr, PADDED_HEAD: tl.constexpr,
+                    ACTUAL_BLOCK_DMODEL: tl.constexpr, SM_SCALE: tl.constexpr, USE_EXP2: tl.constexpr,
+                    RETURN_SCORES: tl.constexpr):
+    if USE_EXP2:
+        RCP_LN2: tl.constexpr = 1.4426950408889634
+
+    # loop over k, v, and update accumulator
+    for start_n in range(block_min, block_max, BLOCK_N):
+        # For padded blocks, we will overrun the tensor size if
+        # we load all BLOCK_N. For others, the blocks are all within range.
+        if MASK_STEPS:
+            k_offs_n = start_n + tl.arange(0, BLOCK_N)
+        else:
+            k_offs_n = None
+        k_offs_k = None if not PADDED_HEAD else tl.arange(0, BLOCK_DMODEL)
+        k = load_fn(k_ptrs, k_offs_k, k_offs_n, ACTUAL_BLOCK_DMODEL, actual_seqlen_k)
+        if PRE_LOAD_V:
+            # We can use the same offsets as k, just with dims transposed.
+            v = load_fn(v_ptrs, k_offs_n, k_offs_k, actual_seqlen_k, ACTUAL_BLOCK_DMODEL)
+        qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
+        # We start from end of seqlen_k so only the first iteration would need
+        # to be checked for padding if it is not a multiple of block_n
+        # TODO: This can be optimized to only be true for the padded block.
+        if MASK_STEPS:
+            # If this is the last block / iteration, we want to
+            # mask if the sequence length is not a multiple of block size
+            # a solution is to always do BLOCK_M // BLOCK_N + 1 steps if not is_modulo_mn.
+            # last step might get wasted but that is okay. check if this masking works For
+            # that case.
+            if (start_n + BLOCK_N == block_max) and (n_extra_tokens != 0):
+                boundary_m = tl.full([BLOCK_M], actual_seqlen_k, dtype=tl.int32)
+                size_n = start_n + OFFS_N[None, :]
+                mask = size_n < boundary_m[:, None]
+                qk = tl.where(mask, qk, float("-inf"))
+
+        # -- compute qk ----
+        qk += tl.dot(q, k)
+        qk_scaled =  qk * SM_SCALE
+        if RETURN_SCORES:
+            score_mask = (OFFS_M[:, None] < actual_seqlen_q) & ((start_n + tl.arange(0, BLOCK_N))[None, :] < actual_seqlen_k)
+            tl.store(score_ptrs, qk_scaled, mask=score_mask)
+
+        if IS_CAUSAL:
+            causal_boundary = start_n + offs_n_causal
+            causal_mask = OFFS_M[:, None] >= causal_boundary[None, :]
+            qk_scaled = tl.where(causal_mask, qk_scaled, float("-inf"))
+        if bias_ptrs is not None:
+            bias_offs_n = start_n + tl.arange(0, BLOCK_N) if MASK_STEPS else None
+            bias = load_fn(bias_ptrs, OFFS_M, bias_offs_n, actual_seqlen_q, actual_seqlen_k)
+            qk_scaled += bias
+
+        if alibi_slope is not None:
+            # Compute the global position of each token within the sequence
+            global_m_positions = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
+            global_n_positions = start_n + tl.arange(0, BLOCK_N)
+            alibi_block = compute_alibi_block(alibi_slope, actual_seqlen_q, actual_seqlen_k, global_m_positions,
+                                              global_n_positions)
+            qk_scaled += alibi_block
+        # get max scores so far
+        m_ij = tl.maximum(m_i, tl.max(qk_scaled, 1))
+
+        # scale and subtract max
+        q_shifted = qk_scaled - m_ij[:, None]
+        if RETURN_SCORES:
+            # NOTE: the returned score is not the same as the reference because we need to adjust as we find new maxes per block. We are not doing that
+            scores_scaled_shifted_mask = (OFFS_M[:, None] < actual_seqlen_q) & ((start_n + tl.arange(0, BLOCK_N))[None, :] < actual_seqlen_k)
+            tl.store(scores_scaled_shifted_ptrs, q_shifted, mask=scores_scaled_shifted_mask)
+
+        # Compute scaled QK and softmax probabilities
+        if USE_EXP2:
+            p = tl.math.exp2(q_shifted * RCP_LN2)
+        else:
+            p = tl.math.exp(q_shifted)
+
+        # CAVEAT: Must update l_ij before applying dropout
+        l_ij = tl.sum(p, 1)
+        if ENABLE_DROPOUT:
+            philox_offset = batch_philox_offset + start_m * BLOCK_M * actual_seqlen_k + start_n - BLOCK_N
+            keep = dropout_mask(philox_seed, philox_offset, dropout_p, BLOCK_M, BLOCK_N, actual_seqlen_k)
+            if RETURN_SCORES:
+                # NOTE: the returned score is not the same as the reference because we need to adjust as we find new maxes per block. We are not doing that
+                exp_score_mask = (OFFS_M[:, None] < actual_seqlen_q) & ((start_n + tl.arange(0, BLOCK_N))[None, :] < actual_seqlen_k)
+                tl.store(exp_scores_ptrs, tl.where(keep, p, -p), mask=exp_score_mask)
+            p = tl.where(keep, p, 0.0)
+        elif RETURN_SCORES:
+            # NOTE: the returned score is not the same as the reference because we need to adjust as we find new maxes per block. We are not doing that
+            exp_score_mask = (OFFS_M[:, None] < actual_seqlen_q) & ((start_n + tl.arange(0, BLOCK_N))[None, :] < actual_seqlen_k)
+            tl.store(exp_scores_ptrs, p, mask=exp_score_mask)
+
+        # -- update output accumulator --
+        # alpha is an adjustment factor for acc and li as we loop and find new maxes
+        # store the diff in maxes to adjust acc and li as we discover new maxes
+        m_diff = m_i - m_ij
+        if USE_EXP2:
+            alpha = tl.math.exp2(m_diff * RCP_LN2)
+        else:
+            alpha = tl.math.exp(m_diff)
+        acc = acc * alpha[:, None]
+        v = None
+        if not PRE_LOAD_V:
+            v = load_fn(v_ptrs, k_offs_n, k_offs_k, actual_seqlen_k, ACTUAL_BLOCK_DMODEL)
+        # -- update m_i and l_i
+        l_i = l_i * alpha + l_ij
+        # update m_i and l_i
+        m_i = m_ij
+        acc += tl.dot(p.to(v.type.element_ty), v)
+        k_ptrs += BLOCK_N * stride_kn
+        v_ptrs += BLOCK_N * stride_vk
+        if bias_ptrs is not None:
+            bias_ptrs += BLOCK_N * stride_bn
+        if RETURN_SCORES:
+            score_ptrs += BLOCK_N
+            scores_scaled_shifted_ptrs += BLOCK_N
+            exp_scores_ptrs += BLOCK_N
+    return acc, l_i, m_i
+
+
+def get_cdna_autotune_configs():
+    return [
+        triton.Config({'BLOCK_M': 128, 'BLOCK_N': 128, 'waves_per_eu': 2, 'PRE_LOAD_V': False}, num_stages=1,
+                      num_warps=4),
+        triton.Config({'BLOCK_M': 128, 'BLOCK_N': 64, 'waves_per_eu': 2, 'PRE_LOAD_V': False}, num_stages=1,
+                      num_warps=4),
+        triton.Config({'BLOCK_M': 128, 'BLOCK_N': 64, 'waves_per_eu': 3, 'PRE_LOAD_V': False}, num_stages=1,
+                      num_warps=4),
+        triton.Config({'BLOCK_M': 128, 'BLOCK_N': 64, 'waves_per_eu': 1, 'PRE_LOAD_V': False}, num_stages=1,
+                      num_warps=4),
+        triton.Config({'BLOCK_M': 128, 'BLOCK_N': 32, 'waves_per_eu': 2, 'PRE_LOAD_V': False}, num_stages=1,
+                      num_warps=4),
+        triton.Config({'BLOCK_M': 64, 'BLOCK_N': 64, 'waves_per_eu': 1, 'PRE_LOAD_V': False}, num_stages=1,
+                      num_warps=4),
+        # Fall-back config.
+        triton.Config({'BLOCK_M': 16, 'BLOCK_N': 16, 'waves_per_eu': 1, 'PRE_LOAD_V': False}, num_stages=1,
+                      num_warps=4),
+    ], ['IS_CAUSAL', 'dropout_p', 'MAX_SEQLENS_Q', 'MAX_SEQLENS_K', 'ACTUAL_BLOCK_DMODEL', 'VARLEN', 'HQ', 'HK']
+
+
+def get_rdna_autotune_configs():
+    return [
+        triton.Config({'BLOCK_M': 32, 'BLOCK_N': 32, 'waves_per_eu': 4, 'PRE_LOAD_V': False}, num_stages=1,
+                      num_warps=2),
+        triton.Config({'BLOCK_M': 32, 'BLOCK_N': 32, 'waves_per_eu': 2, 'PRE_LOAD_V': False}, num_stages=1,
+                      num_warps=2),
+        triton.Config({'BLOCK_M': 32, 'BLOCK_N': 16, 'waves_per_eu': 4, 'PRE_LOAD_V': False}, num_stages=1,
+                      num_warps=2),
+        triton.Config({'BLOCK_M': 32, 'BLOCK_N': 16, 'waves_per_eu': 2, 'PRE_LOAD_V': False}, num_stages=1,
+                      num_warps=2),
+        triton.Config({'BLOCK_M': 16, 'BLOCK_N': 16, 'waves_per_eu': 4, 'PRE_LOAD_V': False}, num_stages=1,
+                      num_warps=2),
+        triton.Config({'BLOCK_M': 16, 'BLOCK_N': 16, 'waves_per_eu': 2, 'PRE_LOAD_V': False}, num_stages=1,
+                      num_warps=2),
+        # Fall-back config.
+        triton.Config({'BLOCK_M': 16, 'BLOCK_N': 16, 'waves_per_eu': 1, 'PRE_LOAD_V': False}, num_stages=1,
+                      num_warps=2),
+    ], ['IS_CAUSAL', 'dropout_p', 'MAX_SEQLENS_Q', 'MAX_SEQLENS_K', 'ACTUAL_BLOCK_DMODEL', 'VARLEN', 'HQ', 'HK']
+
+
+def get_autotune_configs():
+    if AUTOTUNE:
+        if is_rdna():
+            return get_rdna_autotune_configs()
+        elif is_cdna():
+            return get_cdna_autotune_configs()
+        else:
+            raise ValueError("Unknown Device Type")
+    else:
+        return [
+            triton.Config(
+                {"BLOCK_M": 64, "BLOCK_N": 64, "waves_per_eu": 1, "PRE_LOAD_V": False},
+                num_stages=1,
+                num_warps=4,
+            ),
+        ], [
+            "IS_CAUSAL",
+            "dropout_p",
+            "MAX_SEQLENS_Q",
+            "MAX_SEQLENS_K",
+            "ACTUAL_BLOCK_DMODEL",
+            "VARLEN",
+            "HQ",
+            "HK",
+        ]
+
+
+autotune_configs, autotune_keys = get_autotune_configs()
+
+@triton.autotune(
+    configs=autotune_configs,
+    key=autotune_keys,
+    # use_cuda_graph=True,
+)
+@triton.jit
+def attn_fwd(Q, K, V, bias, SM_SCALE: tl.constexpr, LSE, Out, stride_qz, stride_qh, stride_qm, stride_qk,
+             stride_kz, stride_kh, stride_kn, stride_kk, stride_vz, stride_vh, stride_vk, stride_vn,
+             stride_oz, stride_oh, stride_om, stride_on, stride_bz, stride_bh, stride_bm, stride_bn, stride_az, stride_ah, # pylint: disable=unused-argument
+             stride_sz, stride_sh, stride_sm, stride_sn, stride_lse_z, stride_lse_h, stride_lse_m, cu_seqlens_q, cu_seqlens_k,
+             dropout_p, philox_seed, philox_offset_base, scores, scores_scaled_shifted, exp_scores, alibi_slopes,  HQ: tl.constexpr,
+             HK: tl.constexpr, ACTUAL_BLOCK_DMODEL: tl.constexpr, MAX_SEQLENS_Q: tl.constexpr,
+             MAX_SEQLENS_K: tl.constexpr, VARLEN: tl.constexpr, IS_CAUSAL: tl.constexpr, BLOCK_M: tl.constexpr,
+             BLOCK_DMODEL: tl.constexpr, BLOCK_N: tl.constexpr, PRE_LOAD_V: tl.constexpr, USE_BIAS: tl.constexpr,
+             ENABLE_DROPOUT: tl.constexpr, RETURN_SCORES: tl.constexpr, USE_ALIBI: tl.constexpr, USE_EXP2: tl.constexpr):
+    start_m = tl.program_id(0)
+    off_h_q = tl.program_id(1)
+    off_z = tl.program_id(2)
+    offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    offs_n = tl.arange(0, BLOCK_N)
+    offs_d = tl.arange(0, BLOCK_DMODEL)
+    if VARLEN:
+        cu_seqlens_q_start = tl.load(cu_seqlens_q + off_z)
+        cu_seqlens_q_end = tl.load(cu_seqlens_q + off_z + 1)
+        # print("cu_seqlens_q_start:", cu_seqlens_q_start)
+
+        seqlen_q = cu_seqlens_q_end - cu_seqlens_q_start
+        # We have a one-size-fits-all grid in id(0). Some seqlens might be too
+        # small for all start_m so for those we return early.
+        if start_m * BLOCK_M > seqlen_q:
+            return
+        cu_seqlens_k_start = tl.load(cu_seqlens_k + off_z)
+        cu_seqlens_k_end = tl.load(cu_seqlens_k + off_z + 1)
+        seqlen_k = cu_seqlens_k_end - cu_seqlens_k_start
+    else:
+        cu_seqlens_q_start = 0
+        cu_seqlens_k_start = 0
+        seqlen_q = MAX_SEQLENS_Q
+        seqlen_k = MAX_SEQLENS_K
+
+    # Now we compute whether we need to exit early due to causal masking.
+    # This is because for seqlen_q > seqlen_k, M rows of the attn scores
+    # are completely masked, resulting in 0s written to the output, and
+    # inf written to LSE. We don't need to do any GEMMs in this case.
+    # This block of code determines what N is, and if this WG is operating
+    # on those M rows.
+    n_blocks = cdiv_fn(seqlen_k, BLOCK_N)
+    if IS_CAUSAL:
+        # If seqlen_q == seqlen_k, the attn scores are a square matrix.
+        # If seqlen_q != seqlen_k, attn scores are rectangular which means
+        # the causal mask boundary is bottom right aligned, and ends at either
+        # the top edge (seqlen_q < seqlen_k) or left edge.
+        # This captures the decrease in n_blocks if we have a rectangular attn matrix
+        n_blocks_seqlen = cdiv_fn((start_m + 1) * BLOCK_M + seqlen_k - seqlen_q, BLOCK_N)
+        # This is what adjusts the block_max for the current WG, only
+        # if IS_CAUSAL. Otherwise we want to always iterate through all n_blocks
+        n_blocks = min(n_blocks, n_blocks_seqlen)
+        # If we have no blocks after adjusting for seqlen deltas, this WG is part of
+        # the blocks that are all 0. We exit early.
+        if n_blocks <= 0:
+            o_offset = Out + off_z * stride_oz + off_h_q * stride_oh + cu_seqlens_q_start * stride_om
+            o_ptrs = o_offset + offs_m[:, None] * stride_om + offs_d[None, :] * stride_on
+            acc = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=Out.type.element_ty)
+            o_ptrs_mask = offs_m[:, None] < seqlen_q
+            # We still need to write 0s to the result
+            tl.store(o_ptrs, acc, mask=o_ptrs_mask)
+            # The tensor allocated for L is based on MAX_SEQLENS_Q as that is
+            # statically known.
+            l_offset = LSE + off_z * stride_lse_z + off_h_q * stride_lse_h + cu_seqlens_q_start * stride_lse_m
+            l_ptrs = l_offset + offs_m * stride_lse_m 
+
+            l = tl.full([BLOCK_M], value=0.0, dtype=tl.float32)
+
+            # mask_m_offsets = start_m + tl.arange(0, BLOCK_M)
+            # lse_mask = mask_m_offsets < causal_start_idx
+            # softmax_lse = tl.where(lse_mask, 0.0, softmax_lse)
+            l_ptrs_mask = offs_m < MAX_SEQLENS_Q
+            tl.store(l_ptrs, l, mask=l_ptrs_mask)
+            # TODO: Should dropout and return encoded softmax be handled here too?
+            return
+
+    # If MQA / GQA, set the K and V head offsets appropriately.
+    GROUP_SIZE: tl.constexpr = HQ // HK
+    if GROUP_SIZE != 1:
+        off_h_k = off_h_q // GROUP_SIZE
+    else:
+        off_h_k = off_h_q
+
+    n_extra_tokens = 0
+    # print("n_extra_tokens:", n_extra_tokens)
+    # print("seqlen_k:", seqlen_k)
+    # print("BLOCK_N:", BLOCK_N)
+    # return
+    if seqlen_k < BLOCK_N:
+        n_extra_tokens = BLOCK_N - seqlen_k
+    elif seqlen_k % BLOCK_N:
+        n_extra_tokens = seqlen_k % BLOCK_N
+    PADDED_HEAD: tl.constexpr = ACTUAL_BLOCK_DMODEL != BLOCK_DMODEL
+
+    # Compute pointers for all the tensors used in this kernel.
+    q_offset = Q + off_z * stride_qz + off_h_q * stride_qh + cu_seqlens_q_start * stride_qm
+    q_ptrs = q_offset + offs_m[:, None] * stride_qm + offs_d[None, :] * stride_qk
+    k_offset = K + off_z * stride_kz + off_h_k * stride_kh + cu_seqlens_k_start * stride_kn
+    k_ptrs = k_offset + offs_d[:, None] * stride_kk + offs_n[None, :] * stride_kn
+    v_offset = V + off_z * stride_vz + off_h_k * stride_vh + cu_seqlens_k_start * stride_vk
+    v_ptrs = v_offset + offs_n[:, None] * stride_vk + offs_d[None, :] * stride_vn
+    if USE_BIAS:
+        # Note: this might get large enough to overflow on some configs
+        bias_offset = off_h_q * stride_bh
+        bias_ptrs = bias + bias_offset + offs_m[:, None] * stride_bm + offs_n[None, :] * stride_bn
+    else:
+        bias_ptrs = None
+
+    if USE_ALIBI:
+        a_offset = off_z * stride_az + off_h_q * stride_ah
+        alibi_slope = tl.load(alibi_slopes + a_offset)
+    else:
+        alibi_slope = None
+
+    if RETURN_SCORES:
+        scores_offset = scores + off_z * stride_sz + off_h_q * stride_sh + cu_seqlens_q_start * stride_sm
+        score_ptrs = scores_offset + offs_m[:, None] * stride_sm + offs_n[None, :] * stride_sn
+
+        scores_scaled_shifted_offset = scores_scaled_shifted + off_z * stride_sz + off_h_q * stride_sh + cu_seqlens_q_start * stride_sm
+        scores_scaled_shifted_ptrs = scores_scaled_shifted_offset + offs_m[:, None] * stride_sm + offs_n[None, :] * stride_sn
+
+        exp_scores_offset = exp_scores + off_z * stride_sz + off_h_q * stride_sh + cu_seqlens_q_start * stride_sm
+        exp_scores_ptrs = exp_scores_offset + offs_m[:, None] * stride_sm + offs_n[None, :] * stride_sn
+    else:
+        score_ptrs = None
+        scores_scaled_shifted_ptrs = None
+        exp_scores_ptrs = None
+
+    if ENABLE_DROPOUT:
+        off_hz = off_z * HQ + off_h_q
+        batch_philox_offset = philox_offset_base + off_hz * seqlen_q * seqlen_k
+    else:
+        batch_philox_offset = 0
+    # initialize pointer to m and l
+    m_i = tl.full([BLOCK_M], float("-inf"), dtype=tl.float32)
+    l_i = tl.full([BLOCK_M], 1.0, dtype=tl.float32)
+    acc = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32)
+    # Q is loaded once at the beginning and shared by all N blocks.
+    q_ptrs_mask = offs_m[:, None] < seqlen_q
+    if PADDED_HEAD:
+        q_ptrs_mask = q_ptrs_mask & (offs_d[None, :] < ACTUAL_BLOCK_DMODEL)
+    q = tl.load(q_ptrs, mask=q_ptrs_mask, other=0.0)
+
+    # Here we compute how many full and masked blocks we have.
+    padded_block_k = n_extra_tokens != 0
+    is_modulo_mn = not padded_block_k and (seqlen_q % BLOCK_M == 0)
+    if IS_CAUSAL:
+        # There are always at least BLOCK_M // BLOCK_N masked blocks.
+        # Additionally there might be one more due to dissimilar seqlens.
+        masked_blocks = BLOCK_M // BLOCK_N + (not is_modulo_mn)
+    else:
+        # Padding on Q does not need to be masked in the FA loop.
+        masked_blocks = padded_block_k
+    # if IS_CAUSAL, not is_modulo_mn does not always result in an additional block.
+    # In this case we might exceed n_blocks so pick the min.
+    masked_blocks = min(masked_blocks, n_blocks)
+    n_full_blocks = n_blocks - masked_blocks
+    block_min = 0
+    block_max = n_blocks * BLOCK_N
+    # Compute for full blocks. Here we set causal to false regardless of its actual
+    # value because there is no masking. Similarly we do not need padding.
+    if n_full_blocks > 0:
+        block_max = (n_blocks - masked_blocks) * BLOCK_N
+        acc, l_i, m_i = _attn_fwd_inner(acc, l_i, m_i, q, k_ptrs, v_ptrs, bias_ptrs, stride_kn, stride_vk, stride_bn,
+                                        start_m, seqlen_k, seqlen_q, dropout_p, philox_seed, batch_philox_offset,
+                                        exp_scores_ptrs,
+                                        # _, _, offs_n_causal, masked_blocks, n_extra_tokens, _
+                                        block_min, block_max, 0, 0, 0, alibi_slope, score_ptrs, scores_scaled_shifted_ptrs,
+                                        # IS_CAUSAL, ....
+                                        False, BLOCK_M, BLOCK_DMODEL, BLOCK_N, offs_m, offs_n,
+                                        # _, MASK_STEPS, ...
+                                        PRE_LOAD_V, False, ENABLE_DROPOUT, PADDED_HEAD,
+                                        ACTUAL_BLOCK_DMODEL, SM_SCALE,  USE_EXP2=USE_EXP2, RETURN_SCORES=RETURN_SCORES)
+        block_min = block_max
+        block_max = n_blocks * BLOCK_N
+
+    tl.debug_barrier()
+    # Remaining blocks, if any, are full / not masked.
+    if masked_blocks > 0:
+        if IS_CAUSAL:
+            offs_n_causal = offs_n + (seqlen_q - seqlen_k)
+        else:
+            offs_n_causal = 0
+        k_ptrs += n_full_blocks * BLOCK_N * stride_kn
+        v_ptrs += n_full_blocks * BLOCK_N * stride_vk
+        if USE_BIAS:
+            bias_ptrs += n_full_blocks * BLOCK_N * stride_bn
+        if RETURN_SCORES:
+            score_ptrs += n_full_blocks * BLOCK_N
+            scores_scaled_shifted_ptrs += n_full_blocks * BLOCK_N
+            exp_scores_ptrs += n_full_blocks * BLOCK_N
+        acc, l_i, m_i = _attn_fwd_inner(acc, l_i, m_i, q, k_ptrs, v_ptrs, bias_ptrs, stride_kn, stride_vk, stride_bn,
+                                        start_m, seqlen_k, seqlen_q, dropout_p, philox_seed, batch_philox_offset,
+                                        exp_scores_ptrs, block_min, block_max, offs_n_causal, masked_blocks,
+                                        n_extra_tokens, alibi_slope, score_ptrs, scores_scaled_shifted_ptrs,
+                                        IS_CAUSAL, BLOCK_M, BLOCK_DMODEL, BLOCK_N, offs_m, offs_n,
+                                        # _, MASK_STEPS, ...
+                                        PRE_LOAD_V, True, ENABLE_DROPOUT, PADDED_HEAD,
+                                        ACTUAL_BLOCK_DMODEL, SM_SCALE, USE_EXP2=USE_EXP2, RETURN_SCORES=RETURN_SCORES)
+    # epilogue
+    # This helps the compiler do Newton Raphson on l_i vs on acc which is much larger.
+    l_recip = 1 / l_i[:, None]
+    acc = acc * l_recip
+    if ENABLE_DROPOUT:
+        acc = acc / (1 - dropout_p)
+    # If seqlen_q > seqlen_k but the delta is not a multiple of BLOCK_M,
+    # then we have one block with a row of all NaNs which come from computing
+    # softmax over a row of all -infs (-inf - inf = NaN). We check for that here
+    # and store 0s where there are NaNs as these rows should've been zeroed out.
+    end_m_idx = (start_m + 1) * BLOCK_M
+    start_m_idx = start_m * BLOCK_M
+    causal_start_idx = seqlen_q - seqlen_k
+    acc = acc.to(Out.type.element_ty)
+    if IS_CAUSAL:
+        if causal_start_idx > start_m_idx and causal_start_idx < end_m_idx:
+            out_mask_boundary = tl.full((BLOCK_DMODEL, ), causal_start_idx, dtype=tl.int32)
+            mask_m_offsets = start_m_idx + tl.arange(0, BLOCK_M)
+            out_ptrs_mask = mask_m_offsets[:, None] >= out_mask_boundary[None, :]
+            z: tl.tensor = 0.0
+            acc = tl.where(out_ptrs_mask, acc, z.to(acc.type.element_ty))
+
+    # write back LSE(Log Sum Exponents), the log of the normalization constant
+    l_offset = LSE + off_z * stride_lse_z + off_h_q * stride_lse_h + cu_seqlens_q_start * stride_lse_m
+    l_ptrs = l_offset + offs_m * stride_lse_m 
+    if USE_EXP2:
+        RCP_LN2: tl.constexpr = 1.4426950408889634
+        LN2: tl.constexpr = 0.6931471824645996
+        # compute log-sum-exp in base 2 units
+        mi_base2 = m_i * RCP_LN2
+        softmax_lse = mi_base2 + tl.math.log2(l_i)
+        # convert back to natural units
+        softmax_lse *= LN2
+    else:
+        softmax_lse = m_i + tl.math.log(l_i)
+
+    if IS_CAUSAL:
+        # zero out nans caused by -infs when doing causal
+        lse_mask = (start_m_idx + tl.arange(0, BLOCK_M)) < causal_start_idx
+        softmax_lse = tl.where(lse_mask, 0.0, softmax_lse)
+
+    # If seqlen_q not multiple of BLOCK_M, we need to mask out the last few rows.
+    # This is only true for the last M block. For others, overflow_size will be -ve
+    overflow_size = end_m_idx - seqlen_q
+    if overflow_size > 0:
+        boundary = tl.full((BLOCK_M, ), BLOCK_M - overflow_size, dtype=tl.int32)
+        l_ptrs_mask = tl.arange(0, BLOCK_M) < boundary
+        tl.store(l_ptrs, softmax_lse, mask=l_ptrs_mask) # the log of the normalization constant
+    else:
+        tl.store(l_ptrs, softmax_lse) # the log of the normalization constant
+
+    # write back O
+    o_offset = Out + off_z * stride_oz + off_h_q * stride_oh + cu_seqlens_q_start * stride_om
+    o_ptrs = o_offset + offs_m[:, None] * stride_om + offs_d[None, :] * stride_on
+    o_ptrs_mask = tl.full([BLOCK_M, BLOCK_DMODEL], 1, dtype=tl.int1)
+    if overflow_size > 0:
+        o_ptrs_mask = o_ptrs_mask & (offs_m[:, None] < seqlen_q)
+    if PADDED_HEAD:
+        o_ptrs_mask = o_ptrs_mask & (offs_d[None, :] < ACTUAL_BLOCK_DMODEL)
+    tl.store(o_ptrs, acc.to(Out.dtype.element_ty), mask=o_ptrs_mask)
+
+
+def attention_prefill_forward_triton_impl(
+                                        q,
+                                        k,
+                                        v,
+                                        o,
+                                        sm_scale,
+                                        alibi_slopes,
+                                        causal,
+                                        bias,
+                                        dropout_p,
+                                        layout,
+                                        cu_seqlens_q,
+                                        cu_seqlens_k,
+                                        max_seqlens_q,
+                                        max_seqlens_k,
+                                        return_scores,
+                                        use_exp2):
+    # check if varlen
+    is_varlen = layout == "thd"
+
+    # NOTE: a large bias tensor leads to overflow during pointer arithmetic
+    if bias is not None:
+        assert bias.numel() < 2**31
+
+    batch, nheads_q, nheads_k, head_size, seqlen_q, seqlen_k = get_shape_from_layout(q, k, layout, cu_seqlens_q, cu_seqlens_k, max_seqlens_q, max_seqlens_k) # pylint: disable=unused-variable
+    q_strides, k_strides, v_strides, o_strides = get_strides_from_layout(q, k, v, o, layout)
+
+    # Get closest power of 2 over or equal to 32.
+    padded_d_model = 1 << (head_size - 1).bit_length()
+    # Smallest head_dim supported is 16. If smaller, the tile in the
+    # kernel is padded - there is no padding in memory for any dims.
+    padded_d_model = max(padded_d_model, 16)
+
+    grid = lambda META: (triton.cdiv(max_seqlens_q, META['BLOCK_M']), nheads_q, batch) # pylint: disable=unnecessary-lambda-assignment
+
+    if return_scores:
+        scores = torch.zeros((batch, nheads_q, max_seqlens_q, max_seqlens_k), device=q.device,
+                                        dtype=torch.float32)
+        scores_scaled_shifted = torch.zeros((batch, nheads_q, max_seqlens_q, max_seqlens_k), device=q.device,
+                                        dtype=torch.float32)
+        scores_strides = (scores.stride(0), scores.stride(1), scores.stride(2), scores.stride(3))
+    else:
+        scores = None
+        scores_scaled_shifted = None
+        scores_strides = (0, 0 , 0 , 0)
+
+    # exp_scores is used to validate dropout behavior vs the PyTorch SDPA math backend reference.  We zero this out
+    # to give a consistent starting point and then populate it with the output of softmax with the sign bit set according
+    # to the dropout mask. The resulting return allows this mask to be fed into the reference implementation for testing
+    # only.  This return holds no useful output aside from debugging.
+    if return_scores:
+        exp_scores = torch.zeros((batch, nheads_q, max_seqlens_q, max_seqlens_k), device=q.device,
+                                        dtype=torch.float32)
+    else:
+        exp_scores = None
+
+    # stores LSE the log of the normalization constant / sum of expoential score(unnormalzied probablities)
+    if is_varlen:
+        softmax_lse = torch.empty((q.shape[0], nheads_q), device=q.device, dtype=torch.float32)
+        stride_lse_m, stride_lse_h = softmax_lse.stride()
+        stride_lse_z = 0
+    else:
+        softmax_lse = torch.empty((batch, nheads_q, max_seqlens_q), device=q.device, dtype=torch.float32)
+        stride_lse_z, stride_lse_h, stride_lse_m = softmax_lse.stride()
+
+    # Seed the RNG so we get reproducible results for testing.
+    philox_seed = 0x1BF52
+    philox_offset = 0x1D4B42
+
+    if bias is not None:
+        bias_strides = (bias.stride(0), bias.stride(1),bias.stride(2),
+                        bias.stride(3))
+    else:
+        bias_strides = (0, 0, 0, 0)
+
+    if alibi_slopes is not None:
+        alibi_strides = (alibi_slopes.stride(0), alibi_slopes.stride(1))
+    else:
+        alibi_strides = (0, 0)
+
+
+    attn_fwd[grid](q, k, v, bias, sm_scale, softmax_lse, o, *q_strides, *k_strides, *v_strides, *o_strides,
+                    *bias_strides, *alibi_strides, *scores_strides, stride_lse_z, stride_lse_h, stride_lse_m, cu_seqlens_q, cu_seqlens_k,
+                    dropout_p=dropout_p, philox_seed=philox_seed, philox_offset_base=philox_offset, scores=scores,
+                    scores_scaled_shifted=scores_scaled_shifted, exp_scores=exp_scores, alibi_slopes=alibi_slopes,
+                    HQ=nheads_q, HK=nheads_k, ACTUAL_BLOCK_DMODEL=head_size, MAX_SEQLENS_Q=max_seqlens_q,
+                    MAX_SEQLENS_K=max_seqlens_k, IS_CAUSAL=causal, VARLEN=is_varlen,
+                    BLOCK_DMODEL=padded_d_model, USE_BIAS=False if bias is None else True,
+                    USE_ALIBI=False if alibi_slopes is None else True, ENABLE_DROPOUT=dropout_p
+                    > 0.0, USE_EXP2=use_exp2, RETURN_SCORES=return_scores)
+
+    return o, softmax_lse, exp_scores, grid, head_size, philox_seed, philox_offset, scores, scores_scaled_shifted
diff --git a/modules/flash_attn_triton_amd/fwd_ref.py b/modules/flash_attn_triton_amd/fwd_ref.py
new file mode 100644
index 000000000..03e53efde
--- /dev/null
+++ b/modules/flash_attn_triton_amd/fwd_ref.py
@@ -0,0 +1,258 @@
+import math
+import torch
+
+
+def attention_forward_core_ref_impl(q, k, v, sm_scale, causal, use_exp2):
+    # Compute attention scores
+    attention_scores = torch.matmul(q.to(torch.float32), k.transpose(-2, -1).to(torch.float32))
+
+    # Scale scores
+    attention_scaled_scores = sm_scale * attention_scores
+
+    # Apply causal mask if necessary
+    if causal:
+        L_q, L_k = q.shape[1], k.shape[1]
+        row_idx = torch.arange(L_q, device=q.device).unsqueeze(1)
+        col_idx = torch.arange(L_k, device=q.device).unsqueeze(0)
+        col_offset = L_q-L_k
+        causal_mask = row_idx >= (col_offset + col_idx)
+        # set -inf to places the causal mask is false
+        attention_scaled_scores = attention_scaled_scores.masked_fill(
+             torch.logical_not(causal_mask.unsqueeze(0)), float('-inf')
+        )
+
+
+    # Compute max for numerical stability
+    max_scores = torch.max(attention_scaled_scores, dim=-1, keepdim=True)[0]
+    if causal:
+        # Replace -inf in max_scores with zeros to avoid NaN in subtraction
+        max_scores = torch.where(
+            torch.isinf(max_scores), torch.zeros_like(max_scores), max_scores
+        )
+
+    # Shift scores
+    attention_shifted_scaled_scores = attention_scaled_scores - max_scores
+
+    # Exponentiate
+    if use_exp2:
+        RCP_LN = 1 / math.log(2)
+        exp_scores = torch.exp2(RCP_LN * attention_shifted_scaled_scores)
+    else:
+        exp_scores = torch.exp(attention_shifted_scaled_scores)
+
+    # Sum of exponentials
+    sum_exp_scores = torch.sum(exp_scores, dim=-1, keepdim=True)
+    if causal:
+        # if sum of exp scores is 0.0 it means scores where -inf, we cannot compute softmax and softmax_lse. Setting to 1 deals with -inf case cleanly
+        sum_exp_scores = torch.where(
+        sum_exp_scores == 0,
+        torch.ones_like(sum_exp_scores),
+        sum_exp_scores
+        )
+
+    # Compute softmax probabilities
+    softmax = exp_scores / sum_exp_scores
+
+    # Compute log-sum-exp
+    if use_exp2:
+        LN2 = math.log(2)
+        RCP_LN = 1 / math.log(2)
+        max_scores_base2 = max_scores * RCP_LN
+        softmax_lse_base2 = max_scores_base2 + torch.log2(sum_exp_scores)
+        softmax_lse = softmax_lse_base2 * LN2
+        softmax_lse.squeeze_(-1)
+    else:
+        softmax_lse = max_scores + torch.log(sum_exp_scores)
+        softmax_lse = softmax_lse.squeeze(-1)
+
+    # Compute output
+    o = torch.matmul(softmax, v.to(torch.float32)).to(torch.float16)
+
+    return o, softmax_lse, exp_scores, softmax, attention_shifted_scaled_scores, attention_scaled_scores, attention_scores
+
+def attention_vanilla_forward_pytorch_ref_impl(q, k, v, sm_scale, causal, layout, use_exp2):
+    """Compute reference output and softmax_lse using PyTorch's built-in function"""
+
+    # Ensure the layout is 'bhsd'
+    if layout == "bshd":
+        q = q.transpose(1, 2).contiguous()
+        k = k.transpose(1, 2).contiguous()
+        v = v.transpose(1, 2).contiguous()
+    elif layout != "bhsd":
+        raise ValueError(f"Unknown layout {layout}")
+
+    # Prepare tensors in [batch_size * num_heads, seq_len, head_dim] format
+    batch_size, num_heads, seq_len_q, head_dim = q.shape
+    seq_len_k = k.shape[2]
+
+    # Merge batch and heads dimensions
+    q = q.reshape(batch_size * num_heads, seq_len_q, head_dim)
+    k = k.reshape(batch_size * num_heads, seq_len_k, head_dim)
+    v = v.reshape(batch_size * num_heads, seq_len_k, head_dim)
+
+    # Call the core attention function
+    o, softmax_lse, exp_scores, softmax, attention_shifted_scaled_scores, attention_scaled_scores, attention_scores = attention_forward_core_ref_impl(
+        q, k, v, sm_scale, causal, use_exp2
+    )
+
+    # Reshape outputs back to [batch_size, num_heads, seq_len, head_dim]
+    o = o.reshape(batch_size, num_heads, seq_len_q, head_dim)
+    softmax_lse = softmax_lse.reshape(batch_size, num_heads, seq_len_q)
+    exp_scores = exp_scores.reshape(batch_size, num_heads, seq_len_q, seq_len_k)
+    softmax = softmax.reshape(batch_size, num_heads, seq_len_q, seq_len_k)
+    attention_shifted_scaled_scores = attention_shifted_scaled_scores.reshape(batch_size, num_heads, seq_len_q, seq_len_k)
+    attention_scaled_scores = attention_scaled_scores.reshape(batch_size, num_heads, seq_len_q, seq_len_k)
+    attention_scores = attention_scores.reshape(batch_size, num_heads, seq_len_q, seq_len_k)
+
+    # Restore original layout if necessary
+    if layout == "bshd":
+        o = o.transpose(1, 2)
+
+    return o, softmax_lse, exp_scores, softmax, attention_shifted_scaled_scores, attention_scaled_scores, attention_scores
+
+def attention_varlen_forward_pytorch_ref_impl(
+    q,
+    k,
+    v,
+    sm_scale,
+    causal,
+    layout,
+    cu_seqlens_q,
+    cu_seqlens_k,
+    max_seqlen_q, max_seqlen_k, # pylint: disable=unused-argument
+    use_exp2
+):
+    # Ensure the layout is 'thd'
+    if layout != 'thd':
+        raise ValueError(f"Unsupported layout {layout}. Expected 'thd'.")
+
+    batch_size = cu_seqlens_q.shape[0] - 1
+    num_heads = q.shape[1]
+    head_dim = q.shape[2]
+
+    # Pre-allocate outputs
+    total_L_q = q.shape[0]
+    total_L_k = k.shape[0] # pylint: disable=unused-variable
+
+    o = torch.empty((total_L_q, num_heads, head_dim), dtype=q.dtype, device=q.device)
+    softmax_lse = torch.empty((total_L_q, num_heads), dtype=torch.float32, device=q.device)
+
+    for i in range(batch_size):
+        # Get the start and end indices for the current sequence
+        start_q = cu_seqlens_q[i].item()
+        end_q = cu_seqlens_q[i + 1].item()
+        start_k = cu_seqlens_k[i].item()
+        end_k = cu_seqlens_k[i + 1].item()
+
+        # Extract q_i, k_i, v_i
+        q_i = q[start_q:end_q, :, :]  # [L_q_i, num_heads, head_dim]
+        k_i = k[start_k:end_k, :, :]  # [L_k_i, num_heads, head_dim]
+        v_i = v[start_k:end_k, :, :]  # [L_k_i, num_heads, head_dim]
+
+        # Permute to [num_heads, L_q_i, head_dim]
+        q_i = q_i.permute(1, 0, 2)
+        k_i = k_i.permute(1, 0, 2)
+        v_i = v_i.permute(1, 0, 2)
+
+        # Call the core attention function for this sequence
+        (
+            o_i,
+            softmax_lse_i,
+            exp_scores_i,
+            softmax_i,
+            attention_shifted_scaled_scores_i,
+            attention_scaled_scores_i,
+            attention_scores_i,
+        ) = attention_forward_core_ref_impl(q_i, k_i, v_i, sm_scale, causal, use_exp2)
+
+        # Convert back to 'thd' layout and float16
+        o_i = o_i.permute(1, 0, 2).to(torch.float16)  # [L_q_i, num_heads, head_dim]
+
+        # Place outputs in pre-allocated tensors
+        o[start_q:end_q, :, :] = o_i
+        softmax_lse[start_q:end_q, :] = softmax_lse_i.transpose(0, 1)  # Transpose to [L_q_i, num_heads]
+
+        # For variable-sized outputs, map them into the preallocated tensors
+        # exp_scores_i: [num_heads, L_q_i, L_k_i] -> [L_q_i, num_heads, L_k_i]
+        exp_scores_i = exp_scores_i.permute(1, 0, 2)
+        softmax_i = softmax_i.permute(1, 0, 2)
+        attention_shifted_scaled_scores_i = attention_shifted_scaled_scores_i.permute(1, 0, 2)
+        attention_scaled_scores_i = attention_scaled_scores_i.permute(1, 0, 2)
+        attention_scores_i = attention_scores_i.permute(1, 0, 2)
+
+    return (
+        o,
+        softmax_lse,
+        None,
+        None,
+        None,
+        None,
+        None,
+    )
+
+
+def attention_forward_pytorch_ref_impl(
+    q,
+    k,
+    v,
+    sm_scale,
+    causal,
+    layout,
+    cu_seqlens_q,
+    cu_seqlens_k,
+    max_seqlen_q,
+    max_seqlen_k,
+    use_exp2
+    ):
+     # compute reference
+    if layout == "thd":
+        (
+            o_ref,
+            softmax_lse_ref,
+            exp_scores_ref,
+            softmax_ref,
+            attention_shifted_scaled_scores_ref,
+            attention_scaled_scores_ref,
+            attention_scores_ref,
+        ) = attention_varlen_forward_pytorch_ref_impl(
+            q.clone(),
+            k.clone(),
+            v.clone(),
+            sm_scale,
+            causal,
+            layout,
+            cu_seqlens_q,
+            cu_seqlens_k,
+            max_seqlen_q,
+            max_seqlen_k,
+            use_exp2,
+        )
+    else:
+        (
+            o_ref,
+            softmax_lse_ref,
+            exp_scores_ref,
+            softmax_ref,
+            attention_shifted_scaled_scores_ref,
+            attention_scaled_scores_ref,
+            attention_scores_ref,
+        ) = attention_vanilla_forward_pytorch_ref_impl(
+            q.clone(), k.clone(), v.clone(), sm_scale, causal, layout, use_exp2
+        )
+
+    return (
+            o_ref,
+            softmax_lse_ref,
+            exp_scores_ref,
+            softmax_ref,
+            attention_shifted_scaled_scores_ref,
+            attention_scaled_scores_ref,
+            attention_scores_ref,
+    )
+
+
+def compute_alibi_tensor_ref(alibi_slopes, seqlen_q, seqlen_k):
+    q_idx = torch.arange(seqlen_q, dtype=torch.int32, device="cuda").unsqueeze(-1)  # (N_CTX_Q, 1)
+    k_idx = torch.arange(seqlen_k, dtype=torch.int32, device="cuda").unsqueeze(0)  # (1, N_CTX_K)
+    relative_pos = torch.abs(q_idx + seqlen_k - seqlen_q - k_idx)  # (N_CTX_Q, N_CTX_K)
+    return -1 * alibi_slopes.unsqueeze(-1).unsqueeze(-1) * relative_pos  # (Z, H, N_CTX_Q, N_CTX_K)
diff --git a/modules/flash_attn_triton_amd/interface_fa.py b/modules/flash_attn_triton_amd/interface_fa.py
new file mode 100644
index 000000000..72373d35f
--- /dev/null
+++ b/modules/flash_attn_triton_amd/interface_fa.py
@@ -0,0 +1,394 @@
+import os
+import torch
+from modules.flash_attn_triton_amd.fwd_prefill import attention_prefill_forward_triton_impl
+from modules.flash_attn_triton_amd.bwd_prefill import attention_prefill_backward_triton_impl
+from modules.flash_attn_triton_amd.fwd_decode import attention_decode_forward_triton_impl
+from modules.flash_attn_triton_amd.fwd_ref import attention_forward_pytorch_ref_impl
+from modules.flash_attn_triton_amd.bwd_ref import attention_backward_pytorch_ref_impl
+from modules.flash_attn_triton_amd.utils import MetaData, get_shape_from_layout
+
+
+USE_REF = os.environ.get('FLASH_ATTENTION_TRITON_AMD_REF', '0').lower() in ('1', 'true', 'yes')
+
+
+def fwd(q,
+    k,
+    v,
+    o,
+    alibi_slopes,
+    dropout_p,
+    softmax_scale,
+    causal,
+    window_size_left, window_size_right, softcap, # pylint: disable=unused-argument
+    return_softmax,
+    gen_ # pylint: disable=unused-argument
+):
+    if dropout_p != 0.0:
+        raise ValueError("dropout is not supported on AMD's Triton Backend yet")
+
+    if o is None:
+        o = torch.empty_like(q)
+
+    # Setup metadata
+    metadata = MetaData(sm_scale=softmax_scale)
+    metadata.max_seqlens_q = q.shape[1]
+    metadata.max_seqlens_k = k.shape[1]
+    metadata.layout = "bshd"
+    if return_softmax:
+        metadata.return_scores = True
+
+    batch, nheads_q, nheads_k, head_size, _, _ = get_shape_from_layout(q, k, metadata.layout) # pylint: disable=unused-variable
+
+    if causal:
+        metadata.need_causal()
+
+    if alibi_slopes is not None:
+        metadata.need_alibi(alibi_slopes, batch, nheads_q)
+
+    if dropout_p > 0.0:
+        metadata.need_dropout(dropout_p, return_softmax)
+
+    # Check arguments
+    metadata.check_args(q, k, v, o)
+    if USE_REF:
+        (output,
+        softmax_lse,
+        exp_scores,
+        _,
+        _,
+        _,
+        _) = attention_forward_pytorch_ref_impl(
+                                                q,
+                                                k,
+                                                v,
+                                                metadata.sm_scale,
+                                                metadata.causal,
+                                                metadata.layout,
+                                                metadata.cu_seqlens_q,
+                                                metadata.cu_seqlens_k,
+                                                metadata.max_seqlens_q,
+                                                metadata.max_seqlens_k,
+                                                metadata.use_exp2)
+        o.copy_(output)
+    else:
+        (_,
+        softmax_lse,
+        exp_scores,
+        _,
+        _,
+        _,
+        _,
+        _,
+        _) = attention_prefill_forward_triton_impl(
+                                                q,
+                                                k,
+                                                v,
+                                                o,
+                                                metadata.sm_scale,
+                                                metadata.alibi_slopes,
+                                                metadata.causal,
+                                                metadata.bias,
+                                                metadata.dropout_p,
+                                                metadata.layout,
+                                                metadata.cu_seqlens_q,
+                                                metadata.cu_seqlens_k,
+                                                metadata.max_seqlens_q,
+                                                metadata.max_seqlens_k,
+                                                metadata.return_scores,
+                                                metadata.use_exp2)
+
+    return o, softmax_lse, exp_scores, None
+
+
+def bwd(
+    dout,
+    q,
+    k,
+    v,
+    out,
+    softmax_lse,
+    dq,
+    dk,
+    dv,
+    alibi_slopes,
+    dropout_p,
+    softmax_scale,
+    causal,
+    window_size_left, window_size_right, softcap, deterministic, gen_, rng_state, # pylint: disable=unused-argument
+):
+    if dropout_p != 0.0:
+        raise ValueError("dropout is not supported on AMD yet")
+
+    if USE_REF:
+        dq_ref, dk_ref, dv_ref, delta_ref = attention_backward_pytorch_ref_impl(
+            dout,
+            q,
+            k,
+            v,
+            out,
+            softmax_lse,
+            softmax_scale,
+            causal,
+            "bshd",
+            None,
+            None,
+            None,
+            None,
+            False,
+        )
+        dq.copy_(dq_ref)
+        dk.copy_(dk_ref)
+        dv.copy_(dv_ref)
+        delta = delta_ref
+    else:
+        dq_triton, dk_triton, dv_triton, delta_triton, _, _ = attention_prefill_backward_triton_impl( # pylint: disable=unused-variable
+            dout,
+            q,
+            k,
+            v,
+            out,
+            softmax_lse,
+            dq,
+            dk,
+            dv,
+            softmax_scale,
+            alibi_slopes,
+            causal,
+            "bshd",
+            None,
+            None,
+            None,
+            None,
+            False,
+        )
+        delta = delta_triton
+
+    return dq, dk, dv, delta
+
+
+def varlen_fwd(
+    q,
+    k,
+    v,
+    o,
+    cu_seqlens_q,
+    cu_seqlens_k,
+    seqused_k, leftpad_k, block_table_, # pylint: disable=unused-argument
+    alibi_slopes,\
+    max_seqlen_q,
+    max_seqlen_k,
+    dropout_p,
+    softmax_scale,
+    zero_tensors, # pylint: disable=unused-argument
+    causal,
+    window_size_left, window_size_right, softcap, # pylint: disable=unused-argument
+    return_softmax,
+    gen_ # pylint: disable=unused-argument
+):
+    if dropout_p != 0.0:
+        raise ValueError("dropout is not supported on AMD's Triton Backend yet")
+
+    if o is None:
+        o = torch.empty_like(q)
+
+    # Setup metadata
+    metadata = MetaData(sm_scale=softmax_scale)
+    if return_softmax:
+        metadata.return_scores = True
+    metadata.set_varlen_params(cu_seqlens_q, cu_seqlens_k)  # set layout to "thd" and other metdata
+
+    # get shapes
+    batch, nheads_q, nheads_k, head_size , seqlen_q, seqlen_k = get_shape_from_layout(q, k, metadata.layout, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k) # pylint: disable=unused-variable
+
+    if causal:
+        metadata.need_causal()
+
+    if alibi_slopes is not None:
+        metadata.need_alibi(alibi_slopes, batch, nheads_q)
+
+    if dropout_p > 0.0:
+        metadata.need_dropout(dropout_p, return_softmax)
+
+    # Check arguments
+    metadata.check_args(q, k, v, o)
+    if o is None:
+        o = torch.empty_like(q, dtype=v.dtype)
+
+    if USE_REF:
+        (output,
+        softmax_lse,
+        exp_scores,
+        _,
+        _,
+        _,
+        _) = attention_forward_pytorch_ref_impl(
+                                                q,
+                                                k,
+                                                v,
+                                                metadata.sm_scale,
+                                                metadata.causal,
+                                                metadata.layout,
+                                                metadata.cu_seqlens_q,
+                                                metadata.cu_seqlens_k,
+                                                metadata.max_seqlens_q,
+                                                metadata.max_seqlens_k,
+                                                metadata.use_exp2)
+        o.copy_(output)
+    else:
+        (_,
+        softmax_lse,
+        exp_scores,
+        _,
+        _,
+        _,
+        _,
+        _,
+        _) = attention_prefill_forward_triton_impl(
+                                                    q,
+                                                    k,
+                                                    v,
+                                                    o,
+                                                    metadata.sm_scale,
+                                                    metadata.alibi_slopes,
+                                                    metadata.causal,
+                                                    metadata.bias,
+                                                    metadata.dropout_p,
+                                                    metadata.layout,
+                                                    metadata.cu_seqlens_q,
+                                                    metadata.cu_seqlens_k,
+                                                    metadata.max_seqlens_q,
+                                                    metadata.max_seqlens_k,
+                                                    metadata.return_scores,
+                                                    metadata.use_exp2)
+
+    return o, softmax_lse, exp_scores, None
+
+
+def varlen_bwd(
+    dout,
+    q,
+    k,
+    v,
+    out,
+    softmax_lse,
+    dq,
+    dk,
+    dv,
+    cu_seqlens_q,
+    cu_seqlens_k,
+    alibi_slopes,
+    max_seqlen_q,
+    max_seqlen_k,
+    dropout_p,
+    softmax_scale,
+    zero_tensors, # pylint: disable=unused-argument
+    causal,
+    window_size_left, window_size_right, softcap, deterministic, gen_, rng_state, # pylint: disable=unused-argument
+):
+    if dropout_p != 0.0:
+        raise ValueError("dropout is not supported on AMD yet")
+
+    if USE_REF:
+        dq_ref, dk_ref, dv_ref, delta_ref = attention_backward_pytorch_ref_impl(
+            dout,
+            q,
+            k,
+            v,
+            out,
+            softmax_lse,
+            softmax_scale,
+            causal,
+            "thd",
+            cu_seqlens_q,
+            cu_seqlens_k,
+            max_seqlen_q,
+            max_seqlen_k,
+            False,
+        )
+        dq.copy_(dq_ref)
+        dk.copy_(dk_ref)
+        dv.copy_(dv_ref)
+        delta = delta_ref
+    else:
+        dq_triton, dk_triton, dv_triton, delta_triton, _, _ = attention_prefill_backward_triton_impl( # pylint: disable=unused-variable
+            dout,
+            q,
+            k,
+            v,
+            out,
+            softmax_lse,
+            dq,
+            dk,
+            dv,
+            softmax_scale,
+            alibi_slopes,
+            causal,
+            "thd",
+            cu_seqlens_q,
+            cu_seqlens_k,
+            max_seqlen_q,
+            max_seqlen_k,
+            False,
+        )
+        delta = delta_triton
+
+    return dq, dk, dv, delta
+
+
+def fwd_kvcache(
+    q,
+    k_cache,
+    v_cache,
+    k,
+    v,
+    cache_seqlens,
+    rotary_cos, rotary_sin, # pylint: disable=unused-argument
+    cache_batch_idx,
+    cache_leftpad, block_table, # pylint: disable=unused-argument
+    alibi_slopes,
+    out,
+    softmax_scale,
+    causal,
+    window_size_left, window_size_right, softcap, rotary_interleaved, num_splits, # pylint: disable=unused-argument
+):
+    if out is None:
+        out = torch.empty_like(q)
+
+    # fill metadata
+    metadata = MetaData(sm_scale=softmax_scale)
+    metadata.layout = "bshd"
+    metadata.max_seqlens_q = q.shape[1]
+    metadata.max_seqlens_k = k_cache.shape[1]
+    metadata.cache_seqlens = cache_seqlens
+    metadata.cache_batch_idx = cache_batch_idx
+
+    if k is not None and v is not None:
+        metadata.new_kv = True
+        metadata.seqlen_new = k.shape[1]
+        metadata.k_new = k
+        metadata.v_new = v
+
+    if causal:
+        metadata.need_causal()
+
+    if alibi_slopes is not None:
+        batch, _ , nheads_q, _= q.shape
+        metadata.need_alibi(alibi_slopes, batch, nheads_q)
+
+    # launch kernel
+    # TODO: pass output as an arg. Maybe we are copying output which is causing slow down
+    output, softmax_lse = attention_decode_forward_triton_impl(
+        q,
+        k_cache,
+        v_cache,
+        metadata.sm_scale,
+        metadata.causal,
+        metadata.alibi_slopes,
+        metadata.layout,
+        metadata.cache_seqlens,
+        metadata.cache_batch_idx,
+        metadata.new_kv,
+        metadata.k_new,
+        metadata.v_new,
+    )
+    return output, softmax_lse
diff --git a/modules/flash_attn_triton_amd/utils.py b/modules/flash_attn_triton_amd/utils.py
new file mode 100644
index 000000000..77384cff6
--- /dev/null
+++ b/modules/flash_attn_triton_amd/utils.py
@@ -0,0 +1,280 @@
+import os
+import torch
+import triton
+
+
+AUTOTUNE = os.environ.get('FLASH_ATTENTION_TRITON_AMD_AUTOTUNE', '0').lower() in ('1', 'true', 'yes')
+PERF = os.environ.get('FLASH_ATTENTION_TRITON_AMD_PERF', '0').lower() in ('1', 'true', 'yes')
+
+
+class MetaData():
+    cu_seqlens_q = None
+    cu_seqlens_k = None
+    max_seqlens_q = 0
+    max_seqlens_k = 0
+    bias = None
+    alibi_slopes = None
+    causal = False
+    num_contexts = 0
+    varlen = False
+    layout = None
+    cache_seqlens = None
+    cache_batch_idx = None
+    new_kv = False
+    seqlen_new = None
+    k_new = None
+    v_new = None
+    dropout_p, return_scores= 0.0, False
+    # NOTE: scale sm_scale by log_2(e) and use 2^x in the loop as we do not have native e^x support in HW.
+    use_exp2 = False
+
+    def __repr__(self) -> str:
+        return (f"MetaData(\n"
+                f"  sm_scale={self.sm_scale},\n"
+                f"  cu_seqlens_q={self.cu_seqlens_q},\n"
+                f"  cu_seqlens_k={self.cu_seqlens_k},\n"
+                f"  max_seqlens_q={self.max_seqlens_q},\n"
+                f"  max_seqlens_k={self.max_seqlens_k},\n"
+                f"  bias={self.bias},\n"
+                f"  alibi_slopes={self.alibi_slopes},\n"
+                f"  causal={self.causal},\n"
+                f"  num_contexts={self.num_contexts},\n"
+                f"  varlen={self.varlen},\n"
+                f"  layout={self.layout},\n"
+                f"  cache_seqlens={self.cache_seqlens},\n"
+                f"  cache_batch_idx={self.cache_batch_idx},\n"
+                f"  new_kv={self.new_kv},\n"
+                f"  seqlen_new={self.seqlen_new},\n"
+                f"  k_new={self.k_new},\n"
+                f"  v_new={self.v_new},\n"
+                f"  dropout_p={self.dropout_p},\n"
+                f"  return_scores={self.return_scores}\n"
+                f")")
+
+    def __init__(self, sm_scale=1.0):
+        self.sm_scale = sm_scale
+
+    def set_varlen_params(self, cu_seqlens_q, cu_seqlens_k):
+        self.varlen = True
+        self.layout = 'thd'
+        self.cu_seqlens_q = cu_seqlens_q
+        self.cu_seqlens_k = cu_seqlens_k
+        # Without "varlen", there should still be one sequence.
+        assert len(cu_seqlens_q) >= 2
+        assert len(cu_seqlens_q) == len(cu_seqlens_k)
+        self.num_contexts = len(cu_seqlens_q) - 1
+        for i in range(0, self.num_contexts):
+            self.max_seqlens_q = max(cu_seqlens_q[i + 1].item() - cu_seqlens_q[i].item(), self.max_seqlens_q)
+            self.max_seqlens_k = max(cu_seqlens_k[i + 1].item() - cu_seqlens_k[i].item(), self.max_seqlens_k)
+
+    def need_bias(self, bias, batch, nheads, seqlen_q, seqlen_k): # pylint: disable=unused-argument
+        assert bias.is_cuda
+        assert bias.dim() == 4
+        assert bias.shape[0] == 1
+        assert bias.shape[2:] == (seqlen_q, seqlen_k)
+        self.bias = bias
+
+    def need_alibi(self, alibi_slopes, batch, nheads):
+        assert alibi_slopes.is_cuda
+        assert alibi_slopes.dim() == 2
+        assert alibi_slopes.shape[0] == batch
+        assert alibi_slopes.shape[1] == nheads
+        self.alibi_slopes = alibi_slopes
+
+    def need_causal(self):
+        self.causal = True
+
+    def need_dropout(self, dropout_p, return_scores):
+        self.dropout_p = dropout_p
+        self.return_scores = return_scores
+
+    def check_args(self, q, k, v, o):
+        assert q.dim() == k.dim() and q.dim() == v.dim()
+
+        batch, nheads_q, nheads_k, head_size, _, _ = get_shape_from_layout(q, k, self.layout, self.cu_seqlens_q, self.cu_seqlens_k, self.max_seqlens_q, self.max_seqlens_k) # pylint: disable=unused-variable
+        if self.varlen:
+            assert q.dim() == 3
+            assert self.cu_seqlens_q is not None
+            assert self.cu_seqlens_k is not None
+            assert len(self.cu_seqlens_q) == len(self.cu_seqlens_k)
+            # TODO: Remove once bias is supported with varlen
+            assert self.bias is None
+            # TODO:Remove once dropout is supported with varlen
+            assert self.dropout_p == 0.0
+            # assert not self.return_scores
+        else:
+            assert q.dim() == 4
+            assert self.max_seqlens_q > 0 and self.max_seqlens_k > 0
+            assert self.cu_seqlens_q is None and self.cu_seqlens_k is None
+        assert k.shape == v.shape
+        assert q.shape[-1] == k.shape[-1] and q.shape[-1] == v.shape[-1]
+        # TODO: Change assert if we support qkl f8 and v f16
+        assert q.dtype == k.dtype and q.dtype == v.dtype
+        assert head_size <= 256
+        assert o.shape == q.shape
+        assert (nheads_q % nheads_k) == 0
+        assert self.layout is not None
+        assert self.layout == 'thd' or not self.varlen
+
+def input_helper(Z, HQ, HK, N_CTX_Q, N_CTX_K, D_HEAD, dtype, layout, device="cuda", DEBUG_INPUT=False):
+    torch.manual_seed(20)
+
+    # Initialize q, k, v
+    if layout == 'bhsd':
+        q_tensor_shape = (Z, HQ, N_CTX_Q, D_HEAD)
+        k_tensor_shape = (Z, HK, N_CTX_K, D_HEAD)
+    elif layout == 'bshd':
+        q_tensor_shape = (Z, N_CTX_Q, HQ, D_HEAD)
+        k_tensor_shape = (Z, N_CTX_K, HK, D_HEAD)
+    else:
+        assert False, f'Got unsupported tensor layout: {layout}'
+
+    q = None
+    k = None
+    v = None
+
+    if DEBUG_INPUT:
+        if layout == "bhsd":
+            q = torch.arange(N_CTX_Q, dtype=dtype, device=device).view(1, 1, N_CTX_Q, 1).expand(*q_tensor_shape).contiguous().requires_grad_()
+            k = torch.arange(N_CTX_K, dtype=dtype, device=device).view(1, 1, N_CTX_K, 1).expand(*k_tensor_shape).contiguous().requires_grad_()
+            v = torch.arange(N_CTX_K, dtype=dtype, device=device).view(1, 1, N_CTX_K, 1).expand(*k_tensor_shape).contiguous().requires_grad_()
+        elif layout == "bshd":
+            q = torch.arange(N_CTX_Q, dtype=dtype, device=device).view(1, N_CTX_Q, 1, 1).expand(*q_tensor_shape).contiguous().requires_grad_()
+            k = torch.arange(N_CTX_K, dtype=dtype, device=device).view(1, N_CTX_K, 1, 1).expand(*k_tensor_shape).contiguous().requires_grad_()
+            v = torch.arange(N_CTX_K, dtype=dtype, device=device).view(1, N_CTX_K, 1, 1).expand(*k_tensor_shape).contiguous().requires_grad_()
+    else:
+        q = torch.randn(q_tensor_shape, dtype=dtype, device=device, requires_grad=True)
+        k = torch.randn(k_tensor_shape, dtype=dtype, device=device, requires_grad=True)
+        v = torch.randn(k_tensor_shape, dtype=dtype, device=device, requires_grad=True)
+
+    if DEBUG_INPUT:
+        sm_scale = 1
+    else:
+        sm_scale = D_HEAD**-0.5
+    input_metadata = MetaData(sm_scale=sm_scale)
+    input_metadata.max_seqlens_q = N_CTX_Q
+    input_metadata.max_seqlens_k = N_CTX_K
+    input_metadata.layout = layout
+    return q, k, v, input_metadata
+
+
+def varlen_input_helper(Z, HQ, HK, N_CTX_Q, N_CTX_K, D_HEAD, dtype, device="cuda", equal_seqlens=False, DEBUG_INPUT=False):
+    torch.manual_seed(20)
+
+    # Random or equal sequence lengths based on 'equal_seqlens' flag
+    if not equal_seqlens:
+        max_seqlens_q = N_CTX_Q // Z
+        max_seqlens_k = N_CTX_K // Z
+        seqlens_q = torch.randint(1, max_seqlens_q + 1, (Z,), dtype=torch.int32)
+        seqlens_k = torch.randint(1, max_seqlens_k + 1, (Z,), dtype=torch.int32)
+    else:
+        seqlens_q = torch.full((Z,), N_CTX_Q // Z, dtype=torch.int32)
+        seqlens_k = torch.full((Z,), N_CTX_K // Z, dtype=torch.int32)
+
+    # Calculate cumulative sequence lengths
+    cu_seqlens_q = torch.cat([torch.tensor([0], dtype=torch.int32), seqlens_q.cumsum(dim=0)])
+    cu_seqlens_k = torch.cat([torch.tensor([0], dtype=torch.int32), seqlens_k.cumsum(dim=0)])
+    cu_seqlens_q = cu_seqlens_q.to(device=device).to(torch.int32)
+    cu_seqlens_k = cu_seqlens_k.to(device=device).to(torch.int32)
+
+    # Total lengths
+    total_q = cu_seqlens_q[-1].item()
+    total_k = cu_seqlens_k[-1].item()
+
+    if DEBUG_INPUT:
+        # Initialize q, k, v with deterministic values
+        q = torch.arange(total_q, dtype=dtype, device=device).view(total_q, 1, 1)
+        q = q.expand(total_q, HQ, D_HEAD).contiguous().requires_grad_()
+        k = torch.arange(total_k, dtype=dtype, device=device).view(total_k, 1, 1)
+        k = k.expand(total_k, HK, D_HEAD).contiguous().requires_grad_()
+        v = torch.arange(total_k, dtype=dtype, device=device).view(total_k, 1, 1)
+        v = v.expand(total_k, HK, D_HEAD).contiguous().requires_grad_()
+        sm_scale = 1
+    else:
+        # Initialize q, k, v with random values
+        q = torch.randn((total_q, HQ, D_HEAD), dtype=dtype, device=device).requires_grad_()
+        k = torch.randn((total_k, HK, D_HEAD), dtype=dtype, device=device).requires_grad_()
+        v = torch.randn((total_k, HK, D_HEAD), dtype=dtype, device=device).requires_grad_()
+        sm_scale = D_HEAD ** -0.5
+
+    input_metadata = MetaData(sm_scale=sm_scale)
+    input_metadata.set_varlen_params(cu_seqlens_q, cu_seqlens_k)
+    return q, k, v, input_metadata
+
+
+def get_shape_from_layout(q, k, layout, cu_seqlens_q = None, cu_seqlens_k = None, max_seqlen_q=None, max_seqlen_k=None):
+    if layout == 'bhsd':
+        batch_q, nheads_q, max_seqlen_q, head_size_q = q.shape
+        batch_k, nheads_k, max_seqlen_k, head_size_k = k.shape
+    elif layout == 'bshd':
+        batch_q, max_seqlen_q, nheads_q, head_size_q = q.shape
+        batch_k, max_seqlen_k, nheads_k, head_size_k = k.shape
+    elif  layout == 'thd':
+        batch_q, max_seqlen_q, nheads_q, head_size_q = len(cu_seqlens_q) - 1, max_seqlen_q, q.shape[1], q.shape[2] # pylint: disable=self-assigning-variable
+        batch_k, max_seqlen_k, nheads_k, head_size_k = len(cu_seqlens_k) - 1, max_seqlen_k, k.shape[1], k.shape[2] # pylint: disable=self-assigning-variable
+    else:
+        assert False, "Got unsupported layout."
+
+    # assert
+    assert batch_q == batch_k
+    assert head_size_q == head_size_k
+
+    return batch_q, nheads_q, nheads_k, head_size_q, max_seqlen_q, max_seqlen_k
+
+
+def get_strides_from_layout(q, k, v, o, layout):
+    if layout == 'thd':
+        q_strides = (0, q.stride(1), q.stride(0), q.stride(2))
+        k_strides = (0, k.stride(1), k.stride(0), k.stride(2))
+        v_strides = (0, v.stride(1), v.stride(0), v.stride(2))
+        o_strides = (0, o.stride(1), o.stride(0), o.stride(2))
+    elif layout == 'bhsd':
+        q_strides = (q.stride(0), q.stride(1), q.stride(2), q.stride(3))
+        k_strides = (k.stride(0), k.stride(1), k.stride(2), k.stride(3))
+        v_strides = (v.stride(0), v.stride(1), v.stride(2), v.stride(3))
+        o_strides = (o.stride(0), o.stride(1), o.stride(2), o.stride(3))
+    elif layout == 'bshd':
+        q_strides = (q.stride(0), q.stride(2), q.stride(1), q.stride(3))
+        k_strides = (k.stride(0), k.stride(2), k.stride(1), k.stride(3))
+        v_strides = (v.stride(0), v.stride(2), v.stride(1), v.stride(3))
+        o_strides = (o.stride(0), o.stride(2), o.stride(1), o.stride(3))
+    else:
+        assert False, 'Got unsupported layout.'
+    return q_strides, k_strides, v_strides, o_strides
+
+
+def get_padded_headsize(size):
+    # Get closest power of 2 over or equal to 32.
+    padded_d_model = 1 << (size - 1).bit_length()
+    # Smallest head_dim supported is 16. If smaller, the tile in the
+    # kernel is padded - there is no padding in memory for any dims.
+    padded_d_model = max(padded_d_model, 16)
+    return padded_d_model
+
+
+def _strides(x: torch.Tensor, *stride_names: str):
+    if x is None:
+        return {f"stride_{s}": 0 for i, s in enumerate(stride_names)}
+
+    assert x.ndim == len(stride_names)
+    return {f"stride_{s}": x.stride(i) for i, s in enumerate(stride_names)}
+
+
+def get_input_shapes():
+    cases = [(max(1, 2**(16 - i)), 1, 2**i, 16, 1, 128)
+             for i in range(8, 18)] + [(max(1, 2**(16 - i)), 1, 2**i, 16, 2, 128) for i in range(8, 18)]
+    return cases
+
+
+def is_hip():
+    return triton.runtime.driver.active.get_current_target().backend == "hip"
+
+
+def is_cdna():
+    return is_hip() and triton.runtime.driver.active.get_current_target().arch in ('gfx940', 'gfx941', 'gfx942',
+                                                                                   'gfx90a', 'gfx908')
+
+
+def is_rdna():
+    return is_hip() and triton.runtime.driver.active.get_current_target().arch in ("gfx1030", "gfx1100", "gfx1101",
+                                                                                   "gfx1102", "gfx1200", "gfx1201")
diff --git a/modules/generation_parameters_copypaste.py b/modules/generation_parameters_copypaste.py
index 613c09ff2..788edb1e5 100644
--- a/modules/generation_parameters_copypaste.py
+++ b/modules/generation_parameters_copypaste.py
@@ -218,6 +218,9 @@ def connect_paste(button, local_paste_fields, input_comp, override_settings_comp
             else:
                 try:
                     valtype = type(output.value)
+                    if hasattr(output, "step") and type(output.step) == float:
+                        valtype = float
+                    debug(f'Paste: "{key}"="{v}" type={valtype} var={vars(output)}')
                     if valtype == bool and v == "False":
                         val = False
                     else:
diff --git a/modules/img2img.py b/modules/img2img.py
index 75971c608..ca71ff0e7 100644
--- a/modules/img2img.py
+++ b/modules/img2img.py
@@ -6,9 +6,10 @@ from PIL import Image, ImageOps, ImageFilter, ImageEnhance, ImageChops, Unidenti
 import modules.scripts
 from modules import shared, processing, images
 from modules.generation_parameters_copypaste import create_override_settings_dict
-from modules.ui import plaintext_to_html
+from modules.ui_common import plaintext_to_html
 from modules.memstats import memory_stats
 
+
 debug = shared.log.trace if os.environ.get('SD_PROCESS_DEBUG', None) is not None else lambda *args, **kwargs: None
 debug('Trace: PROCESS')
 
diff --git a/modules/infiniteyou/__init__.py b/modules/infiniteyou/__init__.py
new file mode 100644
index 000000000..142921909
--- /dev/null
+++ b/modules/infiniteyou/__init__.py
@@ -0,0 +1,2 @@
+from .pipeline_flux_infusenet import FluxInfuseNetPipeline
+from .pipeline_infu_flux import InfUFluxPipeline
diff --git a/modules/infiniteyou/pipeline_flux_infusenet.py b/modules/infiniteyou/pipeline_flux_infusenet.py
new file mode 100644
index 000000000..38fa186f9
--- /dev/null
+++ b/modules/infiniteyou/pipeline_flux_infusenet.py
@@ -0,0 +1,612 @@
+# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates.
+# Copyright (c) 2024 Black Forest Labs, The HuggingFace Team and The InstantX Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import inspect
+from typing import Any, Callable, Dict, List, Optional, Union
+
+import numpy as np
+import torch
+from diffusers import FluxControlNetPipeline
+from diffusers.models.controlnet_flux import FluxControlNetModel, FluxMultiControlNetModel
+from diffusers.image_processor import PipelineImageInput
+from diffusers.pipelines.flux.pipeline_output import FluxPipelineOutput
+from diffusers.utils import is_torch_xla_available, logging
+
+
+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
+
+
+# Copied from diffusers.pipelines.flux.pipeline_flux.calculate_shift
+def calculate_shift(
+    image_seq_len,
+    base_seq_len: int = 256,
+    max_seq_len: int = 4096,
+    base_shift: float = 0.5,
+    max_shift: float = 1.16,
+):
+    m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
+    b = base_shift - m * base_seq_len
+    mu = image_seq_len * m + b
+    return mu
+
+
+# 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,
+    sigmas: Optional[List[float]] = None,
+    **kwargs,
+):
+    r"""
+    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 override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
+            `num_inference_steps` and `sigmas` must be `None`.
+        sigmas (`List[float]`, *optional*):
+            Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
+            `num_inference_steps` and `timesteps` 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 and sigmas is not None:
+        raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
+    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)
+    elif sigmas is not None:
+        accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+        if not accept_sigmas:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" sigmas schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(sigmas=sigmas, 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 FluxInfuseNetPipeline(FluxControlNetPipeline):
+    @torch.no_grad()
+    def __call__(
+        self,
+        prompt: Union[str, List[str]] = None,
+        prompt_2: Optional[Union[str, List[str]]] = None,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        num_inference_steps: int = 28,
+        timesteps: List[int] = None,
+        guidance_scale: float = 3.5,
+        id_image: PipelineImageInput = None,
+        controlnet_guidance_scale: float = 1.0,
+        control_guidance_start: Union[float, List[float]] = 0.0,
+        control_guidance_end: Union[float, List[float]] = 1.0,
+        control_image: PipelineImageInput = None,
+        control_mode: Optional[Union[int, List[int]]] = None,
+        controlnet_conditioning_scale: Union[float, List[float]] = 1.0,
+        num_images_per_prompt: Optional[int] = 1,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        latents: Optional[torch.FloatTensor] = None,
+        prompt_embeds: Optional[torch.FloatTensor] = None,
+        pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+        output_type: Optional[str] = "pil",
+        return_dict: bool = True,
+        joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+        callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
+        callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+        max_sequence_length: int = 512,
+
+        # ID-specific parameters
+        controlnet_prompt_embeds: Optional[torch.FloatTensor] = None,
+
+        # True CFG parameters
+        true_guidance_scale: float = 1.0,
+        negative_prompt: Optional[Union[str, List[str]]] = None,
+        negative_prompt_2: Optional[Union[str, List[str]]] = None,
+        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+        negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+    ):
+        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 `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
+                will be used instead
+            height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
+                The height in pixels of the generated image. This is set to 1024 by default for the best results.
+            width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
+                The width in pixels of the generated image. This is set to 1024 by default for the best results.
+            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.
+            timesteps (`List[int]`, *optional*):
+                Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument
+                in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is
+                passed will be used. Must be in descending order.
+            guidance_scale (`float`, *optional*, defaults to 7.0):
+                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.
+            controlnet_guidance_scale (`float`, *optional*, defaults to 7.0):
+                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
+                `controlnet_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.
+            control_guidance_start (`float` or `List[float]`, *optional*, defaults to 0.0):
+                The percentage of total steps at which the ControlNet starts applying.
+            control_guidance_end (`float` or `List[float]`, *optional*, defaults to 1.0):
+                The percentage of total steps at which the ControlNet stops applying.
+            control_image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, `List[np.ndarray]`,:
+                    `List[List[torch.Tensor]]`, `List[List[np.ndarray]]` or `List[List[PIL.Image.Image]]`):
+                The ControlNet input condition to provide guidance to the `unet` for generation. If the type is
+                specified as `torch.Tensor`, it is passed to ControlNet as is. `PIL.Image.Image` can also be accepted
+                as an image. The dimensions of the output image defaults to `image`'s dimensions. If height and/or
+                width are passed, `image` is resized accordingly. If multiple ControlNets are specified in `init`,
+                images must be passed as a list such that each element of the list can be correctly batched for input
+                to a single ControlNet.
+            controlnet_conditioning_scale (`float` or `List[float]`, *optional*, defaults to 1.0):
+                The outputs of the ControlNet are multiplied by `controlnet_conditioning_scale` before they are added
+                to the residual in the original `unet`. If multiple ControlNets are specified in `init`, you can set
+                the corresponding scale as a list.
+            control_mode (`int` or `List[int]`,, *optional*, defaults to None):
+                The control mode when applying ControlNet-Union.
+            num_images_per_prompt (`int`, *optional*, defaults to 1):
+                The number of images to generate per prompt.
+            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.FloatTensor`, *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.FloatTensor`, *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.
+            pooled_prompt_embeds (`torch.FloatTensor`, *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.
+            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.flux.FluxPipelineOutput`] instead of a plain tuple.
+            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).
+            callback_on_step_end (`Callable`, *optional*):
+                A function that calls at the end of each denoising steps during the inference. The function is called
+                with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
+                callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
+                `callback_on_step_end_tensor_inputs`.
+            callback_on_step_end_tensor_inputs (`List`, *optional*):
+                The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
+                will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
+                `._callback_tensor_inputs` attribute of your pipeline class.
+            max_sequence_length (`int` defaults to 512): Maximum sequence length to use with the `prompt`.
+            controlnet_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated embeddings for the InfuseNet. Can be used to easily tweak inputs, *e.g.* image embeddings.
+                If not provided, embeddings will be generated from `prompt` or `prompt_embeds` input arguments.
+            true_guidance_scale (`float`, *optional*, defaults to 1.0):
+                True CFG scale as defined in [Classifier-Free Diffusion Guidance]((https://arxiv.org/abs/2207.12598).
+            negative_prompt (`str` or `List[str]`, *optional*):
+                The negative prompt or negative prompts to guide the image generation. If not defined, one has to pass
+                `negative_prompt_embeds`. instead.
+            negative_prompt_2 (`str` or `List[str]`, *optional*):
+                The negative prompt or negative prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined,
+                `negative_prompt` is will be used instead.
+            negative_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+                weighting. If not provided, negative text embeddings will be generated from `negative_prompt` input
+                argument.
+            negative_pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated negative pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+                weighting. If not provided, negative pooled text embeddings will be generated from
+                `negative_prompt` input argument.
+
+        Examples:
+
+        Returns:
+            [`~pipelines.flux.FluxPipelineOutput`] or `tuple`: [`~pipelines.flux.FluxPipelineOutput`] if `return_dict`
+            is True, otherwise a `tuple`. When returning a tuple, the first element is a list with the generated
+            images.
+        """
+
+        height = height or self.default_sample_size * self.vae_scale_factor
+        width = width or self.default_sample_size * self.vae_scale_factor
+
+        if not isinstance(control_guidance_start, list) and isinstance(control_guidance_end, list):
+            control_guidance_start = len(control_guidance_end) * [control_guidance_start]
+        elif not isinstance(control_guidance_end, list) and isinstance(control_guidance_start, list):
+            control_guidance_end = len(control_guidance_start) * [control_guidance_end]
+        elif not isinstance(control_guidance_start, list) and not isinstance(control_guidance_end, list):
+            mult = len(self.controlnet.nets) if isinstance(self.controlnet, FluxMultiControlNetModel) else 1
+            control_guidance_start, control_guidance_end = (
+                mult * [control_guidance_start],
+                mult * [control_guidance_end],
+            )
+
+        # 1. Check inputs. Raise error if not correct
+        self.check_inputs(
+            prompt,
+            prompt_2,
+            height,
+            width,
+            prompt_embeds=prompt_embeds,
+            pooled_prompt_embeds=pooled_prompt_embeds,
+            callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
+            max_sequence_length=max_sequence_length,
+        )
+
+        self._guidance_scale = guidance_scale
+        self._controlnet_guidance_scale = controlnet_guidance_scale
+        self._true_guidance_scale = true_guidance_scale
+        self._joint_attention_kwargs = joint_attention_kwargs
+        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
+        dtype = self.transformer.dtype
+
+        lora_scale = (
+            self.joint_attention_kwargs.get("scale", None) if self.joint_attention_kwargs is not None else None
+        )
+        (
+            prompt_embeds,
+            pooled_prompt_embeds,
+            text_ids,
+        ) = self.encode_prompt(
+            prompt=prompt,
+            prompt_2=prompt_2,
+            prompt_embeds=prompt_embeds,
+            pooled_prompt_embeds=pooled_prompt_embeds,
+            device=device,
+            num_images_per_prompt=num_images_per_prompt,
+            max_sequence_length=max_sequence_length,
+            lora_scale=lora_scale,
+        )
+        if negative_prompt is not None or (negative_prompt_embeds is not None and negative_pooled_prompt_embeds is not None):
+            (
+                negative_prompt_embeds,
+                negative_pooled_prompt_embeds,
+                negative_text_ids,
+            ) = self.encode_prompt(
+                prompt=negative_prompt,
+                prompt_2=negative_prompt_2,
+                prompt_embeds=negative_prompt_embeds,
+                pooled_prompt_embeds=negative_pooled_prompt_embeds,
+                device=device,
+                num_images_per_prompt=num_images_per_prompt,
+                max_sequence_length=max_sequence_length,
+                lora_scale=lora_scale,
+            )
+
+        if controlnet_prompt_embeds is None:
+            controlnet_prompt_embeds = prompt_embeds
+        (
+            controlnet_prompt_embeds,
+            pooled_prompt_embeds,
+            controlnet_text_ids,
+        ) = self.encode_prompt(
+            prompt=prompt,
+            prompt_2=prompt_2,
+            prompt_embeds=controlnet_prompt_embeds,
+            pooled_prompt_embeds=pooled_prompt_embeds,
+            device=device,
+            num_images_per_prompt=num_images_per_prompt,
+            max_sequence_length=max_sequence_length,
+            lora_scale=lora_scale,
+        )
+
+        # 3. Prepare control image
+        num_channels_latents = self.transformer.config.in_channels // 4
+        if isinstance(self.controlnet, FluxControlNetModel) or True:
+            control_image = self.prepare_image(
+                image=control_image,
+                width=width,
+                height=height,
+                batch_size=batch_size * num_images_per_prompt,
+                num_images_per_prompt=num_images_per_prompt,
+                device=device,
+                dtype=self.vae.dtype,
+            )
+            height, width = control_image.shape[-2:]
+
+            # xlab controlnet has a input_hint_block and instantx controlnet does not
+            controlnet_blocks_repeat = False if self.controlnet.input_hint_block is None else True
+            if self.controlnet.input_hint_block is None:
+                # vae encode
+                control_image = self.vae.encode(control_image).latent_dist.sample()
+                control_image = (control_image - self.vae.config.shift_factor) * self.vae.config.scaling_factor
+
+                # pack
+                height_control_image, width_control_image = control_image.shape[2:]
+                control_image = self._pack_latents(
+                    control_image,
+                    batch_size * num_images_per_prompt,
+                    num_channels_latents,
+                    height_control_image,
+                    width_control_image,
+                )
+
+            # Here we ensure that `control_mode` has the same length as the control_image.
+            if control_mode is not None:
+                if not isinstance(control_mode, int):
+                    raise ValueError(" For `FluxControlNet`, `control_mode` should be an `int` or `None`")
+                control_mode = torch.tensor(control_mode).to(device, dtype=torch.long)
+                control_mode = control_mode.view(-1, 1).expand(control_image.shape[0], 1)
+
+        elif isinstance(self.controlnet, FluxMultiControlNetModel):
+            control_images = []
+            # xlab controlnet has a input_hint_block and instantx controlnet does not
+            controlnet_blocks_repeat = False if self.controlnet.nets[0].input_hint_block is None else True
+            for _i, control_image_ in enumerate(control_image):
+                control_image_ = self.prepare_image(
+                    image=control_image_,
+                    width=width,
+                    height=height,
+                    batch_size=batch_size * num_images_per_prompt,
+                    num_images_per_prompt=num_images_per_prompt,
+                    device=device,
+                    dtype=self.vae.dtype,
+                )
+                height, width = control_image_.shape[-2:]
+
+                if self.controlnet.nets[0].input_hint_block is None:
+                    # vae encode
+                    control_image_ = self.vae.encode(control_image_).latent_dist.sample()
+                    control_image_ = (control_image_ - self.vae.config.shift_factor) * self.vae.config.scaling_factor
+
+                    # pack
+                    height_control_image, width_control_image = control_image_.shape[2:]
+                    control_image_ = self._pack_latents(
+                        control_image_,
+                        batch_size * num_images_per_prompt,
+                        num_channels_latents,
+                        height_control_image,
+                        width_control_image,
+                    )
+                control_images.append(control_image_)
+
+            control_image = control_images
+
+            # Here we ensure that `control_mode` has the same length as the control_image.
+            if isinstance(control_mode, list) and len(control_mode) != len(control_image):
+                raise ValueError("For Multi-ControlNet, `control_mode` must be a list of the same length as the number of controlnets (control images) specified")
+            if not isinstance(control_mode, list):
+                control_mode = [control_mode] * len(control_image)
+            # set control mode
+            control_modes = []
+            for cmode in control_mode:
+                if cmode is None:
+                    cmode = -1
+                control_mode = torch.tensor(cmode).expand(control_images[0].shape[0]).to(device, dtype=torch.long)
+                control_modes.append(control_mode)
+            control_mode = control_modes
+
+        # 4. Prepare latent variables
+        num_channels_latents = self.transformer.config.in_channels // 4
+        latents, latent_image_ids = self.prepare_latents(
+            batch_size * num_images_per_prompt,
+            num_channels_latents,
+            height,
+            width,
+            prompt_embeds.dtype,
+            device,
+            generator,
+            latents,
+        )
+
+        # 5. Prepare timesteps
+        sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps)
+        image_seq_len = latents.shape[1]
+        mu = calculate_shift(
+            image_seq_len,
+            self.scheduler.config.base_image_seq_len,
+            self.scheduler.config.max_image_seq_len,
+            self.scheduler.config.base_shift,
+            self.scheduler.config.max_shift,
+        )
+        timesteps, num_inference_steps = retrieve_timesteps(
+            self.scheduler,
+            num_inference_steps,
+            device,
+            timesteps,
+            sigmas,
+            mu=mu,
+        )
+
+        num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
+        self._num_timesteps = len(timesteps)
+
+        # 6. Create tensor stating which controlnets to keep
+        controlnet_keep = []
+        for i in range(len(timesteps)):
+            keeps = [
+                1.0 - float(i / len(timesteps) < s or (i + 1) / len(timesteps) > e)
+                for s, e in zip(control_guidance_start, control_guidance_end)
+            ]
+            controlnet_keep.append(keeps[0] if isinstance(self.controlnet, FluxControlNetModel) else keeps)
+
+        # 7. Denoising loop
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                if self.interrupt:
+                    continue
+
+                # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
+                timestep = t.expand(latents.shape[0]).to(latents.dtype)
+
+                if isinstance(self.controlnet, FluxMultiControlNetModel):
+                    use_guidance = self.controlnet.nets[0].config.guidance_embeds
+                else:
+                    use_guidance = self.controlnet.config.guidance_embeds
+
+                guidance = torch.tensor([controlnet_guidance_scale], device=device) if use_guidance else None
+                guidance = guidance.expand(latents.shape[0]) if guidance is not None else None
+
+                if isinstance(controlnet_keep[i], list):
+                    if not isinstance(controlnet_conditioning_scale, list):
+                        controlnet_conditioning_scale = len(controlnet_keep) * [controlnet_conditioning_scale]
+                    cond_scale = [c * s for c, s in zip(controlnet_conditioning_scale, controlnet_keep[i])]
+                    controlnet_conditioning_scale = controlnet_conditioning_scale[0]
+                else:
+                    controlnet_cond_scale = controlnet_conditioning_scale
+                    if isinstance(controlnet_cond_scale, list):
+                        controlnet_cond_scale = controlnet_cond_scale[0]
+                    cond_scale = controlnet_cond_scale * controlnet_keep[i]
+
+                # controlnet
+                controlnet_block_samples, controlnet_single_block_samples = self.controlnet(
+                    hidden_states=latents,
+                    controlnet_cond=control_image,
+                    controlnet_mode=control_mode,
+                    conditioning_scale=cond_scale[0],
+                    timestep=timestep / 1000,
+                    guidance=guidance,
+                    pooled_projections=pooled_prompt_embeds,
+                    encoder_hidden_states=controlnet_prompt_embeds,
+                    txt_ids=controlnet_text_ids,
+                    img_ids=latent_image_ids,
+                    joint_attention_kwargs=self.joint_attention_kwargs,
+                    return_dict=False,
+                )
+
+                guidance = (
+                    torch.tensor([guidance_scale], device=device) if self.transformer.config.guidance_embeds else None
+                )
+                guidance = guidance.expand(latents.shape[0]) if guidance is not None else None
+
+                noise_pred = self.transformer(
+                    hidden_states=latents,
+                    timestep=timestep / 1000,
+                    guidance=guidance,
+                    pooled_projections=pooled_prompt_embeds,
+                    encoder_hidden_states=prompt_embeds,
+                    controlnet_block_samples=controlnet_block_samples,
+                    controlnet_single_block_samples=controlnet_single_block_samples,
+                    txt_ids=text_ids,
+                    img_ids=latent_image_ids,
+                    joint_attention_kwargs=self.joint_attention_kwargs,
+                    return_dict=False,
+                    controlnet_blocks_repeat=controlnet_blocks_repeat,
+                )[0]
+
+                # perform true CFG
+                if negative_prompt_embeds is not None and negative_pooled_prompt_embeds is not None and negative_text_ids is not None:
+                    noise_pred_uncond = self.transformer(
+                        hidden_states=latents,
+                        timestep=timestep / 1000,
+                        guidance=guidance,
+                        pooled_projections=negative_pooled_prompt_embeds,
+                        encoder_hidden_states=negative_prompt_embeds,
+                        controlnet_block_samples=None,
+                        controlnet_single_block_samples=None,
+                        txt_ids=negative_text_ids,
+                        img_ids=latent_image_ids,
+                        joint_attention_kwargs=self.joint_attention_kwargs,
+                        return_dict=False,
+                        controlnet_blocks_repeat=controlnet_blocks_repeat,
+                    )[0]
+
+                    noise_pred = noise_pred_uncond + true_guidance_scale * (noise_pred - noise_pred_uncond)
+
+                # compute the previous noisy sample x_t -> x_t-1
+                latents_dtype = latents.dtype
+                latents = self.scheduler.step(noise_pred, t, latents, 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)
+
+                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)
+
+                # 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 XLA_AVAILABLE:
+                    xm.mark_step()
+
+        if output_type == "latent":
+            image = latents
+
+        else:
+            latents = self._unpack_latents(latents, height, width, self.vae_scale_factor)
+            latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor
+
+            image = self.vae.decode(latents, return_dict=False)[0]
+            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 FluxPipelineOutput(images=image)
diff --git a/modules/infiniteyou/pipeline_infu_flux.py b/modules/infiniteyou/pipeline_infu_flux.py
new file mode 100644
index 000000000..8ae9f6e95
--- /dev/null
+++ b/modules/infiniteyou/pipeline_infu_flux.py
@@ -0,0 +1,325 @@
+# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates. All rights reserved.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import math
+import os
+import random
+from typing import Optional
+
+import cv2
+import numpy as np
+import torch
+from diffusers.models import FluxControlNetModel
+from facexlib.recognition import init_recognition_model
+from huggingface_hub import snapshot_download
+from insightface.app import FaceAnalysis
+from insightface.utils import face_align
+from PIL import Image
+
+from modules import shared, devices, model_quant
+from .pipeline_flux_infusenet import FluxInfuseNetPipeline
+from .resampler import Resampler
+
+
+def seed_everything(seed, deterministic=False):
+    """Set random seed.
+
+    Args:
+        seed (int): Seed to be used.
+        deterministic (bool): Whether to set the deterministic option for
+            CUDNN backend, i.e., set `torch.backends.cudnn.deterministic`
+            to True and `torch.backends.cudnn.benchmark` to False.
+            Default: False.
+    """
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    os.environ['PYTHONHASHSEED'] = str(seed)
+    if deterministic:
+        torch.backends.cudnn.deterministic = True
+        torch.backends.cudnn.benchmark = False
+
+
+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")
+
+
+# modified from https://github.com/instantX-research/InstantID/blob/main/pipeline_stable_diffusion_xl_instantid.py
+def draw_kps(image_pil, kps, color_list=[(255,0,0), (0,255,0), (0,0,255), (255,255,0), (255,0,255)]):
+    stickwidth = 4
+    limbSeq = np.array([[0, 2], [1, 2], [3, 2], [4, 2]])
+    kps = np.array(kps)
+
+    w, h = image_pil.size
+    out_img = np.zeros([h, w, 3])
+
+    for i in range(len(limbSeq)):
+        index = limbSeq[i]
+        color = color_list[index[0]]
+
+        x = kps[index][:, 0]
+        y = kps[index][:, 1]
+        length = ((x[0] - x[1]) ** 2 + (y[0] - y[1]) ** 2) ** 0.5
+        angle = math.degrees(math.atan2(y[0] - y[1], x[0] - x[1]))
+        polygon = cv2.ellipse2Poly((int(np.mean(x)), int(np.mean(y))), (int(length / 2), stickwidth), int(angle), 0, 360, 1)
+        out_img = cv2.fillConvexPoly(out_img.copy(), polygon, color)
+    out_img = (out_img * 0.6).astype(np.uint8)
+
+    for idx_kp, kp in enumerate(kps):
+        color = color_list[idx_kp]
+        x, y = kp
+        out_img = cv2.circle(out_img.copy(), (int(x), int(y)), 10, color, -1)
+
+    out_img_pil = Image.fromarray(out_img.astype(np.uint8))
+    return out_img_pil
+
+
+def extract_arcface_bgr_embedding(in_image, landmark, arcface_model=None, in_settings=None): # pylint: disable=unused-argument
+    kps = landmark
+    arc_face_image = face_align.norm_crop(in_image, landmark=np.array(kps), image_size=112)
+    arc_face_image = torch.from_numpy(arc_face_image).unsqueeze(0).permute(0,3,1,2) / 255.
+    arc_face_image = 2 * arc_face_image - 1
+    arc_face_image = arc_face_image.cuda().contiguous()
+    if arcface_model is None:
+        arcface_model = init_recognition_model('arcface', device=devices.device)
+    face_emb = arcface_model(arc_face_image)[0] # [512], normalized
+    return face_emb
+
+
+def resize_and_pad_image(source_img, target_img_size):
+    # Get original and target sizes
+    source_img_size = source_img.size
+    target_width, target_height = target_img_size
+
+    # Determine the new size based on the shorter side of target_img
+    if target_width <= target_height:
+        new_width = target_width
+        new_height = int(target_width * (source_img_size[1] / source_img_size[0]))
+    else:
+        new_height = target_height
+        new_width = int(target_height * (source_img_size[0] / source_img_size[1]))
+
+    # Resize the source image using LANCZOS interpolation for high quality
+    resized_source_img = source_img.resize((new_width, new_height), Image.Resampling.LANCZOS)
+
+    # Compute padding to center resized image
+    pad_left = (target_width - new_width) // 2
+    pad_top = (target_height - new_height) // 2
+
+    # Create a new image with white background
+    padded_img = Image.new("RGB", target_img_size, (255, 255, 255))
+    padded_img.paste(resized_source_img, (pad_left, pad_top))
+
+    return padded_img
+
+
+class InfUFluxPipeline:
+    def __init__(
+            self,
+            pipe,
+            image_proj_num_tokens=8,
+            infu_flux_version='v1.0',
+            model_version='aes_stage2',
+        ):
+
+        self.infu_flux_version = infu_flux_version
+        self.model_version = model_version
+
+        # Load pipeline
+        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 = {}
+
+        self.infusenet = FluxControlNetModel.from_pretrained(
+            infusenet_path,
+            torch_dtype=devices.dtype,
+            **quant_args,
+        )
+
+        self.pipe = FluxInfuseNetPipeline(
+                vae=pipe.vae,
+                text_encoder=pipe.text_encoder,
+                text_encoder_2=pipe.text_encoder_2,
+                tokenizer=pipe.tokenizer,
+                tokenizer_2=pipe.tokenizer_2,
+                transformer=pipe.transformer,
+                scheduler=pipe.scheduler,
+                controlnet=self.infusenet,
+            )
+
+        # Load image proj model
+        num_tokens = image_proj_num_tokens
+        image_emb_dim = 512
+        image_proj_model = Resampler(
+            dim=1280,
+            depth=4,
+            dim_head=64,
+            heads=20,
+            num_queries=num_tokens,
+            embedding_dim=image_emb_dim,
+            output_dim=4096,
+            ff_mult=4,
+        )
+        image_proj_model_path = os.path.join(infiniteyou_path, 'image_proj_model.bin')
+        ipm_state_dict = torch.load(image_proj_model_path, map_location="cpu")
+        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
+
+        # Load face encoder
+        insightface_root_path = os.path.join(local_path, 'supports', 'insightface')
+        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)
+        self.app_320.prepare(ctx_id=0, det_size=(320, 320))
+        self.app_160 = FaceAnalysis(name='antelopev2', root=insightface_root_path, providers=devices.onnx)
+        self.app_160.prepare(ctx_id=0, det_size=(160, 160))
+        self.arcface_model = init_recognition_model('arcface', device=devices.device)
+
+    def load_loras(self, loras):
+        names, scales = [],[]
+        for lora_path, lora_name, lora_scale in loras:
+            if lora_path != "":
+                print(f"loading lora {lora_path}")
+                self.pipe.load_lora_weights(lora_path, adapter_name = lora_name)
+                names.append(lora_name)
+                scales.append(lora_scale)
+
+        if len(names) > 0:
+            self.pipe.set_adapters(names, adapter_weights=scales)
+
+    def _detect_face(self, id_image_cv2):
+        face_info = self.app_640.get(id_image_cv2)
+        if len(face_info) > 0:
+            return face_info
+
+        face_info = self.app_320.get(id_image_cv2)
+        if len(face_info) > 0:
+            return face_info
+
+        face_info = self.app_160.get(id_image_cv2)
+        return face_info
+
+    def __call__(
+        self,
+        prompt: str,
+        id_image: Image.Image, # PIL.Image.Image (RGB)
+        negative_prompt = None,
+        control_image: Optional[Image.Image] = None, # PIL.Image.Image (RGB) or None
+        width = 1024,
+        height = 1024,
+        seed = 42,
+        guidance_scale = 3.5,
+        controlnet_guidance_scale = 1.0,
+        num_inference_steps = 30,
+        infusenet_conditioning_scale = 1.0,
+        infusenet_guidance_start = 0.0,
+        infusenet_guidance_end = 1.0,
+        output_type = 'pil',
+        generator = None,
+        *args, **kwargs # pylint: disable=unused-argument
+    ):
+        # Extract ID embeddings
+        id_image_cv2 = cv2.cvtColor(np.array(id_image), cv2.COLOR_RGB2BGR)
+        face_info = self._detect_face(id_image_cv2)
+        if len(face_info) == 0:
+            raise ValueError('No face detected in the input ID image')
+
+        face_info = sorted(face_info, key=lambda x:(x['bbox'][2]-x['bbox'][0])*(x['bbox'][3]-x['bbox'][1]))[-1] # only use the maximum face
+        landmark = face_info['kps']
+        id_embed = extract_arcface_bgr_embedding(id_image_cv2, landmark, self.arcface_model)
+        id_embed = id_embed.clone().unsqueeze(0).float().cuda()
+        id_embed = id_embed.reshape([1, -1, 512])
+        id_embed = id_embed.to(device=devices.device, dtype=devices.dtype)
+        with torch.no_grad():
+            id_embed = self.image_proj_model(id_embed)
+            bs_embed, seq_len, _ = id_embed.shape
+            id_embed = id_embed.repeat(1, 1, 1)
+            id_embed = id_embed.view(bs_embed * 1, seq_len, -1)
+            id_embed = id_embed.to(device=devices.device, dtype=devices.dtype)
+
+        # Load control image
+        if control_image is not None:
+            control_image = control_image.convert("RGB")
+            control_image = resize_and_pad_image(control_image, (width, height))
+            face_info = self._detect_face(cv2.cvtColor(np.array(control_image), cv2.COLOR_RGB2BGR))
+            if len(face_info) == 0:
+                raise ValueError('No face detected in the control image')
+            face_info = sorted(face_info, key=lambda x:(x['bbox'][2]-x['bbox'][0])*(x['bbox'][3]-x['bbox'][1]))[-1] # only use the maximum face
+            control_image = draw_kps(control_image, face_info['kps'])
+        else:
+            out_img = np.zeros([height, width, 3])
+            control_image = Image.fromarray(out_img.astype(np.uint8))
+
+        """
+        control_image = self.pipe.prepare_image(
+            image=control_image,
+            width=width,
+            height=height,
+            batch_size=1,
+            num_images_per_prompt=1,
+            device=devices.device,
+            dtype=devices.dtype,
+        )
+        control_image = retrieve_latents(self.pipe.vae.encode(control_image), generator=generator)
+        control_image = (control_image - self.pipe.vae.config.shift_factor) * self.pipe.vae.config.scaling_factor
+        # pack
+        height_control_image, width_control_image = control_image.shape[2:]
+        num_channels_latents = self.pipe.transformer.config.in_channels // 4
+        control_image = self.pipe._pack_latents(
+            control_image,
+            1,
+            num_channels_latents,
+            height_control_image,
+            width_control_image,
+        )
+        """
+
+        # Perform inference
+        seed_everything(seed)
+        latents = self.pipe(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            controlnet_prompt_embeds=id_embed,
+            control_image=control_image,
+            guidance_scale=guidance_scale,
+            num_inference_steps=num_inference_steps,
+            controlnet_guidance_scale=controlnet_guidance_scale,
+            controlnet_conditioning_scale=infusenet_conditioning_scale,
+            control_guidance_start=infusenet_guidance_start,
+            control_guidance_end=infusenet_guidance_end,
+            height=height,
+            width=width,
+            output_type=output_type,
+            callback_on_step_end=kwargs.get('callback_on_step_end', None),
+            callback_on_step_end_tensor_inputs=kwargs.get('callback_on_step_end_tensor_inputs', None),
+        )
+
+        return latents
diff --git a/modules/infiniteyou/resampler.py b/modules/infiniteyou/resampler.py
new file mode 100644
index 000000000..6d0011e83
--- /dev/null
+++ b/modules/infiniteyou/resampler.py
@@ -0,0 +1,121 @@
+# Modified from https://github.com/mlfoundations/open_flamingo/blob/main/open_flamingo/src/helpers.py
+
+import math
+
+import torch
+import torch.nn as nn
+
+
+# FFN
+def FeedForward(dim, mult=4):
+    inner_dim = int(dim * mult)
+    return nn.Sequential(
+        nn.LayerNorm(dim),
+        nn.Linear(dim, inner_dim, bias=False),
+        nn.GELU(),
+        nn.Linear(inner_dim, dim, bias=False),
+    )
+
+
+def reshape_tensor(x, heads):
+    bs, length, width = x.shape
+    #(bs, length, width) --> (bs, length, n_heads, dim_per_head)
+    x = x.view(bs, length, heads, -1)
+    # (bs, length, n_heads, dim_per_head) --> (bs, n_heads, length, dim_per_head)
+    x = x.transpose(1, 2)
+    # (bs, n_heads, length, dim_per_head) --> (bs*n_heads, length, dim_per_head)
+    x = x.reshape(bs, heads, length, -1)
+    return x
+
+
+class PerceiverAttention(nn.Module):
+    def __init__(self, *, dim, dim_head=64, heads=8):
+        super().__init__()
+        self.scale = dim_head**-0.5
+        self.dim_head = dim_head
+        self.heads = heads
+        inner_dim = dim_head * heads
+
+        self.norm1 = nn.LayerNorm(dim)
+        self.norm2 = nn.LayerNorm(dim)
+
+        self.to_q = nn.Linear(dim, inner_dim, bias=False)
+        self.to_kv = nn.Linear(dim, inner_dim * 2, bias=False)
+        self.to_out = nn.Linear(inner_dim, dim, bias=False)
+
+    def forward(self, x, latents):
+        """
+        Args:
+            x (torch.Tensor): image features
+                shape (b, n1, D)
+            latent (torch.Tensor): latent features
+                shape (b, n2, D)
+        """
+        x = self.norm1(x)
+        latents = self.norm2(latents)
+
+        b, l, _ = latents.shape
+
+        q = self.to_q(latents)
+        kv_input = torch.cat((x, latents), dim=-2)
+        k, v = self.to_kv(kv_input).chunk(2, dim=-1)
+
+        q = reshape_tensor(q, self.heads)
+        k = reshape_tensor(k, self.heads)
+        v = reshape_tensor(v, self.heads)
+
+        # attention
+        scale = 1 / math.sqrt(math.sqrt(self.dim_head))
+        weight = (q * scale) @ (k * scale).transpose(-2, -1) # More stable with f16 than dividing afterwards
+        weight = torch.softmax(weight.float(), dim=-1).type(weight.dtype)
+        out = weight @ v
+
+        out = out.permute(0, 2, 1, 3).reshape(b, l, -1)
+
+        return self.to_out(out)
+
+
+class Resampler(nn.Module):
+    def __init__(
+        self,
+        dim=1024,
+        depth=8,
+        dim_head=64,
+        heads=16,
+        num_queries=8,
+        embedding_dim=768,
+        output_dim=1024,
+        ff_mult=4,
+    ):
+        super().__init__()
+
+        self.latents = nn.Parameter(torch.randn(1, num_queries, dim) / dim**0.5)
+
+        self.proj_in = nn.Linear(embedding_dim, dim)
+
+        self.proj_out = nn.Linear(dim, output_dim)
+        self.norm_out = nn.LayerNorm(output_dim)
+
+        self.layers = nn.ModuleList([])
+        for _ in range(depth):
+            self.layers.append(
+                nn.ModuleList(
+                    [
+                        PerceiverAttention(dim=dim, dim_head=dim_head, heads=heads),
+                        FeedForward(dim=dim, mult=ff_mult),
+                    ]
+                )
+            )
+
+    def forward(self, x):
+
+        latents = self.latents.repeat(x.size(0), 1, 1)
+
+        x = self.proj_in(x)
+
+        for attn, ff in self.layers:
+            latents = attn(x, latents) + latents
+            latents = ff(latents) + latents
+
+        latents = self.proj_out(latents)
+        return self.norm_out(latents)
diff --git a/modules/infotext.py b/modules/infotext.py
index 78c1fd92e..497879d31 100644
--- a/modules/infotext.py
+++ b/modules/infotext.py
@@ -31,7 +31,7 @@ def unquote(text):
 # disabled by default can be enabled if needed
 def check_lora(params):
     try:
-        import modules.lora.networks as networks
+        from modules.lora import lora_load
         from modules.errors import log # pylint: disable=redefined-outer-name
     except Exception:
         return
@@ -39,14 +39,14 @@ def check_lora(params):
     found = []
     missing = []
     for l in loras:
-        lora = networks.available_network_hash_lookup.get(l, None)
+        lora = lora_load.available_network_hash_lookup.get(l, None)
         if lora is not None:
             found.append(lora.name)
         else:
             missing.append(l)
     loras = [s.strip() for s in params.get('LoRA networks', '').split(',')]
     for l in loras:
-        lora = networks.available_network_aliases.get(l, None)
+        lora = lora_load.available_network_aliases.get(l, None)
         if lora is not None:
             found.append(lora.name)
         else:
@@ -54,7 +54,7 @@ def check_lora(params):
     # networks.available_network_aliases.get(name, None)
     loras = re_lora.findall(params.get('Prompt', ''))
     for l in loras:
-        lora = networks.available_network_aliases.get(l, None)
+        lora = lora_load.available_network_aliases.get(l, None)
         if lora is not None:
             found.append(lora.name)
         else:
diff --git a/modules/intel/ipex/attention.py b/modules/intel/ipex/attention.py
index 400b59b66..177f5bc5e 100644
--- a/modules/intel/ipex/attention.py
+++ b/modules/intel/ipex/attention.py
@@ -61,13 +61,13 @@ def dynamic_scaled_dot_product_attention(query, key, value, attn_mask=None, drop
     if query.device.type != "xpu":
         return original_scaled_dot_product_attention(query, key, value, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal, **kwargs)
     is_unsqueezed = False
-    if len(query.shape) == 3:
+    if query.dim() == 3:
         query = query.unsqueeze(0)
         is_unsqueezed = True
-    if len(key.shape) == 3:
-        key = key.unsqueeze(0)
-    if len(value.shape) == 3:
-        value = value.unsqueeze(0)
+        if key.dim() == 3:
+            key = key.unsqueeze(0)
+        if value.dim() == 3:
+            value = value.unsqueeze(0)
     do_batch_split, do_head_split, do_query_split, split_batch_size, split_head_size, split_query_size = find_sdpa_slice_sizes(query.shape, key.shape, query.element_size(), slice_rate=attention_slice_rate, trigger_rate=sdpa_slice_trigger_rate)
 
     # Slice SDPA
@@ -115,5 +115,5 @@ def dynamic_scaled_dot_product_attention(query, key, value, attn_mask=None, drop
     else:
         hidden_states = original_scaled_dot_product_attention(query, key, value, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal, **kwargs)
     if is_unsqueezed:
-        hidden_states.squeeze(0)
+        hidden_states = hidden_states.squeeze(0)
     return hidden_states
diff --git a/modules/intel/ipex/hijacks.py b/modules/intel/ipex/hijacks.py
index e47065a62..b9630afd3 100644
--- a/modules/intel/ipex/hijacks.py
+++ b/modules/intel/ipex/hijacks.py
@@ -118,7 +118,7 @@ original_torch_bmm = torch.bmm
 @wraps(torch.bmm)
 def torch_bmm(input, mat2, *, out=None):
     if input.dtype != mat2.dtype:
-        mat2 = mat2.to(input.dtype)
+        mat2 = mat2.to(dtype=input.dtype)
     return original_torch_bmm(input, mat2, out=out)
 
 # Diffusers FreeU
diff --git a/modules/intel/openvino/__init__.py b/modules/intel/openvino/__init__.py
index 2d2dc0b79..a5584575f 100644
--- a/modules/intel/openvino/__init__.py
+++ b/modules/intel/openvino/__init__.py
@@ -479,8 +479,8 @@ def openvino_fx(subgraph, example_inputs, options=None):
         subgraph_type[3] is torch.nn.modules.linear.Linear):
 
         dont_use_faketensors = True
-        dont_use_nncf = bool("Text Encoder" not in shared.opts.nncf_compress_weights)
-        dont_use_quant = bool("Text Encoder" not in shared.opts.nncf_quantize)
+        dont_use_nncf = bool("TE" not in shared.opts.nncf_compress_weights)
+        dont_use_quant = bool("TE" not in shared.opts.nncf_quantize)
 
     # Create a hash to be used for caching
     shared.compiled_model_state.model_hash_str = ""
diff --git a/modules/interrogate/deepbooru.py b/modules/interrogate/deepbooru.py
index 30227dc21..1e47e6cc8 100644
--- a/modules/interrogate/deepbooru.py
+++ b/modules/interrogate/deepbooru.py
@@ -4,7 +4,7 @@ import threading
 import torch
 import numpy as np
 from PIL import Image
-from modules import modelloader, paths, devices, shared
+from modules import modelloader, paths, devices, shared, sd_models
 
 re_special = re.compile(r'([\\()])')
 load_lock = threading.Lock()
@@ -35,11 +35,11 @@ class DeepDanbooru:
 
     def start(self):
         self.load()
-        self.model.to(devices.device)
+        sd_models.move_model(self.model, devices.device)
 
     def stop(self):
         if shared.opts.interrogate_offload:
-            self.model.to(devices.cpu)
+            sd_models.move_model(self.model, devices.cpu)
         devices.torch_gc()
 
     def tag(self, pil_image):
diff --git a/modules/interrogate/deepseek.py b/modules/interrogate/deepseek.py
index 5138c5693..b2d340248 100644
--- a/modules/interrogate/deepseek.py
+++ b/modules/interrogate/deepseek.py
@@ -12,7 +12,7 @@ import os
 import sys
 import importlib
 from transformers import AutoModelForCausalLM
-from modules import shared, devices, paths
+from modules import shared, devices, paths, sd_models
 
 
 # model_path = "deepseek-ai/deepseek-vl2-small"
@@ -73,7 +73,7 @@ def predict(question, image, repo):
     ).to(device=devices.device, dtype=devices.dtype)
     inputs_embeds = vl_gpt.prepare_inputs_embeds(**prepare_inputs)
     inputs_embeds = inputs_embeds.to(device=devices.device, dtype=devices.dtype)
-    vl_gpt = vl_gpt.to(devices.device)
+    sd_models.move_model(vl_gpt, devices.device)
     with devices.inference_context():
         outputs = vl_gpt.language.generate(
             inputs_embeds=inputs_embeds,
diff --git a/modules/interrogate/interrogate.py b/modules/interrogate/interrogate.py
index f68212d91..ce3f75193 100644
--- a/modules/interrogate/interrogate.py
+++ b/modules/interrogate/interrogate.py
@@ -28,7 +28,7 @@ def interrogate(image):
     elif shared.opts.interrogate_default_type == 'VLM':
         shared.log.info(f'Interrogate: type={shared.opts.interrogate_default_type} vlm="{shared.opts.interrogate_vlm_model}" prompt="{shared.opts.interrogate_vlm_prompt}"')
         from modules.interrogate import vqa
-        prompt = vqa.interrogate(image=image, model_name=shared.opts.interrogate_vlm_model, question=shared.opts.interrogate_vlm_prompt, prompt=None)
+        prompt = vqa.interrogate(image=image, model_name=shared.opts.interrogate_vlm_model, question=shared.opts.interrogate_vlm_prompt, prompt=None, system_prompt=shared.opts.interrogate_vlm_system)
         shared.log.debug(f'Interrogate: time={time.time()-t0:.2f} answer="{prompt}"')
         return prompt
     else:
diff --git a/modules/interrogate/openclip.py b/modules/interrogate/openclip.py
index 761c0fa39..792b4df85 100644
--- a/modules/interrogate/openclip.py
+++ b/modules/interrogate/openclip.py
@@ -10,7 +10,7 @@ import gradio as gr
 from PIL import Image
 from torchvision import transforms
 from torchvision.transforms.functional import InterpolationMode
-from modules import devices, paths, shared, lowvram, errors
+from modules import devices, paths, shared, lowvram, errors, sd_models
 
 
 caption_models = {
@@ -125,7 +125,7 @@ class InterrogateModels:
             else:
                 model, preprocess = clip.load(clip_model_name, download_root=shared.opts.clip_models_path)
             model.eval()
-            model = model.to(devices.device)
+            sd_models.move_model(model, devices.device)
             return model, preprocess
 
     def load(self):
@@ -133,23 +133,23 @@ class InterrogateModels:
             self.blip_model = self.load_blip_model()
             if not shared.opts.no_half and not self.running_on_cpu:
                 self.blip_model = self.blip_model.half()
-        self.blip_model = self.blip_model.to(devices.device)
         if self.clip_model is None:
             self.clip_model, self.clip_preprocess = self.load_clip_model()
             if not shared.opts.no_half and not self.running_on_cpu:
                 self.clip_model = self.clip_model.half()
-        self.clip_model = self.clip_model.to(devices.device)
         self.dtype = next(self.clip_model.parameters()).dtype
+        sd_models.move_model(self.blip_model, devices.device)
+        sd_models.move_model(self.clip_model, devices.device)
 
     def send_clip_to_ram(self):
         if shared.opts.interrogate_offload:
             if self.clip_model is not None:
-                self.clip_model = self.clip_model.to(devices.cpu)
+                sd_models.move_model(self.blip_model, devices.cpu)
 
     def send_blip_to_ram(self):
         if shared.opts.interrogate_offload:
             if self.blip_model is not None:
-                self.blip_model = self.blip_model.to(devices.cpu)
+                sd_models.move_model(self.blip_model, devices.cpu)
 
     def unload(self):
         self.send_clip_to_ram()
@@ -291,8 +291,8 @@ def load_interrogator(clip_model, blip_model):
 
 def unload_clip_model():
     if ci is not None and shared.opts.interrogate_offload:
-        ci.caption_model = ci.caption_model.to(devices.cpu)
-        ci.clip_model = ci.clip_model.to(devices.cpu)
+        sd_models.move_model(ci.caption_model, devices.cpu)
+        sd_models.move_model(ci.clip_model, devices.cpu)
         ci.caption_offloaded = True
         ci.clip_offloaded = True
         devices.torch_gc()
diff --git a/modules/interrogate/vqa.py b/modules/interrogate/vqa.py
index afe5aac09..b7f6aabb3 100644
--- a/modules/interrogate/vqa.py
+++ b/modules/interrogate/vqa.py
@@ -7,7 +7,8 @@ import torch
 import transformers
 import transformers.dynamic_module_utils
 from PIL import Image
-from modules import shared, devices, errors
+from modules import shared, devices, errors, sd_models
+
 
 processor = None
 model = None
@@ -22,28 +23,32 @@ vlm_models = {
     "CogFlorence 2.0 Large": "thwri/CogFlorence-2-Large-Freeze", # 1.6GB
     "CogFlorence 2.2 Large": "thwri/CogFlorence-2.2-Large", # 1.6GB
     "Moondream 2": "vikhyatk/moondream2", # 3.7GB
+    "Google Gemma 3 4B": "google/gemma-3-4b-it",
+    "Google Pix Textcaps": "google/pix2struct-textcaps-base", # 1.1GB
+    "Google PaliGemma 2 3B": "google/paligemma2-3b-pt-224",
     "Alibaba Qwen VL2 2B": "Qwen/Qwen2-VL-2B-Instruct",
     "Huggingface Smol VL2 0.5B": "HuggingFaceTB/SmolVLM-500M-Instruct",
     "Huggingface Smol VL2 2B": "HuggingFaceTB/SmolVLM-Instruct",
     "Salesforce BLIP Base": "Salesforce/blip-vqa-base", # 1.5GB
     "Salesforce BLIP Large": "Salesforce/blip-vqa-capfilt-large", # 1.5GB
-    "Google Pix Textcaps": "google/pix2struct-textcaps-base", # 1.1GB
     "Microsoft GIT TextCaps Base": "microsoft/git-base-textcaps", # 0.7GB
     "Microsoft GIT VQA Base": "microsoft/git-base-vqav2", # 0.7GB
     "Microsoft GIT VQA Large": "microsoft/git-large-vqav2", # 1.6GB
     "ToriiGate 0.4 2B": "Minthy/ToriiGate-v0.4-2B",
     "ToriiGate 0.4 7B": "Minthy/ToriiGate-v0.4-7B",
     "ViLT Base": "dandelin/vilt-b32-finetuned-vqa", # 0.5GB
-    "Google PaliGemma 2 3B": "google/paligemma2-3b-pt-224",
     "JoyCaption": "fancyfeast/llama-joycaption-alpha-two-hf-llava", # 0.7GB
     "JoyTag": "fancyfeast/joytag", # 17.4GB
     "AIDC Ovis2 1B": "AIDC-AI/Ovis2-1B",
     "AIDC Ovis2 2B": "AIDC-AI/Ovis2-2B",
     "AIDC Ovis2 4B": "AIDC-AI/Ovis2-4B",
+    "ByteDance Sa2VA 1B": "ByteDance/Sa2VA-1B",
+    "ByteDance Sa2VA 4B": "ByteDance/Sa2VA-4B",
     # "OpenGVLab InternVL 2.5 1B": "OpenGVLab/InternVL2_5-1B"
     # "DeepSeek VL2 Tiny": "deepseek-ai/deepseek-vl2-tiny", # broken
     # "nVidia Eagle 2 1B": "nvidia/Eagle2-1B", # not compatible with latest transformers
 }
+vlm_system = 'You are image captioning expert, creative, unbiased and uncensored.'
 vlm_prompts = [
     '',
     '',
@@ -70,6 +75,7 @@ def b64(image):
 
 
 def clean(response, question):
+    strip = ['---', '\r', '\t', '**', '"', '“', '”', 'Assistant:', 'Caption:', '<|im_end|>', '']
     if isinstance(response, dict):
         if 'task' in response:
             response = response['task']
@@ -81,12 +87,10 @@ def clean(response, question):
     question = question.replace('<', '').replace('>', '').replace('_', ' ')
     if question in response:
         response = response.split(question, 1)[1]
-    response = response.replace('\n', '').replace('\r', '').replace('\t', '').strip()
-    if response.startswith('"'):
-        response = response[1:]
-    if response.endswith('"'):
-        response = response[:-1]
-    response = response.replace('Assistant:', '').strip()
+    while any(s in response for s in strip):
+        for s in strip:
+            response = response.replace(s, '')
+    response = response.replace('\n\n', '\n').replace('  ', ' ').replace('*  ', '- ').strip()
     return response
 
 
@@ -106,24 +110,26 @@ def get_kwargs():
     return kwargs
 
 
-def qwen(question: str, image: Image.Image, repo: str = None):
+def qwen(question: str, image: Image.Image, repo: str = None, system_prompt: str = None):
     global processor, model, loaded # pylint: disable=global-statement
     if model is None or loaded != repo:
         shared.log.debug(f'Interrogate load: vlm="{repo}"')
+        model = None
         model = transformers.Qwen2VLForConditionalGeneration.from_pretrained(
             repo,
             cache_dir=shared.opts.hfcache_dir
         )
+        model = model.to(devices.device, devices.dtype)
         processor = transformers.AutoProcessor.from_pretrained(repo, cache_dir=shared.opts.hfcache_dir)
         loaded = repo
-    model = model.to(devices.device, devices.dtype)
+        devices.torch_gc()
+    sd_models.move_model(model, devices.device)
     question = question.replace('<', '').replace('>', '').replace('_', ' ')
+    system_prompt = system_prompt or shared.opts.vlm_system
     conversation = [
         {
             "role": "system",
-            "content": [
-                {"type": "text", "text": "You are image captioning expert, creative, unbiased and uncensored."}
-            ],
+            "content": [{"type": "text", "text": system_prompt}],
         },
         {
             "role": "user",
@@ -134,7 +140,6 @@ def qwen(question: str, image: Image.Image, repo: str = None):
         }
     ]
     text_prompt = processor.apply_chat_template(conversation, add_generation_prompt=True)
-    # '<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>Describe this image.<|im_end|>\n<|im_start|>assistant\n'
     inputs = processor(text=[text_prompt], images=[image], padding=True, return_tensors="pt")
     inputs = inputs.to(devices.device, devices.dtype)
     output_ids = model.generate(
@@ -149,18 +154,68 @@ def qwen(question: str, image: Image.Image, repo: str = None):
     return response
 
 
+def gemma(question: str, image: Image.Image, repo: str = None, system_prompt: str = None):
+    global processor, model, loaded # pylint: disable=global-statement
+    if not hasattr(transformers, 'Gemma3ForConditionalGeneration'):
+        shared.log.error(f'Interrogate: vlm="{repo}" gemma is not available')
+        return ''
+    if model is None or loaded != repo:
+        shared.log.debug(f'Interrogate load: vlm="{repo}"')
+        model = None
+        model = transformers.Gemma3ForConditionalGeneration.from_pretrained(repo, cache_dir=shared.opts.hfcache_dir)
+        model = model.to(devices.device, devices.dtype)
+        processor = transformers.AutoProcessor.from_pretrained(repo, cache_dir=shared.opts.hfcache_dir)
+        loaded = repo
+        devices.torch_gc()
+    sd_models.move_model(model, devices.device)
+    question = question.replace('<', '').replace('>', '').replace('_', ' ')
+    system_prompt = system_prompt or shared.opts.vlm_system
+    conversation = [
+        {
+            "role": "system",
+            "content": [{"type": "text", "text": system_prompt}]
+        },
+        {
+            "role": "user",
+            "content": [
+                {"type": "image", "image": b64(image)},
+                {"type": "text", "text": question}
+            ]
+        }
+    ]
+    inputs = processor.apply_chat_template(
+        conversation,
+        add_generation_prompt=True,
+        tokenize=True,
+        return_dict=True,
+        return_tensors="pt",
+    ).to(device=devices.device, dtype=devices.dtype)
+    input_len = inputs["input_ids"].shape[-1]
+    with devices.inference_context():
+        generation = model.generate(
+            **inputs,
+            **get_kwargs(),
+        )
+        generation = generation[0][input_len:]
+    response = processor.decode(generation, skip_special_tokens=True)
+    return response
+
+
 def paligemma(question: str, image: Image.Image, repo: str = None):
     global processor, model, loaded # pylint: disable=global-statement
     if model is None or loaded != repo:
         shared.log.debug(f'Interrogate load: vlm="{repo}"')
         processor = transformers.PaliGemmaProcessor.from_pretrained(repo, cache_dir=shared.opts.hfcache_dir)
+        model = None
         model = transformers.PaliGemmaForConditionalGeneration.from_pretrained(
             repo,
             cache_dir=shared.opts.hfcache_dir,
             torch_dtype=devices.dtype,
         )
+        model = model.to(devices.device, devices.dtype)
         loaded = repo
-    model = model.to(devices.device, devices.dtype)
+        devices.torch_gc()
+    sd_models.move_model(model, devices.device)
     question = question.replace('<', '').replace('>', '').replace('_', ' ')
     model_inputs = processor(text=question, images=image, return_tensors="pt").to(devices.device, devices.dtype)
     input_len = model_inputs["input_ids"].shape[-1]
@@ -183,6 +238,7 @@ def ovis(question: str, image: Image.Image, repo: str = None):
     global model, loaded # pylint: disable=global-statement
     if model is None or loaded != repo:
         shared.log.debug(f'Interrogate load: vlm="{repo}"')
+        model = None
         model = transformers.AutoModelForCausalLM.from_pretrained(
             repo,
             torch_dtype=devices.dtype,
@@ -190,8 +246,10 @@ def ovis(question: str, image: Image.Image, repo: str = None):
             trust_remote_code=True,
             cache_dir=shared.opts.hfcache_dir,
         )
+        model = model.to(devices.device, devices.dtype)
         loaded = repo
-    model = model.to(devices.device, devices.dtype)
+        devices.torch_gc()
+    sd_models.move_model(model, devices.device)
     text_tokenizer = model.get_text_tokenizer()
     visual_tokenizer = model.get_visual_tokenizer()
     max_partition = 9
@@ -219,26 +277,28 @@ def ovis(question: str, image: Image.Image, repo: str = None):
     return response
 
 
-def smol(question: str, image: Image.Image, repo: str = None):
+def smol(question: str, image: Image.Image, repo: str = None, system_prompt: str = None):
     global processor, model, loaded # pylint: disable=global-statement
     if model is None or loaded != repo:
         shared.log.debug(f'Interrogate load: vlm="{repo}"')
+        model = None
         model = transformers.AutoModelForVision2Seq.from_pretrained(
             repo,
             cache_dir=shared.opts.hfcache_dir,
             torch_dtype=devices.dtype,
             _attn_implementation="eager",
             )
+        model.to(devices.device, devices.dtype)
         processor = transformers.AutoProcessor.from_pretrained(repo, cache_dir=shared.opts.hfcache_dir)
         loaded = repo
-    model.to(devices.device, devices.dtype)
+        devices.torch_gc()
+    sd_models.move_model(model, devices.device)
     question = question.replace('<', '').replace('>', '').replace('_', ' ')
+    system_prompt = system_prompt or shared.opts.vlm_system
     conversation = [
         {
             "role": "system",
-            "content": [
-                {"type": "text", "text": "You are image captioning expert, creative, unbiased and uncensored."}
-            ],
+            "content": [{"type": "text", "text": system_prompt}],
         },
         {
             "role": "user",
@@ -249,7 +309,6 @@ def smol(question: str, image: Image.Image, repo: str = None):
         }
     ]
     text_prompt = processor.apply_chat_template(conversation, add_generation_prompt=True)
-    # '<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>Describe this image.<|im_end|>\n<|im_start|>assistant\n'
     inputs = processor(text=text_prompt, images=[image], padding=True, return_tensors="pt")
     inputs = inputs.to(devices.device, devices.dtype)
     output_ids = model.generate(
@@ -264,13 +323,16 @@ def git(question: str, image: Image.Image, repo: str = None):
     global processor, model, loaded # pylint: disable=global-statement
     if model is None or loaded != repo:
         shared.log.debug(f'Interrogate load: vlm="{repo}"')
+        model = None
         model = transformers.GitForCausalLM.from_pretrained(
             repo,
             cache_dir=shared.opts.hfcache_dir,
         )
+        model.to(devices.device, devices.dtype)
         processor = transformers.GitProcessor.from_pretrained(repo, cache_dir=shared.opts.hfcache_dir)
         loaded = repo
-    model.to(devices.device, devices.dtype)
+        devices.torch_gc()
+    sd_models.move_model(model, devices.device)
     pixel_values = processor(images=image, return_tensors="pt").pixel_values
     git_dict = {}
     git_dict['pixel_values'] = pixel_values.to(devices.device, devices.dtype)
@@ -289,13 +351,16 @@ def blip(question: str, image: Image.Image, repo: str = None):
     global processor, model, loaded # pylint: disable=global-statement
     if model is None or loaded != repo:
         shared.log.debug(f'Interrogate load: vlm="{repo}"')
+        model = None
         model = transformers.BlipForQuestionAnswering.from_pretrained(
             repo,
             cache_dir=shared.opts.hfcache_dir,
         )
+        model.to(devices.device, devices.dtype)
         processor = transformers.BlipProcessor.from_pretrained(repo, cache_dir=shared.opts.hfcache_dir)
         loaded = repo
-    model.to(devices.device, devices.dtype)
+        devices.torch_gc()
+    sd_models.move_model(model, devices.device)
     inputs = processor(image, question, return_tensors="pt")
     inputs = inputs.to(devices.device, devices.dtype)
     with devices.inference_context():
@@ -308,13 +373,16 @@ def vilt(question: str, image: Image.Image, repo: str = None):
     global processor, model, loaded # pylint: disable=global-statement
     if model is None or loaded != repo:
         shared.log.debug(f'Interrogate load: vlm="{repo}"')
+        model = None
         model = transformers.ViltForQuestionAnswering.from_pretrained(
             repo,
             cache_dir=shared.opts.hfcache_dir,
         )
+        model.to(devices.device)
         processor = transformers.ViltProcessor.from_pretrained(repo, cache_dir=shared.opts.hfcache_dir)
         loaded = repo
-    model.to(devices.device)
+        devices.torch_gc()
+    sd_models.move_model(model, devices.device)
     inputs = processor(image, question, return_tensors="pt")
     inputs = inputs.to(devices.device)
     with devices.inference_context():
@@ -329,13 +397,16 @@ def pix(question: str, image: Image.Image, repo: str = None):
     global processor, model, loaded # pylint: disable=global-statement
     if model is None or loaded != repo:
         shared.log.debug(f'Interrogate load: vlm="{repo}"')
+        model = None
         model = transformers.Pix2StructForConditionalGeneration.from_pretrained(
             repo,
             cache_dir=shared.opts.hfcache_dir,
         )
+        model.to(devices.device)
         processor = transformers.Pix2StructProcessor.from_pretrained(repo, cache_dir=shared.opts.hfcache_dir)
         loaded = repo
-    model.to(devices.device)
+        devices.torch_gc()
+    sd_models.move_model(model, devices.device)
     if len(question) > 0:
         inputs = processor(images=image, text=question, return_tensors="pt").to(devices.device)
     else:
@@ -350,6 +421,7 @@ def moondream(question: str, image: Image.Image, repo: str = None):
     global processor, model, loaded # pylint: disable=global-statement
     if model is None or loaded != repo:
         shared.log.debug(f'Interrogate load: vlm="{repo}"')
+        model = None
         model = transformers.AutoModelForCausalLM.from_pretrained(
             repo,
             revision="2024-08-26",
@@ -358,8 +430,10 @@ def moondream(question: str, image: Image.Image, repo: str = None):
         )
         processor = transformers.AutoTokenizer.from_pretrained(repo, cache_dir=shared.opts.hfcache_dir)
         loaded = repo
+        model.to(devices.device, devices.dtype)
         model.eval()
-    model.to(devices.device, devices.dtype)
+        devices.torch_gc()
+    sd_models.move_model(model, devices.device)
     question = question.replace('<', '').replace('>', '').replace('_', ' ')
     encoded = model.encode_image(image)
     with devices.inference_context():
@@ -381,6 +455,7 @@ def florence(question: str, image: Image.Image, repo: str = None, revision: str
     if model is None or loaded != repo:
         shared.log.debug(f'Interrogate load: vlm="{repo}" path="{shared.opts.hfcache_dir}"')
         transformers.dynamic_module_utils.get_imports = get_imports
+        model = None
         model = transformers.AutoModelForCausalLM.from_pretrained(
             repo,
             trust_remote_code=True,
@@ -390,8 +465,10 @@ def florence(question: str, image: Image.Image, repo: str = None, revision: str
         processor = transformers.AutoProcessor.from_pretrained(repo, trust_remote_code=True, revision=revision, cache_dir=shared.opts.hfcache_dir)
         transformers.dynamic_module_utils.get_imports = _get_imports
         loaded = repo
+        model.to(devices.device, devices.dtype)
         model.eval()
-    model.to(devices.device, devices.dtype)
+        devices.torch_gc()
+    sd_models.move_model(model, devices.device)
     if question.startswith('<'):
         task = question.split('>', 1)[0] + '>'
     else:
@@ -410,7 +487,43 @@ def florence(question: str, image: Image.Image, repo: str = None, revision: str
     return response
 
 
-def interrogate(question, prompt, image, model_name, quiet:bool=False):
+def sa2(question: str, image: Image.Image, repo: str = None):
+    global processor, model, loaded # pylint: disable=global-statement
+    if model is None or loaded != repo:
+        model = None
+        model = transformers.AutoModel.from_pretrained(
+            repo,
+            torch_dtype=devices.dtype,
+            low_cpu_mem_usage=True,
+            use_flash_attn=False,
+            trust_remote_code=True)
+        model = model.to(devices.device, devices.dtype)
+        model = model.eval()
+        processor = transformers.AutoTokenizer.from_pretrained(
+            repo,
+            trust_remote_code=True,
+            use_fast=False,
+        )
+        loaded = repo
+        devices.torch_gc()
+    sd_models.move_model(model, devices.device)
+    if question.startswith('<'):
+        task = question.split('>', 1)[0] + '>'
+    else:
+        task = ''
+    input_dict = {
+        'image': image,
+        'text': f'{task}',
+        'past_text': '',
+        'mask_prompts': None,
+        'tokenizer': processor,
+        }
+    return_dict = model.predict_forward(**input_dict)
+    response = return_dict["prediction"] # the text format answer
+    return response
+
+
+def interrogate(question, system_prompt, prompt, image, model_name, quiet:bool=False):
     if not quiet:
         shared.state.begin('Interrogate')
     t0 = time.time()
@@ -457,9 +570,9 @@ def interrogate(question, prompt, image, model_name, quiet:bool=False):
         elif 'florence' in vqa_model.lower():
             answer = florence(question, image, vqa_model)
         elif 'qwen' in vqa_model.lower() or 'torii' in vqa_model.lower():
-            answer = qwen(question, image, vqa_model)
+            answer = qwen(question, image, vqa_model, system_prompt)
         elif 'smol' in vqa_model.lower():
-            answer = smol(question, image, vqa_model)
+            answer = smol(question, image, vqa_model, system_prompt)
         elif 'joytag' in vqa_model.lower():
             from modules.interrogate import joytag
             answer = joytag.predict(image)
@@ -471,15 +584,19 @@ def interrogate(question, prompt, image, model_name, quiet:bool=False):
             answer = deepseek.predict(question, image, vqa_model)
         elif 'paligemma' in vqa_model.lower():
             answer = paligemma(question, image, vqa_model)
+        elif 'gemma' in vqa_model.lower():
+            answer = gemma(question, image, vqa_model, system_prompt)
         elif 'ovis' in vqa_model.lower():
             answer = ovis(question, image, vqa_model)
+        elif 'sa2' in vqa_model.lower():
+            answer = sa2(question, image, vqa_model)
         else:
             answer = 'unknown model'
     except Exception as e:
         errors.display(e, 'VQA')
         answer = 'error'
     if shared.opts.interrogate_offload and model is not None:
-        model.to(devices.cpu)
+        sd_models.move_model(model, devices.cpu)
     devices.torch_gc()
     answer = clean(answer, question)
     t1 = time.time()
@@ -489,7 +606,7 @@ def interrogate(question, prompt, image, model_name, quiet:bool=False):
     return answer
 
 
-def batch(model_name, batch_files, batch_folder, batch_str, question, prompt, write, append, recursive):
+def batch(model_name, system_prompt, batch_files, batch_folder, batch_str, question, prompt, write, append, recursive):
     class BatchWriter:
         def __init__(self, folder, mode='w'):
             self.folder = folder
@@ -536,7 +653,7 @@ def batch(model_name, batch_files, batch_folder, batch_str, question, prompt, wr
                 if shared.state.interrupted:
                     break
                 image = Image.open(file)
-                prompt = interrogate(question, prompt, image, model_name, quiet=True)
+                prompt = interrogate(question, system_prompt, prompt, image, model_name, quiet=True)
                 prompts.append(prompt)
                 if write:
                     writer.add(file, prompt)
diff --git a/modules/lora/extra_networks_lora.py b/modules/lora/extra_networks_lora.py
index 357c5291f..0f981d93a 100644
--- a/modules/lora/extra_networks_lora.py
+++ b/modules/lora/extra_networks_lora.py
@@ -2,11 +2,12 @@ from typing import List
 import os
 import re
 import numpy as np
-from modules.lora import networks, network_overrides
+from modules.lora import networks, lora_overrides, lora_load
+from modules.lora import lora_common as l
 from modules import extra_networks, shared
 
 
-debug = os.environ.get('SD_SCRIPT_DEBUG', None) is not None
+debug = os.environ.get('SD_LORA_DEBUG', None) is not None
 debug_log = shared.log.trace if debug else lambda *args, **kwargs: None
 
 
@@ -26,6 +27,7 @@ def get_stepwise(param, step, steps): # from https://github.com/cheald/sd-webui-
             if m[1][-1] <= 1.0:
                 step = step / (max_steps - step_offset) if max_steps > 0 else 1.0
             v = np.interp(step, m[1], m[0])
+            debug_log(f"Network load: type=LoRA step={step} steps={max_steps} v={v}")
             return v
         else:
             return m
@@ -38,7 +40,7 @@ def prompt(p):
     if shared.opts.lora_apply_tags == 0:
         return
     all_tags = []
-    for loaded in networks.loaded_networks:
+    for loaded in l.loaded_networks:
         page = [en for en in shared.extra_networks if en.name == 'lora'][0]
         item = page.create_item(loaded.name)
         tags = (item or {}).get("tags", {})
@@ -52,7 +54,7 @@ def prompt(p):
         all_tags = list(set(all_tags))
         all_tags = [t for t in all_tags if t not in p.prompt]
         if len(all_tags) > 0:
-            shared.log.debug(f"Load network: type=LoRA tags={all_tags} max={shared.opts.lora_apply_tags} apply")
+            shared.log.debug(f"Network load: type=LoRA tags={all_tags} max={shared.opts.lora_apply_tags} apply")
         all_tags = ', '.join(all_tags)
         p.extra_generation_params["LoRA tags"] = all_tags
         if '_tags_' in p.prompt:
@@ -68,12 +70,12 @@ def prompt(p):
 
 
 def infotext(p):
-    names = [i.name for i in networks.loaded_networks]
+    names = [i.name for i in l.loaded_networks]
     if len(names) > 0:
         p.extra_generation_params["LoRA networks"] = ", ".join(names)
     if shared.opts.lora_add_hashes_to_infotext:
         network_hashes = []
-        for item in networks.loaded_networks:
+        for item in l.loaded_networks:
             if not item.network_on_disk.shorthash:
                 continue
             network_hashes.append(item.network_on_disk.shorthash)
@@ -112,6 +114,19 @@ def parse(p, params_list, step=0):
     return names, te_multipliers, unet_multipliers, dyn_dims
 
 
+def unload_diffusers():
+    if hasattr(shared.sd_model, "unfuse_lora"):
+        try:
+            shared.sd_model.unfuse_lora()
+        except Exception:
+            pass
+    if hasattr(shared.sd_model, "unload_lora_weights"):
+        try:
+            shared.sd_model.unload_lora_weights() # fails for non-CLIP models
+        except Exception:
+            pass
+
+
 class ExtraNetworkLora(extra_networks.ExtraNetwork):
 
     def __init__(self):
@@ -129,12 +144,12 @@ class ExtraNetworkLora(extra_networks.ExtraNetwork):
             sd_model.loaded_loras = {}
         key = f'{",".join(include)}:{",".join(exclude)}'
         loaded = sd_model.loaded_loras.get(key, [])
-        # shared.log.trace(f'Load network: type=LoRA key="{key}" requested={requested} loaded={loaded}')
-        if (len(requested) == 0) or (len(requested) != len(loaded)):
+        # shared.log.trace(f'Network load: type=LoRA key="{key}" requested={requested} loaded={loaded}')
+        if len(requested) != len(loaded):
             sd_model.loaded_loras[key] = requested
             return True
-        for r, l in zip(requested, loaded):
-            if r != l:
+        for req, load in zip(requested, loaded):
+            if req != load:
                 sd_model.loaded_loras[key] = requested
                 return True
         return False
@@ -153,46 +168,42 @@ class ExtraNetworkLora(extra_networks.ExtraNetwork):
         if debug:
             import sys
             fn = f'{sys._getframe(2).f_code.co_name}:{sys._getframe(1).f_code.co_name}' # pylint: disable=protected-access
-            debug_log(f'Load network: type=LoRA include={include} exclude={exclude} requested={requested} fn={fn}')
+            debug_log(f'Network load: type=LoRA include={include} exclude={exclude} requested={requested} fn={fn}')
 
-        force_diffusers = network_overrides.check_override()
+        force_diffusers = lora_overrides.check_override()
         if force_diffusers:
             has_changed = False # diffusers handle their own loading
             if len(exclude) == 0:
-                networks.network_load(names, te_multipliers, unet_multipliers, dyn_dims) # load only on first call
+                shared.state.begin('LoRA')
+                lora_load.network_load(names, te_multipliers, unet_multipliers, dyn_dims) # load only on first call
+                shared.state.end()
         else:
-            networks.network_load(names, te_multipliers, unet_multipliers, dyn_dims) # load
+            lora_load.network_load(names, te_multipliers, unet_multipliers, dyn_dims) # load
             has_changed = self.changed(requested, include, exclude)
             if has_changed:
-                networks.network_deactivate(include, exclude)
+                shared.state.begin('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)
-                debug_log(f'Load network: type=LoRA previous={[n.name for n in networks.previously_loaded_networks]} current={[n.name for n in networks.loaded_networks]} changed')
+                if len(exclude) > 0: # only update on last activation
+                    l.previously_loaded_networks = l.loaded_networks.copy()
+                shared.state.end()
+                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(networks.loaded_networks) > 0 and (len(networks.applied_layers) > 0 or force_diffusers) and step == 0:
+        if len(l.loaded_networks) > 0 and (len(networks.applied_layers) > 0 or force_diffusers) and step == 0:
             infotext(p)
             prompt(p)
             if (has_changed or force_diffusers) and len(include) == 0: # print only once
-                shared.log.info(f'Load network: type=LoRA apply={[n.name for n in networks.loaded_networks]} mode={"fuse" if shared.opts.lora_fuse_diffusers else "backup"} te={te_multipliers} unet={unet_multipliers} time={networks.timer.summary}')
+                shared.log.info(f'Network load: type=LoRA apply={[n.name for n in l.loaded_networks]} mode={"fuse" if shared.opts.lora_fuse_diffusers else "backup"} te={te_multipliers} unet={unet_multipliers} time={l.timer.summary}')
 
     def deactivate(self, p):
-        if shared.native:
-            networks.previously_loaded_networks = networks.loaded_networks.copy()
-            debug_log(f'Load network: type=LoRA active={[n.name for n in networks.previously_loaded_networks]} deactivate')
-        if shared.native and len(networks.diffuser_loaded) > 0:
+        if shared.native and len(lora_load.diffuser_loaded) > 0:
             if not (shared.compiled_model_state is not None and shared.compiled_model_state.is_compiled is True):
-                if hasattr(shared.sd_model, "unfuse_lora"):
-                    try:
-                        shared.sd_model.unfuse_lora()
-                    except Exception:
-                        pass
-                if hasattr(shared.sd_model, "unload_lora_weights"):
-                    try:
-                        shared.sd_model.unload_lora_weights() # fails for non-CLIP models
-                    except Exception:
-                        pass
-        if self.active and networks.debug:
-            shared.log.debug(f"Network end: type=LoRA time={networks.timer.summary}")
+                unload_diffusers()
+        if self.active and l.debug:
+            shared.log.debug(f"Network end: type=LoRA time={l.timer.summary}")
         if self.errors:
             for k, v in self.errors.items():
-                shared.log.error(f'LoRA: name="{k}" errors={v}')
+                shared.log.error(f'Network: type=LoRA name="{k}" errors={v}')
             self.errors.clear()
diff --git a/modules/lora/lora_apply.py b/modules/lora/lora_apply.py
new file mode 100644
index 000000000..8bf327151
--- /dev/null
+++ b/modules/lora/lora_apply.py
@@ -0,0 +1,203 @@
+from typing import Union
+import re
+import time
+import torch
+import diffusers.models.lora
+from modules.lora import lora_common as l
+from modules import shared, devices, errors, model_quant
+
+
+bnb = None
+re_network_name = re.compile(r"(.*)\s*\([0-9a-fA-F]+\)")
+
+
+def network_backup_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn.GroupNorm, torch.nn.LayerNorm, diffusers.models.lora.LoRACompatibleLinear, diffusers.models.lora.LoRACompatibleConv], network_layer_name: str, wanted_names: tuple):
+    global bnb # pylint: disable=W0603
+    backup_size = 0
+    if len(l.loaded_networks) > 0 and network_layer_name is not None and any([net.modules.get(network_layer_name, None) for net in l.loaded_networks]): # noqa: C419 # pylint: disable=R1729
+        t0 = time.time()
+
+        weights_backup = getattr(self, "network_weights_backup", None)
+        bias_backup = getattr(self, "network_bias_backup", None)
+        if weights_backup is not None or bias_backup is not None:
+            if (shared.opts.lora_fuse_diffusers and not isinstance(weights_backup, bool)) or (not shared.opts.lora_fuse_diffusers and isinstance(weights_backup, bool)): # invalidate so we can change direct/backup on-the-fly
+                weights_backup = None
+                bias_backup = None
+                self.network_weights_backup = weights_backup
+                self.network_bias_backup = bias_backup
+
+        if weights_backup is None and wanted_names != (): # pylint: disable=C1803
+            weight = getattr(self, 'weight', None)
+            self.network_weights_backup = None
+            if getattr(weight, "quant_type", None) in ['nf4', 'fp4']:
+                if bnb is None:
+                    bnb = model_quant.load_bnb('Network load: type=LoRA', silent=True)
+                if bnb is not None:
+                    if shared.opts.lora_fuse_diffusers:
+                        self.network_weights_backup = True
+                    else:
+                        self.network_weights_backup = bnb.functional.dequantize_4bit(weight, quant_state=weight.quant_state, quant_type=weight.quant_type, blocksize=weight.blocksize,)
+                    self.quant_state, self.quant_type, self.blocksize = weight.quant_state, weight.quant_type, weight.blocksize
+                else:
+                    self.network_weights_backup = weight.clone().to(devices.cpu) if not shared.opts.lora_fuse_diffusers else True
+            else:
+                self.network_weights_backup = weight.clone().to(devices.cpu) if not shared.opts.lora_fuse_diffusers else True
+
+        if bias_backup is None:
+            if getattr(self, 'bias', None) is not None:
+                if shared.opts.lora_fuse_diffusers:
+                    self.network_bias_backup = True
+                else:
+                    bias_backup = self.bias.clone()
+                    bias_backup = bias_backup.to(devices.cpu)
+
+        if getattr(self, 'network_weights_backup', None) is not None:
+            backup_size += self.network_weights_backup.numel() * self.network_weights_backup.element_size() if isinstance(self.network_weights_backup, torch.Tensor) else 0
+        if getattr(self, 'network_bias_backup', None) is not None:
+            backup_size += self.network_bias_backup.numel() * self.network_bias_backup.element_size() if isinstance(self.network_bias_backup, torch.Tensor) else 0
+        l.timer.backup += time.time() - t0
+    return backup_size
+
+
+def network_calc_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn.GroupNorm, torch.nn.LayerNorm, diffusers.models.lora.LoRACompatibleLinear, diffusers.models.lora.LoRACompatibleConv], network_layer_name: str, use_previous: bool = False):
+    if shared.opts.diffusers_offload_mode == "none":
+        try:
+            self.to(devices.device)
+        except Exception:
+            pass
+    batch_updown = None
+    batch_ex_bias = None
+    loaded = l.loaded_networks if not use_previous else l.previously_loaded_networks
+    for net in loaded:
+        module = net.modules.get(network_layer_name, None)
+        if module is None:
+            continue
+        try:
+            t0 = time.time()
+            try:
+                weight = self.weight.to(devices.device)
+            except Exception:
+                weight = self.weight
+            updown, ex_bias = module.calc_updown(weight)
+            del module
+            if updown is not None:
+                if batch_updown is not None:
+                    batch_updown += updown.to(batch_updown.device)
+                else:
+                    batch_updown = updown.to(devices.device)
+            if ex_bias is not None:
+                if batch_ex_bias:
+                    batch_ex_bias += ex_bias.to(batch_ex_bias.device)
+                else:
+                    batch_ex_bias = ex_bias.to(devices.device)
+            l.timer.calc += time.time() - t0
+            if shared.opts.diffusers_offload_mode == "sequential":
+                t0 = time.time()
+                if batch_updown is not None:
+                    batch_updown = batch_updown.to(devices.cpu)
+                if batch_ex_bias is not None:
+                    batch_ex_bias = batch_ex_bias.to(devices.cpu)
+                t1 = time.time()
+                l.timer.move += t1 - t0
+        except RuntimeError as e:
+            l.extra_network_lora.errors[net.name] = l.extra_network_lora.errors.get(net.name, 0) + 1
+            module_name = net.modules.get(network_layer_name, None)
+            shared.log.error(f'Network: type=LoRA name="{net.name}" module="{module_name}" layer="{network_layer_name}" apply weight: {e}')
+            if l.debug:
+                errors.display(e, 'LoRA')
+                raise RuntimeError('LoRA apply weight') from e
+        continue
+    return batch_updown, batch_ex_bias
+
+
+def network_add_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn.GroupNorm, torch.nn.LayerNorm, diffusers.models.lora.LoRACompatibleLinear, diffusers.models.lora.LoRACompatibleConv], model_weights: Union[None, torch.Tensor] = None, lora_weights: torch.Tensor = None, deactivate: bool = False, bias: bool = False):
+    if lora_weights is None:
+        return None
+    if deactivate:
+        lora_weights *= -1
+    if model_weights is None: # weights are used if provided-from-backup else use self.weight
+        model_weights = self.weight
+    # TODO lora: add other quantization types
+    weight = None
+    if self.__class__.__name__ == 'Linear4bit' and bnb is not None:
+        try:
+            dequant_weight = bnb.functional.dequantize_4bit(model_weights.to(devices.device), quant_state=self.quant_state, quant_type=self.quant_type, blocksize=self.blocksize)
+            new_weight = dequant_weight.to(devices.device) + lora_weights.to(devices.device)
+            weight = bnb.nn.Params4bit(new_weight, quant_state=self.quant_state, quant_type=self.quant_type, blocksize=self.blocksize, requires_grad=False)
+            # TODO lora: maybe force imediate quantization
+            # weight._quantize(devices.device) / weight.to(device=device)
+        except Exception as e:
+            shared.log.error(f'Network load: type=LoRA quant=bnb cls={self.__class__.__name__} type={self.quant_type} blocksize={self.blocksize} state={vars(self.quant_state)} weight={self.weight} bias={lora_weights} {e}')
+    else:
+        try:
+            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
+        del model_weights
+        del lora_weights
+        weight = torch.nn.Parameter(new_weight, requires_grad=False)
+        del new_weight # without this its a massive memory leak
+    if weight is not None:
+        if not bias:
+            self.weight = weight
+        else:
+            self.bias = weight
+    return weight
+
+
+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):
+        bias_backup = True
+    if not weights_backup and not bias_backup:
+        return
+    t0 = time.time()
+
+    if weights_backup:
+        if updown is not None and len(self.weight.shape) == 4 and self.weight.shape[1] == 9: # inpainting model so zero pad updown to make channel 4 to 9
+            updown = torch.nn.functional.pad(updown, (0, 0, 0, 0, 0, 5)) # pylint: disable=not-callable
+        if updown is not None:
+            network_add_weights(self, lora_weights=updown, deactivate=deactivate, bias=False)
+
+    if bias_backup:
+        if ex_bias is not None:
+            network_add_weights(self, lora_weights=ex_bias, deactivate=deactivate, bias=True)
+
+    if hasattr(self, "qweight") and hasattr(self, "freeze"):
+        self.freeze()
+
+    l.timer.apply += time.time() - t0
+
+
+def network_apply_weights(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, device: torch.device, deactivate: bool = False):
+    weights_backup = getattr(self, "network_weights_backup", None)
+    bias_backup = getattr(self, "network_bias_backup", None)
+    if weights_backup is None and bias_backup is None:
+        return
+    t0 = time.time()
+
+    if weights_backup is not None:
+        self.weight = None
+        if updown is not None and len(weights_backup.shape) == 4 and weights_backup.shape[1] == 9: # inpainting model. zero pad updown to make channel[1]  4 to 9
+            updown = torch.nn.functional.pad(updown, (0, 0, 0, 0, 0, 5)) # pylint: disable=not-callable
+        if updown is not None:
+            network_add_weights(self, model_weights=weights_backup, lora_weights=updown, deactivate=deactivate, bias=False)
+        else:
+            self.weight = torch.nn.Parameter(weights_backup.to(device), requires_grad=False)
+
+    if bias_backup is not None:
+        self.bias = None
+        if ex_bias is not None:
+            network_add_weights(self, model_weights=bias_backup, lora_weights=ex_bias, deactivate=deactivate, bias=True)
+        else:
+            self.bias = torch.nn.Parameter(bias_backup.to(device), requires_grad=False)
+
+    if hasattr(self, "qweight") and hasattr(self, "freeze"):
+        self.freeze()
+
+    l.timer.apply += time.time() - t0
diff --git a/modules/lora/lora_common.py b/modules/lora/lora_common.py
new file mode 100644
index 000000000..a6b15ae13
--- /dev/null
+++ b/modules/lora/lora_common.py
@@ -0,0 +1,21 @@
+from typing import List
+import os
+from modules.lora import lora_timers
+from modules.lora import network_lora, network_hada, network_ia3, network_oft, network_lokr, network_full, network_norm, network_glora
+
+
+timer = lora_timers.Timer()
+debug = os.environ.get('SD_LORA_DEBUG', None) is not None
+module_types = [
+    network_lora.ModuleTypeLora(),
+    network_hada.ModuleTypeHada(),
+    network_ia3.ModuleTypeIa3(),
+    network_oft.ModuleTypeOFT(),
+    network_lokr.ModuleTypeLokr(),
+    network_full.ModuleTypeFull(),
+    network_norm.ModuleTypeNorm(),
+    network_glora.ModuleTypeGLora(),
+]
+loaded_networks: List = [] # no type due to circular import
+previously_loaded_networks: List = [] # no type due to circular import
+extra_network_lora = None # initialized in extra_networks.py
diff --git a/modules/lora/lora_extract.py b/modules/lora/lora_extract.py
index 21351d9e3..1217b952d 100644
--- a/modules/lora/lora_extract.py
+++ b/modules/lora/lora_extract.py
@@ -5,7 +5,7 @@ import datetime
 import torch
 from safetensors.torch import save_file
 import gradio as gr
-from rich import progress as p
+from rich import progress as rp
 from modules import shared, devices
 from modules.ui_common import create_refresh_button
 from modules.call_queue import wrap_gradio_gpu_call
@@ -134,7 +134,7 @@ def make_lora(fn, maxrank, auto_rank, rank_ratio, modules, overwrite):
     shared.log.debug(f'LoRA extract: modules={modules} maxrank={maxrank} auto={auto_rank} ratio={rank_ratio} fn="{fn}"')
     shared.state.begin('LoRA extract')
 
-    with p.Progress(p.TextColumn('[cyan]LoRA extract'), p.BarColumn(), p.TaskProgressColumn(), p.TimeRemainingColumn(), p.TimeElapsedColumn(), p.TextColumn('[cyan]{task.description}'), console=shared.console) as progress:
+    with rp.Progress(rp.TextColumn('[cyan]LoRA extract'), rp.BarColumn(), rp.TaskProgressColumn(), rp.TimeRemainingColumn(), rp.TimeElapsedColumn(), rp.TextColumn('[cyan]{task.description}'), console=shared.console) as progress:
 
         if 'te' in modules and getattr(shared.sd_model, 'text_encoder', None) is not None:
             modules = shared.sd_model.text_encoder.named_modules()
@@ -182,7 +182,7 @@ def make_lora(fn, maxrank, auto_rank, rank_ratio, modules, overwrite):
             progress.remove_task(task)
         t3 = time.time()
 
-        # TODO: lora make support quantized flux
+        # TODO: lora support pre-quantized flux
         # if 'te' in modules and getattr(shared.sd_model, 'transformer', None) is not None:
         #     for name, module in shared.sd_model.transformer.named_modules():
         #         if "norm" in name and "linear" not in name:
diff --git a/modules/lora/lora_load.py b/modules/lora/lora_load.py
new file mode 100644
index 000000000..110cc0e46
--- /dev/null
+++ b/modules/lora/lora_load.py
@@ -0,0 +1,283 @@
+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.lora import network, lora_overrides, lora_convert
+from modules.lora import lora_common as l
+
+
+diffuser_loaded = []
+diffuser_scales = []
+lora_cache = {}
+available_networks = {}
+available_network_aliases = {}
+forbidden_network_aliases = {}
+available_network_hash_lookup = {}
+
+
+def load_diffusers(name, network_on_disk, lora_scale=shared.opts.extra_networks_default_multiplier) -> Union[network.Network, None]:
+    t0 = time.time()
+    name = name.replace(".", "_")
+    shared.log.debug(f'Network load: type=LoRA name="{name}" file="{network_on_disk.filename}" detected={network_on_disk.sd_version} method=diffusers scale={lora_scale} fuse={shared.opts.lora_fuse_diffusers}')
+    if not shared.native:
+        return None
+    if not hasattr(shared.sd_model, 'load_lora_weights'):
+        shared.log.error(f'Network load: type=LoRA class={shared.sd_model.__class__} does not implement load lora')
+        return None
+    try:
+        shared.sd_model.load_lora_weights(network_on_disk.filename, adapter_name=name)
+    except Exception as e:
+        if 'already in use' in str(e):
+            pass
+        else:
+            if 'The following keys have not been correctly renamed' in str(e):
+                shared.log.error(f'Network load: type=LoRA name="{name}" diffusers unsupported format')
+            else:
+                shared.log.error(f'Network load: type=LoRA name="{name}" {e}')
+            if l.debug:
+                errors.display(e, "LoRA")
+            return None
+    if name not in diffuser_loaded:
+        diffuser_loaded.append(name)
+        diffuser_scales.append(lora_scale)
+    net = network.Network(name, network_on_disk)
+    net.mtime = os.path.getmtime(network_on_disk.filename)
+    l.timer.activate += time.time() - t0
+    return net
+
+
+def load_safetensors(name, network_on_disk) -> Union[network.Network, None]:
+    if not shared.sd_loaded:
+        return None
+
+    cached = lora_cache.get(name, None)
+    if l.debug:
+        shared.log.debug(f'Network load: type=LoRA name="{name}" file="{network_on_disk.filename}" type=lora {"cached" if cached else ""}')
+    if cached is not None:
+        return cached
+    net = network.Network(name, network_on_disk)
+    net.mtime = os.path.getmtime(network_on_disk.filename)
+    sd = sd_models.read_state_dict(network_on_disk.filename, what='network')
+    if shared.sd_model_type == 'f1': # if kohya flux lora, convert state_dict
+        sd = lora_convert._convert_kohya_flux_lora_to_diffusers(sd) or sd # pylint: disable=protected-access
+    if shared.sd_model_type == 'sd3': # if kohya flux lora, convert state_dict
+        try:
+            sd = lora_convert._convert_kohya_sd3_lora_to_diffusers(sd) or sd # pylint: disable=protected-access
+        except ValueError: # EAFP for diffusers PEFT keys
+            pass
+    lora_convert.assign_network_names_to_compvis_modules(shared.sd_model)
+    keys_failed_to_match = {}
+    matched_networks = {}
+    bundle_embeddings = {}
+    dtypes = []
+    convert = lora_convert.KeyConvert()
+    device = devices.device if shared.opts.lora_apply_gpu else devices.cpu
+    for key_network, weight in sd.items():
+        parts = key_network.split('.')
+        if parts[0] == "bundle_emb":
+            emb_name, vec_name = parts[1], key_network.split(".", 2)[-1]
+            emb_dict = bundle_embeddings.get(emb_name, {})
+            emb_dict[vec_name] = weight
+            bundle_embeddings[emb_name] = emb_dict
+            continue
+        if len(parts) > 5: # messy handler for diffusers peft lora
+            key_network_without_network_parts = '_'.join(parts[:-2])
+            if not key_network_without_network_parts.startswith('lora_'):
+                key_network_without_network_parts = 'lora_' + key_network_without_network_parts
+            network_part = '.'.join(parts[-2:]).replace('lora_A', 'lora_down').replace('lora_B', 'lora_up')
+        else:
+            key_network_without_network_parts, network_part = key_network.split(".", 1)
+        key, sd_module = convert(key_network_without_network_parts)
+        if sd_module is None:
+            keys_failed_to_match[key_network] = key
+            continue
+        if key not in matched_networks:
+            matched_networks[key] = network.NetworkWeights(network_key=key_network, sd_key=key, w={}, sd_module=sd_module)
+        matched_networks[key].w[network_part] = weight
+        if weight.dtype not in dtypes:
+            dtypes.append(weight.dtype)
+    network_types = []
+    for key, weights in matched_networks.items():
+        net_module = None
+        for nettype in l.module_types:
+            net_module = nettype.create_module(net, weights)
+            if net_module is not None:
+                network_types.append(nettype.__class__.__name__)
+                break
+        if net_module is None:
+            shared.log.error(f'LoRA unhandled: name={name} key={key} weights={weights.w.keys()}')
+        else:
+            net.modules[key] = net_module
+    if len(keys_failed_to_match) > 0:
+        shared.log.warning(f'Network load: type=LoRA name="{name}" type={set(network_types)} unmatched={len(keys_failed_to_match)} matched={len(matched_networks)}')
+        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}')
+    if len(matched_networks) == 0:
+        return None
+    lora_cache[name] = net
+    net.bundle_embeddings = bundle_embeddings
+    return net
+
+
+def maybe_recompile_model(names, te_multipliers):
+    recompile_model = False
+    skip_lora_load = False
+    if shared.compiled_model_state is not None and shared.compiled_model_state.is_compiled:
+        if len(names) == len(shared.compiled_model_state.lora_model):
+            for i, name in enumerate(names):
+                if shared.compiled_model_state.lora_model[
+                    i] != f"{name}:{te_multipliers[i] if te_multipliers else shared.opts.extra_networks_default_multiplier}":
+                    recompile_model = True
+                    shared.compiled_model_state.lora_model = []
+                    break
+            if not recompile_model:
+                skip_lora_load = True
+                if len(l.loaded_networks) > 0 and l.debug:
+                    shared.log.debug('Model Compile: Skipping LoRa loading')
+                return recompile_model, skip_lora_load
+        else:
+            recompile_model = True
+            shared.compiled_model_state.lora_model = []
+    if recompile_model:
+        backup_cuda_compile = shared.opts.cuda_compile
+        backup_scheduler = getattr(shared.sd_model, "scheduler", None)
+        sd_models.unload_model_weights(op='model')
+        shared.opts.cuda_compile = []
+        sd_models.reload_model_weights(op='model')
+        shared.opts.cuda_compile = backup_cuda_compile
+        if backup_scheduler is not None:
+            shared.sd_model.scheduler = backup_scheduler
+    return recompile_model, skip_lora_load
+
+
+def list_available_networks():
+    t0 = time.time()
+    available_networks.clear()
+    available_network_aliases.clear()
+    forbidden_network_aliases.clear()
+    available_network_hash_lookup.clear()
+    forbidden_network_aliases.update({"none": 1, "Addams": 1})
+    if not os.path.exists(shared.cmd_opts.lora_dir):
+        shared.log.warning(f'LoRA directory not found: path="{shared.cmd_opts.lora_dir}"')
+
+    def add_network(filename):
+        if not os.path.isfile(filename):
+            return
+        name = os.path.splitext(os.path.basename(filename))[0]
+        name = name.replace('.', '_')
+        try:
+            entry = network.NetworkOnDisk(name, filename)
+            available_networks[entry.name] = entry
+            if entry.alias in available_network_aliases:
+                forbidden_network_aliases[entry.alias.lower()] = 1
+            if shared.opts.lora_preferred_name == 'filename':
+                available_network_aliases[entry.name] = entry
+            else:
+                available_network_aliases[entry.alias] = entry
+            if entry.shorthash:
+                available_network_hash_lookup[entry.shorthash] = entry
+        except OSError as e: # should catch FileNotFoundError and PermissionError etc.
+            shared.log.error(f'LoRA: filename="{filename}" {e}')
+
+    candidates = sorted(files_cache.list_files(shared.cmd_opts.lora_dir, ext_filter=[".pt", ".ckpt", ".safetensors"]))
+    with concurrent.futures.ThreadPoolExecutor(max_workers=shared.max_workers) as executor:
+        for fn in candidates:
+            executor.submit(add_network, fn)
+    t1 = time.time()
+    l.timer.list = t1 - t0
+    shared.log.info(f'Available LoRAs: path="{shared.cmd_opts.lora_dir}" items={len(available_networks)} folders={len(forbidden_network_aliases)} time={t1 - t0:.2f}')
+
+
+def network_download(name):
+    from huggingface_hub import hf_hub_download
+    if os.path.exists(name):
+        return network.NetworkOnDisk(name, name)
+    parts = name.split('/')
+    if len(parts) >= 5 and parts[1] == 'huggingface.co':
+        repo_id = f'{parts[2]}/{parts[3]}'
+        filename = '/'.join(parts[4:])
+        fn = hf_hub_download(repo_id=repo_id, filename=filename, cache_dir=shared.opts.hfcache_dir)
+        return network.NetworkOnDisk(name, fn)
+    return None
+
+
+def network_load(names, te_multipliers=None, unet_multipliers=None, dyn_dims=None):
+    networks_on_disk: list[network.NetworkOnDisk] = [available_network_aliases.get(name, None) for name in names]
+    if any(x is None for x in networks_on_disk):
+        list_available_networks()
+        networks_on_disk: list[network.NetworkOnDisk] = [available_network_aliases.get(name, None) for name in names]
+    for i in range(len(names)):
+        if names[i].startswith('/'):
+            networks_on_disk[i] = network_download(names[i])
+    failed_to_load_networks = []
+    recompile_model, skip_lora_load = maybe_recompile_model(names, te_multipliers)
+
+    l.loaded_networks.clear()
+    diffuser_loaded.clear()
+    diffuser_scales.clear()
+    t0 = time.time()
+
+    for i, (network_on_disk, name) in enumerate(zip(networks_on_disk, names)):
+        net = None
+        if network_on_disk is not None:
+            shorthash = getattr(network_on_disk, 'shorthash', '').lower()
+            if l.debug:
+                shared.log.debug(f'Network load: type=LoRA name="{name}" file="{network_on_disk.filename}" hash="{shorthash}"')
+            try:
+                if recompile_model:
+                    shared.compiled_model_state.lora_model.append(f"{name}:{te_multipliers[i] if te_multipliers else shared.opts.extra_networks_default_multiplier}")
+                if shared.opts.lora_force_diffusers or lora_overrides.check_override(shorthash): # OpenVINO only works with Diffusers LoRa loading
+                    net = load_diffusers(name, network_on_disk, lora_scale=te_multipliers[i] if te_multipliers else shared.opts.extra_networks_default_multiplier)
+                else:
+                    net = load_safetensors(name, network_on_disk)
+                if net is not None:
+                    net.mentioned_name = name
+                    network_on_disk.read_hash()
+            except Exception as e:
+                shared.log.error(f'Network load: type=LoRA file="{network_on_disk.filename}" {e}')
+                if l.debug:
+                    errors.display(e, 'LoRA')
+                continue
+        if net is None:
+            failed_to_load_networks.append(name)
+            shared.log.error(f'Network load: type=LoRA name="{name}" detected={network_on_disk.sd_version if network_on_disk is not None else None} failed')
+            continue
+        if hasattr(shared.sd_model, 'embedding_db'):
+            shared.sd_model.embedding_db.load_diffusers_embedding(None, net.bundle_embeddings)
+        net.te_multiplier = te_multipliers[i] if te_multipliers else shared.opts.extra_networks_default_multiplier
+        net.unet_multiplier = unet_multipliers[i] if unet_multipliers else shared.opts.extra_networks_default_multiplier
+        net.dyn_dim = dyn_dims[i] if dyn_dims else shared.opts.extra_networks_default_multiplier
+        l.loaded_networks.append(net)
+
+    while len(lora_cache) > shared.opts.lora_in_memory_limit:
+        name = next(iter(lora_cache))
+        lora_cache.pop(name, None)
+
+    if not skip_lora_load and len(diffuser_loaded) > 0:
+        shared.log.debug(f'Network load: type=LoRA loaded={diffuser_loaded} available={shared.sd_model.get_list_adapters()} active={shared.sd_model.get_active_adapters()} scales={diffuser_scales}')
+        try:
+            t0 = time.time()
+            shared.sd_model.set_adapters(adapter_names=diffuser_loaded, adapter_weights=diffuser_scales)
+            if shared.opts.lora_fuse_diffusers and not lora_overrides.check_fuse():
+                shared.sd_model.fuse_lora(adapter_names=diffuser_loaded, lora_scale=1.0, fuse_unet=True, fuse_text_encoder=True) # diffusers with fuse uses fixed scale since later apply does the scaling
+                shared.sd_model.unload_lora_weights()
+            l.timer.activate += time.time() - t0
+        except Exception as e:
+            shared.log.error(f'Network load: type=LoRA {e}')
+            if l.debug:
+                errors.display(e, 'LoRA')
+
+    if len(l.loaded_networks) > 0 and l.debug:
+        shared.log.debug(f'Network load: type=LoRA loaded={[n.name for n in l.loaded_networks]} cache={list(lora_cache)}')
+
+    if recompile_model:
+        shared.log.info("Network load: type=LoRA recompiling model")
+        backup_lora_model = shared.compiled_model_state.lora_model
+        if 'Model' in shared.opts.cuda_compile:
+            shared.sd_model = sd_models_compile.compile_diffusers(shared.sd_model)
+        shared.compiled_model_state.lora_model = backup_lora_model
+
+    l.timer.load = time.time() - t0
diff --git a/modules/lora/network_overrides.py b/modules/lora/lora_overrides.py
similarity index 94%
rename from modules/lora/network_overrides.py
rename to modules/lora/lora_overrides.py
index 65448ef2d..22d251c47 100644
--- a/modules/lora/network_overrides.py
+++ b/modules/lora/lora_overrides.py
@@ -30,8 +30,14 @@ force_models = [ # forced always
     'sc',
     'kandinsky',
     'hunyuandit',
-    'hunyuanvideo',
     'auraflow',
+    # video models
+    'hunyuanvideo',
+    'cogvideo',
+    'wanvideo',
+    'ltxvideo',
+    'mochivideo',
+    'allegrovideo',
 ]
 
 force_classes = [ # forced always
diff --git a/modules/lora/network.py b/modules/lora/network.py
index 97feb76f1..f6d93009c 100644
--- a/modules/lora/network.py
+++ b/modules/lora/network.py
@@ -17,6 +17,7 @@ class SdVersion(enum.Enum):
     SDXL = 4
     SC = 5
     F1 = 6
+    HV = 7
 
 
 class NetworkOnDisk:
@@ -56,6 +57,8 @@ class NetworkOnDisk:
             return 'sd3'
         if base.startswith("flux"):
             return 'f1'
+        if base.startswith("hunyuan_video"):
+            return 'hv'
 
         if arch.startswith("stable-diffusion-v1"):
             return 'sd1'
@@ -65,6 +68,8 @@ class NetworkOnDisk:
             return 'sc'
         if arch.startswith("flux"):
             return 'f1'
+        if arch.startswith("hunyuan-video"):
+            return 'hv'
 
         if "v1-5" in str(self.metadata.get('ss_sd_model_name', "")):
             return 'sd1'
@@ -86,8 +91,10 @@ class NetworkOnDisk:
             self.set_hash(hashes.sha256(self.filename, "lora/" + self.name, use_addnet_hash=self.is_safetensors) or '')
 
     def get_alias(self):
-        import modules.lora.networks as networks
-        return self.name if shared.opts.lora_preferred_name == "filename" or self.alias.lower() in networks.forbidden_network_aliases else self.alias
+        if shared.opts.lora_preferred_name == "filename":
+            return self.name
+        else:
+            return self.alias
 
 
 class Network:  # LoraModule
diff --git a/modules/lora/networks.py b/modules/lora/networks.py
index 9e981a234..a36ba5631 100644
--- a/modules/lora/networks.py
+++ b/modules/lora/networks.py
@@ -1,570 +1,17 @@
-from typing import Union, List
 from contextlib import nullcontext
-import os
-import re
 import time
-import concurrent
-import torch
-import diffusers.models.lora
 import rich.progress as rp
-
-from modules.lora import lora_timers, network, lora_convert, network_overrides
-from modules.lora import network_lora, network_hada, network_ia3, network_oft, network_lokr, network_full, network_norm, network_glora
-from modules.lora.extra_networks_lora import ExtraNetworkLora
-from modules import shared, devices, sd_models, sd_models_compile, errors, files_cache, model_quant
+from modules.lora import lora_common as l
+from modules.lora.lora_apply import network_apply_weights, network_apply_direct, network_backup_weights, network_calc_weights
+from modules import shared, devices, sd_models
 
 
-debug = os.environ.get('SD_LORA_DEBUG', None) is not None
-extra_network_lora = ExtraNetworkLora()
-available_networks = {}
-available_network_aliases = {}
-loaded_networks: List[network.Network] = []
-previously_loaded_networks: List[network.Network] = []
 applied_layers: list[str] = []
-bnb = None
-lora_cache = {}
-diffuser_loaded = []
-diffuser_scales = []
-available_network_hash_lookup = {}
-forbidden_network_aliases = {}
-re_network_name = re.compile(r"(.*)\s*\([0-9a-fA-F]+\)")
-timer = lora_timers.Timer()
-module_types = [
-    network_lora.ModuleTypeLora(),
-    network_hada.ModuleTypeHada(),
-    network_ia3.ModuleTypeIa3(),
-    network_oft.ModuleTypeOFT(),
-    network_lokr.ModuleTypeLokr(),
-    network_full.ModuleTypeFull(),
-    network_norm.ModuleTypeNorm(),
-    network_glora.ModuleTypeGLora(),
-]
-
-# section: load networks from disk
-
-def load_diffusers(name, network_on_disk, lora_scale=shared.opts.extra_networks_default_multiplier) -> Union[network.Network, None]:
-    t0 = time.time()
-    name = name.replace(".", "_")
-    shared.log.debug(f'Load network: type=LoRA name="{name}" file="{network_on_disk.filename}" detected={network_on_disk.sd_version} method=diffusers scale={lora_scale} fuse={shared.opts.lora_fuse_diffusers}')
-    if not shared.native:
-        return None
-    if not hasattr(shared.sd_model, 'load_lora_weights'):
-        shared.log.error(f'Load network: type=LoRA class={shared.sd_model.__class__} does not implement load lora')
-        return None
-    try:
-        shared.sd_model.load_lora_weights(network_on_disk.filename, adapter_name=name)
-    except Exception as e:
-        if 'already in use' in str(e):
-            pass
-        else:
-            if 'The following keys have not been correctly renamed' in str(e):
-                shared.log.error(f'Load network: type=LoRA name="{name}" diffusers unsupported format')
-            else:
-                shared.log.error(f'Load network: type=LoRA name="{name}" {e}')
-            if debug:
-                errors.display(e, "LoRA")
-            return None
-    if name not in diffuser_loaded:
-        diffuser_loaded.append(name)
-        diffuser_scales.append(lora_scale)
-    net = network.Network(name, network_on_disk)
-    net.mtime = os.path.getmtime(network_on_disk.filename)
-    timer.activate += time.time() - t0
-    return net
-
-
-def load_safetensors(name, network_on_disk) -> Union[network.Network, None]:
-    if not shared.sd_loaded:
-        return None
-
-    cached = lora_cache.get(name, None)
-    if debug:
-        shared.log.debug(f'Load network: type=LoRA name="{name}" file="{network_on_disk.filename}" type=lora {"cached" if cached else ""}')
-    if cached is not None:
-        return cached
-    net = network.Network(name, network_on_disk)
-    net.mtime = os.path.getmtime(network_on_disk.filename)
-    sd = sd_models.read_state_dict(network_on_disk.filename, what='network')
-    if shared.sd_model_type == 'f1':  # if kohya flux lora, convert state_dict
-        sd = lora_convert._convert_kohya_flux_lora_to_diffusers(sd) or sd  # pylint: disable=protected-access
-    if shared.sd_model_type == 'sd3':  # if kohya flux lora, convert state_dict
-        try:
-            sd = lora_convert._convert_kohya_sd3_lora_to_diffusers(sd) or sd  # pylint: disable=protected-access
-        except ValueError:  # EAFP for diffusers PEFT keys
-            pass
-    lora_convert.assign_network_names_to_compvis_modules(shared.sd_model)
-    keys_failed_to_match = {}
-    matched_networks = {}
-    bundle_embeddings = {}
-    convert = lora_convert.KeyConvert()
-    for key_network, weight in sd.items():
-        parts = key_network.split('.')
-        if parts[0] == "bundle_emb":
-            emb_name, vec_name = parts[1], key_network.split(".", 2)[-1]
-            emb_dict = bundle_embeddings.get(emb_name, {})
-            emb_dict[vec_name] = weight
-            bundle_embeddings[emb_name] = emb_dict
-            continue
-        if len(parts) > 5: # messy handler for diffusers peft lora
-            key_network_without_network_parts = '_'.join(parts[:-2])
-            if not key_network_without_network_parts.startswith('lora_'):
-                key_network_without_network_parts = 'lora_' + key_network_without_network_parts
-            network_part = '.'.join(parts[-2:]).replace('lora_A', 'lora_down').replace('lora_B', 'lora_up')
-        else:
-            key_network_without_network_parts, network_part = key_network.split(".", 1)
-        key, sd_module = convert(key_network_without_network_parts)
-        if sd_module is None:
-            keys_failed_to_match[key_network] = key
-            continue
-        if key not in matched_networks:
-            matched_networks[key] = network.NetworkWeights(network_key=key_network, sd_key=key, w={}, sd_module=sd_module)
-        matched_networks[key].w[network_part] = weight
-    network_types = []
-    for key, weights in matched_networks.items():
-        net_module = None
-        for nettype in module_types:
-            net_module = nettype.create_module(net, weights)
-            if net_module is not None:
-                network_types.append(nettype.__class__.__name__)
-                break
-        if net_module is None:
-            shared.log.error(f'LoRA unhandled: name={name} key={key} weights={weights.w.keys()}')
-        else:
-            net.modules[key] = net_module
-    if len(keys_failed_to_match) > 0:
-        shared.log.warning(f'Load network: type=LoRA name="{name}" type={set(network_types)} unmatched={len(keys_failed_to_match)} matched={len(matched_networks)}')
-        if debug:
-            shared.log.debug(f'Load network: type=LoRA name="{name}" unmatched={keys_failed_to_match}')
-    else:
-        shared.log.debug(f'Load network: type=LoRA name="{name}" type={set(network_types)} keys={len(matched_networks)} direct={shared.opts.lora_fuse_diffusers}')
-    if len(matched_networks) == 0:
-        return None
-    lora_cache[name] = net
-    net.bundle_embeddings = bundle_embeddings
-    return net
-
-
-def maybe_recompile_model(names, te_multipliers):
-    recompile_model = False
-    skip_lora_load = False
-    if shared.compiled_model_state is not None and shared.compiled_model_state.is_compiled:
-        if len(names) == len(shared.compiled_model_state.lora_model):
-            for i, name in enumerate(names):
-                if shared.compiled_model_state.lora_model[
-                    i] != f"{name}:{te_multipliers[i] if te_multipliers else shared.opts.extra_networks_default_multiplier}":
-                    recompile_model = True
-                    shared.compiled_model_state.lora_model = []
-                    break
-            if not recompile_model:
-                skip_lora_load = True
-                if len(loaded_networks) > 0 and debug:
-                    shared.log.debug('Model Compile: Skipping LoRa loading')
-                return recompile_model, skip_lora_load
-        else:
-            recompile_model = True
-            shared.compiled_model_state.lora_model = []
-    if recompile_model:
-        backup_cuda_compile = shared.opts.cuda_compile
-        backup_scheduler = getattr(shared.sd_model, "scheduler", None)
-        sd_models.unload_model_weights(op='model')
-        shared.opts.cuda_compile = []
-        sd_models.reload_model_weights(op='model')
-        shared.opts.cuda_compile = backup_cuda_compile
-        if backup_scheduler is not None:
-            shared.sd_model.scheduler = backup_scheduler
-    return recompile_model, skip_lora_load
-
-
-def list_available_networks():
-    t0 = time.time()
-    available_networks.clear()
-    available_network_aliases.clear()
-    forbidden_network_aliases.clear()
-    available_network_hash_lookup.clear()
-    forbidden_network_aliases.update({"none": 1, "Addams": 1})
-    if not os.path.exists(shared.cmd_opts.lora_dir):
-        shared.log.warning(f'LoRA directory not found: path="{shared.cmd_opts.lora_dir}"')
-
-    def add_network(filename):
-        if not os.path.isfile(filename):
-            return
-        name = os.path.splitext(os.path.basename(filename))[0]
-        name = name.replace('.', '_')
-        try:
-            entry = network.NetworkOnDisk(name, filename)
-            available_networks[entry.name] = entry
-            if entry.alias in available_network_aliases:
-                forbidden_network_aliases[entry.alias.lower()] = 1
-            if shared.opts.lora_preferred_name == 'filename':
-                available_network_aliases[entry.name] = entry
-            else:
-                available_network_aliases[entry.alias] = entry
-            if entry.shorthash:
-                available_network_hash_lookup[entry.shorthash] = entry
-        except OSError as e:  # should catch FileNotFoundError and PermissionError etc.
-            shared.log.error(f'LoRA: filename="{filename}" {e}')
-
-    candidates = sorted(files_cache.list_files(shared.cmd_opts.lora_dir, ext_filter=[".pt", ".ckpt", ".safetensors"]))
-    with concurrent.futures.ThreadPoolExecutor(max_workers=shared.max_workers) as executor:
-        for fn in candidates:
-            executor.submit(add_network, fn)
-    t1 = time.time()
-    timer.list = t1 - t0
-    shared.log.info(f'Available LoRAs: path="{shared.cmd_opts.lora_dir}" items={len(available_networks)} folders={len(forbidden_network_aliases)} time={t1 - t0:.2f}')
-
-
-def network_download(name):
-    from huggingface_hub import hf_hub_download
-    if os.path.exists(name):
-        return network.NetworkOnDisk(name, name)
-    parts = name.split('/')
-    if len(parts) >= 5 and parts[1] == 'huggingface.co':
-        repo_id = f'{parts[2]}/{parts[3]}'
-        filename = '/'.join(parts[4:])
-        fn = hf_hub_download(repo_id=repo_id, filename=filename, cache_dir=shared.opts.hfcache_dir)
-        return network.NetworkOnDisk(name, fn)
-    return None
-
-
-def network_load(names, te_multipliers=None, unet_multipliers=None, dyn_dims=None):
-    networks_on_disk: list[network.NetworkOnDisk] = [available_network_aliases.get(name, None) for name in names]
-    if any(x is None for x in networks_on_disk):
-        list_available_networks()
-        networks_on_disk: list[network.NetworkOnDisk] = [available_network_aliases.get(name, None) for name in names]
-    for i in range(len(names)):
-        if names[i].startswith('/'):
-            networks_on_disk[i] = network_download(names[i])
-    failed_to_load_networks = []
-    recompile_model, skip_lora_load = maybe_recompile_model(names, te_multipliers)
-
-    loaded_networks.clear()
-    diffuser_loaded.clear()
-    diffuser_scales.clear()
-    t0 = time.time()
-
-    for i, (network_on_disk, name) in enumerate(zip(networks_on_disk, names)):
-        net = None
-        if network_on_disk is not None:
-            shorthash = getattr(network_on_disk, 'shorthash', '').lower()
-            if debug:
-                shared.log.debug(f'Load network: type=LoRA name="{name}" file="{network_on_disk.filename}" hash="{shorthash}"')
-            try:
-                if recompile_model:
-                    shared.compiled_model_state.lora_model.append(f"{name}:{te_multipliers[i] if te_multipliers else shared.opts.extra_networks_default_multiplier}")
-                if shared.opts.lora_force_diffusers or network_overrides.check_override(shorthash): # OpenVINO only works with Diffusers LoRa loading
-                    net = load_diffusers(name, network_on_disk, lora_scale=te_multipliers[i] if te_multipliers else shared.opts.extra_networks_default_multiplier)
-                else:
-                    net = load_safetensors(name, network_on_disk)
-                if net is not None:
-                    net.mentioned_name = name
-                    network_on_disk.read_hash()
-            except Exception as e:
-                shared.log.error(f'Load network: type=LoRA file="{network_on_disk.filename}" {e}')
-                if debug:
-                    errors.display(e, 'LoRA')
-                continue
-        if net is None:
-            failed_to_load_networks.append(name)
-            shared.log.error(f'Load network: type=LoRA name="{name}" detected={network_on_disk.sd_version if network_on_disk is not None else None} failed')
-            continue
-        if hasattr(shared.sd_model, 'embedding_db'):
-            shared.sd_model.embedding_db.load_diffusers_embedding(None, net.bundle_embeddings)
-        net.te_multiplier = te_multipliers[i] if te_multipliers else shared.opts.extra_networks_default_multiplier
-        net.unet_multiplier = unet_multipliers[i] if unet_multipliers else shared.opts.extra_networks_default_multiplier
-        net.dyn_dim = dyn_dims[i] if dyn_dims else shared.opts.extra_networks_default_multiplier
-        loaded_networks.append(net)
-
-    while len(lora_cache) > shared.opts.lora_in_memory_limit:
-        name = next(iter(lora_cache))
-        lora_cache.pop(name, None)
-
-    if not skip_lora_load and len(diffuser_loaded) > 0:
-        shared.log.debug(f'Load network: type=LoRA loaded={diffuser_loaded} available={shared.sd_model.get_list_adapters()} active={shared.sd_model.get_active_adapters()} scales={diffuser_scales}')
-        try:
-            t0 = time.time()
-            shared.sd_model.set_adapters(adapter_names=diffuser_loaded, adapter_weights=diffuser_scales)
-            if shared.opts.lora_fuse_diffusers and not network_overrides.check_fuse():
-                shared.sd_model.fuse_lora(adapter_names=diffuser_loaded, lora_scale=1.0, fuse_unet=True, fuse_text_encoder=True) # fuse uses fixed scale since later apply does the scaling
-                shared.sd_model.unload_lora_weights()
-            timer.activate += time.time() - t0
-        except Exception as e:
-            shared.log.error(f'Load network: type=LoRA {e}')
-            if debug:
-                errors.display(e, 'LoRA')
-
-    if len(loaded_networks) > 0 and debug:
-        shared.log.debug(f'Load network: type=LoRA loaded={[n.name for n in loaded_networks]} cache={list(lora_cache)}')
-
-    if recompile_model:
-        shared.log.info("Load network: type=LoRA recompiling model")
-        backup_lora_model = shared.compiled_model_state.lora_model
-        if 'Model' in shared.opts.cuda_compile:
-            shared.sd_model = sd_models_compile.compile_diffusers(shared.sd_model)
-        shared.compiled_model_state.lora_model = backup_lora_model
-
-    if len(loaded_networks) > 0:
-        devices.torch_gc()
-
-    timer.load = time.time() - t0
-
-
-# section: process loaded networks
-
-def network_backup_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn.GroupNorm, torch.nn.LayerNorm, diffusers.models.lora.LoRACompatibleLinear, diffusers.models.lora.LoRACompatibleConv], network_layer_name: str, wanted_names: tuple):
-    global bnb # pylint: disable=W0603
-    backup_size = 0
-    if len(loaded_networks) > 0 and network_layer_name is not None and any([net.modules.get(network_layer_name, None) for net in loaded_networks]): # noqa: C419 # pylint: disable=R1729
-        t0 = time.time()
-
-        weights_backup = getattr(self, "network_weights_backup", None)
-        bias_backup = getattr(self, "network_bias_backup", None)
-        if weights_backup is not None or bias_backup is not None:
-            if (shared.opts.lora_fuse_diffusers and not isinstance(weights_backup, bool)) or (not shared.opts.lora_fuse_diffusers and isinstance(weights_backup, bool)): # invalidate so we can change direct/backup on-the-fly
-                weights_backup = None
-                bias_backup = None
-                self.network_weights_backup = weights_backup
-                self.network_bias_backup = bias_backup
-
-        if weights_backup is None and wanted_names != (): # pylint: disable=C1803
-            weight = getattr(self, 'weight', None)
-            self.network_weights_backup = None
-            if getattr(weight, "quant_type", None) in ['nf4', 'fp4']:
-                if bnb is None:
-                    bnb = model_quant.load_bnb('Load network: type=LoRA', silent=True)
-                if bnb is not None:
-                    with devices.inference_context():
-                        if shared.opts.lora_fuse_diffusers:
-                            self.network_weights_backup = True
-                        else:
-                            self.network_weights_backup = bnb.functional.dequantize_4bit(weight, quant_state=weight.quant_state, quant_type=weight.quant_type, blocksize=weight.blocksize,)
-                        self.quant_state = weight.quant_state
-                        self.quant_type = weight.quant_type
-                        self.blocksize = weight.blocksize
-                else:
-                    if shared.opts.lora_fuse_diffusers:
-                        self.network_weights_backup = True
-                    else:
-                        weights_backup = weight.clone()
-                        self.network_weights_backup = weights_backup.to(devices.cpu)
-            else:
-                if shared.opts.lora_fuse_diffusers:
-                    self.network_weights_backup = True
-                else:
-                    self.network_weights_backup = weight.clone().to(devices.cpu)
-
-        if bias_backup is None:
-            if getattr(self, 'bias', None) is not None:
-                if shared.opts.lora_fuse_diffusers:
-                    self.network_bias_backup = True
-                else:
-                    bias_backup = self.bias.clone()
-                    bias_backup = bias_backup.to(devices.cpu)
-
-        if getattr(self, 'network_weights_backup', None) is not None:
-            backup_size += self.network_weights_backup.numel() * self.network_weights_backup.element_size() if isinstance(self.network_weights_backup, torch.Tensor) else 0
-        if getattr(self, 'network_bias_backup', None) is not None:
-            backup_size += self.network_bias_backup.numel() * self.network_bias_backup.element_size() if isinstance(self.network_bias_backup, torch.Tensor) else 0
-        timer.backup += time.time() - t0
-    return backup_size
-
-
-def network_calc_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn.GroupNorm, torch.nn.LayerNorm, diffusers.models.lora.LoRACompatibleLinear, diffusers.models.lora.LoRACompatibleConv], network_layer_name: str, use_previous: bool = False):
-    if shared.opts.diffusers_offload_mode == "none":
-        try:
-            self.to(devices.device)
-        except Exception:
-            pass
-    batch_updown = None
-    batch_ex_bias = None
-    loaded = loaded_networks if not use_previous else previously_loaded_networks
-    for net in loaded:
-        module = net.modules.get(network_layer_name, None)
-        if module is None:
-            continue
-        try:
-            t0 = time.time()
-            try:
-                weight = self.weight.to(devices.device)
-            except Exception:
-                weight = self.weight
-            updown, ex_bias = module.calc_updown(weight)
-            if batch_updown is not None and updown is not None:
-                batch_updown += updown.to(batch_updown.device)
-            else:
-                batch_updown = updown
-            if batch_ex_bias is not None and ex_bias is not None:
-                batch_ex_bias += ex_bias.to(batch_ex_bias.device)
-            else:
-                batch_ex_bias = ex_bias
-            timer.calc += time.time() - t0
-            if shared.opts.diffusers_offload_mode == "sequential":
-                t0 = time.time()
-                if batch_updown is not None:
-                    batch_updown = batch_updown.to(devices.cpu)
-                if batch_ex_bias is not None:
-                    batch_ex_bias = batch_ex_bias.to(devices.cpu)
-                t1 = time.time()
-                timer.move += t1 - t0
-        except RuntimeError as e:
-            extra_network_lora.errors[net.name] = extra_network_lora.errors.get(net.name, 0) + 1
-            if debug:
-                module_name = net.modules.get(network_layer_name, None)
-                shared.log.error(f'LoRA apply weight name="{net.name}" module="{module_name}" layer="{network_layer_name}" {e}')
-                errors.display(e, 'LoRA')
-                raise RuntimeError('LoRA apply weight') from e
-        continue
-    return batch_updown, batch_ex_bias
-
-
-def network_add_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn.GroupNorm, torch.nn.LayerNorm, diffusers.models.lora.LoRACompatibleLinear, diffusers.models.lora.LoRACompatibleConv], model_weights: Union[None, torch.Tensor] = None, lora_weights: torch.Tensor = None, deactivate: bool = False):
-    if lora_weights is None:
-        return self.weight
-    if deactivate:
-        lora_weights *= -1
-    if model_weights is None: # weights are used if provided-from-backup else use self.weight
-        model_weights = self.weight
-    # TODO lora: add other quantization types
-    if self.__class__.__name__ == 'Linear4bit' and bnb is not None:
-        try:
-            dequant_weight = bnb.functional.dequantize_4bit(model_weights.to(devices.device), quant_state=self.quant_state, quant_type=self.quant_type, blocksize=self.blocksize)
-            new_weight = dequant_weight.to(devices.device) + lora_weights.to(devices.device)
-            self.weight = bnb.nn.Params4bit(new_weight, quant_state=self.quant_state, quant_type=self.quant_type, blocksize=self.blocksize)
-        except Exception as e:
-            shared.log.error(f'Load network: type=LoRA quant=bnb cls={self.__class__.__name__} type={self.quant_type} blocksize={self.blocksize} state={vars(self.quant_state)} weight={self.weight} bias={lora_weights} {e}')
-    else:
-        try:
-            new_weight = model_weights.to(devices.device) + lora_weights.to(devices.device)
-        except Exception:
-            new_weight = model_weights + lora_weights # try without device cast
-        self.weight = torch.nn.Parameter(new_weight, requires_grad=False)
-    try:
-        self.weight = self.weight.to(device=devices.device) # required since quantization happens only during .to call, not during params creation
-    except Exception:
-        pass # may fail if weights is meta tensor
-    return self.weight
-
-
-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):
-    weights_backup = getattr(self, "network_weights_backup", False)
-    bias_backup = getattr(self, "network_bias_backup", False)
-    if not isinstance(weights_backup, bool): # remove previous backup if we switched settings
-        weights_backup = True
-    if not isinstance(bias_backup, bool):
-        bias_backup = True
-    if not weights_backup and not bias_backup:
-        return None, None
-    t0 = time.time()
-
-    if weights_backup:
-        if updown is not None and len(self.weight.shape) == 4 and self.weight.shape[1] == 9: # inpainting model. zero pad updown to make channel[1]  4 to 9
-            updown = torch.nn.functional.pad(updown, (0, 0, 0, 0, 0, 5))  # pylint: disable=not-callable
-        if updown is not None:
-            self.weight = network_add_weights(self, lora_weights=updown, deactivate=deactivate)
-
-    if bias_backup:
-        if ex_bias is not None:
-            self.bias = network_add_weights(self, lora_weights=ex_bias, deactivate=deactivate)
-
-    if hasattr(self, "qweight") and hasattr(self, "freeze"):
-        self.freeze()
-
-    timer.apply += time.time() - t0
-    return self.weight.device, self.weight.dtype
-
-
-def network_apply_weights(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, orig_device: torch.device, deactivate: bool = False):
-    weights_backup = getattr(self, "network_weights_backup", None)
-    bias_backup = getattr(self, "network_bias_backup", None)
-    if weights_backup is None and bias_backup is None:
-        return None, None
-    t0 = time.time()
-
-    if weights_backup is not None:
-        self.weight = None
-        if updown is not None and len(weights_backup.shape) == 4 and weights_backup.shape[1] == 9: # inpainting model. zero pad updown to make channel[1]  4 to 9
-            updown = torch.nn.functional.pad(updown, (0, 0, 0, 0, 0, 5))  # pylint: disable=not-callable
-        if updown is not None:
-            self.weight = network_add_weights(self, model_weights=weights_backup, lora_weights=updown, deactivate=deactivate)
-        else:
-            self.weight = torch.nn.Parameter(weights_backup.to(device=orig_device), requires_grad=False)
-
-    if bias_backup is not None:
-        self.bias = None
-        if ex_bias is not None:
-            self.weight = network_add_weights(self, model_weights=weights_backup, lora_weights=ex_bias, deactivate=deactivate)
-        else:
-            self.bias = torch.nn.Parameter(bias_backup.to(device=orig_device), requires_grad=False)
-
-    if hasattr(self, "qweight") and hasattr(self, "freeze"):
-        self.freeze()
-
-    timer.apply += time.time() - t0
-    return self.weight.device, self.weight.dtype
-
-
-def network_deactivate(include=[], exclude=[]):
-    if not shared.opts.lora_fuse_diffusers or shared.opts.lora_force_diffusers:
-        return
-    t0 = time.time()
-    sd_model = getattr(shared.sd_model, "pipe", shared.sd_model)  # wrapped model compatiblility
-    if shared.opts.diffusers_offload_mode == "sequential":
-        sd_models.disable_offload(sd_model)
-        sd_models.move_model(sd_model, device=devices.cpu)
-    modules = {}
-
-    components = include if len(include) > 0 else ['text_encoder', 'text_encoder_2', 'text_encoder_3', 'unet', 'transformer']
-    components = [x for x in components if x not in exclude]
-    active_components = []
-    for name in components:
-        component = getattr(sd_model, name, None)
-        if component is not None and hasattr(component, 'named_modules'):
-            modules[name] = list(component.named_modules())
-            active_components.append(name)
-    total = sum(len(x) for x in modules.values())
-    if len(previously_loaded_networks) > 0 and 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)
-    else:
-        task = None
-        pbar = nullcontext()
-    with devices.inference_context(), pbar:
-        applied_layers.clear()
-        weights_devices = []
-        weights_dtypes = []
-        for component in modules.keys():
-            orig_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:
-                    if task is not None:
-                        pbar.update(task, advance=1)
-                    continue
-                batch_updown, batch_ex_bias = network_calc_weights(module, network_layer_name, use_previous=True)
-                if shared.opts.lora_fuse_diffusers:
-                    weights_device, weights_dtype = network_apply_direct(module, batch_updown, batch_ex_bias, deactivate=True)
-                else:
-                    weights_device, weights_dtype = network_apply_weights(module, batch_updown, batch_ex_bias, orig_device, deactivate=True)
-                weights_devices.append(weights_device)
-                weights_dtypes.append(weights_dtype)
-                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
-                module.network_current_names = ()
-                if task is not None:
-                    pbar.update(task, advance=1, description=f'networks={len(previously_loaded_networks)} modules={active_components} layers={total} unapply={len(applied_layers)}')
-
-    timer.deactivate = time.time() - t0
-    if debug and len(previously_loaded_networks) > 0:
-        weights_devices, weights_dtypes = list(set([x for x in weights_devices if x is not None])), list(set([x for x in weights_dtypes if x is not None]))  # noqa: C403 # pylint: disable=R1718
-        shared.log.debug(f'Deactivate network: type=LoRA networks={[n.name for n in previously_loaded_networks]} modules={active_components} layers={total} apply={len(applied_layers)} device={weights_devices} dtype={weights_dtypes} fuse={shared.opts.lora_fuse_diffusers} time={timer.summary}')
-    modules.clear()
-    if shared.opts.diffusers_offload_mode == "sequential":
-        sd_models.set_diffuser_offload(sd_model, op="model")
 
 
 def network_activate(include=[], exclude=[]):
     t0 = time.time()
-    sd_model = getattr(shared.sd_model, "pipe", shared.sd_model)  # wrapped model compatiblility
+    sd_model = getattr(shared.sd_model, "pipe", shared.sd_model)
     if shared.opts.diffusers_offload_mode == "sequential":
         sd_models.disable_offload(sd_model)
         sd_models.move_model(sd_model, device=devices.cpu)
@@ -578,18 +25,19 @@ def network_activate(include=[], exclude=[]):
             active_components.append(name)
             modules[name] = list(component.named_modules())
     total = sum(len(x) for x in modules.values())
-    if len(loaded_networks) > 0:
+    if len(l.loaded_networks) > 0:
         pbar = rp.Progress(rp.TextColumn('[cyan]Network: type=LoRA action=activate'), rp.BarColumn(), rp.TaskProgressColumn(), rp.TimeRemainingColumn(), rp.TimeElapsedColumn(), rp.TextColumn('[cyan]{task.description}'), console=shared.console)
         task = pbar.add_task(description='' , total=total)
     else:
         task = None
         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 loaded_networks) if len(loaded_networks) > 0 else ()
+        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
-        weights_devices = []
-        weights_dtypes = []
         for component in modules.keys():
             orig_device = getattr(sd_model, component, None).device
             for _, module in modules[component]:
@@ -602,24 +50,83 @@ def network_activate(include=[], exclude=[]):
                 backup_size += network_backup_weights(module, network_layer_name, wanted_names)
                 batch_updown, batch_ex_bias = network_calc_weights(module, network_layer_name)
                 if shared.opts.lora_fuse_diffusers:
-                    weights_device, weights_dtype = network_apply_direct(module, batch_updown, batch_ex_bias)
+                    network_apply_direct(module, batch_updown, batch_ex_bias, device=device)
                 else:
-                    weights_device, weights_dtype = network_apply_weights(module, batch_updown, batch_ex_bias, orig_device)
-                weights_devices.append(weights_device)
-                weights_dtypes.append(weights_dtype)
+                    network_apply_weights(module, batch_updown, batch_ex_bias, device=orig_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
+                    applied_bias += 1 if batch_ex_bias is not None else 0
+                batch_updown, batch_ex_bias = None, None
                 del batch_updown, batch_ex_bias
                 module.network_current_names = wanted_names
                 if task is not None:
-                    pbar.update(task, advance=1, description=f'networks={len(loaded_networks)} modules={active_components} layers={total} apply={len(applied_layers)} backup={backup_size}')
+                    bs = round(backup_size/1024/1024/1024, 2) if backup_size > 0 else None
+                    pbar.update(task, advance=1, description=f'networks={len(l.loaded_networks)} modules={active_components} layers={total} weights={applied_weight} bias={applied_bias} backup={bs} device={device}')
 
         if task is not None and len(applied_layers) == 0:
             pbar.remove_task(task) # hide progress bar for no action
-    timer.activate += time.time() - t0
-    if debug and len(loaded_networks) > 0:
-        weights_devices, weights_dtypes = list(set([x for x in weights_devices if x is not None])), list(set([x for x in weights_dtypes if x is not None])) # noqa: C403 # pylint: disable=R1718
-        shared.log.debug(f'Load network: type=LoRA networks={[n.name for n in loaded_networks]} modules={active_components} layers={total} apply={len(applied_layers)} device={weights_devices} dtype={weights_dtypes} backup={backup_size} fuse={shared.opts.lora_fuse_diffusers} time={timer.summary}')
+    l.timer.activate += time.time() - t0
+    if l.debug and len(l.loaded_networks) > 0:
+        shared.log.debug(f'Network load: type=LoRA networks={[n.name for n in l.loaded_networks]} modules={active_components} layers={total} weights={applied_weight} bias={applied_bias} backup={round(backup_size/1024/1024/1024, 2)} fuse={shared.opts.lora_fuse_diffusers} device={device} time={l.timer.summary}')
     modules.clear()
-    if shared.opts.diffusers_offload_mode == "sequential":
+    if len(applied_layers) > 0 or shared.opts.diffusers_offload_mode == "sequential":
+        sd_models.set_diffuser_offload(sd_model, op="model")
+
+
+def network_deactivate(include=[], exclude=[]):
+    if not shared.opts.lora_fuse_diffusers or shared.opts.lora_force_diffusers:
+        return
+    if len(l.previously_loaded_networks) == 0:
+        return
+    t0 = time.time()
+    sd_model = getattr(shared.sd_model, "pipe", shared.sd_model)
+    if shared.opts.diffusers_offload_mode == "sequential":
+        sd_models.disable_offload(sd_model)
+        sd_models.move_model(sd_model, device=devices.cpu)
+    modules = {}
+
+    components = include if len(include) > 0 else ['text_encoder', 'text_encoder_2', 'text_encoder_3', 'unet', 'transformer']
+    components = [x for x in components if x not in exclude]
+    active_components = []
+    for name in components:
+        component = getattr(sd_model, name, None)
+        if component is not None and hasattr(component, 'named_modules'):
+            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)
+    else:
+        task = None
+        pbar = nullcontext()
+    with devices.inference_context(), pbar:
+        applied_layers.clear()
+        for component in modules.keys():
+            orig_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:
+                    if task is not None:
+                        pbar.update(task, advance=1)
+                    continue
+                batch_updown, batch_ex_bias = network_calc_weights(module, network_layer_name, use_previous=True)
+                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)
+                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
+                module.network_current_names = ()
+                if task is not None:
+                    pbar.update(task, advance=1, description=f'networks={len(l.previously_loaded_networks)} modules={active_components} layers={total} unapply={len(applied_layers)}')
+
+    l.timer.deactivate = time.time() - t0
+    if l.debug and len(l.previously_loaded_networks) > 0:
+        shared.log.debug(f'Network deactivate: type=LoRA networks={[n.name for n in l.previously_loaded_networks]} modules={active_components} layers={total} apply={len(applied_layers)} fuse={shared.opts.lora_fuse_diffusers} time={l.timer.summary}')
+    modules.clear()
+    if len(applied_layers) > 0 or shared.opts.diffusers_offload_mode == "sequential":
         sd_models.set_diffuser_offload(sd_model, op="model")
diff --git a/modules/memstats.py b/modules/memstats.py
index 160492069..90512e870 100644
--- a/modules/memstats.py
+++ b/modules/memstats.py
@@ -56,15 +56,18 @@ def memory_stats():
             fail_once = True
         mem.update({ 'ram': { 'error': str(e) } })
     try:
-        s = torch.cuda.mem_get_info()
-        gpu = { 'used': gb(s[1] - s[0]), 'total': gb(s[1]) }
-        s = dict(torch.cuda.memory_stats())
-        if s.get('num_ooms', 0) > 0:
+        free, total = torch.cuda.mem_get_info()
+        gpu = { 'used': gb(total - free), 'total': gb(total) }
+        stats = dict(torch.cuda.memory_stats())
+        if stats.get('num_ooms', 0) > 0:
             shared.state.oom = True
         mem.update({
+            'job': shared.state.job,
             'gpu': gpu,
-            'retries': s.get('num_alloc_retries', 0),
-            'oom': s.get('num_ooms', 0)
+            'active': gb(stats.get('active_bytes.all.current', 0)),
+            'peak': gb(stats.get('active_bytes.all.peak', 0)),
+            'retries': stats.get('num_alloc_retries', 0),
+            'oom': stats.get('num_ooms', 0),
         })
         return mem
     except Exception:
@@ -72,6 +75,13 @@ def memory_stats():
     return mem
 
 
+def reset_stats():
+    try:
+        torch.cuda.reset_memory_stats()
+    except Exception:
+        pass
+
+
 def memory_cache():
     return mem
 
diff --git a/modules/model_auraflow.py b/modules/model_auraflow.py
index 83320040b..790095cae 100644
--- a/modules/model_auraflow.py
+++ b/modules/model_auraflow.py
@@ -17,5 +17,5 @@ def load_auraflow(checkpoint_info, diffusers_load_config={}):
         cache_dir = shared.opts.diffusers_dir,
         **diffusers_load_config,
     )
-    devices.torch_gc()
+    devices.torch_gc(force=True)
     return pipe
diff --git a/modules/model_cogview.py b/modules/model_cogview.py
new file mode 100644
index 000000000..a37dd88ca
--- /dev/null
+++ b/modules/model_cogview.py
@@ -0,0 +1,95 @@
+import transformers
+import diffusers
+from modules import shared, devices, sd_models
+
+
+def load_common(diffusers_load_config={}, module=None):
+    from modules import model_quant, modelloader
+    modelloader.hf_login()
+
+    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']
+
+    quant_args = model_quant.create_config(module=module)
+    if quant_args:
+        shared.log.debug(f'Load model: type=CogView quantization module="{module}" {quant_args}')
+
+    return diffusers_load_config, quant_args
+
+
+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')
+    transformer = diffusers.CogView3PlusTransformer2DModel.from_pretrained(
+        repo_id,
+        subfolder="transformer",
+        cache_dir=shared.opts.diffusers_dir,
+        **diffusers_load_config,
+        **quant_args,
+    )
+
+    diffusers_load_config, quant_args = load_common(diffusers_load_config, module='TE')
+    text_encoder = transformers.T5EncoderModel.from_pretrained(
+        repo_id,
+        subfolder="text_encoder",
+        cache_dir=shared.opts.diffusers_dir,
+        **diffusers_load_config,
+        **quant_args,
+    )
+
+    pipe = diffusers.CogView3PlusPipeline.from_pretrained(
+        repo_id,
+        text_encoder=text_encoder,
+        transformer=transformer,
+        cache_dir=shared.opts.diffusers_dir,
+        **diffusers_load_config,
+    )
+    devices.torch_gc()
+    return pipe
+
+
+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')
+    transformer = diffusers.CogView4Transformer2DModel.from_pretrained(
+        repo_id,
+        subfolder="transformer",
+        cache_dir=shared.opts.diffusers_dir,
+        **diffusers_load_config,
+        **quant_args,
+    )
+
+    diffusers_load_config, 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,
+        **quant_args,
+    )
+
+    pipe = diffusers.CogView4Pipeline.from_pretrained(
+        repo_id,
+        text_encoder=text_encoder,
+        transformer=transformer,
+        cache_dir=shared.opts.diffusers_dir,
+        **diffusers_load_config,
+    )
+    if shared.opts.diffusers_eval:
+        pipe.text_encoder.eval()
+        pipe.transformer.eval()
+    pipe.enable_model_cpu_offload() # TODO cogview4: balanced offload does not work for GlmModel
+    devices.torch_gc()
+    return pipe
diff --git a/modules/model_flux.py b/modules/model_flux.py
index bdb42037b..3dfe83ff2 100644
--- a/modules/model_flux.py
+++ b/modules/model_flux.py
@@ -5,7 +5,7 @@ import diffusers
 import transformers
 from safetensors.torch import load_file
 from huggingface_hub import hf_hub_download
-from modules import shared, devices, modelloader, sd_models, sd_unet, model_te, model_quant
+from modules import shared, errors, devices, modelloader, sd_models, sd_unet, model_te, model_quant
 
 
 debug = shared.log.trace if os.environ.get('SD_LOAD_DEBUG', None) is not None else lambda *args, **kwargs: None
@@ -44,7 +44,6 @@ def load_flux_quanto(checkpoint_info):
     except Exception as e:
         shared.log.error(f"Load model: type=FLUX failed to load Quanto transformer: {e}")
         if debug:
-            from modules import errors
             errors.display(e, 'FLUX Quanto:')
 
     try:
@@ -72,7 +71,6 @@ def load_flux_quanto(checkpoint_info):
     except Exception as e:
         shared.log.error(f"Load model: type=FLUX failed to load Quanto text encoder: {e}")
         if debug:
-            from modules import errors
             errors.display(e, 'FLUX Quanto:')
 
     return transformer, text_encoder_2
@@ -105,66 +103,23 @@ def load_flux_bnb(checkpoint_info, diffusers_load_config): # pylint: disable=unu
         shared.log.error(f"Load model: type=FLUX failed to load BnB transformer: {e}")
         transformer, text_encoder_2 = None, None
         if debug:
-            from modules import errors
             errors.display(e, 'FLUX:')
     return transformer, text_encoder_2
 
 
-"""
-def quant_flux_bnb(checkpoint_info, transformer, text_encoder_2):
-    repo_id = sd_models.path_to_repo(checkpoint_info.name)
-    cache_dir=shared.opts.diffusers_dir
-    if len(shared.opts.bnb_quantization) > 0 and (transformer is None or text_encoder_2 is None):
-        from modules.model_quant import load_bnb
-        load_bnb('Load model: type=FLUX')
-        try:
-            bnb_config = diffusers.BitsAndBytesConfig(
-                load_in_8bit=shared.opts.bnb_quantization_type in ['fp8'],
-                load_in_4bit=shared.opts.bnb_quantization_type in ['nf4', 'fp4'],
-                bnb_4bit_quant_storage=shared.opts.bnb_quantization_storage,
-                bnb_4bit_quant_type=shared.opts.bnb_quantization_type,
-                bnb_4bit_compute_dtype=devices.dtype
-            )
-            if ('Model' in shared.opts.bnb_quantization) and (transformer is None):
-                transformer = diffusers.FluxTransformer2DModel.from_pretrained(repo_id, subfolder="transformer", cache_dir=cache_dir, quantization_config=bnb_config, torch_dtype=devices.dtype)
-                shared.log.debug(f'Quantization: module=transformer type=bnb dtype={shared.opts.bnb_quantization_type} storage={shared.opts.bnb_quantization_storage}')
-            if ('Text Encoder' in shared.opts.bnb_quantization) and (text_encoder_2 is None):
-                if repo_id == 'sayakpaul/flux.1-dev-nf4':
-                    repo_id = 'black-forest-labs/FLUX.1-dev' # workaround since sayakpaul model is missing model_index.json
-                text_encoder_2 = transformers.T5EncoderModel.from_pretrained(repo_id, subfolder="text_encoder_2", cache_dir=cache_dir, quantization_config=bnb_config, torch_dtype=devices.dtype)
-                shared.log.debug(f'Quantization: module=t5 type=bnb dtype={shared.opts.bnb_quantization_type} storage={shared.opts.bnb_quantization_storage}')
-        except Exception as e:
-            shared.log.error(f"Load model: type=FLUX failed quantize using BnB: {e}")
-            transformer, text_encoder_2 = None, None
-            if debug:
-                from modules import errors
-                errors.display(e, 'FLUX:')
-    return transformer, text_encoder_2
-"""
-
-
 def load_quants(kwargs, repo_id, cache_dir, allow_quant):
     try:
-        if not allow_quant:
-            return kwargs
-        quant_args = {}
-        quant_args = model_quant.create_bnb_config(quant_args)
-        if quant_args:
-            model_quant.load_bnb(f'Load model: type=FLUX quant={quant_args}')
-        if not quant_args:
-            quant_args = model_quant.create_ao_config(quant_args)
+        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(allow=allow_quant, module='Transformer')
             if quant_args:
-                model_quant.load_torchao(f'Load model: type=FLUX quant={quant_args}')
-        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):
-            kwargs['transformer'] = diffusers.FluxTransformer2DModel.from_pretrained(repo_id, subfolder="transformer", cache_dir=cache_dir, torch_dtype=devices.dtype, **quant_args)
-            shared.log.debug(f'Quantization: module=transformer type=bnb dtype={shared.opts.bnb_quantization_type} storage={shared.opts.bnb_quantization_storage}')
-        if 'text_encoder_2' not in kwargs and ('Text Encoder' in shared.opts.bnb_quantization or 'Text Encoder' in shared.opts.torchao_quantization):
-            kwargs['text_encoder_2'] = transformers.T5EncoderModel.from_pretrained(repo_id, subfolder="text_encoder_2", cache_dir=cache_dir, torch_dtype=devices.dtype, **quant_args)
-            shared.log.debug(f'Quantization: module=t5 type=bnb dtype={shared.opts.bnb_quantization_type} storage={shared.opts.bnb_quantization_storage}')
+                kwargs['transformer'] = diffusers.FluxTransformer2DModel.from_pretrained(repo_id, subfolder="transformer", cache_dir=cache_dir, torch_dtype=devices.dtype, **quant_args)
+        if 'text_encoder_2' 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(allow=allow_quant, module='TE')
+            if quant_args:
+                kwargs['text_encoder_2'] = transformers.T5EncoderModel.from_pretrained(repo_id, subfolder="text_encoder_2", cache_dir=cache_dir, torch_dtype=devices.dtype, **quant_args)
     except Exception as e:
         shared.log.error(f'Quantization: {e}')
+        errors.display(e, 'Quantization:')
     return kwargs
 
 
@@ -230,15 +185,13 @@ def load_transformer(file_path): # triggered by opts.sd_unet change
     else:
         quant_args = model_quant.create_bnb_config({})
         if quant_args:
-            model_quant.load_bnb(f'Load model: type=FLUX quant={quant_args}')
             shared.log.info(f'Load module: type=UNet/Transformer file="{file_path}" offload={shared.opts.diffusers_offload_mode} quant=bnb dtype={devices.dtype}')
             from modules.model_flux_nf4 import load_flux_nf4
             transformer, _text_encoder_2 = load_flux_nf4(file_path, prequantized=False)
             if transformer is not None:
                 return transformer
-        quant_args = model_quant.create_ao_config({})
+        quant_args = model_quant.create_config(module='Transformer')
         if quant_args:
-            model_quant.load_torchao(f'Load model: type=FLUX quant={quant_args}')
             shared.log.info(f'Load module: type=UNet/Transformer file="{file_path}" offload={shared.opts.diffusers_offload_mode} quant=torchao dtype={devices.dtype}')
             transformer = diffusers.FluxTransformer2DModel.from_single_file(file_path, **diffusers_load_config, **quant_args)
             if transformer is not None:
@@ -249,14 +202,14 @@ def load_transformer(file_path): # triggered by opts.sd_unet change
         transformer = diffusers.FluxTransformer2DModel.from_single_file(file_path, **diffusers_load_config)
     if transformer is None:
         shared.log.error('Failed to load UNet model')
-        shared.opts.sd_unet = 'None'
+        shared.opts.sd_unet = 'Default'
     return transformer
 
 
 def load_flux(checkpoint_info, diffusers_load_config): # triggered by opts.sd_checkpoint change
-    quant = model_quant.get_quant(checkpoint_info.path)
+    prequantized = model_quant.get_quant(checkpoint_info.path)
     repo_id = sd_models.path_to_repo(checkpoint_info.name)
-    shared.log.debug(f'Load model: type=FLUX model="{checkpoint_info.name}" repo="{repo_id}" unet="{shared.opts.sd_unet}" te="{shared.opts.sd_text_encoder}" vae="{shared.opts.sd_vae}" quant={quant} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype}')
+    shared.log.debug(f'Load model: type=FLUX model="{checkpoint_info.name}" repo="{repo_id}" unet="{shared.opts.sd_unet}" te="{shared.opts.sd_text_encoder}" vae="{shared.opts.sd_vae}" quant={prequantized} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype}')
     debug(f'Load model: type=FLUX config={diffusers_load_config}')
     modelloader.hf_login()
 
@@ -271,20 +224,19 @@ def load_flux(checkpoint_info, diffusers_load_config): # triggered by opts.sd_ch
     devices.torch_gc(force=True)
 
     # load overrides if any
-    if shared.opts.sd_unet != 'None':
+    if shared.opts.sd_unet != 'Default':
         try:
             debug(f'Load model: type=FLUX unet="{shared.opts.sd_unet}"')
             transformer = load_transformer(sd_unet.unet_dict[shared.opts.sd_unet])
             if transformer is None:
-                shared.opts.sd_unet = 'None'
+                shared.opts.sd_unet = 'Default'
                 sd_unet.failed_unet.append(shared.opts.sd_unet)
         except Exception as e:
             shared.log.error(f"Load model: type=FLUX failed to load UNet: {e}")
-            shared.opts.sd_unet = 'None'
+            shared.opts.sd_unet = 'Default'
             if debug:
-                from modules import errors
                 errors.display(e, 'FLUX UNet:')
-    if shared.opts.sd_text_encoder != 'None':
+    if shared.opts.sd_text_encoder != 'Default':
         try:
             debug(f'Load model: type=FLUX te="{shared.opts.sd_text_encoder}"')
             from modules.model_te import load_t5, load_vit_l
@@ -294,11 +246,10 @@ def load_flux(checkpoint_info, diffusers_load_config): # triggered by opts.sd_ch
                 text_encoder_2 = load_t5(name=shared.opts.sd_text_encoder, cache_dir=shared.opts.diffusers_dir)
         except Exception as e:
             shared.log.error(f"Load model: type=FLUX failed to load T5: {e}")
-            shared.opts.sd_text_encoder = 'None'
+            shared.opts.sd_text_encoder = 'Default'
             if debug:
-                from modules import errors
                 errors.display(e, 'FLUX T5:')
-    if shared.opts.sd_vae != 'None' and shared.opts.sd_vae != 'Automatic':
+    if shared.opts.sd_vae != 'Default' and shared.opts.sd_vae != 'Automatic':
         try:
             debug(f'Load model: type=FLUX vae="{shared.opts.sd_vae}"')
             from modules import sd_vae
@@ -309,13 +260,12 @@ def load_flux(checkpoint_info, diffusers_load_config): # triggered by opts.sd_ch
                 vae = diffusers.AutoencoderKL.from_single_file(vae_file, config=vae_config, **diffusers_load_config)
         except Exception as e:
             shared.log.error(f"Load model: type=FLUX failed to load VAE: {e}")
-            shared.opts.sd_vae = 'None'
+            shared.opts.sd_vae = 'Default'
             if debug:
-                from modules import errors
                 errors.display(e, 'FLUX VAE:')
 
     # load quantized components if any
-    if quant == 'nf4':
+    if prequantized == 'nf4':
         try:
             from modules.model_flux_nf4 import load_flux_nf4
             _transformer, _text_encoder = load_flux_nf4(checkpoint_info)
@@ -326,9 +276,8 @@ def load_flux(checkpoint_info, diffusers_load_config): # triggered by opts.sd_ch
         except Exception as e:
             shared.log.error(f"Load model: type=FLUX failed to load NF4 components: {e}")
             if debug:
-                from modules import errors
                 errors.display(e, 'FLUX NF4:')
-    if quant == 'qint8' or quant == 'qint4':
+    if prequantized == 'qint8' or prequantized == 'qint4':
         try:
             _transformer, _text_encoder = load_flux_quanto(checkpoint_info)
             if _transformer is not None:
@@ -338,7 +287,6 @@ def load_flux(checkpoint_info, diffusers_load_config): # triggered by opts.sd_ch
         except Exception as e:
             shared.log.error(f"Load model: type=FLUX failed to load Quanto components: {e}")
             if debug:
-                from modules import errors
                 errors.display(e, 'FLUX Quanto:')
 
     # initialize pipeline with pre-loaded components
@@ -375,12 +323,11 @@ def load_flux(checkpoint_info, diffusers_load_config): # triggered by opts.sd_ch
             except Exception:
                 pass
 
-    allow_quant = 'gguf' not in (sd_unet.loaded_unet or '') and (quant is None or quant == 'none')
+    allow_quant = 'gguf' not in (sd_unet.loaded_unet or '') and (prequantized is None or prequantized == 'none')
     fn = checkpoint_info.path
     if (fn is None) or (not os.path.exists(fn) or os.path.isdir(fn)):
         kwargs = load_quants(kwargs, repo_id, cache_dir=shared.opts.diffusers_dir, allow_quant=allow_quant)
-    kwargs = model_quant.create_bnb_config(kwargs, allow_quant)
-    kwargs = model_quant.create_ao_config(kwargs, allow_quant)
+    # kwargs = model_quant.create_config(kwargs, allow_quant)
     if fn.endswith('.safetensors') and os.path.isfile(fn):
         pipe = diffusers.FluxPipeline.from_single_file(fn, cache_dir=shared.opts.diffusers_dir, **kwargs, **diffusers_load_config)
     else:
@@ -393,6 +340,5 @@ def load_flux(checkpoint_info, diffusers_load_config): # triggered by opts.sd_ch
     vae = None
     for k in kwargs.keys():
         kwargs[k] = None
-    devices.torch_gc()
-
+    devices.torch_gc(force=True)
     return pipe
diff --git a/modules/model_flux_nf4.py b/modules/model_flux_nf4.py
index b00c3320e..2290bacba 100644
--- a/modules/model_flux_nf4.py
+++ b/modules/model_flux_nf4.py
@@ -47,7 +47,7 @@ def _replace_with_bnb_linear(
                         in_features,
                         out_features,
                         module.bias is not None,
-                        compute_dtype=torch.bfloat16,
+                        compute_dtype=devices.dtype,
                         compress_statistics=False,
                         quant_type="nf4",
                     )
diff --git a/modules/model_kolors.py b/modules/model_kolors.py
index 932763b4b..b6c35c85a 100644
--- a/modules/model_kolors.py
+++ b/modules/model_kolors.py
@@ -23,5 +23,5 @@ def load_kolors(_checkpoint_info, diffusers_load_config={}):
         **diffusers_load_config,
     )
     pipe.vae.config.force_upcast = True
-    devices.torch_gc()
+    devices.torch_gc(force=True)
     return pipe
diff --git a/modules/model_lumina.py b/modules/model_lumina.py
index 9dff4dccd..f19fcd7da 100644
--- a/modules/model_lumina.py
+++ b/modules/model_lumina.py
@@ -21,7 +21,7 @@ def load_lumina(_checkpoint_info, diffusers_load_config={}):
         cache_dir = shared.opts.diffusers_dir,
         **diffusers_load_config,
     )
-    devices.torch_gc()
+    devices.torch_gc(force=True)
     return pipe
 
 
@@ -32,14 +32,13 @@ def load_lumina2(checkpoint_info, diffusers_load_config={}):
     if quant_args:
         model_quant.load_bnb(f'Load model: type=Lumina quant={quant_args}')
     if not quant_args:
-        quant_args = model_quant.create_ao_config(quant_args)
-        if quant_args:
-            model_quant.load_torchao(f'Load model: type=Lumina quant={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):
+    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 ('Text Encoder' in shared.opts.bnb_quantization or 'Text Encoder' in shared.opts.torchao_quantization):
+    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)
+    devices.torch_gc(force=True)
     return sd_model
diff --git a/modules/model_meissonic.py b/modules/model_meissonic.py
index 69ceab458..d705a32d9 100644
--- a/modules/model_meissonic.py
+++ b/modules/model_meissonic.py
@@ -33,5 +33,5 @@ def load_meissonic(checkpoint_info, diffusers_load_config={}):
     diffusers.pipelines.auto_pipeline.AUTO_TEXT2IMAGE_PIPELINES_MAPPING["meissonic"] = PipelineMeissonic
     diffusers.pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING["meissonic"] = PipelineMeissonicImg2Img
     diffusers.pipelines.auto_pipeline.AUTO_INPAINT_PIPELINES_MAPPING["meissonic"] = PipelineMeissonicInpaint
-    devices.torch_gc()
+    devices.torch_gc(force=True)
     return pipe
diff --git a/modules/model_omnigen.py b/modules/model_omnigen.py
index a08ad4ed5..b7b6e3546 100644
--- a/modules/model_omnigen.py
+++ b/modules/model_omnigen.py
@@ -20,12 +20,5 @@ def load_omnigen(checkpoint_info, diffusers_load_config={}): # pylint: disable=u
     if shared.opts.diffusers_eval:
         pipe.model.eval()
     pipe.vae.to(devices.device, dtype=devices.dtype)
-    devices.torch_gc()
-
-    # register
-    # from diffusers import pipelines
-    # pipelines.auto_pipeline.AUTO_TEXT2IMAGE_PIPELINES_MAPPING["omnigen"] = pipe.__class__
-    # pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING["omnigen"] = pipe.__class__
-    # pipelines.auto_pipeline.AUTO_INPAINT_PIPELINES_MAPPING["omnigen"] = pipe.__class__
-
+    devices.torch_gc(force=True)
     return pipe
diff --git a/modules/model_pixart.py b/modules/model_pixart.py
index c017cc468..0757a1216 100644
--- a/modules/model_pixart.py
+++ b/modules/model_pixart.py
@@ -26,5 +26,5 @@ def load_pixart(checkpoint_info, diffusers_load_config={}):
         **kwargs,
         **diffusers_load_config,
     )
-    devices.torch_gc()
+    devices.torch_gc(force=True)
     return pipe
diff --git a/modules/model_quant.py b/modules/model_quant.py
index 23cc37aec..40a2ce3f8 100644
--- a/modules/model_quant.py
+++ b/modules/model_quant.py
@@ -1,17 +1,18 @@
+import os
 import sys
 import copy
 import time
 import diffusers
-from installer import install, log, setup_logging
+from installer import installed, install, log, setup_logging
 
 
 ao = None
 bnb = None
 intel_nncf = None
 optimum_quanto = None
-
 quant_last_model_name = None
 quant_last_model_device = None
+debug = os.environ.get('SD_QUANT_DEBUG', None) is not None
 
 
 def get_quant(name):
@@ -30,10 +31,10 @@ def get_quant(name):
     return 'none'
 
 
-def create_bnb_config(kwargs = None, allow_bnb: bool = True):
+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:
+        if 'Model' in shared.opts.bnb_quantization or (module is not None and module in shared.opts.bnb_quantization):
             load_bnb()
             if bnb is None:
                 return kwargs
@@ -44,7 +45,7 @@ def create_bnb_config(kwargs = None, allow_bnb: bool = True):
                 bnb_4bit_quant_type=shared.opts.bnb_quantization_type,
                 bnb_4bit_compute_dtype=devices.dtype
             )
-            shared.log.debug(f'Quantization: module=all type=bnb dtype={shared.opts.bnb_quantization_type} storage={shared.opts.bnb_quantization_storage}')
+            log.debug(f'Quantization: module="{module}" type=bnb dtype={shared.opts.bnb_quantization_type} storage={shared.opts.bnb_quantization_storage}')
             if kwargs is None:
                 return bnb_config
             else:
@@ -53,16 +54,15 @@ def create_bnb_config(kwargs = None, allow_bnb: bool = True):
     return kwargs
 
 
-def create_ao_config(kwargs = None, allow_ao: bool = True):
+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:
+        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:
                 return kwargs
-            diffusers.utils.import_utils.is_torchao_available = lambda: True
             ao_config = diffusers.TorchAoConfig(shared.opts.torchao_quantization_type)
-            shared.log.debug(f'Quantization: module=all type=torchao dtype={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
             else:
@@ -71,16 +71,62 @@ def create_ao_config(kwargs = None, allow_ao: bool = True):
     return kwargs
 
 
+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):
+            load_quanto(silent=True)
+            if optimum_quanto is None:
+                return kwargs
+            quanto_config = diffusers.QuantoConfig(
+                weights_dtype=shared.opts.quanto_quantization_type,
+            )
+            quanto_config.activations = None # patch so it works with transformers
+            log.debug(f'Quantization: module="{module}" type=quanto dtype={shared.opts.quanto_quantization_type}')
+            if kwargs is None:
+                return quanto_config
+            else:
+                kwargs['quantization_config'] = quanto_config
+                return kwargs
+    return kwargs
+
+
+def create_config(kwargs = None, allow: bool = True, module: str = 'Model'):
+    if kwargs is None:
+        kwargs = {}
+    kwargs = create_bnb_config(kwargs, allow_bnb=allow, module=module)
+    if kwargs is not None and 'quantization_config' in kwargs:
+        if debug:
+            log.trace(f'Quantization: type=bnb config={kwargs.get("quantization_config", None)}')
+        return kwargs
+    kwargs = create_ao_config(kwargs, allow_ao=allow, module=module)
+    if kwargs is not None and 'quantization_config' in kwargs:
+        if debug:
+            log.trace(f'Quantization: type=torchao config={kwargs.get("quantization_config", None)}')
+        return kwargs
+    kwargs = create_quanto_config(kwargs, allow_quanto=allow, module=module)
+    if kwargs is not None and 'quantization_config' in kwargs:
+        if debug:
+            log.trace(f'Quantization: type=quanto config={kwargs.get("quantization_config", None)}')
+        return kwargs
+    return kwargs
+
+
 def load_torchao(msg='', silent=False):
     global ao # pylint: disable=global-statement
     if ao is not None:
         return ao
-    install('torchao==0.8.0', quiet=True)
+    if not installed('torchao'):
+        install('torchao==0.8.0', quiet=True)
+        log.warning('Quantization: torchao installed please restart')
     try:
         import torchao
         ao = torchao
         fn = f'{sys._getframe(2).f_code.co_name}:{sys._getframe(1).f_code.co_name}' # pylint: disable=protected-access
         log.debug(f'Quantization: type=torchao version={ao.__version__} fn={fn}') # pylint: disable=protected-access
+        from diffusers.utils import import_utils
+        import_utils.is_torchao_available = lambda: True
+        import_utils._torchao_available = True # pylint: disable=protected-access
         return ao
     except Exception as e:
         if len(msg) > 0:
@@ -96,15 +142,18 @@ def load_bnb(msg='', silent=False):
     global bnb # pylint: disable=global-statement
     if bnb is not None:
         return bnb
-    if devices.backend == 'cuda':
-        # forcing a version will uninstall the multi-backend-refactor branch of bnb
-        install('bitsandbytes==0.45.1', quiet=True)
+    if not installed('bitsandbytes'):
+        if devices.backend == 'cuda':
+            # forcing a version will uninstall the multi-backend-refactor branch of bnb
+            install('bitsandbytes==0.45.1', quiet=True)
+            log.warning('Quantization: bitsandbytes installed please restart')
     try:
         import bitsandbytes
         bnb = bitsandbytes
-        diffusers.utils.import_utils._bitsandbytes_available = True # pylint: disable=protected-access
-        diffusers.utils.import_utils._bitsandbytes_version = '0.43.3' # pylint: disable=protected-access
-        fn = f'{sys._getframe(2).f_code.co_name}:{sys._getframe(1).f_code.co_name}' # pylint: disable=protected-access
+        from diffusers.utils import import_utils
+        import_utils._bitsandbytes_available = True # pylint: disable=protected-access
+        import_utils._bitsandbytes_version = '0.43.3' # pylint: disable=protected-access
+        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=bitsandbytes version={bnb.__version__} fn={fn}') # pylint: disable=protected-access
         return bnb
     except Exception as e:
@@ -117,18 +166,22 @@ def load_bnb(msg='', silent=False):
 
 
 def load_quanto(msg='', silent=False):
-    from modules import shared
     global optimum_quanto # pylint: disable=global-statement
     if optimum_quanto is not None:
         return optimum_quanto
-    install('optimum-quanto==0.2.6', quiet=True)
+    if not installed('optimum-quanto'):
+        install('optimum-quanto==0.2.7', quiet=True)
+        log.warning('Quantization: optimum-quanto installed please restart')
     try:
         from optimum import quanto # pylint: disable=no-name-in-module
         optimum_quanto = quanto
-        fn = f'{sys._getframe(2).f_code.co_name}:{sys._getframe(1).f_code.co_name}' # pylint: disable=protected-access
+        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
-        if shared.opts.diffusers_offload_mode in {'balanced', 'sequential'}:
-            shared.log.error(f'Quantization: type=quanto offload={shared.opts.diffusers_offload_mode} not supported')
+        from diffusers.utils import import_utils
+        import_utils.is_optimum_quanto_available = lambda: True
+        import_utils._optimum_quanto_available = True # pylint: disable=protected-access
+        import_utils._optimum_quanto_version = quanto.__version__ # pylint: disable=protected-access
+        import_utils._replace_with_quanto_layers = diffusers.quantizers.quanto.utils._replace_with_quanto_layers # pylint: disable=protected-access
         return optimum_quanto
     except Exception as e:
         if len(msg) > 0:
@@ -143,7 +196,9 @@ def load_nncf(msg='', silent=False):
     global intel_nncf # pylint: disable=global-statement
     if intel_nncf is not None:
         return intel_nncf
-    install('nncf==2.7.0', quiet=True)
+    if not installed('nncf'):
+        install('nncf==2.7.0', quiet=True)
+        log.warning('Quantization: nncf installed please restart')
     try:
         import nncf
         intel_nncf = nncf
@@ -169,7 +224,7 @@ def apply_layerwise(sd_model, quiet:bool=False):
         storage_dtype = torch.float8_e5m2
     else:
         storage_dtype = None
-        shared.log.warning(f'Quantization: type=layerwise storage={shared.opts.layerwise_quantization_storage} not supported')
+        log.warning(f'Quantization: type=layerwise storage={shared.opts.layerwise_quantization_storage} not supported')
         return
     non_blocking = False
     if not hasattr(quantization_config.QuantizationMethod, 'LAYERWISE'):
@@ -191,14 +246,14 @@ def apply_layerwise(sd_model, quiet:bool=False):
                     m.enable_layerwise_casting(compute_dtype=devices.dtype, storage_dtype=storage_dtype, non_blocking=non_blocking)
                     m.quantization_method = 'LayerWise'
                     log.quiet(quiet, f'Quantization: type=layerwise module={module} cls={cls} storage={storage_dtype} compute={devices.dtype} blocking={not non_blocking}')
-            if module.startswith('text_encoder') and ('Model' in shared.opts.layerwise_quantization or 'Text Encoder' in shared.opts.layerwise_quantization) and ('clip' not in cls.lower()):
+            if module.startswith('text_encoder') and ('Model' in shared.opts.layerwise_quantization or 'TE' in shared.opts.layerwise_quantization) and ('clip' not in cls.lower()):
                 m = getattr(sd_model, module)
                 if hasattr(m, 'enable_layerwise_casting'):
                     m.enable_layerwise_casting(compute_dtype=devices.dtype, storage_dtype=storage_dtype, non_blocking=non_blocking)
                     m.quantization_method = quantization_config.QuantizationMethod.LAYERWISE # pylint: disable=no-member
                     log.quiet(quiet, f'Quantization: type=layerwise module={module} cls={cls} storage={storage_dtype} compute={devices.dtype} blocking={not non_blocking}')
         except Exception as e:
-            shared.log.error(f'Quantization: type=layerwise {e}')
+            log.error(f'Quantization: type=layerwise {e}')
 
 
 def nncf_send_to_device(model, device):
@@ -244,7 +299,7 @@ def nncf_compress_weights(sd_model):
     try:
         t0 = time.time()
         from modules import shared, devices, sd_models
-        shared.log.info(f"Quantization: type=NNCF modules={shared.opts.nncf_compress_weights}")
+        log.info(f"Quantization: type=NNCF modules={shared.opts.nncf_compress_weights}")
         global quant_last_model_name, quant_last_model_device # pylint: disable=global-statement
 
         sd_model = sd_models.apply_function_to_model(sd_model, nncf_compress_model, shared.opts.nncf_compress_weights, op="nncf")
@@ -259,9 +314,9 @@ def nncf_compress_weights(sd_model):
         quant_last_model_device = None
 
         t1 = time.time()
-        shared.log.info(f"Quantization: type=NNCF time={t1-t0:.2f}")
+        log.info(f"Quantization: type=NNCF time={t1-t0:.2f}")
     except Exception as e:
-        shared.log.warning(f"Quantization: type=NNCF {e}")
+        log.warning(f"Quantization: type=NNCF {e}")
     return sd_model
 
 
@@ -312,9 +367,9 @@ def optimum_quanto_weights(sd_model):
         t0 = time.time()
         from modules import shared, devices, sd_models
         if shared.opts.diffusers_offload_mode in {"balanced", "sequential"}:
-            shared.log.warning(f"Quantization: type=Optimum.quanto offload={shared.opts.diffusers_offload_mode} not compatible")
+            log.warning(f"Quantization: type=Optimum.quanto offload={shared.opts.diffusers_offload_mode} not compatible")
             return sd_model
-        shared.log.info(f"Quantization: type=Optimum.quanto: modules={shared.opts.optimum_quanto_weights}")
+        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)
@@ -361,9 +416,9 @@ def optimum_quanto_weights(sd_model):
             devices.torch_gc(force=True)
 
         t1 = time.time()
-        shared.log.info(f"Quantization: type=Optimum.quanto time={t1-t0:.2f}")
+        log.info(f"Quantization: type=Optimum.quanto time={t1-t0:.2f}")
     except Exception as e:
-        shared.log.warning(f"Quantization: type=Optimum.quanto {e}")
+        log.warning(f"Quantization: type=Optimum.quanto {e}")
     return sd_model
 
 
@@ -374,19 +429,19 @@ def torchao_quantization(sd_model):
 
     fn = getattr(q, shared.opts.torchao_quantization_type, None)
     if fn is None:
-        shared.log.error(f"Quantization: type=TorchAO type={shared.opts.torchao_quantization_type} not supported")
+        log.error(f"Quantization: type=TorchAO type={shared.opts.torchao_quantization_type} not supported")
         return sd_model
     def torchao_model(model, op=None, sd_model=None): # pylint: disable=unused-argument
         q.quantize_(model, fn(), device=devices.device)
         return model
 
-    shared.log.info(f"Quantization: type=TorchAO pipe={sd_model.__class__.__name__} quant={shared.opts.torchao_quantization_type} fn={fn} targets={shared.opts.torchao_quantization}")
+    log.info(f"Quantization: type=TorchAO pipe={sd_model.__class__.__name__} quant={shared.opts.torchao_quantization_type} fn={fn} targets={shared.opts.torchao_quantization}")
     try:
         t0 = time.time()
         sd_models.apply_function_to_model(sd_model, torchao_model, shared.opts.torchao_quantization, op="torchao")
         t1 = time.time()
-        shared.log.info(f"Quantization: type=TorchAO time={t1-t0:.2f}")
+        log.info(f"Quantization: type=TorchAO time={t1-t0:.2f}")
     except Exception as e:
-        shared.log.error(f"Quantization: type=TorchAO {e}")
+        log.error(f"Quantization: type=TorchAO {e}")
     setup_logging() # torchao uses dynamo which messes with logging so reset is needed
     return sd_model
diff --git a/modules/model_sana.py b/modules/model_sana.py
index 79a13592d..7f39f17e0 100644
--- a/modules/model_sana.py
+++ b/modules/model_sana.py
@@ -1,4 +1,3 @@
-import os
 import time
 import torch
 import diffusers
@@ -8,22 +7,14 @@ from modules import shared, sd_models, devices, modelloader, model_quant
 
 def load_quants(kwargs, repo_id, cache_dir):
     quant_args = {}
-    quant_args = model_quant.create_bnb_config(quant_args)
-    if quant_args:
-        model_quant.load_bnb(f'Load model: type=Sana quant={quant_args}')
-    if not quant_args:
-        quant_args = model_quant.create_ao_config(quant_args)
-        if quant_args:
-            model_quant.load_torchao(f'Load model: type=Sana quant={quant_args}')
+    quant_args = model_quant.create_config()
     if not quant_args:
         return kwargs
     load_args = kwargs.copy()
-    if 'transformer' not in kwargs and ('Model' in shared.opts.bnb_quantization or 'Model' in shared.opts.torchao_quantization):
+    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)):
         kwargs['transformer'] = diffusers.models.SanaTransformer2DModel.from_pretrained(repo_id, subfolder="transformer", cache_dir=cache_dir, **load_args, **quant_args)
-        shared.log.debug(f'Quantization: module=transformer type=bnb dtype={shared.opts.bnb_quantization_type} storage={shared.opts.bnb_quantization_storage}')
-    if 'text_encoder' not in kwargs and ('Text Encoder' in shared.opts.bnb_quantization or 'Text Encoder' in shared.opts.torchao_quantization):
+    if 'text_encoder' not in kwargs and ('TE' in shared.opts.bnb_quantization or 'TE' in shared.opts.torchao_quantization or 'TE' in shared.opts.quanto_quantization):
         kwargs['text_encoder'] = transformers.AutoModelForCausalLM.from_pretrained(repo_id, subfolder="text_encoder", cache_dir=cache_dir, **load_args, **quant_args)
-        shared.log.debug(f'Quantization: module=t5 type=bnb dtype={shared.opts.bnb_quantization_type} storage={shared.opts.bnb_quantization_storage}')
     return kwargs
 
 
@@ -39,8 +30,6 @@ def load_sana(checkpoint_info, kwargs={}):
 
     if not repo_id.endswith('_diffusers'):
         repo_id = f'{repo_id}_diffusers'
-    if devices.dtype == torch.bfloat16 and 'BF16' not in repo_id:
-        repo_id = repo_id.replace('_diffusers', '_BF16_diffusers')
 
     if 'Sana_1600M' in repo_id:
         if devices.dtype == torch.bfloat16 or 'BF16' in repo_id:
@@ -53,13 +42,20 @@ def load_sana(checkpoint_info, kwargs={}):
     if 'Sana_600M' in repo_id:
         kwargs['variant'] = 'fp16'
 
-    if (fn is None) or (not os.path.exists(fn) or os.path.isdir(fn)):
-        # TODO sana: fails when quantized
-        # kwargs = load_quants(kwargs, repo_id, cache_dir=shared.opts.diffusers_dir)
-        pass
+    kwargs = load_quants(kwargs, repo_id, cache_dir=shared.opts.diffusers_dir)
     shared.log.debug(f'Load model: type=Sana repo="{repo_id}" args={list(kwargs)}')
     t0 = time.time()
-    pipe = diffusers.SanaPipeline.from_pretrained(repo_id, cache_dir=shared.opts.diffusers_dir, **kwargs)
+    if devices.dtype == torch.bfloat16 or devices.dtype == torch.float32:
+        kwargs['torch_dtype'] = devices.dtype
+    if 'Sprint' in repo_id:
+        cls = diffusers.SanaSprintPipeline
+    else:
+        cls = diffusers.SanaPipeline
+    pipe = cls.from_pretrained(
+        repo_id,
+        cache_dir=shared.opts.diffusers_dir,
+        **kwargs,
+    )
     if devices.dtype == torch.bfloat16 or devices.dtype == torch.float32:
         if 'transformer' not in kwargs:
             pipe.transformer = pipe.transformer.to(dtype=devices.dtype)
@@ -77,6 +73,5 @@ def load_sana(checkpoint_info, kwargs={}):
         pipe.transformer.eval()
     t1 = time.time()
     shared.log.debug(f'Load model: type=Sana target={devices.dtype} te={pipe.text_encoder.dtype} transformer={pipe.transformer.dtype} vae={pipe.vae.dtype} time={t1-t0:.2f}')
-
-    devices.torch_gc()
+    devices.torch_gc(force=True)
     return pipe
diff --git a/modules/model_sd3.py b/modules/model_sd3.py
index bf8644284..e3774b291 100644
--- a/modules/model_sd3.py
+++ b/modules/model_sd3.py
@@ -5,7 +5,7 @@ from modules import shared, devices, sd_models, sd_unet, model_quant, model_tool
 
 
 def load_overrides(kwargs, cache_dir):
-    if shared.opts.sd_unet != 'None':
+    if shared.opts.sd_unet != 'Default':
         try:
             fn = sd_unet.unet_dict[shared.opts.sd_unet]
             if fn.endswith('.safetensors'):
@@ -20,9 +20,9 @@ def load_overrides(kwargs, cache_dir):
                 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}")
-            shared.opts.sd_unet = 'None'
+            shared.opts.sd_unet = 'Default'
             sd_unet.failed_unet.append(shared.opts.sd_unet)
-    if shared.opts.sd_text_encoder != 'None':
+    if shared.opts.sd_text_encoder != 'Default':
         try:
             from modules.model_te import load_t5, load_vit_l, load_vit_g
             if 'vit-l' in shared.opts.sd_text_encoder.lower():
@@ -36,8 +36,8 @@ 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}")
-            shared.opts.sd_text_encoder = 'None'
-    if shared.opts.sd_vae != 'None' and shared.opts.sd_vae != 'Automatic':
+            shared.opts.sd_text_encoder = 'Default'
+    if shared.opts.sd_vae != 'Default' and shared.opts.sd_vae != 'Automatic':
         try:
             from modules import sd_vae
             vae_file = sd_vae.vae_dict[shared.opts.sd_vae]
@@ -47,27 +47,18 @@ 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}")
-            shared.opts.sd_vae = 'None'
+            shared.opts.sd_vae = 'Default'
     return kwargs
 
 
 def load_quants(kwargs, repo_id, cache_dir):
-    quant_args = {}
-    quant_args = model_quant.create_bnb_config(quant_args)
-    if quant_args:
-        model_quant.load_bnb(f'Load model: type=SD3 quant={quant_args}')
-    if not quant_args:
-        quant_args = model_quant.create_ao_config(quant_args)
-        if quant_args:
-            model_quant.load_torchao(f'Load model: type=SD3 quant={quant_args}')
+    quant_args = model_quant.create_config()
     if not quant_args:
         return kwargs
-    if 'Model' in shared.opts.bnb_quantization and 'transformer' not in 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)):
         kwargs['transformer'] = diffusers.SD3Transformer2DModel.from_pretrained(repo_id, subfolder="transformer", cache_dir=cache_dir, torch_dtype=devices.dtype, **quant_args)
-        shared.log.debug(f'Quantization: module=transformer type=bnb dtype={shared.opts.bnb_quantization_type} storage={shared.opts.bnb_quantization_storage}')
-    if 'text_encoder_3' not in kwargs and ('Text Encoder' in shared.opts.bnb_quantization or 'Text Encoder' in shared.opts.torchao_quantization):
+    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):
         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)
-        shared.log.debug(f'Quantization: module=t5 type=bnb dtype={shared.opts.bnb_quantization_type} storage={shared.opts.bnb_quantization_storage}')
     return kwargs
 
 
@@ -157,8 +148,7 @@ 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_bnb_config(kwargs)
-    kwargs = model_quant.create_ao_config(kwargs)
+    kwargs = model_quant.create_config(kwargs)
     pipe = loader(
         repo_id,
         torch_dtype=devices.dtype,
@@ -166,5 +156,5 @@ def load_sd3(checkpoint_info, cache_dir=None, config=None):
         config=config,
         **kwargs,
     )
-    devices.torch_gc()
+    devices.torch_gc(force=True)
     return pipe
diff --git a/modules/model_stablecascade.py b/modules/model_stablecascade.py
index 3c3339dca..fc143e8e7 100644
--- a/modules/model_stablecascade.py
+++ b/modules/model_stablecascade.py
@@ -93,7 +93,7 @@ def load_cascade_combined(checkpoint_info, diffusers_load_config):
     if 'cascade' in checkpoint_info.name.lower():
         diffusers_load_config["variant"] = 'bf16'
 
-    if shared.opts.sd_unet != "None" or 'stabilityai' in checkpoint_info.name.lower():
+    if shared.opts.sd_unet != "Default" or 'stabilityai' in checkpoint_info.name.lower():
         if 'cascade' in checkpoint_info.name and ('lite' in checkpoint_info.name or (checkpoint_info.hash is not None and 'abc818bb0d' in checkpoint_info.hash)):
             decoder_folder = 'decoder_lite'
             prior_folder = 'prior_lite'
@@ -107,7 +107,7 @@ def load_cascade_combined(checkpoint_info, diffusers_load_config):
             decoder = StableCascadeDecoderPipeline.from_pretrained(checkpoint_info.path, cache_dir=shared.opts.diffusers_dir, text_encoder=None, **diffusers_load_config)
         # shared.log.debug(f'StableCascade {decoder_folder}: scale={decoder.latent_dim_scale}')
         prior_text_encoder = None
-        if shared.opts.sd_unet != "None":
+        if shared.opts.sd_unet != "Default":
             prior_unet, prior_text_encoder = load_prior(unet_dict[shared.opts.sd_unet])
         else:
             prior_unet = StableCascadeUNet.from_pretrained("stabilityai/stable-cascade-prior", subfolder=prior_folder, cache_dir=shared.opts.diffusers_dir, **diffusers_load_config)
@@ -155,6 +155,7 @@ def load_cascade_combined(checkpoint_info, diffusers_load_config):
         latent_dim_scale=sd_model.decoder_pipe.config.latent_dim_scale,
     )
 
+    devices.torch_gc(force=True)
     shared.log.debug(f'StableCascade combined: {sd_model.__class__.__name__}')
     return sd_model
 
diff --git a/modules/model_tools.py b/modules/model_tools.py
index fdeda5c2a..8e473ba23 100644
--- a/modules/model_tools.py
+++ b/modules/model_tools.py
@@ -69,13 +69,11 @@ def load_modules(repo_id: str, params: dict):
             subfolder = 'text_encoder_2'
         if cls == transformers.T5EncoderModel: # t5-xxl
             subfolder = 'text_encoder_3'
-            kwargs = model_quant.create_bnb_config(kwargs)
-            kwargs = model_quant.create_ao_config(kwargs)
+            kwargs = model_quant.create_config(kwargs)
             kwargs['variant'] = 'fp16'
         if cls == diffusers.SD3Transformer2DModel:
             subfolder = 'transformer'
-            kwargs = model_quant.create_bnb_config(kwargs)
-            kwargs = model_quant.create_ao_config(kwargs)
+            kwargs = model_quant.create_config(kwargs)
         if subfolder is None:
             continue
         shared.log.debug(f'Load: module={name} class={cls.__name__} repo={repo_id} location={subfolder}')
diff --git a/modules/modeldata.py b/modules/modeldata.py
index 0105771e5..f066eaa68 100644
--- a/modules/modeldata.py
+++ b/modules/modeldata.py
@@ -29,20 +29,31 @@ def get_model_type(pipe):
         model_type = 'auraflow'
     elif "Flux" in name:
         model_type = 'f1'
-    elif "Mochi" in name:
-        model_type = 'mochi'
     elif "Lumina2" in name:
         model_type = 'lumina2'
     elif "Lumina" in name:
         model_type = 'lumina'
     elif "OmniGen" in name:
         model_type = 'omnigen'
-    elif "CogVideo" in name:
-        model_type = 'cogvideox'
+    elif "CogView3" in name:
+        model_type = 'cogview3'
+    elif "CogView4" in name:
+        model_type = 'cogview4'
     elif "Sana" in name:
         model_type = 'sana'
-    elif 'HunyuanVideoPipeline' in name:
+    # video models
+    elif "CogVideo" in name:
+        model_type = 'cogvideo'
+    elif 'HunyuanVideoPipeline' in name or 'HunyuanSkyreels' in name:
         model_type = 'hunyuanvideo'
+    elif 'Wan' in name:
+        model_type = 'wanvideo'
+    elif 'LTX' in name:
+        model_type = 'ltxvideo'
+    elif "Mochi" in name:
+        model_type = 'mochivideo'
+    elif "Allegro" in name:
+        model_type = 'allegrovideo'
     else:
         model_type = name
     return model_type
diff --git a/modules/modelloader.py b/modules/modelloader.py
index f3621fdc3..a132b765a 100644
--- a/modules/modelloader.py
+++ b/modules/modelloader.py
@@ -18,6 +18,7 @@ from modules.paths import script_path, models_path
 loggedin = None
 diffuser_repos = []
 debug = shared.log.trace if os.environ.get('SD_DOWNLOAD_DEBUG', None) is not None else lambda *args, **kwargs: None
+pbar = None
 
 
 def hf_login(token=None):
@@ -61,9 +62,35 @@ def download_civit_meta(model_path: str, model_id):
     return f'CivitAI download error: id={model_id} url={url} code={r.status_code}'
 
 
+def save_video_frame(filepath: str):
+    from modules import video
+    try:
+        frames, fps, duration, w, h, codec, frame = video.get_video_params(filepath, capture=True)
+    except Exception as e:
+        shared.log.error(f'Video: file={filepath} {e}')
+        return None
+    if frame is not None:
+        basename = os.path.splitext(filepath)
+        thumb = f'{basename[0]}.thumb.jpg'
+        shared.log.debug(f'Video: file={filepath} frames={frames} fps={fps} size={w}x{h} codec={codec} duration={duration} thumb={thumb}')
+        frame.save(thumb)
+    else:
+        shared.log.error(f'Video: file={filepath} no frames found')
+    return frame
+
+
 def download_civit_preview(model_path: str, preview_url: str):
+    global pbar # pylint: disable=global-statement
+    if model_path is None:
+        pbar = None
+        return ''
     ext = os.path.splitext(preview_url)[1]
     preview_file = os.path.splitext(model_path)[0] + ext
+    is_video = preview_file.lower().endswith('.mp4')
+    is_json = preview_file.lower().endswith('.json')
+    if is_json:
+        shared.log.warning(f'CivitAI download: url="{preview_url}" skip json')
+        return 'CivitAI download error: JSON file'
     if os.path.exists(preview_file):
         return ''
     res = f'CivitAI download: url={preview_url} file="{preview_file}"'
@@ -73,20 +100,25 @@ def download_civit_preview(model_path: str, preview_url: str):
     written = 0
     img = None
     shared.state.begin('CivitAI')
+    if pbar is None:
+        pbar = p.Progress(p.TextColumn('[cyan]Download'), p.DownloadColumn(), p.BarColumn(), p.TaskProgressColumn(), p.TimeRemainingColumn(), p.TimeElapsedColumn(), p.TransferSpeedColumn(), p.TextColumn('[yellow]{task.description}'), console=shared.console)
     try:
         with open(preview_file, 'wb') as f:
-            with p.Progress(p.TextColumn('[cyan]{task.description}'), p.DownloadColumn(), p.BarColumn(), p.TaskProgressColumn(), p.TimeRemainingColumn(), p.TimeElapsedColumn(), p.TransferSpeedColumn(), console=shared.console) as progress:
-                task = progress.add_task(description="Download starting", total=total_size)
+            with pbar:
+                task = pbar.add_task(description=preview_file, total=total_size)
                 for data in r.iter_content(block_size):
                     written = written + len(data)
                     f.write(data)
-                    progress.update(task, advance=block_size, description="Downloading")
+                    pbar.update(task, advance=block_size)
         if written < 1024: # min threshold
             os.remove(preview_file)
             raise ValueError(f'removed invalid download: bytes={written}')
-        img = Image.open(preview_file)
+        if is_video:
+            img = save_video_frame(preview_file)
+        else:
+            img = Image.open(preview_file)
     except Exception as e:
-        os.remove(preview_file)
+        # os.remove(preview_file)
         res += f' error={e}'
         shared.log.error(f'CivitAI download error: url={preview_url} file="{preview_file}" written={written} {e}')
     shared.state.end()
diff --git a/modules/omnigen/utils.py b/modules/omnigen/utils.py
index 8ace4fab6..bf0a6de62 100644
--- a/modules/omnigen/utils.py
+++ b/modules/omnigen/utils.py
@@ -28,8 +28,6 @@ def update_ema(ema_model, model, decay=0.9999):
         ema_params[name].mul_(decay).add_(param.data, alpha=1 - decay)
 
 
-
-
 def requires_grad(model, flag=True):
     """
     Set requires_grad flag for all parameters in a model.
@@ -59,7 +57,6 @@ def center_crop_arr(pil_image, image_size):
     return Image.fromarray(arr[crop_y: crop_y + image_size, crop_x: crop_x + image_size])
 
 
-
 def crop_arr(pil_image, max_image_size):
     while min(*pil_image.size) >= 2 * max_image_size:
         pil_image = pil_image.resize(
@@ -89,7 +86,6 @@ def crop_arr(pil_image, max_image_size):
     return Image.fromarray(arr)
 
 
-
 def vae_encode(vae, x, weight_dtype):
     if x is not None:
         if vae.config.shift_factor is not None:
@@ -100,6 +96,7 @@ def vae_encode(vae, x, weight_dtype):
         x = x.to(weight_dtype)
     return x
 
+
 def vae_encode_list(vae, x, weight_dtype):
     latents = []
     for img in x:
diff --git a/modules/onnx_impl/execution_providers.py b/modules/onnx_impl/execution_providers.py
index e38199d0f..25e0056b3 100644
--- a/modules/onnx_impl/execution_providers.py
+++ b/modules/onnx_impl/execution_providers.py
@@ -107,10 +107,14 @@ def install_execution_provider(ep: ExecutionProvider):
     elif ep == ExecutionProvider.OpenVINO:
         packages.append("openvino")
         packages.append("onnxruntime-openvino")
+    log.info(f'ONNX install: {packages}')
     for package in packages:
         res += install(package)
     res += '

' res += 'Server restart required' log.info("Server restart required") - importlib.reload(ort) + try: + importlib.reload(ort) + except Exception: + pass return res diff --git a/modules/onnx_impl/pipelines/__init__.py b/modules/onnx_impl/pipelines/__init__.py index a11b07fc7..62db80cd5 100644 --- a/modules/onnx_impl/pipelines/__init__.py +++ b/modules/onnx_impl/pipelines/__init__.py @@ -368,7 +368,7 @@ class OnnxRawPipeline(PipelineBase): if shared.opts.cuda_compile_backend == "olive-ai": submodels_for_olive = [] - if "Text Encoder" in shared.opts.cuda_compile: + if "TE" in shared.opts.cuda_compile: if not self.is_refiner: submodels_for_olive.append("text_encoder") if self._is_sdxl: diff --git a/modules/pixelsmith/pixelsmith_pipeline.py b/modules/pixelsmith/pixelsmith_pipeline.py index 4e04b2d92..702ee67f6 100644 --- a/modules/pixelsmith/pixelsmith_pipeline.py +++ b/modules/pixelsmith/pixelsmith_pipeline.py @@ -133,7 +133,6 @@ class PAGIdentitySelfAttnProcessor: value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) # the output of sdp = (batch, num_heads, seq_len, head_dim) - # TODO: add support for attn.scale when we move to Torch 2.1 hidden_states_org = F.scaled_dot_product_attention( query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False ) @@ -248,7 +247,6 @@ class PAGCFGIdentitySelfAttnProcessor: value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) # the output of sdp = (batch, num_heads, seq_len, head_dim) - # TODO: add support for attn.scale when we move to Torch 2.1 hidden_states_org = F.scaled_dot_product_attention( query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False ) diff --git a/modules/postprocess/sdupscaler_model.py b/modules/postprocess/sdupscaler_model.py index 5ec7168d3..7b8c7a8ca 100644 --- a/modules/postprocess/sdupscaler_model.py +++ b/modules/postprocess/sdupscaler_model.py @@ -24,15 +24,19 @@ class UpscalerDiffusion(Upscaler): def load_model(self, path: str): from modules.sd_models import set_diffuser_options - scaler: UpscalerData = [x for x in self.scalers if x.data_path == path][0] + scaler: UpscalerData = [x for x in self.scalers if x.data_path == path or x.name == path] + if len(scaler) == 0: + shared.log.error(f"Upscaler cannot match model: type={self.name} model={path}") + return None + scaler = scaler[0] if self.models.get(path, None) is not None: shared.log.debug(f"Upscaler cached: type={scaler.name} model={path}") return self.models[path] else: - model = diffusers.DiffusionPipeline.from_pretrained(path, cache_dir=shared.opts.diffusers_dir, torch_dtype=devices.dtype) + model = diffusers.DiffusionPipeline.from_pretrained(scaler.data_path, cache_dir=shared.opts.diffusers_dir, torch_dtype=devices.dtype) if hasattr(model, "set_progress_bar_config"): model.set_progress_bar_config(bar_format='Progress {rate_fmt}{postfix} {bar} {percentage:3.0f}% {n_fmt}/{total_fmt} {elapsed} {remaining} ' + '\x1b[38;5;71m' + 'Upscale', ncols=80, colour='#327fba') - set_diffuser_options(scaler.model, vae=None, op='upscaler') + set_diffuser_options(model, vae=None, op='upscaler') self.models[path] = model return self.models[path] diff --git a/modules/postprocess/yolo.py b/modules/postprocess/yolo.py index 4aa80c613..8054e03ea 100644 --- a/modules/postprocess/yolo.py +++ b/modules/postprocess/yolo.py @@ -265,6 +265,7 @@ class YoloRestorer(Detailer): 'inpaint_full_res_padding': shared.opts.detailer_padding, 'width': resolution, 'height': resolution, + 'vae_type': orig_p.get('vae_type', 'Full'), } if args['denoising_strength'] == 0: shared.log.debug(f'Detailer: model="{name}" strength=0 skip') diff --git a/modules/processing.py b/modules/processing.py index 3a49f99ba..b1a8bd548 100644 --- a/modules/processing.py +++ b/modules/processing.py @@ -6,7 +6,7 @@ import numpy as np from PIL import Image, ImageOps from modules import shared, devices, errors, images, scripts, memstats, lowvram, script_callbacks, extra_networks, detailer, sd_models, sd_checkpoint, sd_vae, processing_helpers, timer, face_restoration, token_merge from modules.sd_hijack_hypertile import context_hypertile_vae, context_hypertile_unet -from modules.processing_class import StableDiffusionProcessing, StableDiffusionProcessingTxt2Img, StableDiffusionProcessingImg2Img, StableDiffusionProcessingControl # pylint: disable=unused-import +from modules.processing_class import StableDiffusionProcessing, StableDiffusionProcessingTxt2Img, StableDiffusionProcessingImg2Img, StableDiffusionProcessingControl, StableDiffusionProcessingVideo # pylint: disable=unused-import from modules.processing_info import create_infotext from modules.modeldata import model_data from modules import pag @@ -81,6 +81,7 @@ class Processed: self.all_seeds = all_seeds or p.all_seeds or [self.seed] self.all_subseeds = all_subseeds or p.all_subseeds or [self.subseed] self.infotexts = infotexts or [self.info] + memstats.reset_stats() def js(self): obj = { diff --git a/modules/processing_args.py b/modules/processing_args.py index ff693a03e..a839c9995 100644 --- a/modules/processing_args.py +++ b/modules/processing_args.py @@ -26,7 +26,7 @@ def task_specific_kwargs(p, model): p.init_images = [helpers.decode_base64_to_image(i, quiet=True) for i in p.init_images] if isinstance(p.init_images[0], Image.Image): p.init_images = [i.convert('RGB') if i.mode != 'RGB' else i for i in p.init_images if i is not None] - if (sd_models.get_diffusers_task(model) == sd_models.DiffusersTaskType.TEXT_2_IMAGE or len(getattr(p, 'init_images', [])) == 0) and not is_img2img_model: + if (sd_models.get_diffusers_task(model) == sd_models.DiffusersTaskType.TEXT_2_IMAGE or len(getattr(p, 'init_images', [])) == 0) and not is_img2img_model and 'video' not in p.ops: p.ops.append('txt2img') if hasattr(p, 'width') and hasattr(p, 'height'): task_args = { @@ -38,13 +38,20 @@ def task_specific_kwargs(p, model): model.register_to_config(requires_aesthetics_score = False) if 'hires' not in p.ops: p.ops.append('img2img') + if p.vae_type == 'Remote': + from modules.sd_vae_remote import remote_encode + p.init_images = remote_encode(p.init_images) task_args = { 'image': p.init_images, 'strength': p.denoising_strength, } if model.__class__.__name__ == 'FluxImg2ImgPipeline': # needs explicit width/height - p.width = 8 * math.ceil(p.init_images[0].width / 8) - p.height = 8 * math.ceil(p.init_images[0].height / 8) + if torch.is_tensor(p.init_images[0]): + p.width = p.init_images[0].shape[-1] * 16 + p.height = p.init_images[0].shape[-2] * 16 + else: + p.width = 8 * math.ceil(p.init_images[0].width / 8) + p.height = 8 * math.ceil(p.init_images[0].height / 8) task_args['width'], task_args['height'] = p.width, p.height if model.__class__.__name__ == 'OmniGenPipeline': p.width = 16 * math.ceil(p.init_images[0].width / 16) @@ -70,9 +77,14 @@ def task_specific_kwargs(p, model): else: p.ops.append('inpaint') width, height = processing_helpers.resize_init_images(p) + mask_image = p.task_args.get('image_mask', None) or getattr(p, 'image_mask', None) or getattr(p, 'mask', None) + if p.vae_type == 'Remote': + from modules.sd_vae_remote import remote_encode + p.init_images = remote_encode(p.init_images) + # mask_image = remote_encode(mask_image) task_args = { 'image': p.init_images, - 'mask_image': p.task_args.get('image_mask', None) or getattr(p, 'image_mask', None) or getattr(p, 'mask', None), + 'mask_image': mask_image, 'strength': p.denoising_strength, 'height': height, 'width': width, @@ -107,16 +119,16 @@ def set_pipeline_args(p, model, prompts:list, negative_prompts:list, prompts_2:t t0 = time.time() shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model) apply_circular(p.tiling, model) - if hasattr(model, "set_progress_bar_config"): - if disable_pbar: - model.set_progress_bar_config(bar_format='Progress {rate_fmt}{postfix} {bar} {percentage:3.0f}% {n_fmt}/{total_fmt} {elapsed} {remaining} ' + '\x1b[38;5;71m' + desc, ncols=80, colour='#327fba', disable=disable_pbar) - else: - model.set_progress_bar_config(bar_format='Progress {rate_fmt}{postfix} {bar} {percentage:3.0f}% {n_fmt}/{total_fmt} {elapsed} {remaining} ' + '\x1b[38;5;71m' + desc, ncols=80, colour='#327fba') args = {} has_vae = hasattr(model, 'vae') or (hasattr(model, 'pipe') and hasattr(model.pipe, 'vae')) if hasattr(model, 'pipe') and not hasattr(model, 'no_recurse'): # recurse model = model.pipe has_vae = has_vae or hasattr(model, 'vae') + if hasattr(model, "set_progress_bar_config"): + if disable_pbar: + model.set_progress_bar_config(bar_format='Progress {rate_fmt}{postfix} {bar} {percentage:3.0f}% {n_fmt}/{total_fmt} {elapsed} {remaining} ' + '\x1b[38;5;71m' + desc, ncols=80, colour='#327fba', disable=disable_pbar) + else: + model.set_progress_bar_config(bar_format='Progress {rate_fmt}{postfix} {bar} {percentage:3.0f}% {n_fmt}/{total_fmt} {elapsed} {remaining} ' + '\x1b[38;5;71m' + desc, ncols=80, colour='#327fba') signature = inspect.signature(type(model).__call__, follow_wrapped=True) possible = list(signature.parameters) @@ -226,13 +238,13 @@ def set_pipeline_args(p, model, prompts:list, negative_prompts:list, prompts_2:t if hasattr(model, 'scheduler') and hasattr(model.scheduler, 'noise_sampler_seed') and hasattr(model.scheduler, 'noise_sampler'): model.scheduler.noise_sampler = None # noise needs to be reset instead of using cached values model.scheduler.noise_sampler_seed = p.seeds # some schedulers have internal noise generator and do not use pipeline generator - if 'seed' in possible: + if 'seed' in possible and p.seed is not None: args['seed'] = p.seed - if 'noise_sampler_seed' in possible: + if 'noise_sampler_seed' in possible and p.seeds is not None: args['noise_sampler_seed'] = p.seeds - if 'guidance_scale' in possible: + if 'guidance_scale' in possible and p.cfg_scale is not None and p.cfg_scale > 0: args['guidance_scale'] = p.cfg_scale - if 'img_guidance_scale' in possible and hasattr(p, 'image_cfg_scale'): + if 'img_guidance_scale' in possible and hasattr(p, 'image_cfg_scale') and p.image_cfg_scale is not None and p.image_cfg_scale > 0: args['img_guidance_scale'] = p.image_cfg_scale if 'generator' in possible: generator = get_generator(p) @@ -286,13 +298,16 @@ def set_pipeline_args(p, model, prompts:list, negative_prompts:list, prompts_2:t p.init_images = kwargs['image'] if isinstance(kwargs['image'], Image.Image): p.init_images = [kwargs['image']] + if isinstance(kwargs['image'], torch.Tensor): + p.init_images = kwargs['image'] # handle remaining args for arg in kwargs: if arg in possible: # add kwargs + if type(kwargs[arg]) == float or type(kwargs[arg]) == int: + if kwargs[arg] <= -1: # skip -1 as default value + continue args[arg] = kwargs[arg] - else: - pass task_kwargs = task_specific_kwargs(p, model) for arg in task_kwargs: @@ -321,6 +336,9 @@ def set_pipeline_args(p, model, prompts:list, negative_prompts:list, prompts_2:t if isinstance(args['image'], torch.Tensor) or isinstance(args['image'], np.ndarray): args['width'] = 8 * args['image'].shape[-1] args['height'] = 8 * args['image'].shape[-2] + elif isinstance(args['image'], Image.Image): + args['width'] = args['image'].width + args['height'] = args['image'].height elif isinstance(args['image'][0], torch.Tensor) or isinstance(args['image'][0], np.ndarray): args['width'] = 8 * args['image'][0].shape[-1] args['height'] = 8 * args['image'][0].shape[-2] diff --git a/modules/processing_callbacks.py b/modules/processing_callbacks.py index b78b1e6a1..0ab91baa6 100644 --- a/modules/processing_callbacks.py +++ b/modules/processing_callbacks.py @@ -59,8 +59,6 @@ def diffusers_callback(pipe, step: int = 0, timestep: int = 0, kwargs: dict = {} if debug: debug_callback(f'Callback: step={step} timestep={timestep} latents={latents.shape if latents is not None else None} kwargs={list(kwargs)}') shared.state.step() - # order = getattr(pipe.scheduler, "order", 1) if hasattr(pipe, 'scheduler') else 1 - # shared.state.sampling_step = step // order if shared.state.interrupted or shared.state.skipped: raise AssertionError('Interrupted...') if shared.state.paused: @@ -80,10 +78,13 @@ def diffusers_callback(pipe, step: int = 0, timestep: int = 0, kwargs: dict = {} ip_adapter_starts = list(p.ip_adapter_starts) ip_adapter_ends = list(p.ip_adapter_ends) if any(end != 1 for end in ip_adapter_ends) or any(start != 0 for start in ip_adapter_starts): - for i in range(len(ip_adapter_scales)): - ip_adapter_scales[i] *= float(step >= pipe.num_timesteps * ip_adapter_starts[i]) - ip_adapter_scales[i] *= float(step <= pipe.num_timesteps * ip_adapter_ends[i]) - debug_callback(f"Callback: IP Adapter scales={ip_adapter_scales}") + if 'Flux' in pipe.__class__.__name__: + ip_adapter_scales = [(ip_adapter_starts[0] + (ip_adapter_ends[0] - ip_adapter_starts[0]) * (i / (19 - 1))) for i in range(19)] + else: + for i in range(len(ip_adapter_scales)): + ip_adapter_scales[i] *= float(step >= pipe.num_timesteps * ip_adapter_starts[i]) + ip_adapter_scales[i] *= float(step <= pipe.num_timesteps * ip_adapter_ends[i]) + debug_callback(f"Callback: IP Adapter scales={ip_adapter_scales}") pipe.set_ip_adapter_scale(ip_adapter_scales) if step != getattr(pipe, 'num_timesteps', 0): kwargs = processing_correction.correction_callback(p, timestep, kwargs, initial=step == 0) @@ -122,6 +123,8 @@ def diffusers_callback(pipe, step: int = 0, timestep: int = 0, kwargs: dict = {} try: shared.state.current_sigma = pipe.scheduler.sigmas[pipe.scheduler.step_index-1] shared.state.current_sigma_next = pipe.scheduler.sigmas[pipe.scheduler.step_index] + if (shared.opts.schedulers_sigma_adjust != 1.0) and (timestep > 1000 * shared.opts.schedulers_sigma_adjust_min) and (timestep < 1000 * shared.opts.schedulers_sigma_adjust_max): + pipe.scheduler.sigmas[pipe.scheduler.step_index+1] = pipe.scheduler.sigmas[pipe.scheduler.step_index+1] * shared.opts.schedulers_sigma_adjust except Exception: pass except Exception as e: diff --git a/modules/processing_class.py b/modules/processing_class.py index 9b61c4539..e38a44fd9 100644 --- a/modules/processing_class.py +++ b/modules/processing_class.py @@ -111,6 +111,8 @@ class StableDiffusionProcessing: refiner_prompt: str = '', refiner_negative: str = '', hr_refiner_start: float = 0, + # prompt enhancer + enhance_prompt: bool = False, # save options outpath_samples=None, outpath_grids=None, @@ -145,6 +147,7 @@ class StableDiffusionProcessing: self.is_refiner_pass = False self.is_api = False self.scheduled_prompt = False + self.enhance_prompt = enhance_prompt self.prompt_embeds = [] self.positive_pooleds = [] self.negative_embeds = [] @@ -236,9 +239,13 @@ class StableDiffusionProcessing: self.height = firstphase_height self.sampler_name = sampler_name or processing_helpers.get_sampler_name(sampler_index, img=True) self.hr_sampler_name: str = hr_sampler_name if hr_sampler_name != 'Same as primary' else self.sampler_name - self.override_settings = {k: v for k, v in (override_settings or {}).items() if k not in shared.restricted_opts} self.inpaint_full_res = inpaint_full_res if isinstance(inpaint_full_res, bool) else self.inpaint_full_res self.inpaint_full_res = inpaint_full_res != 0 if isinstance(inpaint_full_res, int) else self.inpaint_full_res + try: + self.override_settings = {k: v for k, v in (override_settings or {}).items() if k not in shared.restricted_opts} + except Exception as e: + shared.log.error(f'Override: {override_settings} {e}') + self.override_settings = {} # null items initialized later self.prompts = None @@ -360,6 +367,15 @@ class StableDiffusionProcessing: self.scripts = None +class StableDiffusionProcessingVideo(StableDiffusionProcessing): + def __init__(self, **kwargs): + self.prompt_template: str = None + self.frames: int = 1 + self.scheduler_shift: float = 0.0 + self.vae_tile_frames: int = 0 + debug(f'Process init: mode={self.__class__.__name__} kwargs={kwargs}') # pylint: disable=protected-access + super().__init__(**kwargs) + class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing): def __init__(self, **kwargs): debug(f'Process init: mode={self.__class__.__name__} kwargs={kwargs}') # pylint: disable=protected-access @@ -592,9 +608,6 @@ class StableDiffusionProcessingControl(StableDiffusionProcessingImg2Img): self.hr_upscale_to_x, self.hr_upscale_to_y = 8 * int(self.width * scale / 8), 8 * int(self.height * scale / 8) else: self.hr_upscale_to_x, self.hr_upscale_to_y = self.hr_resize_x, self.hr_resize_y - # hypertile_set(self, hr=True) - # shared.state.job_count = 2 * self.n_iter - # shared.log.debug(f'Control refine: upscaler="{self.hr_upscaler}" scale={scale} fixed={not use_scale} size={self.hr_upscale_to_x}x{self.hr_upscale_to_y}') def switch_class(p: StableDiffusionProcessing, new_class: type, dct: dict = None): diff --git a/modules/processing_correction.py b/modules/processing_correction.py index 73d83d3cd..e06adae15 100644 --- a/modules/processing_correction.py +++ b/modules/processing_correction.py @@ -12,6 +12,14 @@ debug_enabled = os.environ.get('SD_HDR_DEBUG', None) is not None debug = shared.log.trace if debug_enabled else lambda *args, **kwargs: None debug('Trace: HDR') skip_correction = False +warned = False + + +def warn_once(message): + global warned # pylint: disable=global-statement + if not warned: + shared.log.warning(f'VAE: {message}') + warned = True def sharpen_tensor(tensor, ratio=0): @@ -121,6 +129,9 @@ def correction_callback(p, timestep, kwargs, initial: bool = False): return kwargs latents = kwargs["latents"] # debug(f'HDR correction: latents={latents.shape}') + if len(latents.shape) <= 3: # packed latent + warn_once(f'HDR correction: shape={latents.shape} packed latent') + return kwargs if len(latents.shape) == 4: # standard batched latent for i in range(latents.shape[0]): latents[i] = correction(p, timestep, latents[i]) @@ -135,6 +146,6 @@ def correction_callback(p, timestep, kwargs, initial: bool = False): latents[i] = correction(p, timestep, latents[i]) latents = latents.permute(1, 0, 2, 3).unsqueeze(0) else: - shared.log.debug(f'HDR correction: unknown latent shape {latents.shape}') + warn_once(f'HDR correction: shape={latents.shape} unknown latent') kwargs["latents"] = latents return kwargs diff --git a/modules/processing_diffusers.py b/modules/processing_diffusers.py index acea34874..644e0b276 100644 --- a/modules/processing_diffusers.py +++ b/modules/processing_diffusers.py @@ -5,11 +5,11 @@ import numpy as np import torch import torchvision.transforms.functional as TF from PIL import Image -from modules import shared, devices, processing, sd_models, errors, sd_hijack_hypertile, processing_vae, sd_models_compile, hidiffusion, timer, modelstats, extra_networks, ras +from modules import shared, devices, processing, sd_models, errors, sd_hijack_hypertile, processing_vae, sd_models_compile, hidiffusion, timer, modelstats, extra_networks, ras, transformer_cache from modules.processing_helpers import resize_hires, calculate_base_steps, calculate_hires_steps, calculate_refiner_steps, save_intermediate, update_sampler, is_txt2img, is_refiner_enabled, get_job_name from modules.processing_args import set_pipeline_args from modules.onnx_impl import preprocess_pipeline as preprocess_onnx_pipeline, check_parameters_changed as olive_check_parameters_changed -from modules.lora import networks +from modules.lora import lora_common debug = shared.log.trace if os.environ.get('SD_DIFFUSERS_DEBUG', None) is not None else lambda *args, **kwargs: None @@ -71,8 +71,10 @@ def process_base(p: processing.StableDiffusionProcessing): eta=shared.opts.scheduler_eta, guidance_scale=p.cfg_scale, guidance_rescale=p.diffusers_guidance_rescale, + true_cfg_scale=p.pag_scale, denoising_start=0 if use_refiner_start else p.refiner_start if use_denoise_start else None, denoising_end=p.refiner_start if use_refiner_start else 1 if use_denoise_start else None, + num_frames=getattr(p, 'frames', 1), output_type='latent', clip_skip=p.clip_skip, desc='Base', @@ -85,6 +87,7 @@ def process_base(p: processing.StableDiffusionProcessing): try: t0 = time.time() sd_models_compile.check_deepcache(enable=True) + transformer_cache.set_cache() shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model) sd_models.move_model(shared.sd_model, devices.device) if hasattr(shared.sd_model, 'unet'): @@ -107,10 +110,7 @@ def process_base(p: processing.StableDiffusionProcessing): if hasattr(output, 'images'): shared.history.add(output.images, info=processing.create_infotext(p), ops=p.ops) timer.process.record('pipeline') - ras.unapply(shared.sd_model) - hidiffusion.unapply() sd_models_compile.openvino_post_compile(op="base") # only executes on compiled vino models - sd_models_compile.check_deepcache(enable=False) if shared.cmd_opts.profile: t1 = time.time() shared.log.debug(f'Profile: pipeline call: {t1-t0:.2f}') @@ -142,6 +142,10 @@ def process_base(p: processing.StableDiffusionProcessing): shared.log.error(f'Processing: step=base args={err_args} {e}') errors.display(e, 'Processing') modelstats.analyze() + finally: + ras.unapply(shared.sd_model) + hidiffusion.unapply() + sd_models_compile.check_deepcache(enable=False) if hasattr(shared.sd_model, 'embedding_db') and len(shared.sd_model.embedding_db.embeddings_used) > 0: # register used embeddings p.extra_generation_params['Embeddings'] = ', '.join(shared.sd_model.embedding_db.embeddings_used) @@ -363,7 +367,7 @@ def process_decode(p: processing.StableDiffusionProcessing, output): else: width = getattr(p, 'width', 0) height = getattr(p, 'height', 0) - frames = p.task_args.get('num_frames', None) + frames = p.task_args.get('num_frames', None) or getattr(p, 'frames', None) if isinstance(output.images, list): results = [] for i in range(len(output.images)): @@ -474,8 +478,8 @@ def process_diffusers(p: processing.StableDiffusionProcessing): return results extra_networks.deactivate(p) - timer.process.add('lora', networks.timer.total) - networks.timer.clear(complete=True) + timer.process.add('lora', lora_common.timer.total) + lora_common.timer.clear(complete=True) results = process_decode(p, output) timer.process.record('decode') diff --git a/modules/processing_helpers.py b/modules/processing_helpers.py index a5abba6a2..8fffa8313 100644 --- a/modules/processing_helpers.py +++ b/modules/processing_helpers.py @@ -401,24 +401,24 @@ def resize_init_images(p): def resize_hires(p, latents): # input=latents output=pil if not latent_upscaler else latent if not torch.is_tensor(latents): shared.log.warning('Hires: input is not tensor') - first_pass_images = processing_vae.vae_decode(latents=latents, model=shared.sd_model, vae_type=p.vae_type, output_type='pil', width=p.width, height=p.height) - return first_pass_images + decoded = processing_vae.vae_decode(latents=latents, model=shared.sd_model, vae_type=p.vae_type, output_type='pil', width=p.width, height=p.height) + return decoded if (p.hr_upscale_to_x == 0 or p.hr_upscale_to_y == 0) and hasattr(p, 'init_hr'): shared.log.error('Hires: missing upscaling dimensions') - return first_pass_images + return decoded if p.hr_upscaler.lower().startswith('latent'): - resized_image = images.resize_image(p.hr_resize_mode, latents, p.hr_upscale_to_x, p.hr_upscale_to_y, upscaler_name=p.hr_upscaler, context=p.hr_resize_context) - return resized_image + resized = images.resize_image(p.hr_resize_mode, latents, p.hr_upscale_to_x, p.hr_upscale_to_y, upscaler_name=p.hr_upscaler, context=p.hr_resize_context) + return resized - first_pass_images = processing_vae.vae_decode(latents=latents, model=shared.sd_model, vae_type=p.vae_type, output_type='pil', width=p.width, height=p.height) - resized_images = [] - for img in first_pass_images: - resized_image = images.resize_image(p.hr_resize_mode, img, p.hr_upscale_to_x, p.hr_upscale_to_y, upscaler_name=p.hr_upscaler, context=p.hr_resize_context) - resized_images.append(resized_image) + decoded = processing_vae.vae_decode(latents=latents, model=shared.sd_model, vae_type=p.vae_type, output_type='pil', width=p.width, height=p.height) + resized = [] + for image in decoded: + resize = images.resize_image(p.hr_resize_mode, image, p.hr_upscale_to_x, p.hr_upscale_to_y, upscaler_name=p.hr_upscaler, context=p.hr_resize_context) + resized.append(resize) devices.torch_gc() - return resized_images + return resized def fix_prompts(p, prompts, negative_prompts, prompts_2, negative_prompts_2): diff --git a/modules/processing_info.py b/modules/processing_info.py index fa084a2fb..4b57d859d 100644 --- a/modules/processing_info.py +++ b/modules/processing_info.py @@ -50,10 +50,11 @@ def create_infotext(p: StableDiffusionProcessing, all_prompts=None, all_seeds=No "Size": f"{p.width}x{p.height}" if hasattr(p, 'width') and hasattr(p, 'height') else None, "Sampler": p.sampler_name if p.sampler_name != 'Default' else None, "Seed": all_seeds[index], - "Seed resize from": None if p.seed_resize_from_w == 0 or p.seed_resize_from_h == 0 else f"{p.seed_resize_from_w}x{p.seed_resize_from_h}", + "Seed resize from": None if p.seed_resize_from_w <= 0 or p.seed_resize_from_h <= 0 else f"{p.seed_resize_from_w}x{p.seed_resize_from_h}", "CFG scale": p.cfg_scale if p.cfg_scale > 1.0 else None, "CFG rescale": p.diffusers_guidance_rescale if p.diffusers_guidance_rescale > 0 else None, "CFG end": p.cfg_end if p.cfg_end < 1.0 else None, + "CFG true": p.pag_scale if p.pag_scale > 1 else None, "Clip skip": p.clip_skip if p.clip_skip > 1 else None, "Batch": f'{p.n_iter}x{p.batch_size}' if p.n_iter > 1 or p.batch_size > 1 else None, "Model": None if (not shared.opts.add_model_name_to_info) or (not shared.sd_model.sd_checkpoint_info.model_name) else shared.sd_model.sd_checkpoint_info.model_name.replace(',', '').replace(':', ''), @@ -87,8 +88,8 @@ def create_infotext(p: StableDiffusionProcessing, all_prompts=None, all_seeds=No args['Grid'] = grid if shared.native: args['Pipeline'] = shared.sd_model.__class__.__name__ - args['TE'] = None if (not shared.opts.add_model_name_to_info or shared.opts.sd_text_encoder is None or shared.opts.sd_text_encoder == 'None') else shared.opts.sd_text_encoder - args['UNet'] = None if (not shared.opts.add_model_name_to_info or shared.opts.sd_unet is None or shared.opts.sd_unet == 'None') else shared.opts.sd_unet + args['TE'] = None if (not shared.opts.add_model_name_to_info or shared.opts.sd_text_encoder is None or shared.opts.sd_text_encoder == 'Default') else shared.opts.sd_text_encoder + args['UNet'] = None if (not shared.opts.add_model_name_to_info or shared.opts.sd_unet is None or shared.opts.sd_unet == 'Default') else shared.opts.sd_unet if 'txt2img' in p.ops: args["Variation seed"] = all_subseeds[index] if p.subseed_strength > 0 else None args["Variation strength"] = p.subseed_strength if p.subseed_strength > 0 else None @@ -159,7 +160,7 @@ def create_infotext(p: StableDiffusionProcessing, all_prompts=None, all_seeds=No args["Embeddings"] = ', '.join(sd_hijack.model_hijack.embedding_db.embeddings_used) # samplers - if getattr(p, 'sampler_name', None) is not None: + if getattr(p, 'sampler_name', None) is not None and p.sampler_name.lower() != 'default': args["Sampler eta delta"] = shared.opts.eta_noise_seed_delta if shared.opts.eta_noise_seed_delta != 0 and sd_samplers_common.is_sampler_using_eta_noise_seed_delta(p) else None args["Sampler eta multiplier"] = p.initial_noise_multiplier if getattr(p, 'initial_noise_multiplier', 1.0) != 1.0 else None args['Sampler timesteps'] = shared.opts.schedulers_timesteps if shared.opts.schedulers_timesteps != shared.opts.data_labels.get('schedulers_timesteps').default else None @@ -186,6 +187,9 @@ def create_infotext(p: StableDiffusionProcessing, all_prompts=None, all_seeds=No for k, v in args.copy().items(): if v is None: del args[k] + if type(v) is float or type(v) is int: + if v <= -1: + del args[k] if isinstance(v, str): if len(v) == 0 or v == '0x0': del args[k] diff --git a/modules/processing_vae.py b/modules/processing_vae.py index 54dafb940..00eecc092 100644 --- a/modules/processing_vae.py +++ b/modules/processing_vae.py @@ -139,7 +139,7 @@ def full_vae_decode(latents, model): if latents_mean and latents_std: latents_mean = (torch.tensor(latents_mean).view(1, 4, 1, 1).to(latents.device, latents.dtype)) latents_std = (torch.tensor(latents_std).view(1, 4, 1, 1).to(latents.device, latents.dtype)) - latents = latents * latents_std / scaling_factor + latents_mean + latents = ((latents * latents_std) / scaling_factor) + latents_mean else: latents = latents / scaling_factor if shift_factor: @@ -239,6 +239,8 @@ def vae_postprocess(tensor, model, output_type='np'): if len(tensor.shape) == 3 and tensor.shape[0] == 3: tensor = tensor.unsqueeze(0) if hasattr(model, 'video_processor'): + if len(tensor.shape) == 6 and tensor.shape[1] == 1: + tensor = tensor.squeeze(0) images = model.video_processor.postprocess_video(tensor, output_type='pil') elif hasattr(model, 'image_processor'): images = model.image_processor.postprocess(tensor, output_type=output_type) @@ -293,9 +295,6 @@ def vae_decode(latents, model, output_type='np', vae_type='Full', width=None, he latents = model._unpack_latents(latents, height, width, model.vae_scale_factor) # pylint: disable=protected-access if len(latents.shape) == 3: # lost a batch dim in hires latents = latents.unsqueeze(0) - if latents.shape[0] == 4 and latents.shape[1] != 4: # likely animatediff latent - latents = latents.permute(1, 0, 2, 3) - if latents.shape[-1] <= 4: # not a latent, likely an image decoded = latents.float().cpu().numpy() elif vae_type == 'Full' and hasattr(model, "vae"): diff --git a/modules/prompt_parser_diffusers.py b/modules/prompt_parser_diffusers.py index fc6b2af52..4aaf49a39 100644 --- a/modules/prompt_parser_diffusers.py +++ b/modules/prompt_parser_diffusers.py @@ -10,7 +10,7 @@ from modules import shared, prompt_parser, devices, sd_models from modules.prompt_parser_xhinker import get_weighted_text_embeddings_sd15, get_weighted_text_embeddings_sdxl_2p, get_weighted_text_embeddings_sd3, get_weighted_text_embeddings_flux1 debug_enabled = os.environ.get('SD_PROMPT_DEBUG', None) -debug = shared.log.trace if os.environ.get('SD_PROMPT_DEBUG', None) is not None else lambda *args, **kwargs: None +debug = shared.log.trace if debug_enabled else lambda *args, **kwargs: None debug('Trace: PROMPT') orig_encode_token_ids_to_embeddings = EmbeddingsProvider._encode_token_ids_to_embeddings # pylint: disable=protected-access token_dict = None # used by helper get_tokens diff --git a/modules/rife/__init__.py b/modules/rife/__init__.py index 2a636eb2f..b8e9fbc6f 100644 --- a/modules/rife/__init__.py +++ b/modules/rife/__init__.py @@ -25,7 +25,7 @@ def load(model_path: str = 'rife/flownet-v46.pkl'): from modules import modelloader model_dir = os.path.join(shared.models_path, 'RIFE') model_path = modelloader.load_file_from_url(url=model_url, model_dir=model_dir, file_name='flownet-v46.pkl') - shared.log.debug(f'RIFE load model: file="{model_path}"') + shared.log.debug(f'Video interpolate: model="{model_path}"') model = RifeModel() model.load_model(model_path, -1) model.eval() @@ -46,7 +46,6 @@ def interpolate(images: list, count: int = 2, scale: float = 1.0, pad: int = 1, item = buffer.get() while item is not None: img = item[:, :, ::-1] - # image = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB)) image = Image.fromarray(img) item = buffer.get() interpolated.append(image) @@ -76,6 +75,7 @@ def interpolate(images: list, count: int = 2, scale: float = 1.0, pad: int = 1, pw = ((w - 1) // tmp + 1) * tmp padding = (0, pw - w, 0, ph - h) buffer = Queue(maxsize=8192) + duplicate = 0 _thread.start_new_thread(write, (buffer,)) frame = cv2.cvtColor(np.array(images[0]), cv2.COLOR_RGB2BGR) @@ -93,7 +93,8 @@ def interpolate(images: list, count: int = 2, scale: float = 1.0, pad: int = 1, I1_small = F.interpolate(I1, (32, 32), mode='bilinear', align_corners=False).to(torch.float32) ssim = ssim_matlab(I0_small[:, :3], I1_small[:, :3]) if ssim > 0.99: # skip duplicate frames - continue + duplicate += 1 + # continue if ssim < change: output = [] for _i in range(pad): # fill frames if change rate is above threshold @@ -110,8 +111,8 @@ def interpolate(images: list, count: int = 2, scale: float = 1.0, pad: int = 1, for _i in range(pad): # fill ending frames buffer.put(frame) - while not buffer.empty(): + while not buffer.qsize() > 0: time.sleep(0.1) t1 = time.time() - shared.log.info(f'RIFE interpolate: input={len(images)} frames={len(interpolated)} width={w} height={h} interpolate={count} scale={scale} pad={pad} change={change} time={round(t1 - t0, 2)}') + shared.log.info(f'Video interpolate: input={len(images)} frames={len(interpolated)} buffer={buffer.qsize()} duplicate={duplicate} width={w} height={h} interpolate={count} scale={scale} pad={pad} change={change} time={round(t1 - t0, 2)}') return interpolated diff --git a/modules/rocm.py b/modules/rocm.py index b06e92ef6..a742556a2 100644 --- a/modules/rocm.py +++ b/modules/rocm.py @@ -172,12 +172,12 @@ else: return f'{arr[0]}.{arr[1]}' if len(arr) >= 2 else None def get_agents() -> List[Agent]: - if is_wsl: # WSL does not have 'rocm_agent_enumerator' - agents = spawn("rocminfo").split("\n") - agents = [x.strip().split(" ")[-1] for x in agents if x.startswith(' Name:') and "CPU" not in x] - else: + try: agents = spawn("rocm_agent_enumerator").split("\n") agents = [x for x in agents if x and x != 'gfx000'] + except Exception: # old version of ROCm WSL doesn't have rocm_agent_enumerator + agents = spawn("rocminfo").split("\n") + agents = [x.strip().split(" ")[-1] for x in agents if x.startswith(' Name:') and "CPU" not in x] return [Agent(x) for x in agents] def load_hsa_runtime() -> None: diff --git a/modules/schedulers/scheduler_dpm_flowmatch.py b/modules/schedulers/scheduler_dpm_flowmatch.py index 69452aca9..c1f045e8a 100644 --- a/modules/schedulers/scheduler_dpm_flowmatch.py +++ b/modules/schedulers/scheduler_dpm_flowmatch.py @@ -230,6 +230,7 @@ class FlowMatchDPMSolverMultistepScheduler(SchedulerMixin, ConfigMixin): device: Union[str, torch.device] = None, sigmas: Optional[List[float]] = None, mu: Optional[float] = None, + timesteps: Optional[torch.Tensor] = None, ): """ Sets the discrete timesteps used for the diffusion chain (to be run before inference). @@ -355,12 +356,12 @@ class FlowMatchDPMSolverMultistepScheduler(SchedulerMixin, ConfigMixin): sigma_min = sigmas[-1] sigmas = np.linspace(1.0, sigma_min, num_inference_steps) sigmas = torch.from_numpy(sigmas).to(dtype=torch.float64, device=device) - + if self.config.use_dynamic_shifting: sigmas = self.time_shift(mu, 1.0, sigmas) else: sigmas = self.config.shift * sigmas / (1 + (self.config.shift - 1) * sigmas) - + timesteps = sigmas * self.config.num_train_timesteps self.timesteps = timesteps.to(device=device) self.sigmas = torch.cat([sigmas, torch.zeros(1, device=sigmas.device)]) @@ -508,7 +509,10 @@ class FlowMatchDPMSolverMultistepScheduler(SchedulerMixin, ConfigMixin): def t_fn(_sigma: torch.Tensor) -> torch.Tensor: return _sigma.log().neg() sigma = self.sigmas[self.step_index] - sigma_next = self.sigmas[self.step_index + 1] + try: + sigma_next = self.sigmas[self.step_index + 1] + except Exception: + sigma_next = self.sigmas[-1] sigma_prev = self.sigmas[self.step_index - 1] if self.config.algorithm_type == "dpmsolver2": if self.config.solver_order == 2: @@ -517,7 +521,7 @@ class FlowMatchDPMSolverMultistepScheduler(SchedulerMixin, ConfigMixin): model_output = sample - sigma * model_output d = (sample - model_output) / sigma dt = sigma_next - sigma - sample = sample + d * dt + sample = sample + d * dt else: # DPM-Solver2 sigma_mid = sigma.log().lerp(sigma_next.log(), 0.5).exp() @@ -596,7 +600,7 @@ class FlowMatchDPMSolverMultistepScheduler(SchedulerMixin, ConfigMixin): elif self.config.algorithm_type == "dpmsolver++2M": if self.config.solver_order == 2: t, t_next = t_fn(sigma), t_fn(sigma_next) - h = t_next - t + h = t_next - t if self.model_outputs[-2] is None or sigma_next == 0: sample = (sigma_fn(t_next) / sigma_fn(t)) * sample - (-h).expm1() * model_output else: @@ -703,7 +707,7 @@ class FlowMatchDPMSolverMultistepScheduler(SchedulerMixin, ConfigMixin): if self.config.use_noise_sampler: sample = sample + self.noise_sampler(sigma_fn(t), sigma_fn(t_next)) * self.config.s_noise * su else: - sample = sample + noise * self.config.s_noise * su + sample = sample + noise * self.config.s_noise * su del x_2 del denoised_2 del d @@ -745,13 +749,13 @@ class FlowMatchDPMSolverMultistepScheduler(SchedulerMixin, ConfigMixin): t, s = -sigma.log(), -sigma_next.log() h = s - t h_eta = h * 2 - + # 3. Delta timestep dt = sigma_next - sigma sample = sample + model_output * dt sample = torch.exp(-h_eta) * sample + (-h_eta).expm1().neg() * model_output - + if self.h_2 is not None: r0 = self.h_1 / h r1 = self.h_2 / h @@ -780,13 +784,13 @@ class FlowMatchDPMSolverMultistepScheduler(SchedulerMixin, ConfigMixin): sample = sample + self.noise_sampler(sigma, sigma_next) * sigma_next * (-2 * h).expm1().neg().sqrt() * self.config.s_noise else: sample = sample + noise * sigma_next * (-2 * h).expm1().neg().sqrt() * self.config.s_noise - + self.h_2 = self.h_1 self.h_1 = h if not self.config.use_noise_sampler and noise is not None: del noise prev_sample = sample - + # Cast sample back to expected dtype prev_sample = prev_sample.to(model_output.dtype) diff --git a/modules/scripts_postprocessing.py b/modules/scripts_postprocessing.py index 3afaf68d3..da65ee551 100644 --- a/modules/scripts_postprocessing.py +++ b/modules/scripts_postprocessing.py @@ -65,7 +65,7 @@ class ScriptPostprocessingRunner: script.args_from = len(inputs) script.args_to = len(inputs) script.controls = wrap_call(script.ui, script.filename, "ui") - for control in script.controls.values(): + for control in script.controls.values() if script.controls is not None else []: control.custom_script_source = os.path.basename(script.filename) inputs += list(script.controls.values()) script.args_to = len(inputs) diff --git a/modules/sd_detect.py b/modules/sd_detect.py index f011fd77a..642043f29 100644 --- a/modules/sd_detect.py +++ b/modules/sd_detect.py @@ -75,8 +75,10 @@ def detect_pipeline(f: str, op: str = 'model', warning=True, quiet=False): guess = 'Kolors' if 'auraflow' in f.lower(): guess = 'AuraFlow' - if 'cogview' in f.lower(): - guess = 'CogView' + if 'cogview3' in f.lower(): + guess = 'CogView3' + if 'cogview4' in f.lower(): + guess = 'CogView4' if 'meissonic' in f.lower(): guess = 'Meissonic' pipeline = 'custom' diff --git a/modules/sd_hijack.py b/modules/sd_hijack.py index d4456945a..57573a493 100644 --- a/modules/sd_hijack.py +++ b/modules/sd_hijack.py @@ -375,11 +375,11 @@ if devices.backend != "ipex": # disable_compile for AutoencoderKLOutput is the only change if torch.__version__.startswith("2.6"): from dataclasses import dataclass - from torch.compiler import disable as disable_compile - import diffusers.models.autoencoders.autoencoder_kl + from torch.compiler import disable as disable_compile # pylint: disable=ungrouped-imports + import diffusers.models.autoencoders.autoencoder_kl # pylint: disable=ungrouped-imports @dataclass @disable_compile class AutoencoderKLOutput(diffusers.utils.BaseOutput): - latent_dist: "DiagonalGaussianDistribution" # noqa: F821 + latent_dist: "DiagonalGaussianDistribution" # noqa: F821 diffusers.models.autoencoders.autoencoder_kl.AutoencoderKLOutput = AutoencoderKLOutput diff --git a/modules/sd_hijack_dynamic_atten.py b/modules/sd_hijack_dynamic_atten.py index 6c3e69e3a..7a394772a 100644 --- a/modules/sd_hijack_dynamic_atten.py +++ b/modules/sd_hijack_dynamic_atten.py @@ -56,13 +56,13 @@ if devices.sdpa_pre_dyanmic_atten is None: @wraps(devices.sdpa_pre_dyanmic_atten) def dynamic_scaled_dot_product_attention(query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False, **kwargs): is_unsqueezed = False - if len(query.shape) == 3: + if query.dim() == 3: query = query.unsqueeze(0) is_unsqueezed = True - if len(key.shape) == 3: - key = key.unsqueeze(0) - if len(value.shape) == 3: - value = value.unsqueeze(0) + if key.dim() == 3: + key = key.unsqueeze(0) + if value.dim() == 3: + value = value.unsqueeze(0) do_batch_split, do_head_split, do_query_split, split_batch_size, split_head_size, split_query_size = find_sdpa_slice_sizes(query.shape, key.shape, query.element_size(), slice_rate=shared.opts.dynamic_attention_slice_rate, trigger_rate=shared.opts.dynamic_attention_trigger_rate) # Slice SDPA @@ -111,7 +111,7 @@ def dynamic_scaled_dot_product_attention(query, key, value, attn_mask=None, drop else: hidden_states = devices.sdpa_pre_dyanmic_atten(query, key, value, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal, **kwargs) if is_unsqueezed: - hidden_states.squeeze(0) + hidden_states = hidden_states.squeeze(0) return hidden_states diff --git a/modules/sd_hijack_optimizations.py b/modules/sd_hijack_optimizations.py index ce6537925..f89465e36 100644 --- a/modules/sd_hijack_optimizations.py +++ b/modules/sd_hijack_optimizations.py @@ -315,7 +315,7 @@ def get_xformers_flash_attention_op(q, k, v): if 'Flash attention' not in shared.opts.xformers_options: return None try: - flash_attention_op = xformers.ops.MemoryEfficientAttentionFlashAttentionOp # pylint: disable=used-before-assignment + flash_attention_op = xformers.ops.MemoryEfficientAttentionFlashAttentionOp # pylint: disable=possibly-used-before-assignment, used-before-assignment fw, _bw = flash_attention_op if fw.supports(xformers.ops.fmha.Inputs(query=q, key=k, value=v, attn_bias=None)): return flash_attention_op diff --git a/modules/sd_models.py b/modules/sd_models.py index bb5b672f7..4e85ec74c 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -130,7 +130,7 @@ def set_diffuser_options(sd_model, vae=None, op:str='model', offload:bool=True, model.requires_grad_(False) model.eval() return model - sd_model = apply_function_to_model(sd_model, eval_model, ["Model", "VAE", "Text Encoder"], op="eval") + sd_model = apply_function_to_model(sd_model, eval_model, ["Model", "VAE", "TE"], op="eval") if len(shared.opts.torchao_quantization) > 0 and shared.opts.torchao_quantization_mode == 'post': sd_model = model_quant.torchao_quantization(sd_model) @@ -295,9 +295,13 @@ def load_diffuser_force(model_type, checkpoint_info, diffusers_load_config, op=' sd_model = load_lumina2(checkpoint_info, diffusers_load_config) elif model_type in ['Stable Diffusion 3']: from modules.model_sd3 import load_sd3 - shared.log.debug(f'Load {op}: model="Stable Diffusion 3"') - shared.opts.scheduler = 'Default' 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 + from modules.model_cogview import load_cogview3 + sd_model = load_cogview3(checkpoint_info, diffusers_load_config) + elif model_type in ['CogView4']: # 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 from modules.model_meissonic import load_meissonic sd_model = load_meissonic(checkpoint_info, diffusers_load_config) @@ -571,7 +575,7 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No prompt_parser_diffusers.cache.clear() set_diffuser_options(sd_model, vae, op, offload=False) - if shared.opts.nncf_compress_weights and not ('Model' in shared.opts.cuda_compile and shared.opts.cuda_compile_backend == "openvino_fx"): + if 'Model' in shared.opts.nncf_compress_weights and not ('Model' in shared.opts.cuda_compile and shared.opts.cuda_compile_backend == "openvino_fx"): sd_model = model_quant.nncf_compress_weights(sd_model) # run this before move model so it can be compressed in CPU if shared.opts.optimum_quanto_weights: sd_model = model_quant.optimum_quanto_weights(sd_model) # run this before move model so it can be compressed in CPU @@ -627,9 +631,12 @@ class DiffusersTaskType(Enum): def get_diffusers_task(pipe: diffusers.DiffusionPipeline) -> DiffusersTaskType: - if pipe.__class__.__name__ in ["StableVideoDiffusionPipeline", "LEditsPPPipelineStableDiffusion", "LEditsPPPipelineStableDiffusionXL", "OmniGenPipeline"]: + cls = pipe.__class__.__name__ + if cls in ["LEditsPPPipelineStableDiffusion", "LEditsPPPipelineStableDiffusionXL", "OmniGenPipeline"]: # special case return DiffusersTaskType.IMAGE_2_IMAGE - elif pipe.__class__.__name__ == "StableDiffusionXLInstructPix2PixPipeline": + elif 'ImageToVideo' in cls or cls in ['LTXConditionPipeline', 'StableVideoDiffusionPipeline']: # i2v pipelines + return DiffusersTaskType.IMAGE_2_IMAGE + elif 'Instruct' in cls: return DiffusersTaskType.INSTRUCT elif pipe.__class__ in diffusers.pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING.values(): return DiffusersTaskType.IMAGE_2_IMAGE @@ -757,7 +764,6 @@ def set_diffuser_pipe(pipe, new_pipe_type): 'InstantIRPipeline', 'FluxFillPipeline', 'FluxControlPipeline', - 'StableVideoDiffusionPipeline', 'PixelSmithXLPipeline', 'PhotoMakerStableDiffusionXLPipeline', 'StableDiffusionXLInstantIDPipeline', @@ -774,6 +780,8 @@ def set_diffuser_pipe(pipe, new_pipe_type): cls = pipe.__class__.__name__ if cls in exclude: return pipe + if 'Video' in cls: + return pipe if 'Onnx' in cls: return pipe @@ -944,7 +952,7 @@ def get_native(pipe: diffusers.DiffusionPipeline): def reload_text_encoder(initial=False): - if initial and (shared.opts.sd_text_encoder is None or shared.opts.sd_text_encoder == 'None'): + if initial and (shared.opts.sd_text_encoder is None or shared.opts.sd_text_encoder == 'Default'): return # dont unload signature = get_signature(shared.sd_model) t5 = [k for k, v in signature.items() if 'T5EncoderModel' in str(v)] @@ -1043,14 +1051,14 @@ def reload_model_weights(sd_model=None, info=None, reuse_dict=False, op='model', def clear_caches(): - # shared.log.debug('Cache clear') if not shared.opts.lora_legacy: - from modules.lora import networks - networks.loaded_networks.clear() - networks.previously_loaded_networks.clear() - networks.lora_cache.clear() - from modules import prompt_parser_diffusers + from modules.lora import lora_common, lora_load + lora_common.loaded_networks.clear() + lora_common.previously_loaded_networks.clear() + lora_load.lora_cache.clear() + from modules import prompt_parser_diffusers, memstats prompt_parser_diffusers.cache.clear() + memstats.reset_stats() def unload_model_weights(op='model'): diff --git a/modules/sd_models_utils.py b/modules/sd_models_utils.py index 93367197c..0b2794bac 100644 --- a/modules/sd_models_utils.py +++ b/modules/sd_models_utils.py @@ -67,19 +67,15 @@ def read_state_dict(checkpoint_file, map_location=None, what:str='model'): # pyl return None if shared.opts.stream_load: if extension.lower() == ".safetensors": - # shared.log.debug('Model weights loading: type=safetensors mode=buffered') buffer = f.read() pl_sd = safetensors.torch.load(buffer) else: - # shared.log.debug('Model weights loading: type=checkpoint mode=buffered') buffer = io.BytesIO(f.read()) pl_sd = torch.load(buffer, map_location='cpu') else: if extension.lower() == ".safetensors": - # shared.log.debug('Model weights loading: type=safetensors mode=mmap') pl_sd = safetensors.torch.load_file(checkpoint_file, device='cpu') else: - # shared.log.debug('Model weights loading: type=checkpoint mode=direct') pl_sd = torch.load(f, map_location='cpu') sd = get_state_dict_from_checkpoint(pl_sd) del pl_sd @@ -168,7 +164,7 @@ def apply_function_to_model(sd_model, function, options, op=None): sd_model.prior_pipe.prior = function(sd_model.prior_pipe.prior, op="prior_pipe.prior", sd_model=sd_model) if op == "nncf" and "StableCascade" in sd_model.__class__.__name__: sd_model.prior_pipe.prior.clip_txt_pooled_mapper = backup_clip_txt_pooled_mapper - if "Text Encoder" in options: + if "TE" in options: if hasattr(sd_model, 'text_encoder') and hasattr(sd_model.text_encoder, 'config'): if hasattr(sd_model, 'decoder_pipe') and hasattr(sd_model.decoder_pipe, 'text_encoder') and hasattr(sd_model.decoder_pipe.text_encoder, 'config'): sd_model.decoder_pipe.text_encoder = function(sd_model.decoder_pipe.text_encoder, op="decoder_pipe.text_encoder", sd_model=sd_model) diff --git a/modules/sd_modules.py b/modules/sd_modules.py new file mode 100644 index 000000000..781287f9a --- /dev/null +++ b/modules/sd_modules.py @@ -0,0 +1,73 @@ +from dataclasses import dataclass +import inspect +import torch + + +@dataclass +class ModuleStats: + module: str + cls: str + params: float + size: float + quant: str + dtype: str + + def __init__(self, module: str, cls: str, params: float, size: float, quant: str, dtype: str): + self.module = module + self.cls = cls + self.params = params + self.size = size + self.quant = quant + self.dtype = dtype + + def __str__(self): + return f'module="{self.module}" cls={self.cls} params={self.params:.3f} size={self.size:.3f} quant={self.quant} dtype={self.dtype}' + + +def get_signature(cls): + signature = inspect.signature(cls.__init__, follow_wrapped=True) + return signature.parameters + + +def get_module_stats(name, module): + if not isinstance(module, torch.nn.Module): + return None + try: + module_size = sum(p.numel() * p.element_size() for p in module.parameters(recurse=True)) / 1024 / 1024 / 1024 + param_num = sum(p.numel() for p in module.parameters(recurse=True)) / 1024 / 1024 / 1024 + except Exception: + module_size = 0 + param_num = 0 + cls = module.__class__.__name__ + quant = getattr(module, "quantization_method", None) + module_stats = ModuleStats(name, cls, param_num, module_size, quant, module.dtype) + return module_stats + + +def get_model_stats(model, exclude=None): + # from transformers import Gemma3ForCausalLM + modules = [] + + if isinstance(model, torch.nn.Module): + module_stats = get_module_stats(model.__class__.__name__, model) + if module_stats is not None: + modules.append(module_stats) + return modules + + if hasattr(model, "_internal_dict"): + modules_names = model._internal_dict.keys() # pylint: disable=protected-access + else: + modules_names = get_signature(model).keys() + + if modules_names is None or not isinstance(modules_names, list) or len(modules_names) == 0: + return modules + + modules_names = [m for m in modules_names if m is not None and m not in exclude and not m.startswith('_')] + for module_name in modules_names: + module = getattr(model, module_name, None) + if module is not None: + module_stats = get_module_stats(module_name, module) + if module_stats is not None: + modules.append(module_stats) + + return modules diff --git a/modules/sd_offload.py b/modules/sd_offload.py index 33ce1b2c0..e4b24c17b 100644 --- a/modules/sd_offload.py +++ b/modules/sd_offload.py @@ -3,15 +3,15 @@ import sys import time import inspect import torch -import diffusers import accelerate.hooks from modules import shared, devices, errors, model_quant from modules.timer import process as process_timer debug_move = shared.log.trace if os.environ.get('SD_MOVE_DEBUG', None) is not None else lambda *args, **kwargs: None -should_offload = ['sc', 'sd3', 'f1', 'hunyuandit', 'auraflow', 'omnigen', 'hunyuanvideo', 'cogvideox', 'mochi'] +should_offload = ['sc', 'sd3', 'f1', 'hunyuandit', 'auraflow', 'omnigen', 'cogview4'] offload_hook_instance = None +balanced_offload_exclude = ['OmniGenPipeline', 'CogView4Pipeline'] def get_signature(cls): @@ -65,7 +65,7 @@ def set_diffuser_offload(sd_model, op:str='model', quiet:bool=False): if not (hasattr(sd_model, "has_accelerate") and sd_model.has_accelerate): sd_model.has_accelerate = False if shared.opts.diffusers_offload_mode == "none": - if shared.sd_model_type in should_offload: + if shared.sd_model_type in should_offload or 'video' in shared.sd_model_type: shared.log.warning(f'Setting {op}: offload={shared.opts.diffusers_offload_mode} type={shared.sd_model.__class__.__name__} large model') else: shared.log.quiet(quiet, f'Setting {op}: offload={shared.opts.diffusers_offload_mode} limit={shared.opts.cuda_mem_fraction}') @@ -126,7 +126,7 @@ class OffloadHook(accelerate.hooks.ModelHook): self.cpu = int(shared.cpu_memory * shared.opts.diffusers_offload_max_cpu_memory * 1024*1024*1024) self.offload_map = {} self.param_map = {} - gpu = f'{shared.gpu_memory * shared.opts.diffusers_offload_min_gpu_memory:.3f}-{shared.gpu_memory * shared.opts.diffusers_offload_max_gpu_memory}:{shared.gpu_memory}' + gpu = f'{(shared.gpu_memory * shared.opts.diffusers_offload_min_gpu_memory):.2f}-{(shared.gpu_memory * shared.opts.diffusers_offload_max_gpu_memory):.2f}:{shared.gpu_memory:.2f}' shared.log.info(f'Offload: type=balanced op=init watermark={self.min_watermark}-{self.max_watermark} gpu={gpu} cpu={shared.cpu_memory:.3f} limit={shared.opts.cuda_mem_fraction:.2f}') self.validate() super().__init__() @@ -166,9 +166,7 @@ class OffloadHook(accelerate.hooks.ModelHook): keys = device_map.keys() for v in keys: if isinstance(device_map[v], int): - # int implies CUDA or XPU device, but it will break DirectML backend. - # Therefore, the type of device should be added. - device_map[v] = f"{devices.device.type}:{device_map[v]}" + device_map[v] = f"{devices.device.type}:{device_map[v]}" # int implies CUDA or XPU device, but it will break DirectML backend so we add type module = accelerate.dispatch_model(module, device_map=device_map, offload_dir=offload_dir) module._hf_hook.execution_device = torch.device(devices.device) # pylint: disable=protected-access module.balanced_offload_device_map = device_map @@ -182,7 +180,7 @@ class OffloadHook(accelerate.hooks.ModelHook): return module -def apply_balanced_offload(sd_model, exclude=[]): +def apply_balanced_offload(sd_model=None, exclude=[]): global offload_hook_instance # pylint: disable=global-statement if shared.opts.diffusers_offload_mode != "balanced": return sd_model @@ -193,8 +191,7 @@ def apply_balanced_offload(sd_model, exclude=[]): if sd_model is None: return sd_model t0 = time.time() - excluded = ['OmniGenPipeline'] - if sd_model.__class__.__name__ in excluded: + if sd_model.__class__.__name__ in balanced_offload_exclude: return sd_model cached = True checkpoint_name = sd_model.sd_checkpoint_info.name if getattr(sd_model, "sd_checkpoint_info", None) is not None else None @@ -285,22 +282,8 @@ def apply_balanced_offload(sd_model, exclude=[]): apply_balanced_offload_to_module(sd_model.prior_pipe) if hasattr(sd_model, "decoder_pipe"): apply_balanced_offload_to_module(sd_model.decoder_pipe) - if shared.opts.layerwise_quantization: model_quant.apply_layerwise(sd_model, quiet=True) # need to reapply since hooks were removed/readded - if shared.opts.pab_enabled and hasattr(sd_model, 'transformer'): - pab_config = diffusers.PyramidAttentionBroadcastConfig( - spatial_attention_block_skip_range=shared.opts.pab_block_skip_range, - spatial_attention_timestep_skip_range=(int(100 * shared.opts.pab_timestep_skip_start), int(100 * shared.opts.pab_timestep_skip_end)), - current_timestep_callback=lambda: sd_model.current_timestep, # pylint: disable=protected-access - ) - try: - diffusers.apply_pyramid_attention_broadcast(sd_model.transformer, pab_config) - except Exception: # hook may already exist - pass - if not cached: - shared.log.info(f'Applying PAB: cls={sd_model.transformer.__class__.__name__} block={shared.opts.pab_block_skip_range} start={shared.opts.pab_timestep_skip_start} end={shared.opts.pab_timestep_skip_end}') - set_accelerate(sd_model) t = time.time() - t0 process_timer.add('offload', t) diff --git a/modules/sd_samplers.py b/modules/sd_samplers.py index accd6b0ed..be375dbff 100644 --- a/modules/sd_samplers.py +++ b/modules/sd_samplers.py @@ -12,6 +12,8 @@ samplers = all_samplers samplers_for_img2img = all_samplers samplers_map = {} loaded_config = None +flow_models = ['Flux', 'StableDiffusion3', 'Lumina', 'AuraFlow', 'Sana', 'CogView4'] +flow_models += ['Hunyuan', 'LTX', 'Mochi'] def list_samplers(): @@ -79,10 +81,9 @@ def create_sampler(name, model): shared.log.debug(f'Sampler: "{name}" config={config.options}') return sampler elif shared.native: - FlowModels = ['Flux', 'StableDiffusion3', 'Lumina', 'AuraFlow', 'Sana', 'HunyuanVideoPipeline'] if 'KDiffusion' in model.__class__.__name__: return None - if not any(x in model.__class__.__name__ for x in FlowModels) and 'FlowMatch' in name: + if not any(x in model.__class__.__name__ for x in flow_models) and 'FlowMatch' in name: shared.log.warning(f'Sampler: default={current} target="{name}" class={model.__class__.__name__} flow-match scheduler unsupported') return None sampler = config.constructor(model) diff --git a/modules/sd_samplers_common.py b/modules/sd_samplers_common.py index 0b50c5c6c..5eb4a5f94 100644 --- a/modules/sd_samplers_common.py +++ b/modules/sd_samplers_common.py @@ -9,7 +9,7 @@ from modules import shared, devices, processing, images, sd_vae_approx, sd_vae_t SamplerData = namedtuple('SamplerData', ['name', 'constructor', 'aliases', 'options']) approximation_indexes = { "Simple": 0, "Approximate": 1, "TAESD": 2, "Full VAE": 3 } -flow_models = ['f1', 'sd3', 'lumina', 'auraflow', 'sana', 'lumina2'] +flow_models = ['f1', 'sd3', 'lumina', 'auraflow', 'sana', 'lumina2', 'cogview4'] warned = False queue_lock = threading.Lock() diff --git a/modules/sd_samplers_diffusers.py b/modules/sd_samplers_diffusers.py index 667683eb0..644379f1b 100644 --- a/modules/sd_samplers_diffusers.py +++ b/modules/sd_samplers_diffusers.py @@ -50,6 +50,10 @@ try: from modules.schedulers.scheduler_bdia import BDIA_DDIMScheduler # pylint: disable=ungrouped-imports from modules.schedulers.scheduler_ufogen import UFOGenScheduler # pylint: disable=ungrouped-imports from modules.perflow import PeRFlowScheduler # pylint: disable=ungrouped-imports + # from modules.schedulers.scheduler_kohaku import KohakuLoNyuYogScheduler # pylint: disable=ungrouped-imports + # from modules.schedulers.scheduler_smea import SMEAScheduler # pylint: disable=ungrouped-imports + # from modules.schedulers.scheduler_dy import DYScheduler # pylint: disable=ungrouped-imports + # from modules.schedulers.scheduler_negative import EulerNegativeScheduler # pylint: disable=ungrouped-imports except Exception as e: shared.log.error(f'Sampler import: version={diffusers.__version__} error: {e}') if os.environ.get('SD_SAMPLER_DEBUG', None) is not None: @@ -68,6 +72,10 @@ config = { 'Euler SGM': { 'steps_offset': 0, 'interpolation_type': "linear", 'rescale_betas_zero_snr': False, 'final_sigmas_type': 'zero', 'timestep_spacing': 'trailing', 'use_beta_sigmas': False, 'use_exponential_sigmas': False, 'use_karras_sigmas': False, 'prediction_type': "sample" }, 'Euler EDM': { 'sigma_schedule': "karras" }, 'Euler FlowMatch': { 'timestep_spacing': "linspace", 'shift': 1, 'use_dynamic_shifting': False, 'use_karras_sigmas': False, 'use_exponential_sigmas': False, 'use_beta_sigmas': False }, + # 'Euler SMEA': {}, + # 'Euler DY': {}, + # 'Euler Negative': {}, + # 'Kohaku LoNyu': {}, 'DPM++': { 'thresholding': False, 'sample_max_value': 1.0, 'algorithm_type': "dpmsolver++", 'solver_type': "midpoint", 'lower_order_final': True, 'use_karras_sigmas': False, 'use_exponential_sigmas': False, 'use_flow_sigmas': False, 'use_beta_sigmas': False, 'use_lu_lambdas': False, 'final_sigmas_type': 'zero', 'timestep_spacing': 'linspace', 'solver_order': 1 }, 'DPM++ 2M': { 'thresholding': False, 'sample_max_value': 1.0, 'algorithm_type': "dpmsolver++", 'solver_type': "midpoint", 'lower_order_final': True, 'use_karras_sigmas': False, 'use_exponential_sigmas': False, 'use_flow_sigmas': False, 'use_beta_sigmas': False, 'use_lu_lambdas': False, 'final_sigmas_type': 'zero', 'timestep_spacing': 'linspace', 'solver_order': 2 }, @@ -124,6 +132,9 @@ samplers_data_diffusers = [ SamplerData('Euler SGM', lambda model: DiffusionSampler('Euler SGM', EulerDiscreteScheduler, model), [], {}), SamplerData('Euler EDM', lambda model: DiffusionSampler('Euler EDM', EDMEulerScheduler, model), [], {}), SamplerData('Euler FlowMatch', lambda model: DiffusionSampler('Euler FlowMatch', FlowMatchEulerDiscreteScheduler, model), [], {}), + # SamplerData('Euler SMEA', lambda model: DiffusionSampler('Euler SMEA', SMEAScheduler, model), [], {}), + # SamplerData('Euler DY', lambda model: DiffusionSampler('Euler DY', DYScheduler, model), [], {}), + # SamplerData('Euler Negative', lambda model: DiffusionSampler('Euler Negative', EulerNegativeScheduler, model), [], {}), SamplerData('DPM++', lambda model: DiffusionSampler('DPM++', DPMSolverMultistepScheduler, model), [], {}), SamplerData('DPM++ 2M', lambda model: DiffusionSampler('DPM++ 2M', DPMSolverMultistepScheduler, model), [], {}), @@ -169,6 +180,7 @@ samplers_data_diffusers = [ SamplerData('TDD', lambda model: DiffusionSampler('TDD', TDDScheduler, model), [], {}), SamplerData('PeRFlow', lambda model: DiffusionSampler('PeRFlow', PeRFlowScheduler, model), [], {}), SamplerData('UFOGen', lambda model: DiffusionSampler('UFOGen', UFOGenScheduler, model), [], {}), + # SamplerData('Kohaku LoNyu', lambda model: DiffusionSampler('Kohaku LoNyu', KohakuLoNyuYogScheduler, model), [], {}), SamplerData('Same as primary', None, [], {}), ] @@ -181,8 +193,7 @@ class DiffusionSampler: self.name = name self.config = {} self.sampler = None - # if not hasattr(model, 'scheduler'): - # return + if getattr(model, "default_scheduler", None) is None and (model is not None): # sanity check model.default_scheduler = copy.deepcopy(model.scheduler) for key, value in config.get('All', {}).items(): # apply global defaults @@ -205,6 +216,7 @@ class DiffusionSampler: for key, value in kwargs.items(): # apply user args, if any if key in self.config: self.config[key] = value + # finally apply user preferences if shared.opts.schedulers_prediction_type != 'default': self.config['prediction_type'] = shared.opts.schedulers_prediction_type @@ -271,6 +283,7 @@ class DiffusionSampler: del self.config['prediction_type'] if 'SGM' in name: self.config['timestep_spacing'] = 'trailing' + # validate all config params signature = inspect.signature(constructor, follow_wrapped=True) possible = signature.parameters.keys() @@ -281,7 +294,8 @@ class DiffusionSampler: debug_log(f'Sampler: name="{name}"') debug_log(f'Sampler: config={self.config}') debug_log(f'Sampler: signature={possible}') - # shared.log.debug_log(f'Sampler: sampler="{name}" config={self.config}') + + # finally create the new sampler try: sampler = constructor(**self.config) except Exception as e: @@ -290,21 +304,25 @@ class DiffusionSampler: errors.display(e, 'Samplers') self.sampler = None return - accept_sigmas = "sigmas" in set(inspect.signature(sampler.set_timesteps).parameters.keys()) - accepts_timesteps = "timesteps" in set(inspect.signature(sampler.set_timesteps).parameters.keys()) - accept_scale_noise = hasattr(sampler, "scale_noise") - debug_log(f'Sampler: sampler="{name}" sigmas={accept_sigmas} timesteps={accepts_timesteps}') - if ('Flux' in model.__class__.__name__) and (not accept_sigmas): - shared.log.warning(f'Sampler: sampler="{name}" does not accept sigmas') - self.sampler = None - return - if ('StableDiffusion3' in model.__class__.__name__) and (not accept_scale_noise): - shared.log.warning(f'Sampler: sampler="{name}" does not implement scale noise') - self.sampler = None - return + + if hasattr(sampler, 'set_timesteps'): + accept_sigmas = "sigmas" in set(inspect.signature(sampler.set_timesteps).parameters.keys()) + accepts_timesteps = "timesteps" in set(inspect.signature(sampler.set_timesteps).parameters.keys()) + accept_scale_noise = hasattr(sampler, "scale_noise") + debug_log(f'Sampler: sampler="{name}" sigmas={accept_sigmas} timesteps={accepts_timesteps}') + if ('Flux' in model.__class__.__name__) and (not accept_sigmas): + shared.log.warning(f'Sampler: sampler="{name}" does not accept sigmas') + self.sampler = None + return + if ('StableDiffusion3' in model.__class__.__name__) and (not accept_scale_noise): + shared.log.warning(f'Sampler: sampler="{name}" does not implement scale noise') + self.sampler = None + return + self.sampler = sampler if name == 'DC Solver': if not hasattr(self.sampler, 'dc_ratios'): pass + # shared.log.debug_log(f'Sampler: class="{self.sampler.__class__.__name__}" config={self.sampler.config}') self.sampler.name = name diff --git a/modules/sd_unet.py b/modules/sd_unet.py index deb0b24b0..cfba470a1 100644 --- a/modules/sd_unet.py +++ b/modules/sd_unet.py @@ -10,7 +10,7 @@ debug = os.environ.get('SD_LOAD_DEBUG', None) is not None def load_unet(model): global loaded_unet # pylint: disable=global-statement - if shared.opts.sd_unet == 'None': + if shared.opts.sd_unet == 'Default': return if shared.opts.sd_unet not in list(unet_dict): shared.log.error(f'UNet model not found: {shared.opts.sd_unet}') diff --git a/modules/sd_vae.py b/modules/sd_vae.py index 8c34e29a8..17faf12d8 100644 --- a/modules/sd_vae.py +++ b/modules/sd_vae.py @@ -117,7 +117,7 @@ def resolve_vae(checkpoint_file): return None, None if shared.cmd_opts.vae is not None: # 1st return shared.cmd_opts.vae, 'forced' - if shared.opts.sd_vae == "None": # 2nd + if shared.opts.sd_vae == "Default": # 2nd return None, None vae_near_checkpoint = find_vae_near_checkpoint(checkpoint_file) if vae_near_checkpoint is not None: # 3rd diff --git a/modules/sd_vae_remote.py b/modules/sd_vae_remote.py index 7d2645b9a..3c1846ea8 100644 --- a/modules/sd_vae_remote.py +++ b/modules/sd_vae_remote.py @@ -1,3 +1,4 @@ +from typing import List import io import time import json @@ -7,12 +8,17 @@ from PIL import Image from safetensors.torch import _tobytes -hf_endpoints = { +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', 'hunyuanvideo': 'https://o7ywnmrahorts457.us-east-1.aws.endpoints.huggingface.cloud', } +hf_encode_endpoints = { + 'sd': 'https://qc6479g0aac6qwy9.us-east-1.aws.endpoints.huggingface.cloud', + 'sdxl': 'https://xjqqhmyn62rog84g.us-east-1.aws.endpoints.huggingface.cloud', + 'f1': 'https://ptccx55jz97f9zgo.us-east-1.aws.endpoints.huggingface.cloud', +} dtypes = { "float16": torch.float16, "float32": torch.float32, @@ -26,24 +32,34 @@ def remote_decode(latents: torch.Tensor, width: int = 0, height: int = 0, model_ tensors = [] content = 0 model_type = model_type or shared.sd_model_type - url = hf_endpoints.get(model_type, None) + url = hf_decode_endpoints.get(model_type, None) if url is None: shared.log.error(f'Decode: type="remote" type={model_type} unsuppported') return tensors t0 = time.time() modelloader.hf_login() - latents = latents.unsqueeze(0) if len(latents.shape) == 3 else latents - for i in range(latents.shape[0]): + latent_copy = latents.detach().clone().to(device=devices.cpu, dtype=devices.dtype) + latent_copy = latents.unsqueeze(0) if len(latents.shape) == 3 else latents + if model_type == 'hunyuanvideo': + latent_copy = latent_copy.unsqueeze(0) + + for i in range(latent_copy.shape[0]): try: - latent = latents[i].detach().clone().to(device=devices.cpu, dtype=devices.dtype).unsqueeze(0) + latent = latent_copy[i] + if model_type != 'f1': + latent = latent.unsqueeze(0) params = { - "do_scaling": True, "input_tensor_type": "binary", "shape": list(latent.shape), "dtype": str(latent.dtype).split(".", maxsplit=1)[-1], } headers = { "Content-Type": "tensor/binary" } - if shared.opts.remote_vae_type == 'png': + if 'video' in model_type: + params["partial_postprocess"] = False + params["output_type"] = "pt" + params["output_tensor_type"] = "binary" + headers["Accept"] = "tensor/binary" + elif shared.opts.remote_vae_type == 'png': params["image_format"] = "png" params["output_type"] = "pil" headers["Accept"] = "image/png" @@ -59,6 +75,9 @@ def remote_decode(latents: torch.Tensor, width: int = 0, height: int = 0, model_ if (model_type == 'f1') 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: + params['scaling_factor'] = shared.sd_model.vae.config.get("scaling_factor", None) + params['shift_factor'] = shared.sd_model.vae.config.get("shift_factor", None) response = requests.post( url=url, headers=headers, @@ -67,10 +86,10 @@ def remote_decode(latents: torch.Tensor, width: int = 0, height: int = 0, model_ timeout=300, ) if not response.ok: - shared.log.error(f'Decode: type="remote" model={model_type} code={response.status_code} headers={response.headers} {response.json()}') + shared.log.error(f'Decode: type="remote" model={model_type} code={response.status_code} shape={latent.shape} url="{url}" args={params} headers={response.headers} response={response.json()}') else: content += len(response.content) - if shared.opts.remote_vae_type == 'raw': + if shared.opts.remote_vae_type == 'raw' or 'video' in model_type: shape = json.loads(response.headers["shape"]) dtype = response.headers["dtype"] tensor = torch.frombuffer(bytearray(response.content), dtype=dtypes[dtype]).reshape(shape) @@ -86,3 +105,42 @@ def remote_decode(latents: torch.Tensor, width: int = 0, height: int = 0, model_ t1 = time.time() shared.log.debug(f'Decode: type="remote" model={model_type} mode={shared.opts.remote_vae_type} args={params} bytes={content} time={t1-t0:.3f}s') return tensors + + +def remote_encode(images: List[Image.Image], model_type: str = None) -> torch.Tensor: + from diffusers.utils import remote_utils + from modules import devices, shared, errors, modelloader + if not shared.opts.remote_vae_encode: + return images + tensors = [] + model_type = model_type or shared.sd_model_type + url = hf_encode_endpoints.get(model_type, None) + if url is None: + shared.log.error(f'Decode: type="remote" type={model_type} unsuppported') + return images + t0 = time.time() + modelloader.hf_login() + + if isinstance(images, Image.Image): + images = [images] + for init_image in images: + try: + init_latent = remote_utils.remote_encode( + endpoint=url, + image=init_image, + scaling_factor = shared.sd_model.vae.config.get("scaling_factor", None), + shift_factor = shared.sd_model.vae.config.get("shift_factor", None), + ) + tensors.append(init_latent) + except Exception as e: + shared.log.error(f'Encode: type="remote" model={model_type} {e}') + errors.display(e, 'VAE') + + if len(tensors) > 0 and torch.is_tensor(tensors[0]): + tensors = torch.cat(tensors, dim=0) + tensors = tensors.to(dtype=devices.dtype) + else: + return images + t1 = time.time() + shared.log.debug(f'Encode: type="remote" model={model_type} mode={shared.opts.remote_vae_type} image={images} latent={tensors.shape} time={t1-t0:.3f}s') + return tensors diff --git a/modules/sd_vae_taesd.py b/modules/sd_vae_taesd.py index c8a1b882f..a2a447a3e 100644 --- a/modules/sd_vae_taesd.py +++ b/modules/sd_vae_taesd.py @@ -16,6 +16,9 @@ TAESD_MODELS = { 'TAESD 1.2 Chocolate-Dipped Shortbread': { 'fn': 'taesd_12_', 'uri': 'https://github.com/madebyollin/taesd/raw/8909b44e3befaa0efa79c5791e4fe1c4d4f7884e', 'model': None }, 'TAESD 1.1 Fruit Loops': { 'fn': 'taesd_11_', 'uri': 'https://github.com/madebyollin/taesd/raw/3e8a8a2ab4ad4079db60c1c7dc1379b4cc0c6b31', 'model': None }, 'TAESD 1.0': { 'fn': 'taesd_10_', 'uri': 'https://github.com/madebyollin/taesd/raw/88012e67cf0454e6d90f98911fe9d4aef62add86', 'model': None }, + 'TAE HunyuanVideo': { 'fn': 'taehv.pth', 'uri': 'https://github.com/madebyollin/taehv/raw/refs/heads/main/taehv.pth', 'model': None }, + 'TAE WanVideo': { 'fn': 'taew1.pth', 'uri': 'https://github.com/madebyollin/taehv/raw/refs/heads/main/taew2_1.pth', 'model': None }, + 'TAE MochiVideo': { 'fn': 'taem1.pth', 'uri': 'https://github.com/madebyollin/taem1/raw/refs/heads/main/taem1.pth', 'model': None }, } CQYAN_MODELS = { 'Hybrid-Tiny SD': { @@ -35,49 +38,63 @@ prev_model = '' lock = threading.Lock() -def warn_once(msg): +def warn_once(msg, variant=None): from modules import shared + variant = variant or shared.opts.taesd_variant global prev_warnings # pylint: disable=global-statement if not prev_warnings: prev_warnings = True - shared.log.error(f'Decode: type="taesd" variant="{shared.opts.taesd_variant}": {msg}') + shared.log.error(f'Decode: type="taesd" variant="{variant}": {msg}') return Image.new('RGB', (8, 8), color = (0, 0, 0)) -def get_model(model_type = 'decoder'): +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': cls = 'sd' + variant = variant or shared.opts.taesd_variant folder = os.path.join(paths.models_path, "TAESD") os.makedirs(folder, exist_ok=True) - if 'sd' not in cls and 'f1' not in cls: + if 'video' in cls: + return None + if ('sd' not in cls) and ('f1' not in cls): warn_once(f'cls={shared.sd_model.__class__.__name__} type={cls} unsuppported') return None - if shared.opts.taesd_variant.startswith('TAESD'): - cfg = TAESD_MODELS[shared.opts.taesd_variant] - if (cls == prev_cls) and (model_type == prev_type) and (shared.opts.taesd_variant == prev_model) and (cfg['model'] is not None): + if variant.startswith('TAESD'): + cfg = TAESD_MODELS[variant] + if (cls == prev_cls) and (model_type == prev_type) and (variant == prev_model) and (cfg['model'] is not None): return cfg['model'] fn = os.path.join(folder, cfg['fn'] + cls + '_' + model_type + '.pth') if not os.path.exists(fn): uri = cfg['uri'] + '/tae' + cls + '_' + model_type + '.pth' try: - shared.log.info(f'Decode: type="taesd" variant="{shared.opts.taesd_variant}": uri="{uri}" fn="{fn}" download') + shared.log.info(f'Decode: type="taesd" variant="{variant}": uri="{uri}" fn="{fn}" download') torch.hub.download_url_to_file(uri, fn) except Exception as e: warn_once(f'download uri={uri} {e}') if os.path.exists(fn): prev_cls = cls prev_type = model_type - prev_model = shared.opts.taesd_variant - shared.log.debug(f'Decode: type="taesd" variant="{shared.opts.taesd_variant}" fn="{fn}" load') - from modules.taesd.taesd import TAESD - TAESD_MODELS[shared.opts.taesd_variant]['model'] = TAESD(decoder_path=fn if model_type=='decoder' else None, encoder_path=fn if model_type=='encoder' else None) - return TAESD_MODELS[shared.opts.taesd_variant]['model'] - elif shared.opts.taesd_variant.startswith('Hybrid'): - cfg = CQYAN_MODELS[shared.opts.taesd_variant].get(cls, None) - if (cls == prev_cls) and (model_type == prev_type) and (shared.opts.taesd_variant == prev_model) and (cfg['model'] is not None): + prev_model = variant + shared.log.debug(f'Decode: type="taesd" variant="{variant}" fn="{fn}" load') + if 'TAEHV' in variant: + from modules.taesd.taehv import TAEHV + TAESD_MODELS[variant]['model'] = TAEHV(checkpoint_path=fn) + if 'TAEW2' in variant: + from modules.taesd.taehv import TAEHV + TAESD_MODELS[variant]['model'] = TAEHV(checkpoint_path=fn) + elif 'TAEM1' in variant: + from modules.taesd.taem1 import TAEM1 + TAESD_MODELS[variant]['model'] = TAEM1(checkpoint_path=fn) + else: + from modules.taesd.taesd import TAESD + TAESD_MODELS[variant]['model'] = TAESD(decoder_path=fn if model_type=='decoder' else None, encoder_path=fn if model_type=='encoder' else None) + return TAESD_MODELS[variant]['model'] + elif variant.startswith('Hybrid'): + cfg = CQYAN_MODELS[variant].get(cls, None) + if (cls == prev_cls) and (model_type == prev_type) and (variant == prev_model) and (cfg['model'] is not None): return cfg['model'] if cfg is None: warn_once(f'cls={shared.sd_model.__class__.__name__} type={cls} unsuppported') @@ -85,8 +102,8 @@ def get_model(model_type = 'decoder'): repo = cfg['repo'] prev_cls = cls prev_type = model_type - prev_model = shared.opts.taesd_variant - shared.log.debug(f'Decode: type="taesd" variant="{shared.opts.taesd_variant}" id="{repo}" load') + prev_model = variant + shared.log.debug(f'Decode: type="taesd" variant="{variant}" id="{repo}" load') dtype = devices.dtype_vae if devices.dtype_vae != torch.bfloat16 else torch.float16 # taesd does not support bf16 if 'tiny' in repo: from diffusers.models import AutoencoderTiny @@ -95,7 +112,7 @@ def get_model(model_type = 'decoder'): from modules.taesd.hybrid_small import AutoencoderSmall vae = AutoencoderSmall.from_pretrained(repo, cache_dir=shared.opts.hfcache_dir, torch_dtype=dtype) vae = vae.to(devices.device, dtype=dtype) - CQYAN_MODELS[shared.opts.taesd_variant][cls]['model'] = vae + CQYAN_MODELS[variant][cls]['model'] = vae return vae else: warn_once(f'cls={shared.sd_model.__class__.__name__} type={cls} unsuppported') diff --git a/modules/shared.py b/modules/shared.py index 962bf05d0..63d01446c 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -19,7 +19,7 @@ from modules.dml import memory_providers, default_memory_provider, directml_do_h from modules.onnx_impl import initialize_onnx, execution_providers from modules.memstats import memory_stats, ram_stats # pylint: disable=unused-import from modules.interrogate.openclip import caption_models, caption_types, get_clip_models, refresh_clip_models, category_types -from modules.interrogate.vqa import vlm_models, vlm_prompts +from modules.interrogate.vqa import vlm_models, vlm_prompts, vlm_system from modules.ui_components import DropdownEditable from modules.options import OptionInfo import modules.memmon @@ -46,7 +46,7 @@ tab_names = [] extra_networks = [] options_templates = {} hypernetworks = {} -settings_components = None +settings_components = {} restricted_opts = { "samples_filename_pattern", "directories_filename_pattern", @@ -389,7 +389,7 @@ options_templates.update(options_section(('sd', "Models & Loading"), { "diffusers_pipeline": OptionInfo('Autodetect', 'Model pipeline', gr.Dropdown, lambda: {"choices": list(shared_items.get_pipelines()), "visible": native}), "sd_model_checkpoint": OptionInfo(default_checkpoint, "Base model", DropdownEditable, lambda: {"choices": list_checkpoint_titles()}, refresh=refresh_checkpoints), "sd_model_refiner": OptionInfo('None', "Refiner model", gr.Dropdown, lambda: {"choices": ['None'] + list_checkpoint_titles()}, refresh=refresh_checkpoints), - "sd_unet": OptionInfo("None", "UNET model", gr.Dropdown, lambda: {"choices": shared_items.sd_unet_items()}, refresh=shared_items.refresh_unet_list), + "sd_unet": OptionInfo("Default", "UNET model", gr.Dropdown, lambda: {"choices": shared_items.sd_unet_items()}, refresh=shared_items.refresh_unet_list), "latent_history": OptionInfo(16, "Latent history size", gr.Slider, {"minimum": 0, "maximum": 100, "step": 1}), "offload_sep": OptionInfo("

Model Offloading

", "", gr.HTML), @@ -407,9 +407,9 @@ options_templates.update(options_section(('sd', "Models & Loading"), { "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, "Load model directly to GPU"), + "diffusers_to_gpu": OptionInfo(False, "Model Load model direct to GPU"), "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()}, refresh=refresh_checkpoints), + "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 }), })) @@ -422,13 +422,14 @@ options_templates.update(options_section(('vae_encoder', "Variable Auto Encoder" "diffusers_vae_tile_size": OptionInfo(0, "VAE tile size", gr.Slider, {"minimum": 0, "maximum": 4096, "step": 8 }), "diffusers_vae_tile_overlap": OptionInfo(0.25, "VAE tile overlap", gr.Slider, {"minimum": 0, "maximum": 0.95, "step": 0.05 }), "sd_vae_sliced_encode": OptionInfo(False, "VAE sliced encode", gr.Checkbox, {"visible": not native}), - "nan_skip": OptionInfo(False, "Skip Generation if NaN found in latents", gr.Checkbox), + "nan_skip": OptionInfo(False, "Skip Generation if NaN found in latents", gr.Checkbox, {"visible": False}), "remote_vae_type": OptionInfo('raw', "Remote VAE image type", gr.Dropdown, {"choices": ['raw', 'jpg', 'png']}), + "remote_vae_encode": OptionInfo(False, "Remote VAE for encode"), "rollback_vae": OptionInfo(False, "Attempt VAE roll back for NaN values", gr.Checkbox, {"visible": not native}), })) options_templates.update(options_section(('text_encoder', "Text Encoder"), { - "sd_text_encoder": OptionInfo('None', "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", gr.Dropdown, 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}), @@ -493,7 +494,7 @@ options_templates.update(options_section(('backends', "Backend Settings"), { "olive_cache_optimized": OptionInfo(True, 'Olive cache optimized models'), "ipex_sep": OptionInfo("

IPEX

", "", gr.HTML, {"visible": devices.backend == "ipex"}), - "ipex_optimize": OptionInfo([], "IPEX Optimize", gr.CheckboxGroup, {"choices": ["Model", "VAE", "Text Encoder", "Upscaler"], "visible": devices.backend == "ipex"}), + "ipex_optimize": OptionInfo([], "IPEX Optimize", gr.CheckboxGroup, {"choices": ["Model", "VAE", "TE", "Upscaler"], "visible": devices.backend == "ipex"}), "openvino_sep": OptionInfo("

OpenVINO

", "", gr.HTML, {"visible": cmd_opts.use_openvino}), "openvino_devices": OptionInfo([], "OpenVINO devices to use", gr.CheckboxGroup, {"choices": get_openvino_device_list() if cmd_opts.use_openvino else [], "visible": cmd_opts.use_openvino}), # pylint: disable=E0606 @@ -508,32 +509,36 @@ 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", "VAE", "Text Encoder"], "visible": native}), + "bnb_quantization": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "Transformer", "VAE", "TE", "Video", "LLM"], "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}), - "optimum_quanto_sep": OptionInfo("

Optimum Quanto

", "", gr.HTML), - "optimum_quanto_weights": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "VAE", "Text Encoder", "ControlNet"], "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_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), + "optimum_quanto_weights": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "Transformer", "VAE", "TE", "Video", "LLM", "ControlNet"], "visible": native}), "optimum_quanto_weights_type": OptionInfo("qint8", "Quantization weights type", gr.Dropdown, {"choices": ['qint8', 'qfloat8_e4m3fn', 'qfloat8_e5m2', 'qint4', 'qint2'], "visible": native}), "optimum_quanto_activations_type": OptionInfo("none", "Quantization activations type ", gr.Dropdown, {"choices": ['none', 'qint8', 'qfloat8_e4m3fn', 'qfloat8_e5m2'], "visible": native}), "optimum_quanto_shuffle_weights": OptionInfo(False, "Shuffle weights", gr.Checkbox, {"visible": native}), "torchao_sep": OptionInfo("

TorchAO

", "", gr.HTML), - "torchao_quantization": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "VAE", "Text Encoder"], "visible": native}), + "torchao_quantization": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "Transformer", "VAE", "TE", "Video", "LLM"], "visible": native}), "torchao_quantization_mode": OptionInfo("pre", "Quantization mode", gr.Dropdown, {"choices": ['pre', 'post'], "visible": native}), "torchao_quantization_type": OptionInfo("int8_weight_only", "Quantization type", gr.Dropdown, {"choices": ['int4_weight_only', 'int8_dynamic_activation_int4_weight', 'int8_weight_only', 'int8_dynamic_activation_int8_weight', 'float8_weight_only', 'float8_dynamic_activation_float8_weight', 'float8_static_activation_float8_weight'], "visible": native}), - "nncf_compress_sep": OptionInfo("

NNCF

", "", gr.HTML), - "nncf_compress_weights": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "VAE", "Text Encoder", "ControlNet"], "visible": native}), + "nncf_compress_sep": OptionInfo("

NNCF: Neural Network Compression Framework

", "", gr.HTML), + "nncf_compress_weights": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "Transformer", "VAE", "TE", "Video", "LLM", "ControlNet"], "visible": native}), "nncf_compress_weights_mode": OptionInfo("INT8", "Quantization type", gr.Dropdown, {"choices": ['INT8', 'INT8_SYM', 'INT4_ASYM', 'INT4_SYM', 'NF4'] if cmd_opts.use_openvino else ['INT8']}), "nncf_compress_weights_raito": OptionInfo(0, "Compress ratio", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.01, "visible": cmd_opts.use_openvino}), "nncf_compress_weights_group_size": OptionInfo(0, "Group size", gr.Slider, {"minimum": -1, "maximum": 512, "step": 1, "visible": cmd_opts.use_openvino}), - "nncf_quantize": OptionInfo([], "OpenVINO enabled", gr.CheckboxGroup, {"choices": ["Model", "VAE", "Text Encoder"], "visible": cmd_opts.use_openvino}), + "nncf_quantize": OptionInfo([], "OpenVINO enabled", gr.CheckboxGroup, {"choices": ["Model", "VAE", "TE"], "visible": cmd_opts.use_openvino}), "nncf_quantize_mode": OptionInfo("INT8", "OpenVINO mode", gr.Dropdown, {"choices": ['INT8', 'FP8_E4M3', 'FP8_E5M2'], "visible": cmd_opts.use_openvino}), "nncf_quantize_shuffle_weights": OptionInfo(False, "Shuffle weights", gr.Checkbox, {"visible": native}), "layerwise_quantization_sep": OptionInfo("

Layerwise Casting

", "", gr.HTML), - "layerwise_quantization": OptionInfo([], "Layerwise casting enabled", gr.CheckboxGroup, {"choices": ["Model", "Transformer", "Text Encoder"], "visible": native}), + "layerwise_quantization": OptionInfo([], "Layerwise casting enabled", gr.CheckboxGroup, {"choices": ["Model", "Transformer", "TE"], "visible": native}), "layerwise_quantization_storage": OptionInfo("float8_e4m3fn", "Layerwise casting storage", gr.Dropdown, {"choices": ["float8_e4m3fn", "float8_e5m2"], "visible": native}), "layerwise_quantization_nonblocking": OptionInfo(False, "Layerwise non-blocking operations", gr.Checkbox, {"visible": native}), })) @@ -556,9 +561,21 @@ options_templates.update(options_section(('advanced', "Pipeline Modifiers"), { "pab_sep": OptionInfo("

PAB: Pyramid attention broadcast

", "", gr.HTML), "pab_enabled": OptionInfo(False, "Attention cache enabled"), - "pab_block_skip_range": OptionInfo(2, "Block skip range", gr.Slider, {"minimum": 1, "maximum": 4, "step": 1}), - "pab_timestep_skip_start": OptionInfo(0.1, "Timestep skip start", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.05}), - "pab_timestep_skip_end": OptionInfo(0.8, "Timestep skip end", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.05}), + "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}), + + "faster_cache__sep": OptionInfo("

Faster Cache

", "", gr.HTML), + "faster_cache_enabled": OptionInfo(False, "Faster 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}), + "fc_uncond_skip_range": OptionInfo(5, "FC uncond skip range", gr.Slider, {"minimum": 1, "maximum": 4, "step": 1}), + "fc_uncond_skip_start": OptionInfo(0, "FC uncond skip start", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01}), + "fc_uncond_skip_end": OptionInfo(781, "FC uncond skip end", gr.Slider, {"minimum": 0, "maximum": 1, "step": 1}), + "fc_attention_weight": OptionInfo(0.5, "FC spacial skip range", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.05}), + "fc_tensor_format": OptionInfo("BCFHW", "FC tensor format", gr.Radio, {"choices": ["BCFHW", "BFCHW", "BCHW"]}), + "fc_guidance_distilled": OptionInfo(False, "FC guidance distilled", gr.Checkbox), "para_sep": OptionInfo("

Para-attention

", "", gr.HTML), "para_cache_enabled": OptionInfo(False, "First-block cache enabled"), @@ -595,7 +612,7 @@ options_templates.update(options_section(('advanced', "Pipeline Modifiers"), { options_templates.update(options_section(('compile', "Model Compile"), { "cuda_compile_sep": OptionInfo("

Model Compile

", "", gr.HTML), - "cuda_compile": OptionInfo([] if not cmd_opts.use_openvino else ["Model", "VAE", "Upscaler"], "Compile Model", gr.CheckboxGroup, {"choices": ["Model", "VAE", "Text Encoder", "Upscaler"]}), + "cuda_compile": OptionInfo([] if not cmd_opts.use_openvino else ["Model", "VAE", "Upscaler"], "Compile Model", gr.CheckboxGroup, {"choices": ["Model", "VAE", "TE", "Upscaler"]}), "cuda_compile_backend": OptionInfo("none" if not cmd_opts.use_openvino else "openvino_fx", "Model compile backend", gr.Radio, {"choices": ['none', 'inductor', 'cudagraphs', 'aot_ts_nvfuser', 'hidet', 'migraphx', 'ipex', 'onediff', 'stable-fast', 'deep-cache', 'olive-ai', 'openvino_fx']}), "cuda_compile_mode": OptionInfo("default", "Model compile mode", gr.Radio, {"choices": ['default', 'reduce-overhead', 'max-autotune', 'max-autotune-no-cudagraphs']}), "cuda_compile_fullgraph": OptionInfo(True if not cmd_opts.use_openvino else False, "Model compile fullgraph"), @@ -606,7 +623,6 @@ options_templates.update(options_section(('compile', "Model Compile"), { })) options_templates.update(options_section(('system-paths', "System Paths"), { - "clean_temp_dir_at_start": OptionInfo(True, "Cleanup temporary folder on startup"), "models_paths_sep_options": OptionInfo("

Models Paths

", "", gr.HTML), "models_dir": OptionInfo('models', "Root model folder", folder=True), "model_paths_sep_options": OptionInfo("

Paths for specific models

", "", gr.HTML), @@ -618,7 +634,7 @@ options_templates.update(options_section(('system-paths', "System Paths"), { "unet_dir": OptionInfo(os.path.join(paths.models_path, 'UNET'), "Folder with UNET files", folder=True), "te_dir": OptionInfo(os.path.join(paths.models_path, 'Text-encoder'), "Folder with Text encoder files", folder=True), "lora_dir": OptionInfo(os.path.join(paths.models_path, 'Lora'), "Folder with LoRA network(s)", folder=True), - "styles_dir": OptionInfo(os.path.join(paths.data_path, 'styles.csv'), "File or Folder with user-defined styles", folder=True), + "styles_dir": OptionInfo(os.path.join(paths.models_path, 'styles'), "File or Folder with user-defined styles", folder=True), "wildcards_dir": OptionInfo(os.path.join(paths.models_path, 'wildcards'), "Folder with user-defined wildcards", folder=True), "embeddings_dir": OptionInfo(os.path.join(paths.models_path, 'embeddings'), "Folder with textual inversion embeddings", folder=True), "hypernetwork_dir": OptionInfo(os.path.join(paths.models_path, 'hypernetworks'), "Folder with Hypernetwork models", folder=True), @@ -634,6 +650,7 @@ options_templates.update(options_section(('system-paths', "System Paths"), { "ldsr_models_path": OptionInfo(os.path.join(paths.models_path, 'LDSR'), "Folder with LDSR models", folder=True), "clip_models_path": OptionInfo(os.path.join(paths.models_path, 'CLIP'), "Folder with CLIP models", folder=True), "other_paths_sep_options": OptionInfo("

Cache folders

", "", gr.HTML), + "clean_temp_dir_at_start": OptionInfo(True, "Cleanup temporary folder on startup"), "temp_dir": OptionInfo("", "Directory for temporary images; leave empty for default", folder=True), "accelerate_offload_path": OptionInfo('cache/accelerate', "Folder for disk offload", folder=True), "openvino_cache_path": OptionInfo('cache', "Folder for OpenVINO cache", folder=True), @@ -642,15 +659,15 @@ options_templates.update(options_section(('system-paths', "System Paths"), { })) options_templates.update(options_section(('saving-images', "Image Options"), { - "keep_incomplete": OptionInfo(True, "Keep incomplete images"), "samples_save": OptionInfo(True, "Save all generated images"), + "keep_incomplete": OptionInfo(False, "Keep incomplete images"), "samples_format": OptionInfo('jpg', 'File format', gr.Dropdown, {"choices": ["jpg", "png", "webp", "tiff", "jp2", "jxl"]}), "jpeg_quality": OptionInfo(90, "Image quality", gr.Slider, {"minimum": 1, "maximum": 100, "step": 1}), "img_max_size_mp": OptionInfo(1000, "Maximum image size (MP)", gr.Slider, {"minimum": 100, "maximum": 2000, "step": 1}), "webp_lossless": OptionInfo(False, "WebP lossless compression"), - "save_selected_only": OptionInfo(True, "Save only saves selected image"), + "save_selected_only": OptionInfo(True, "UI save only saves selected image"), "include_mask": OptionInfo(False, "Include mask in outputs"), - "samples_save_zip": OptionInfo(True, "Create ZIP archive"), + "samples_save_zip": OptionInfo(False, "Create ZIP archive for multiple images"), "image_background": OptionInfo("#000000", "Resize background color", gr.ColorPicker, {}), "image_sep_metadata": OptionInfo("

Metadata/Logging

", "", gr.HTML), @@ -774,7 +791,10 @@ options_templates.update(options_section(('sampler-params', "Sampler Settings"), 'schedulers_beta_end': OptionInfo(0, "Beta end", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.00001, "visible": native}), 'schedulers_timesteps_range': OptionInfo(1000, "Timesteps range", gr.Slider, {"minimum": 250, "maximum": 4000, "step": 1, "visible": native}), 'schedulers_shift': OptionInfo(3, "Sampler shift", gr.Slider, {"minimum": 0.1, "maximum": 10, "step": 0.1, "visible": False}), - 'schedulers_dynamic_shift': OptionInfo(True, "Sampler dynamic shift", gr.Checkbox, {"visible": False}), + 'schedulers_dynamic_shift': OptionInfo(False, "Sampler dynamic shift", gr.Checkbox, {"visible": False}), + 'schedulers_sigma_adjust': OptionInfo(1.0, "Sigma adjust", gr.Slider, {"minimum": 0.5, "maximum": 1.5, "step": 0.01, "visible": False}), + 'schedulers_sigma_adjust_min': OptionInfo(0.2, "Sigma adjust start", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01, "visible": False}), + 'schedulers_sigma_adjust_max': OptionInfo(0.8, "Sigma adjust end", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01, "visible": False}), # managed from ui.py for backend original k-diffusion "always_batch_cond_uncond": OptionInfo(False, "Disable conditional batching", gr.Checkbox, {"visible": not native}), @@ -855,6 +875,7 @@ options_templates.update(options_section(('interrogate', "Interrogate"), { "interrogate_vlm_sep": OptionInfo("

VLM

", "", gr.HTML), "interrogate_vlm_model": OptionInfo(list(vlm_models)[0], "VLM: default model", gr.Dropdown, {"choices": list(vlm_models)}), "interrogate_vlm_prompt": OptionInfo(vlm_prompts[2], "VLM: default prompt", DropdownEditable, {"choices": vlm_prompts }), + "interrogate_vlm_system": OptionInfo(vlm_system, "VLM: default prompt"), "interrogate_vlm_num_beams": OptionInfo(3, "VLM: num beams", gr.Slider, {"minimum": 1, "maximum": 16, "step": 1, "visible": False}), "interrogate_vlm_max_length": OptionInfo(512, "VLM: max length", gr.Slider, {"minimum": 1, "maximum": 4096, "step": 1, "visible": False}), "interrogate_vlm_do_sample": OptionInfo(False, "VLM: use sample method"), @@ -905,11 +926,12 @@ 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_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}), - "lora_in_memory_limit": OptionInfo(0, "LoRA memory cache", gr.Slider, {"minimum": 0, "maximum": 24, "step": 1}), + "lora_in_memory_limit": OptionInfo(1, "LoRA memory cache", gr.Slider, {"minimum": 0, "maximum": 32, "step": 1}), "lora_quant": OptionInfo("NF4","LoRA precision when quantized", gr.Radio, {"choices": ["NF4", "FP4"]}), "extra_networks_styles_sep": OptionInfo("

Styles

", "", gr.HTML), @@ -925,7 +947,7 @@ options_templates.update(options_section(('extra_networks', "Networks"), { "wildcards_enabled": OptionInfo(True, "Enable file wildcards support"), })) -options_templates.update(options_section((None, "Internal options"), { +options_templates.update(options_section((None, "Hidden options"), { "diffusers_version": OptionInfo("", "Diffusers version", gr.Textbox, {"visible": False}), "disabled_extensions": OptionInfo([], "Disable these extensions"), "sd_checkpoint_hash": OptionInfo("", "SHA256 hash of the current checkpoint"), @@ -1228,7 +1250,7 @@ total_tqdm = TotalTQDM() def restart_server(restart=True): if demo is None: return - log.warning('Server shutdown requested') + log.critical('Server shutdown requested') try: sys.tracebacklimit = 0 stdout = io.StringIO() diff --git a/modules/shared_items.py b/modules/shared_items.py index 5c1e3aebb..91299850e 100644 --- a/modules/shared_items.py +++ b/modules/shared_items.py @@ -5,7 +5,7 @@ def postprocessing_scripts(): def sd_vae_items(): import modules.sd_vae - return ["Automatic", "None"] + list(modules.sd_vae.vae_dict) + return ["Automatic", "Default"] + list(modules.sd_vae.vae_dict) def sd_taesd_items(): @@ -89,7 +89,8 @@ def get_pipelines(): 'SegMoE': getattr(diffusers, 'StableDiffusionPipeline', None), # dynamically redefined and loaded in sd_models.load_diffuser 'Kolors': getattr(diffusers, 'KolorsPipeline', None), 'AuraFlow': getattr(diffusers, 'AuraFlowPipeline', None), - 'CogView': getattr(diffusers, 'CogView3PlusPipeline', 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), diff --git a/modules/styles.py b/modules/styles.py index 41812222c..2de41c2a8 100644 --- a/modules/styles.py +++ b/modules/styles.py @@ -154,6 +154,7 @@ def apply_styles_to_extra(p, style: Style): p.original_negative = p.negative_prompt style_extra = apply_wildcards_to_prompt(style.extra, [style.wildcards], silent=True) + style_extra = ' ' + style_extra.lower() extra.update(infotext.parse(style_extra)) extra.pop('Prompt', None) extra.pop('Negative prompt', None) @@ -162,14 +163,15 @@ def apply_styles_to_extra(p, style: Style): for k, v in extra.items(): k = k.lower() k = k.replace(' ', '_') - if k in name_exclude: # exclude some fields - continue if k in name_map: # rename some fields k = name_map[k] + if k in name_exclude: # exclude some fields + continue if hasattr(p, k): orig = getattr(p, k) - if type(orig) != type(v) and orig is not None: - v = type(orig)(v) + if (type(orig) != type(v)) and (orig is not None): + if not (type(orig) == int and type(v) == float): # dont convert float to int + v = type(orig)(v) setattr(p, k, v) fields.append(f'{k}={v}') else: @@ -260,6 +262,7 @@ class StyleDatabase: for future in concurrent.futures.as_completed(future_items): future.result() + self.built_in = shared.opts.extra_networks_styles list_folder(self.path) t1 = time.time() shared.log.info(f'Available Styles: path="{self.path}" items={len(self.styles.keys())} time={t1-t0:.2f}') diff --git a/modules/taesd/taehv.py b/modules/taesd/taehv.py new file mode 100644 index 000000000..4a424f137 --- /dev/null +++ b/modules/taesd/taehv.py @@ -0,0 +1,284 @@ +#!/usr/bin/env python3 +""" +Tiny AutoEncoder for Hunyuan Video +(DNN for encoding / decoding videos to Hunyuan Video's latent space) +""" +from collections import namedtuple +import torch +import torch.nn as nn +import torch.nn.functional as F +from tqdm.auto import tqdm + +DecoderResult = namedtuple("DecoderResult", ("frame", "memory")) +TWorkItem = namedtuple("TWorkItem", ("input_tensor", "block_index")) + +def conv(n_in, n_out, **kwargs): + return nn.Conv2d(n_in, n_out, 3, padding=1, **kwargs) + +class Clamp(nn.Module): + def forward(self, x): + return torch.tanh(x / 3) * 3 + +class MemBlock(nn.Module): + def __init__(self, n_in, n_out): + super().__init__() + self.conv = nn.Sequential(conv(n_in * 2, n_out), nn.ReLU(inplace=True), conv(n_out, n_out), nn.ReLU(inplace=True), conv(n_out, n_out)) + self.skip = nn.Conv2d(n_in, n_out, 1, bias=False) if n_in != n_out else nn.Identity() + self.act = nn.ReLU(inplace=True) + def forward(self, x, past): + return self.act(self.conv(torch.cat([x, past], 1)) + self.skip(x)) + +class TPool(nn.Module): + def __init__(self, n_f, stride): + super().__init__() + self.stride = stride + self.conv = nn.Conv2d(n_f*stride,n_f, 1, bias=False) + def forward(self, x): + _NT, C, H, W = x.shape + return self.conv(x.reshape(-1, self.stride * C, H, W)) + +class TGrow(nn.Module): + def __init__(self, n_f, stride): + super().__init__() + self.stride = stride + self.conv = nn.Conv2d(n_f, n_f*stride, 1, bias=False) + def forward(self, x): + _NT, C, H, W = x.shape + x = self.conv(x) + return x.reshape(-1, C, H, W) + +def apply_model_with_memblocks(model, x, parallel, show_progress_bar): + """ + Apply a sequential model with memblocks to the given input. + Args: + - model: nn.Sequential of blocks to apply + - x: input data, of dimensions NTCHW + - parallel: if True, parallelize over timesteps (fast but uses O(T) memory) + if False, each timestep will be processed sequentially (slow but uses O(1) memory) + - show_progress_bar: if True, enables tqdm progressbar display + + Returns NTCHW tensor of output data. + """ + assert x.ndim == 5, f"TAEHV operates on NTCHW tensors, but got {x.ndim}-dim tensor" + N, T, C, H, W = x.shape + if parallel: + x = x.reshape(N*T, C, H, W) + # parallel over input timesteps, iterate over blocks + for b in tqdm(model, disable=not show_progress_bar): + if isinstance(b, MemBlock): + NT, C, H, W = x.shape + T = NT // N + _x = x.reshape(N, T, C, H, W) + mem = F.pad(_x, (0,0,0,0,0,0,1,0), value=0)[:,:T].reshape(x.shape) + x = b(x, mem) + else: + x = b(x) + NT, C, H, W = x.shape + T = NT // N + x = x.view(N, T, C, H, W) + else: + # TODO(oboerbohan): at least on macos this still gradually uses more memory during decode... + # need to fix :( + out = [] + # iterate over input timesteps and also iterate over blocks. + # because of the cursed TPool/TGrow blocks, this is not a nested loop, + # it's actually a ***graph traversal*** problem! so let's make a queue + work_queue = [TWorkItem(xt, 0) for t, xt in enumerate(x.reshape(N, T * C, H, W).chunk(T, dim=1))] + # in addition to manually managing our queue, we also need to manually manage our progressbar. + # we'll update it for every source node that we consume. + progress_bar = tqdm(range(T), disable=not show_progress_bar) + # we'll also need a separate addressable memory per node as well + mem = [None] * len(model) + while work_queue: + xt, i = work_queue.pop(0) + if i == 0: + # new source node consumed + progress_bar.update(1) + if i == len(model): + # reached end of the graph, append result to output list + out.append(xt) + else: + # fetch the block to process + b = model[i] + if isinstance(b, MemBlock): + # mem blocks are simple since we're visiting the graph in causal order + if mem[i] is None: + xt_new = b(xt, xt * 0) + mem[i] = xt + else: + xt_new = b(xt, mem[i]) + mem[i].copy_(xt) # inplace might reduce mysterious pytorch memory allocations? doesn't help though + # add successor to work queue + work_queue.insert(0, TWorkItem(xt_new, i+1)) + elif isinstance(b, TPool): + # pool blocks are miserable + if mem[i] is None: + mem[i] = [] # pool memory is itself a queue of inputs to pool + mem[i].append(xt) + if len(mem[i]) > b.stride: + # pool mem is in invalid state, we should have pooled before this + raise ValueError("???") + elif len(mem[i]) < b.stride: + # pool mem is not yet full, go back to processing the work queue + pass + else: + # pool mem is ready, run the pool block + N, C, H, W = xt.shape + xt = b(torch.cat(mem[i], 1).view(N*b.stride, C, H, W)) + # reset the pool mem + mem[i] = [] + # add successor to work queue + work_queue.insert(0, TWorkItem(xt, i+1)) + elif isinstance(b, TGrow): + xt = b(xt) + NT, C, H, W = xt.shape + # each tgrow has multiple successor nodes + for xt_next in reversed(xt.view(N, b.stride*C, H, W).chunk(b.stride, 1)): + # add successor to work queue + work_queue.insert(0, TWorkItem(xt_next, i+1)) + else: + # normal block with no funny business + xt = b(xt) + # add successor to work queue + work_queue.insert(0, TWorkItem(xt, i+1)) + progress_bar.close() + x = torch.stack(out, 1) + return x + +class TAEHV(nn.Module): + latent_channels = 16 + image_channels = 3 + def __init__(self, checkpoint_path="taehv.pth", decoder_time_upscale=(True, True), decoder_space_upscale=(True, True, True)): + """Initialize pretrained TAEHV from the given checkpoint. + + Arg: + checkpoint_path: path to weight file to load. taehv.pth for Hunyuan, taew2_1.pth for Wan 2.1. + decoder_time_upscale: whether temporal upsampling is enabled for each block. upsampling can be disabled for a cheaper preview. + decoder_space_upscale: whether spatial upsampling is enabled for each block. upsampling can be disabled for a cheaper preview. + """ + super().__init__() + self.encoder = nn.Sequential( + conv(TAEHV.image_channels, 64), nn.ReLU(inplace=True), + TPool(64, 2), conv(64, 64, stride=2, bias=False), MemBlock(64, 64), MemBlock(64, 64), MemBlock(64, 64), + TPool(64, 2), conv(64, 64, stride=2, bias=False), MemBlock(64, 64), MemBlock(64, 64), MemBlock(64, 64), + TPool(64, 1), conv(64, 64, stride=2, bias=False), MemBlock(64, 64), MemBlock(64, 64), MemBlock(64, 64), + conv(64, TAEHV.latent_channels), + ) + n_f = [256, 128, 64, 64] + self.frames_to_trim = 2**sum(decoder_time_upscale) - 1 + self.decoder = nn.Sequential( + Clamp(), conv(TAEHV.latent_channels, n_f[0]), nn.ReLU(inplace=True), + MemBlock(n_f[0], n_f[0]), MemBlock(n_f[0], n_f[0]), MemBlock(n_f[0], n_f[0]), nn.Upsample(scale_factor=2 if decoder_space_upscale[0] else 1), TGrow(n_f[0], 1), conv(n_f[0], n_f[1], bias=False), + MemBlock(n_f[1], n_f[1]), MemBlock(n_f[1], n_f[1]), MemBlock(n_f[1], n_f[1]), nn.Upsample(scale_factor=2 if decoder_space_upscale[1] else 1), TGrow(n_f[1], 2 if decoder_time_upscale[0] else 1), conv(n_f[1], n_f[2], bias=False), + MemBlock(n_f[2], n_f[2]), MemBlock(n_f[2], n_f[2]), MemBlock(n_f[2], n_f[2]), nn.Upsample(scale_factor=2 if decoder_space_upscale[2] else 1), TGrow(n_f[2], 2 if decoder_time_upscale[1] else 1), conv(n_f[2], n_f[3], bias=False), + nn.ReLU(inplace=True), conv(n_f[3], TAEHV.image_channels), + ) + if checkpoint_path is not None: + self.load_state_dict(self.patch_tgrow_layers(torch.load(checkpoint_path, map_location="cpu", weights_only=True))) + + def patch_tgrow_layers(self, sd): + """Patch TGrow layers to use a smaller kernel if needed. + + Args: + sd: state dict to patch + """ + new_sd = self.state_dict() + for i, layer in enumerate(self.decoder): + if isinstance(layer, TGrow): + key = f"decoder.{i}.conv.weight" + if sd[key].shape[0] > new_sd[key].shape[0]: + # take the last-timestep output channels + sd[key] = sd[key][-new_sd[key].shape[0]:] + return sd + + def encode_video(self, x, parallel=True, show_progress_bar=True): + """Encode a sequence of frames. + + Args: + x: input NTCHW RGB (C=3) tensor with values in [0, 1]. + parallel: if True, all frames will be processed at once. + (this is faster but may require more memory). + if False, frames will be processed sequentially. + Returns NTCHW latent tensor with ~Gaussian values. + """ + return apply_model_with_memblocks(self.encoder, x, parallel, show_progress_bar) + + def decode_video(self, x, parallel=True, show_progress_bar=True): + """Decode a sequence of frames. + + Args: + x: input NTCHW latent (C=12) tensor with ~Gaussian values. + parallel: if True, all frames will be processed at once. + (this is faster but may require more memory). + if False, frames will be processed sequentially. + Returns NTCHW RGB tensor with ~[0, 1] values. + """ + x = apply_model_with_memblocks(self.decoder, x, parallel, show_progress_bar) + return x[:, self.frames_to_trim:] + + def forward(self, x): + return self.c(x) + +@torch.no_grad() +def main(): + """Run TAEHV roundtrip reconstruction on the given video paths.""" + import sys + import cv2 # no highly esteemed deed is commemorated here + + class VideoTensorReader: + def __init__(self, video_file_path): + self.cap = cv2.VideoCapture(video_file_path) + assert self.cap.isOpened(), f"Could not load {video_file_path}" + self.fps = self.cap.get(cv2.CAP_PROP_FPS) + def __iter__(self): + return self + def __next__(self): + ret, frame = self.cap.read() + if not ret: + self.cap.release() + raise StopIteration # End of video or error + return torch.from_numpy(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)).permute(2, 0, 1) # BGR HWC -> RGB CHW + + class VideoTensorWriter: + def __init__(self, video_file_path, width_height, fps=30): + self.writer = cv2.VideoWriter(video_file_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, width_height) + assert self.writer.isOpened(), f"Could not create writer for {video_file_path}" + def write(self, frame_tensor): + assert frame_tensor.ndim == 3 and frame_tensor.shape[0] == 3, f"{frame_tensor.shape}??" + self.writer.write(cv2.cvtColor(frame_tensor.permute(1, 2, 0).numpy(), cv2.COLOR_RGB2BGR)) # RGB CHW -> BGR HWC + def __del__(self): + if hasattr(self, 'writer'): + self.writer.release() + + dev = torch.device("cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu") + dtype = torch.float16 + print("Using device", dev, "and dtype", dtype) + taehv = TAEHV().to(dev, dtype) + for video_path in sys.argv[1:]: + print(f"Processing {video_path}...") + video_in = VideoTensorReader(video_path) + video = torch.stack(list(video_in), 0)[None] + vid_dev = video.to(dev, dtype).div_(255.0) + # convert to device tensor + if video.numel() < 100_000_000: + print(f" {video_path} seems small enough, will process all frames in parallel") + # convert to device tensor + vid_enc = taehv.encode_video(vid_dev) + print(f" Encoded {video_path} -> {vid_enc.shape}. Decoding...") + vid_dec = taehv.decode_video(vid_enc) + print(f" Decoded {video_path} -> {vid_dec.shape}") + else: + print(f" {video_path} seems large, will process each frame sequentially") + # convert to device tensor + vid_enc = taehv.encode_video(vid_dev, parallel=False) + print(f" Encoded {video_path} -> {vid_enc.shape}. Decoding...") + vid_dec = taehv.decode_video(vid_enc, parallel=False) + print(f" Decoded {video_path} -> {vid_dec.shape}") + video_out_path = video_path + ".reconstructed_by_taehv.mp4" + video_out = VideoTensorWriter(video_out_path, (vid_dec.shape[-1], vid_dec.shape[-2]), fps=int(round(video_in.fps))) + for frame in vid_dec.clamp_(0, 1).mul_(255).round_().byte().cpu()[0]: + video_out.write(frame) + print(f" Saved to {video_out_path}") + +if __name__ == "__main__": + main() diff --git a/modules/taesd/taem1.py b/modules/taesd/taem1.py new file mode 100644 index 000000000..7d59ca2b6 --- /dev/null +++ b/modules/taesd/taem1.py @@ -0,0 +1,272 @@ +#!/usr/bin/env python3 +""" +Tiny AutoEncoder for Mochi 1 +(DNN for encoding / decoding videos to Mochi 1's latent space) +""" +from collections import namedtuple +import torch +import torch.nn as nn +import torch.nn.functional as F +from tqdm.auto import tqdm + +DecoderResult = namedtuple("DecoderResult", ("frame", "memory")) +TWorkItem = namedtuple("TWorkItem", ("input_tensor", "block_index")) + +def conv(n_in, n_out, **kwargs): + return nn.Conv2d(n_in, n_out, 3, padding=1, **kwargs) + +class Clamp(nn.Module): + def forward(self, x): + return torch.tanh(x / 3) * 3 + +class MemBlock(nn.Module): + def __init__(self, n_in, n_out): + super().__init__() + self.conv = nn.Sequential(conv(n_in * 2, n_out), nn.ReLU(inplace=True), conv(n_out, n_out), nn.ReLU(inplace=True), conv(n_out, n_out)) + self.skip = nn.Conv2d(n_in, n_out, 1, bias=False) if n_in != n_out else nn.Identity() + self.act = nn.ReLU(inplace=True) + def forward(self, x, past): + return self.act(self.conv(torch.cat([x, past], 1)) + self.skip(x)) + +class TPool(nn.Module): + def __init__(self, n_f, stride): + super().__init__() + self.stride = stride + self.conv = nn.Conv2d(n_f*stride,n_f, 1, bias=False) + def forward(self, x): + _NT, C, H, W = x.shape + return self.conv(x.reshape(-1, self.stride * C, H, W)) + +class TGrow(nn.Module): + def __init__(self, n_f, stride): + super().__init__() + self.stride = stride + self.conv = nn.Conv2d(n_f, n_f*stride, 1, bias=False) + def forward(self, x): + _NT, C, H, W = x.shape + x = self.conv(x) + return x.reshape(-1, C, H, W) + +def apply_model_with_memblocks(model, x, parallel, show_progress_bar): + """ + Apply a sequential model with memblocks to the given input. + Args: + - model: nn.Sequential of blocks to apply + - x: input data, of dimensions NTCHW + - parallel: if True, parallelize over timesteps (fast but uses O(T) memory) + if False, each timestep will be processed sequentially (slow but uses O(1) memory) + - show_progress_bar: if True, enables tqdm progressbar display + + Returns NTCHW tensor of output data. + """ + assert x.ndim == 5, f"TAEM1 operates on NTCHW tensors, but got {x.ndim}-dim tensor" + N, T, C, H, W = x.shape + if parallel: + x = x.reshape(N*T, C, H, W) + # parallel over input timesteps, iterate over blocks + for b in tqdm(model, disable=not show_progress_bar): + if isinstance(b, MemBlock): + NT, C, H, W = x.shape + T = NT // N + _x = x.reshape(N, T, C, H, W) + mem = F.pad(_x, (0,0,0,0,0,0,1,0), value=0)[:,:T].reshape(x.shape) + x = b(x, mem) + else: + x = b(x) + NT, C, H, W = x.shape + T = NT // N + x = x.view(N, T, C, H, W) + else: + # TODO(oboerbohan): at least on macos this still gradually uses more memory during decode... + # need to fix :( + out = [] + # iterate over input timesteps and also iterate over blocks. + # because of the cursed TPool/TGrow blocks, this is not a nested loop, + # it's actually a ***graph traversal*** problem! so let's make a queue + work_queue = [TWorkItem(xt, 0) for t, xt in enumerate(x.reshape(N, T * C, H, W).chunk(T, dim=1))] + # in addition to manually managing our queue, we also need to manually manage our progressbar. + # we'll update it for every source node that we consume. + progress_bar = tqdm(range(T), disable=not show_progress_bar) + # we'll also need a separate addressable memory per node as well + mem = [None] * len(model) + while work_queue: + xt, i = work_queue.pop(0) + if i == 0: + # new source node consumed + progress_bar.update(1) + if i == len(model): + # reached end of the graph, append result to output list + out.append(xt) + else: + # fetch the block to process + b = model[i] + if isinstance(b, MemBlock): + # mem blocks are simple since we're visiting the graph in causal order + if mem[i] is None: + xt_new = b(xt, xt * 0) + mem[i] = xt + else: + xt_new = b(xt, mem[i]) + mem[i].copy_(xt) # inplace might reduce mysterious pytorch memory allocations? doesn't help though + # add successor to work queue + work_queue.insert(0, TWorkItem(xt_new, i+1)) + elif isinstance(b, TPool): + # pool blocks are miserable + if mem[i] is None: + mem[i] = [] # pool memory is itself a queue of inputs to pool + mem[i].append(xt) + if len(mem[i]) > b.stride: + # pool mem is in invalid state, we should have pooled before this + raise ValueError("???") + elif len(mem[i]) < b.stride: + # pool mem is not yet full, go back to processing the work queue + pass + else: + # pool mem is ready, run the pool block + N, C, H, W = xt.shape + xt = b(torch.cat(mem[i], 1).view(N*b.stride, C, H, W)) + # reset the pool mem + mem[i] = [] + # add successor to work queue + work_queue.insert(0, TWorkItem(xt, i+1)) + elif isinstance(b, TGrow): + xt = b(xt) + NT, C, H, W = xt.shape + # each tgrow has multiple successor nodes + for xt_next in reversed(xt.view(N, b.stride*C, H, W).chunk(b.stride, 1)): + # add successor to work queue + work_queue.insert(0, TWorkItem(xt_next, i+1)) + else: + # normal block with no funny business + xt = b(xt) + # add successor to work queue + work_queue.insert(0, TWorkItem(xt, i+1)) + progress_bar.close() + x = torch.stack(out, 1) + return x + +class TAEM1(nn.Module): + latent_channels = 12 + image_channels = 3 + def __init__(self, checkpoint_path="taem1.pth"): + """Initialize pretrained TAEM1 from the given checkpoints.""" + super().__init__() + self.encoder = nn.Sequential( + conv(TAEM1.image_channels, 64), nn.ReLU(inplace=True), + TPool(64, 3), conv(64, 64, stride=2, bias=False), MemBlock(64, 64), MemBlock(64, 64), MemBlock(64, 64), + TPool(64, 2), conv(64, 64, stride=2, bias=False), MemBlock(64, 64), MemBlock(64, 64), MemBlock(64, 64), + TPool(64, 1), conv(64, 64, stride=2, bias=False), MemBlock(64, 64), MemBlock(64, 64), MemBlock(64, 64), + conv(64, TAEM1.latent_channels), + ) + n_f = [256, 128, 64, 64] + self.decoder = nn.Sequential( + Clamp(), conv(TAEM1.latent_channels, n_f[0]), nn.ReLU(inplace=True), + MemBlock(n_f[0], n_f[0]), MemBlock(n_f[0], n_f[0]), MemBlock(n_f[0], n_f[0]), nn.Upsample(scale_factor=2), TGrow(n_f[0], 1), conv(n_f[0], n_f[1], bias=False), + MemBlock(n_f[1], n_f[1]), MemBlock(n_f[1], n_f[1]), MemBlock(n_f[1], n_f[1]), nn.Upsample(scale_factor=2), TGrow(n_f[1], 2), conv(n_f[1], n_f[2], bias=False), + MemBlock(n_f[2], n_f[2]), MemBlock(n_f[2], n_f[2]), MemBlock(n_f[2], n_f[2]), nn.Upsample(scale_factor=2), TGrow(n_f[2], 3), conv(n_f[2], n_f[3], bias=False), + nn.ReLU(inplace=True), conv(n_f[3], TAEM1.image_channels), + ) + if checkpoint_path is not None: + self.load_state_dict(torch.load(checkpoint_path, map_location="cpu", weights_only=True)) + + def encode_video(self, x, parallel=True, show_progress_bar=True): + """Encode a sequence of frames. + + Args: + x: input NTCHW RGB (C=3) tensor with values in [0, 1]. + parallel: if True, all frames will be processed at once. + (this is faster but may require more memory). + if False, frames will be processed sequentially. + Returns NTCHW latent tensor with ~Gaussian values. + """ + return apply_model_with_memblocks(self.encoder, x, parallel, show_progress_bar) + + def decode_video(self, x, parallel=True, show_progress_bar=True): + """Decode a sequence of frames. + + Args: + x: input NTCHW latent (C=12) tensor with ~Gaussian values. + parallel: if True, all frames will be processed at once. + (this is faster but may require more memory). + if False, frames will be processed sequentially. + Returns NTCHW RGB tensor with ~[0, 1] values. + """ + x = apply_model_with_memblocks(self.decoder, x, parallel, show_progress_bar) + # NOTE: + # the Mochi VAE does not preserve shape along the time axis; + # videos are encoded to floor((n_in - 1)/6)+1 latent frames + # (which makes sense, it's stride 6, so 12 -> 2 and 13->3) + # but then they're decoded to only the *minimal* number + # of input frames (3 latents get decoded to 13 frames, not 18) + # in order to achieve the intended causal structure... + # anyway, that's why we have to remove some frames here. + # mochi-VAE does the slicing at each TGrow (save compute/mem?) + # but I think it's basically the same + return x[:, 5:] + + def forward(self, x): + return self.c(x) + +@torch.no_grad() +def main(): + """Run TAEM1 roundtrip reconstruction on the given video paths.""" + import sys + import cv2 # no highly esteemed deed is commemorated here + + class VideoTensorReader: + def __init__(self, video_file_path): + self.cap = cv2.VideoCapture(video_file_path) + assert self.cap.isOpened(), f"Could not load {video_file_path}" + self.fps = self.cap.get(cv2.CAP_PROP_FPS) + def __iter__(self): + return self + def __next__(self): + ret, frame = self.cap.read() + if not ret: + self.cap.release() + raise StopIteration # End of video or error + return torch.from_numpy(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)).permute(2, 0, 1) # BGR HWC -> RGB CHW + + class VideoTensorWriter: + def __init__(self, video_file_path, width_height, fps=30): + self.writer = cv2.VideoWriter(video_file_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, width_height) + assert self.writer.isOpened(), f"Could not create writer for {video_file_path}" + def write(self, frame_tensor): + assert frame_tensor.ndim == 3 and frame_tensor.shape[0] == 3, f"{frame_tensor.shape}??" + self.writer.write(cv2.cvtColor(frame_tensor.permute(1, 2, 0).numpy(), cv2.COLOR_RGB2BGR)) # RGB CHW -> BGR HWC + def __del__(self): + if hasattr(self, 'writer'): + self.writer.release() + + dev = torch.device("cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu") + dtype = torch.float16 + print("Using device", dev, "and dtype", dtype) + taem1 = TAEM1().to(dev, dtype) + for video_path in sys.argv[1:]: + print(f"Processing {video_path}...") + video_in = VideoTensorReader(video_path) + video = torch.stack(list(video_in), 0)[None] + vid_dev = video.to(dev, dtype).div_(255.0) + # convert to device tensor + if video.numel() < 100_000_000: + print(f" {video_path} seems small enough, will process all frames in parallel") + # convert to device tensor + vid_enc = taem1.encode_video(vid_dev) + print(f" Encoded {video_path} -> {vid_enc.shape}. Decoding...") + vid_dec = taem1.decode_video(vid_enc) + print(f" Decoded {video_path} -> {vid_dec.shape}") + else: + print(f" {video_path} seems large, will process each frame sequentially") + # convert to device tensor + vid_enc = taem1.encode_video(vid_dev, parallel=False) + print(f" Encoded {video_path} -> {vid_enc.shape}. Decoding...") + vid_dec = taem1.decode_video(vid_enc, parallel=False) + print(f" Decoded {video_path} -> {vid_dec.shape}") + video_out_path = video_path + ".reconstructed_by_taem1.mp4" + video_out = VideoTensorWriter(video_out_path, (vid_dec.shape[-1], vid_dec.shape[-2]), fps=int(round(video_in.fps))) + for frame in vid_dec.clamp_(0, 1).mul_(255).round_().byte().cpu()[0]: + video_out.write(frame) + print(f" Saved to {video_out_path}") + +if __name__ == "__main__": + main() diff --git a/modules/textual_inversion/textual_inversion.py b/modules/textual_inversion/textual_inversion.py index 86c7cd260..a3f16ab9c 100644 --- a/modules/textual_inversion/textual_inversion.py +++ b/modules/textual_inversion/textual_inversion.py @@ -421,7 +421,7 @@ class EmbeddingDatabase: if self.previously_displayed_embeddings != displayed_embeddings and shared.opts.diffusers_enable_embed: self.previously_displayed_embeddings = displayed_embeddings t1 = time.time() - shared.log.info(f"Load network: type=embeddings loaded={len(self.word_embeddings)} skipped={len(self.skipped_embeddings)} time={t1-t0:.2f}") + shared.log.info(f"Network load: type=embeddings loaded={len(self.word_embeddings)} skipped={len(self.skipped_embeddings)} time={t1-t0:.2f}") def find_embedding_at_position(self, tokens, offset): diff --git a/modules/transformer_cache.py b/modules/transformer_cache.py new file mode 100644 index 000000000..c8615c534 --- /dev/null +++ b/modules/transformer_cache.py @@ -0,0 +1,51 @@ +import os +import diffusers +from modules import shared, errors + + +debug = shared.log.trace if os.environ.get('SD_PROCESS_DEBUG', None) is not None else lambda *args, **kwargs: None + + +def set_cache(faster_cache=None, pyramid_attention_broadcast=None): + if not shared.sd_loaded or not hasattr(shared.sd_model, 'transformer'): + return + faster_cache = faster_cache if faster_cache is not None else shared.opts.faster_cache_enabled + pyramid_attention_broadcast = pyramid_attention_broadcast if pyramid_attention_broadcast is not None else shared.opts.pab_enabled + if (not faster_cache) and (not pyramid_attention_broadcast): + return + if (not hasattr(shared.sd_model.transformer, 'enable_cache')) or (not hasattr(shared.sd_model.transformer, 'disable_cache')): + shared.log.debug(f'Transformer cache: cls={shared.sd_model.transformer.__class__.__name__} fc={faster_cache} pab={pyramid_attention_broadcast} not supported') + return + try: + if faster_cache: # https://github.com/huggingface/diffusers/pull/10163 + distilled = shared.opts.fc_guidance_distilled or shared.sd_model_type == 'f1' + config = diffusers.FasterCacheConfig( + spatial_attention_block_skip_range=shared.opts.fc_spacial_skip_range, + spatial_attention_timestep_skip_range=(int(shared.opts.fc_spacial_skip_start), int(shared.opts.fc_spacial_skip_end)), + unconditional_batch_skip_range=shared.opts.fc_uncond_skip_range, + unconditional_batch_timestep_skip_range=(int(shared.opts.fc_uncond_skip_start), int(shared.opts.fc_uncond_skip_end)), + attention_weight_callback=lambda _: shared.opts.fc_attention_weight, + tensor_format=shared.opts.fc_tensor_format, # TODO fc: autodetect tensor format based on model + is_guidance_distilled=distilled, # TODO fc: autodetect distilled based on model + current_timestep_callback=lambda: shared.sd_model.current_timestep, + ) + shared.sd_model.transformer.disable_cache() + shared.sd_model.transformer.enable_cache(config) + shared.log.debug(f'Transformer cache: type={config.__class__.__name__}') + debug(f'Transformer cache: {vars(config)}') + elif pyramid_attention_broadcast: # https://github.com/huggingface/diffusers/pull/9562 + config = diffusers.PyramidAttentionBroadcastConfig( + spatial_attention_block_skip_range=shared.opts.pab_spacial_skip_range, + spatial_attention_timestep_skip_range=(int(shared.opts.pab_spacial_skip_start), int(shared.opts.pab_spacial_skip_end)), + current_timestep_callback=lambda: shared.sd_model.current_timestep, + ) + shared.sd_model.transformer.disable_cache() + shared.sd_model.transformer.enable_cache(config) + shared.log.debug(f'Transformer cache: type={config.__class__.__name__}') + debug(f'Transformer cache: {vars(config)}') + else: + debug('Transformer cache: not enabled') + shared.sd_model.transformer.disable_cache() + except Exception as e: + shared.log.error(f'Transformer cache: {e}') + errors.display(e, 'Transformer cache') diff --git a/modules/txt2img.py b/modules/txt2img.py index 08184fe43..151110d07 100644 --- a/modules/txt2img.py +++ b/modules/txt2img.py @@ -1,7 +1,7 @@ import os from modules import shared, processing, scripts from modules.generation_parameters_copypaste import create_override_settings_dict -from modules.ui import plaintext_to_html +from modules.ui_common import plaintext_to_html debug = shared.log.trace if os.environ.get('SD_PROCESS_DEBUG', None) is not None else lambda *args, **kwargs: None diff --git a/modules/ui.py b/modules/ui.py index 6d52187c8..1c4b7ec0f 100644 --- a/modules/ui.py +++ b/modules/ui.py @@ -1,28 +1,17 @@ -import os import mimetypes import gradio as gr import gradio.routes import gradio.utils -from modules.call_queue import wrap_gradio_call, wrap_gradio_gpu_call # pylint: disable=unused-import -from modules import timer, gr_hijack, shared, theme, sd_models, script_callbacks, modelloader, ui_common, ui_loadsave, ui_symbols, ui_javascript, ui_sections, generation_parameters_copypaste, call_queue +from modules import errors, timer, gr_hijack, shared, script_callbacks, ui_common, ui_symbols, ui_javascript, ui_sections, generation_parameters_copypaste, call_queue, scripts from modules.paths import script_path, data_path # pylint: disable=unused-import -from modules.dml import directml_override_opts -from modules.onnx_impl import install_olive -import modules.scripts -import modules.errors -modules.errors.install() +errors.install() mimetypes.init() mimetypes.add_type('application/javascript', '.js') mimetypes.add_type('image/webp', '.webp') mimetypes.add_type('image/jxl', '.jxl') -log = shared.log -opts = shared.opts -cmd_opts = shared.cmd_opts -ui_system_tabs = None -paste_function = None -wrap_queued_call = call_queue.wrap_queued_call +gr_hijack.init() switch_values_symbol = ui_symbols.switch detect_image_size_symbol = ui_symbols.detect paste_symbol = ui_symbols.paste @@ -32,11 +21,16 @@ folder_symbol = ui_symbols.folder extra_networks_symbol = ui_symbols.networks apply_style_symbol = ui_symbols.apply save_style_symbol = ui_symbols.save -gr_hijack.init() +wrap_queued_call = call_queue.wrap_queued_call # compatibility item +wrap_gradio_call = call_queue.wrap_gradio_call # compatibility item +wrap_gradio_gpu_call = call_queue.wrap_gradio_gpu_call # compatibility item +plaintext_to_html = ui_common.plaintext_to_html # compatibility item +infotext_to_html = ui_common.infotext_to_html # compatibility item create_sampler_and_steps_selection = ui_sections.create_sampler_and_steps_selection # compatibility item +ui_system_tabs = None # required for system-info -if not cmd_opts.share and not cmd_opts.listen: +if not shared.cmd_opts.share and not shared.cmd_opts.listen: # fix gradio phoning home gradio.utils.version_check = lambda: None gradio.utils.get_local_ip_address = lambda: '127.0.0.1' @@ -55,14 +49,6 @@ def create_output_panel(tabname, outdir): # pylint: disable=unused-argument # ou return a, b, c, e -def plaintext_to_html(text): # may be referenced by extensions - return ui_common.plaintext_to_html(text) - - -def infotext_to_html(text): # may be referenced by extensions - return ui_common.infotext_to_html(text) - - def send_gradio_gallery_to_image(x): if len(x) == 0: return None @@ -81,43 +67,6 @@ def setup_progressbar(*args, **kwargs): # pylint: disable=unused-argument pass -def apply_setting(key, value): - if value is None: - return gr.update() - if shared.cmd_opts.freeze: - return gr.update() - if key == 'sd_backend': - return gr.update() - if shared.opts.disable_weights_auto_swap and key in ['sd_model_checkpoint', 'sd_model_refiner', 'sd_model_dict', 'sd_vae', 'sd_unet', 'sd_text_encoder']: - return gr.update() - if key == "sd_model_checkpoint": - ckpt_info = sd_models.get_closet_checkpoint_match(value) - if ckpt_info is not None: - value = ckpt_info.title - else: - return gr.update() - comp_args = opts.data_labels[key].component_args - if comp_args and isinstance(comp_args, dict) and comp_args.get('visible') is False: - return gr.update() - valtype = type(opts.data_labels[key].default) - oldval = opts.data.get(key, None) - opts.data[key] = valtype(value) if valtype != type(None) else value - if oldval != value and opts.data_labels[key].onchange is not None: - opts.data_labels[key].onchange() - opts.save(shared.config_filename) - return getattr(opts, key) - - -def get_value_for_setting(key): - value = getattr(opts, key) - info = opts.data_labels[key] - args = info.component_args() if callable(info.component_args) else info.component_args or {} - args = {k: v for k, v in args.items() if k not in {'precision', 'multiselect', 'visible'}} - # if not args: - # return gr.update() - return gr.update(value=value, **args) - - def ordered_ui_categories(): return ['dimensions', 'sampler', 'seed', 'denoising', 'cfg', 'checkboxes', 'accordions', 'override_settings', 'scripts'] # a1111 compatibility item, not implemented @@ -127,6 +76,7 @@ def create_ui(startup_timer = None): timer.startup = timer.Timer() ui_javascript.reload_javascript() generation_parameters_copypaste.reset() + scripts.scripts_current = None with gr.Blocks(analytics_enabled=False) as txt2img_interface: from modules import ui_txt2img @@ -138,14 +88,18 @@ def create_ui(startup_timer = None): ui_img2img.create_ui() timer.startup.record("ui-img2img") - modules.scripts.scripts_current = None - with gr.Blocks(analytics_enabled=False) as control_interface: if shared.native: from modules import ui_control ui_control.create_ui() timer.startup.record("ui-control") + with gr.Blocks(analytics_enabled=False) as video_interface: + if shared.native: + from modules import ui_video + ui_video.create_ui() + timer.startup.record("ui-video") + with gr.Blocks(analytics_enabled=False) as extras_interface: from modules import ui_postprocessing ui_postprocessing.create_ui() @@ -166,224 +120,13 @@ def create_ui(startup_timer = None): ui_gallery.create_ui() timer.startup.record("ui-gallery") - def create_setting_component(key, is_quicksettings=False): - def fun(): - return opts.data[key] if key in opts.data else opts.data_labels[key].default - - info = opts.data_labels[key] - t = type(info.default) - args = (info.component_args() if callable(info.component_args) else info.component_args) or {} - if info.component is not None: - comp = info.component - elif t == str: - comp = gr.Textbox - elif t == int: - comp = gr.Number - elif t == bool: - comp = gr.Checkbox - else: - raise ValueError(f'bad options item type: {t} for key {key}') - elem_id = f"setting_{key}" - dirty_indicator = None - - if not is_quicksettings: - dirtyable_setting = gr.Group(elem_classes="dirtyable", visible=args.get("visible", True)) - dirtyable_setting.__enter__() - dirty_indicator = gr.Button("", elem_classes="modification-indicator", elem_id="modification_indicator_" + key) - - if info.refresh is not None: - if is_quicksettings: - res = comp(label=info.label, value=fun(), elem_id=elem_id, **args) - ui_common.create_refresh_button(res, info.refresh, info.component_args, f"refresh_{key}") - else: - with gr.Row(): - res = comp(label=info.label, value=fun(), elem_id=elem_id, **args) - ui_common.create_refresh_button(res, info.refresh, info.component_args, f"refresh_{key}") - elif info.folder is not None: - with gr.Row(): - res = comp(label=info.label, value=fun(), elem_id=elem_id, elem_classes="folder-selector", **args) - # ui_common.create_browse_button(res, f"folder_{key}") - else: - try: - res = comp(label=info.label, value=fun(), elem_id=elem_id, **args) - except Exception as e: - log.error(f'Error creating setting: {key} {e}') - res = None - - if res is not None and not is_quicksettings: - res.change(fn=None, inputs=res, _js=f'(val) => markIfModified("{key}", val)') - if dirty_indicator is not None: - dirty_indicator.click(fn=lambda: getattr(opts, key), outputs=res, show_progress=False) - dirtyable_setting.__exit__() - - return res - - def create_dirty_indicator(key, keys_to_reset, **kwargs): - def get_opt_values(): - return [getattr(opts, _key) for _key in keys_to_reset] - - elements_to_reset = [component_dict[_key] for _key in keys_to_reset if component_dict[_key] is not None] - indicator = gr.Button("", elem_classes="modification-indicator", elem_id=f"modification_indicator_{key}", **kwargs) - indicator.click(fn=get_opt_values, outputs=elements_to_reset, show_progress=False) - return indicator - - loadsave = ui_loadsave.UiLoadsave(cmd_opts.ui_config) - components = [] - component_dict = {} - shared.settings_components = component_dict - - script_callbacks.ui_settings_callback() - opts.reorder() - - def run_settings(*args): - changed = [] - for key, value, comp in zip(opts.data_labels.keys(), args, components): - if comp == dummy_component or value=='dummy': - continue - if getattr(comp, 'visible', True) is False: - continue - if not opts.same_type(value, opts.data_labels[key].default): - log.error(f'Setting bad value: {key}={value} expecting={type(opts.data_labels[key].default).__name__}') - continue - if opts.set(key, value): - changed.append(key) - if shared.opts.cuda_compile_backend == "olive-ai": - install_olive() - if cmd_opts.use_directml: - directml_override_opts() - if cmd_opts.use_openvino: - if "Model" not in shared.opts.cuda_compile: - shared.log.warning("OpenVINO: Enabling Torch Compile Model") - shared.opts.cuda_compile.append("Model") - if shared.opts.cuda_compile_backend != "openvino_fx": - shared.log.warning("OpenVINO: Setting Torch Compiler backend to OpenVINO FX") - shared.opts.cuda_compile_backend = "openvino_fx" - if shared.opts.sd_backend != "diffusers": - shared.log.warning("OpenVINO: Setting backend to Diffusers") - shared.opts.sd_backend = "diffusers" - try: - if len(changed) > 0: - opts.save(shared.config_filename) - log.info(f'Settings: changed={len(changed)} {changed}') - except RuntimeError: - log.error(f'Settings failed: change={len(changed)} {changed}') - return opts.dumpjson(), f'{len(changed)} Settings changed without save: {", ".join(changed)}' - return opts.dumpjson(), f'{len(changed)} Settings changed{": " if len(changed) > 0 else ""}{", ".join(changed)}' - - def run_settings_single(value, key, progress=False): - if not opts.same_type(value, opts.data_labels[key].default): - return gr.update(visible=True), opts.dumpjson() - if not opts.set(key, value): - return gr.update(value=getattr(opts, key)), opts.dumpjson() - if key == "cuda_compile_backend" and value == "olive-ai": - install_olive() - if cmd_opts.use_directml: - directml_override_opts() - opts.save(shared.config_filename) - log.debug(f'Setting changed: {key}={value} progress={progress}') - return get_value_for_setting(key), opts.dumpjson() - with gr.Blocks(analytics_enabled=False) as settings_interface: - with gr.Row(elem_id="system_row"): - restart_submit = gr.Button(value="Restart server", variant='primary', elem_id="restart_submit") - shutdown_submit = gr.Button(value="Shutdown server", variant='primary', elem_id="shutdown_submit") - unload_sd_model = gr.Button(value='Unload model', variant='primary', elem_id="sett_unload_sd_model") - reload_sd_model = gr.Button(value='Reload model', variant='primary', elem_id="sett_reload_sd_model") - enable_profiling = gr.Button(value='Start profiling', variant='primary', elem_id="enable_profiling") - - with gr.Tabs(elem_id="system") as system_tabs: - global ui_system_tabs # pylint: disable=global-statement - ui_system_tabs = system_tabs - with gr.TabItem("Settings", id="system_settings", elem_id="tab_settings"): - with gr.Row(elem_id="settings_row"): - settings_submit = gr.Button(value="Apply settings", variant='primary', elem_id="settings_submit") - preview_theme = gr.Button(value="Preview theme", variant='primary', elem_id="settings_preview_theme") - defaults_submit = gr.Button(value="Restore defaults", variant='primary', elem_id="defaults_submit") - with gr.Row(): - _settings_search = gr.Text(label="Search", elem_id="settings_search") - - result = gr.HTML(elem_id="settings_result") - quicksettings_names = opts.quicksettings_list - quicksettings_names = {x: i for i, x in enumerate(quicksettings_names) if x != 'quicksettings'} - quicksettings_list = [] - - previous_section = [] - tab_item_keys = [] - current_tab = None - current_row = None - dummy_component = gr.Textbox(visible=False, value='dummy') - with gr.Tabs(elem_id="settings"): - for i, (k, item) in enumerate(opts.data_labels.items()): - section_must_be_skipped = item.section[0] is None - if previous_section != item.section and not section_must_be_skipped: - if len(item.section) == 2: - elem_id, text = item.section - elif len(item.section) == 3: - _category, elem_id, text = item.section - else: - shared.log.error(f'Settings: section={item.section} invalid') - continue - if current_tab is not None and len(previous_section) > 0: - create_dirty_indicator(previous_section[0], tab_item_keys) - tab_item_keys = [] - current_row.__exit__() - current_tab.__exit__() - current_tab = gr.TabItem(elem_id=f"settings_{elem_id}", label=text) - current_tab.__enter__() - current_row = gr.Column(variant='compact') - current_row.__enter__() - previous_section = item.section - if k in quicksettings_names and not shared.cmd_opts.freeze: - quicksettings_list.append((i, k, item)) - components.append(dummy_component) - elif section_must_be_skipped: - components.append(dummy_component) - else: - component = create_setting_component(k) - component_dict[k] = component - tab_item_keys.append(k) - components.append(component) - if current_tab is not None and len(previous_section) > 0: - create_dirty_indicator(previous_section[0], tab_item_keys) - tab_item_keys = [] - current_row.__exit__() - current_tab.__exit__() - - request_notifications = gr.Button(value='Request browser notifications', elem_id="request_notifications", visible=False) - with gr.TabItem("Show all pages", elem_id="settings_show_all_pages"): - create_dirty_indicator("show_all_pages", [], interactive=False) - - with gr.TabItem("Update", id="system_update", elem_id="tab_update"): - from modules import update - update.create_ui() - - with gr.TabItem("User interface", id="system_config", elem_id="tab_config"): - loadsave.create_ui() - create_dirty_indicator("tab_defaults", [], interactive=False) - - with gr.TabItem("ONNX", id="onnx_config", elem_id="tab_onnx"): - from modules.onnx_impl import ui as ui_onnx - ui_onnx.create_ui() - - def unload_sd_weights(): - modules.sd_models.unload_model_weights(op='model') - modules.sd_models.unload_model_weights(op='refiner') - - def reload_sd_weights(): - modules.sd_models.reload_model_weights(force=True) - - def switch_profiling(): - shared.cmd_opts.profile = not shared.cmd_opts.profile - shared.log.warning(f'Profiling: {shared.cmd_opts.profile}') - return 'Stop profiling' if shared.cmd_opts.profile else 'Start profiling' - - unload_sd_model.click(fn=unload_sd_weights, inputs=[], outputs=[]) - reload_sd_model.click(fn=reload_sd_weights, inputs=[], outputs=[]) - enable_profiling.click(fn=switch_profiling, inputs=[], outputs=[enable_profiling]) - request_notifications.click(fn=lambda: None, inputs=[], outputs=[], _js='function(){}') - preview_theme.click(fn=None, _js='previewTheme', inputs=[], outputs=[]) - - timer.startup.record("ui-settings") + from modules import ui_settings + ui_settings.create_ui() + global ui_system_tabs # pylint: disable=global-statement + ui_system_tabs = ui_settings.ui_system_tabs + shared.opts.reorder() + timer.startup.record("ui-extensions") with gr.Blocks(analytics_enabled=False) as info_interface: with gr.Tabs(elem_id="tabs_info"): @@ -395,10 +138,18 @@ def create_ui(startup_timer = None): from modules import ui_docs ui_docs.create_ui_wiki() + with gr.Blocks(analytics_enabled=False) as extensions_interface: + from modules import ui_extensions + ui_extensions.create_ui() + timer.startup.record("ui-extensions") + interfaces = [] interfaces += [(txt2img_interface, "Text", "txt2img")] interfaces += [(img2img_interface, "Image", "img2img")] - interfaces += [(control_interface, "Control", "control")] if control_interface is not None else [] + if control_interface is not None: + interfaces += [(control_interface, "Control", "control")] + if video_interface is not None: + interfaces += [(video_interface, "Video", "video")] interfaces += [(extras_interface, "Process", "process")] interfaces += [(caption_interface, "Caption", "caption")] interfaces += [(gallery_interface, "Gallery", "gallery")] @@ -406,123 +157,12 @@ def create_ui(startup_timer = None): interfaces += script_callbacks.ui_tabs_callback() interfaces += [(settings_interface, "System", "system")] interfaces += [(info_interface, "Info", "info")] - - from modules import ui_extensions - extensions_interface = ui_extensions.create_ui() interfaces += [(extensions_interface, "Extensions", "extensions")] - timer.startup.record("ui-extensions") + + ui_app = ui_settings.create_quicksettings(interfaces) shared.tab_names = [] for _interface, label, _ifid in interfaces: shared.tab_names.append(label) - with gr.Blocks(theme=theme.gradio_theme, analytics_enabled=False, title="SD.Next") as ui_app: - with gr.Row(elem_id="quicksettings", variant="compact"): - for _i, k, _item in sorted(quicksettings_list, key=lambda x: quicksettings_names.get(x[1], x[0])): - component = create_setting_component(k, is_quicksettings=True) - component_dict[k] = component - - generation_parameters_copypaste.connect_paste_params_buttons() - - with gr.Tabs(elem_id="tabs") as tabs: - for interface, label, ifid in interfaces: - if interface is None: - continue - # if label in shared.opts.hidden_tabs or label == '': - # continue - with gr.TabItem(label, id=ifid, elem_id=f"tab_{ifid}"): - # log.debug(f'UI render: id={ifid}') - interface.render() - for interface, _label, ifid in interfaces: - if interface is None: - continue - if ifid in ["extensions", "system"]: - continue - loadsave.add_block(interface, ifid) - loadsave.add_component(f"webui/Tabs@{tabs.elem_id}", tabs) - loadsave.setup_ui() - - if opts.notification_audio_enable and os.path.exists(os.path.join(script_path, opts.notification_audio_path)): - gr.Audio(interactive=False, value=os.path.join(script_path, opts.notification_audio_path), elem_id="audio_notification", visible=False) - - text_settings = gr.Textbox(elem_id="settings_json", value=lambda: opts.dumpjson(), visible=False) - components = [c for c in components if c is not None] - settings_submit.click( - fn=wrap_gradio_call(run_settings, extra_outputs=[gr.update()]), - inputs=components, - outputs=[text_settings, result], - ) - defaults_submit.click(fn=lambda: shared.restore_defaults(restart=True), _js="restartReload") - restart_submit.click(fn=lambda: shared.restart_server(restart=True), _js="restartReload") - shutdown_submit.click(fn=lambda: shared.restart_server(restart=False), _js="restartReload") - - for _i, k, _item in quicksettings_list: - component = component_dict[k] - info = opts.data_labels[k] - if isinstance(component, gr.components.Textbox): - change_handlers = [component.blur, component.submit] - else: - change_handlers = [component.release if hasattr(component, 'release') else component.change] - for change_handler in change_handlers: - change_handler( - fn=lambda value, k=k, progress=info.refresh is not None: run_settings_single(value, key=k, progress=progress), - inputs=[component], - outputs=[component, text_settings], - show_progress=info.refresh is not None, - ) - - dummy_component = gr.Textbox(visible=False, value='dummy') - button_set_checkpoint = gr.Button('Change model', elem_id='change_checkpoint', visible=False) - button_set_checkpoint.click( - fn=lambda value, _: run_settings_single(value, key='sd_model_checkpoint'), - _js="function(v){ var res = desiredCheckpointName; desiredCheckpointName = ''; return [res || v, null]; }", - inputs=[component_dict['sd_model_checkpoint'], dummy_component], - outputs=[component_dict['sd_model_checkpoint'], text_settings], - ) - button_set_refiner = gr.Button('Change refiner', elem_id='change_refiner', visible=False) - button_set_refiner.click( - fn=lambda value, _: run_settings_single(value, key='sd_model_checkpoint'), - _js="function(v){ var res = desiredCheckpointName; desiredCheckpointName = ''; return [res || v, null]; }", - inputs=[component_dict['sd_model_refiner'], dummy_component], - outputs=[component_dict['sd_model_refiner'], text_settings], - ) - button_set_vae = gr.Button('Change VAE', elem_id='change_vae', visible=False) - button_set_vae.click( - fn=lambda value, _: run_settings_single(value, key='sd_vae'), - _js="function(v){ var res = desiredVAEName; desiredVAEName = ''; return [res || v, null]; }", - inputs=[component_dict['sd_vae'], dummy_component], - outputs=[component_dict['sd_vae'], text_settings], - ) - - def reference_submit(model): - if '@' not in model: # diffusers - loaded = modelloader.load_reference(model) - return model if loaded else opts.sd_model_checkpoint - else: # civitai - model, url = model.split('@') - loaded = modelloader.load_civitai(model, url) - return loaded if loaded is not None else opts.sd_model_checkpoint - - button_set_reference = gr.Button('Change reference', elem_id='change_reference', visible=False) - button_set_reference.click( - fn=reference_submit, - _js="function(v){ return desiredCheckpointName; }", - inputs=[component_dict['sd_model_checkpoint']], - outputs=[component_dict['sd_model_checkpoint']], - ) - component_keys = [k for k in opts.data_labels.keys() if k in component_dict] - - def get_settings_values(): - return [get_value_for_setting(key) for key in component_keys] - - ui_app.load( - fn=get_settings_values, - inputs=[], - outputs=[component_dict[k] for k in component_keys if component_dict[k] is not None], - queue=False, - ) - - timer.startup.record("ui-defaults") - loadsave.dump_defaults() - ui_app.ui_loadsave = loadsave return ui_app diff --git a/modules/ui_caption.py b/modules/ui_caption.py index 04a16702c..ef8d2d6e2 100644 --- a/modules/ui_caption.py +++ b/modules/ui_caption.py @@ -35,6 +35,8 @@ def create_ui(): with gr.Tabs(elem_id="mode_caption"): with gr.Tab("VLM Caption", elem_id="tab_vlm_caption"): from modules.interrogate import vqa + with gr.Row(): + vlm_system = gr.Textbox(label="System prompt", value=vqa.vlm_system, lines=1, elem_id='vlm_system') with gr.Row(): vlm_question = gr.Dropdown(label="Predefined question", allow_custom_value=False, choices=vqa.vlm_prompts, value=vqa.vlm_prompts[2], elem_id='vlm_question') with gr.Row(): @@ -43,9 +45,9 @@ def create_ui(): vlm_model = gr.Dropdown(list(vqa.vlm_models), value=list(vqa.vlm_models)[0], label='VLM Model', elem_id='vlm_model') with gr.Accordion(label='Advanced options', open=False, visible=True): with gr.Row(): - vlm_max_tokens = gr.Slider(label='Max tokens', value=shared.opts.interrogate_vlm_max_length, minimum=16, maximum=4096, step=1, elem_id='vlm_max_tokens') - vlm_num_beams = gr.Slider(label='Num beams', value=shared.opts.interrogate_vlm_num_beams, minimum=1, maximum=16, step=1, elem_id='vlm_num_beams') - vlm_temperature = gr.Slider(label='Temperature', value=shared.opts.interrogate_vlm_temperature, minimum=0.1, maximum=1.0, step=0.01, elem_id='vlm_temperature') + vlm_max_tokens = gr.Slider(label='VLM max tokens', value=shared.opts.interrogate_vlm_max_length, minimum=16, maximum=4096, step=1, elem_id='vlm_max_tokens') + vlm_num_beams = gr.Slider(label='VLM num beams', value=shared.opts.interrogate_vlm_num_beams, minimum=1, maximum=16, step=1, elem_id='vlm_num_beams') + vlm_temperature = gr.Slider(label='VLM Temperature', value=shared.opts.interrogate_vlm_temperature, minimum=0.1, maximum=1.0, step=0.01, elem_id='vlm_temperature') with gr.Row(): vlm_top_k = gr.Slider(label='Top-K', value=shared.opts.interrogate_vlm_top_k, minimum=0, maximum=99, step=1, elem_id='vlm_top_k') vlm_top_p = gr.Slider(label='Top-P', value=shared.opts.interrogate_vlm_top_p, minimum=0.0, maximum=1.0, step=0.01, elem_id='vlm_top_p') @@ -88,7 +90,7 @@ def create_ui(): clip_max_flavors = gr.Slider(label='Max flavors', value=shared.opts.interrogate_clip_max_flavors, minimum=1, maximum=64, step=1, elem_id='clip_max_flavors') clip_flavor_count = gr.Slider(label='Intermediates', value=shared.opts.interrogate_clip_flavor_count, minimum=256, maximum=4096, step=8, elem_id='clip_flavor_intermediate_count') with gr.Row(): - clip_num_beams = gr.Slider(label='Num beams', value=shared.opts.interrogate_clip_num_beams, minimum=1, maximum=16, step=1, elem_id='clip_num_beams') + clip_num_beams = gr.Slider(label='CLiP num beams', value=shared.opts.interrogate_clip_num_beams, minimum=1, maximum=16, step=1, elem_id='clip_num_beams') clip_min_length.change(fn=update_clip_params, inputs=[clip_min_length, clip_max_length, clip_chunk_size, clip_min_flavors, clip_max_flavors, clip_flavor_count, clip_num_beams], outputs=[]) clip_max_length.change(fn=update_clip_params, inputs=[clip_min_length, clip_max_length, clip_chunk_size, clip_min_flavors, clip_max_flavors, clip_flavor_count, clip_num_beams], outputs=[]) clip_chunk_size.change(fn=update_clip_params, inputs=[clip_min_length, clip_max_length, clip_chunk_size, clip_min_flavors, clip_max_flavors, clip_flavor_count, clip_num_beams], outputs=[]) @@ -114,7 +116,7 @@ def create_ui(): btn_clip_analyze_img = gr.Button("Analyze", variant='primary', elem_id="btn_clip_analyze_img") with gr.Column(variant='compact', elem_id='interrogate_output'): with gr.Row(elem_id='interrogate_output_prompt'): - prompt = gr.Textbox(label="Answer", lines=8, placeholder="ai generated image description") + prompt = gr.Textbox(label="Answer", lines=12, placeholder="ai generated image description") with gr.Row(elem_id='interrogate_output_classes'): medium = gr.Label(elem_id="interrogate_label_medium", label="Medium", num_top_classes=5, visible=False) artist = gr.Label(elem_id="interrogate_label_artist", label="Artist", num_top_classes=5, visible=False) @@ -127,8 +129,8 @@ def create_ui(): btn_clip_interrogate_img.click(openclip.interrogate_image, inputs=[image, clip_model, blip_model, clip_mode], outputs=[prompt]) btn_clip_analyze_img.click(openclip.analyze_image, inputs=[image, clip_model, blip_model], outputs=[medium, artist, movement, trending, flavor]) btn_clip_interrogate_batch.click(fn=openclip.interrogate_batch, inputs=[clip_batch_files, clip_batch_folder, clip_batch_str, clip_model, blip_model, clip_mode, clip_save_output, clip_save_append, clip_folder_recursive], outputs=[prompt]) - btn_vlm_caption.click(fn=vqa.interrogate, inputs=[vlm_question, vlm_prompt, image, vlm_model], outputs=[prompt]) - btn_vlm_caption_batch.click(fn=vqa.batch, inputs=[vlm_model, vlm_batch_files, vlm_batch_folder, vlm_batch_str, vlm_question, vlm_prompt, vlm_save_output, vlm_save_append, vlm_folder_recursive], outputs=[prompt]) + btn_vlm_caption.click(fn=vqa.interrogate, inputs=[vlm_question, vlm_system, vlm_prompt, image, vlm_model], outputs=[prompt]) + btn_vlm_caption_batch.click(fn=vqa.batch, inputs=[vlm_model, vlm_system, vlm_batch_files, vlm_batch_folder, vlm_batch_str, vlm_question, vlm_prompt, vlm_save_output, vlm_save_append, vlm_folder_recursive], outputs=[prompt]) for tabname, button in copy_interrogate_buttons.items(): generation_parameters_copypaste.register_paste_params_button(generation_parameters_copypaste.ParamBinding(paste_button=button, tabname=tabname, source_text_component=prompt, source_image_component=image,)) diff --git a/modules/ui_common.py b/modules/ui_common.py index a8167e457..f5260d10c 100644 --- a/modules/ui_common.py +++ b/modules/ui_common.py @@ -237,8 +237,8 @@ def interrogate_booru(image): # legacy function return gr.update() if prompt is None else prompt -def create_output_panel(tabname, preview=True, prompt=None, height=None): - with gr.Column(variant='panel', elem_id=f"{tabname}_results"): +def create_output_panel(tabname, preview=True, prompt=None, height=None, transfer=True, scale=1): + with gr.Column(variant='panel', elem_id=f"{tabname}_results", scale=scale): with gr.Group(elem_id=f"{tabname}_gallery_container"): if tabname == "txt2img": gr.HTML(value="", elem_id="main_info", visible=False, elem_classes=["main-info"]) @@ -270,10 +270,13 @@ def create_output_panel(tabname, preview=True, prompt=None, height=None): clip_files.click(fn=None, _js='clip_gallery_urls', inputs=[result_gallery], outputs=[]) save = gr.Button('Save', elem_id=f'save_{tabname}') delete = gr.Button('Delete', elem_id=f'delete_{tabname}') - if not shared.native: - buttons = generation_parameters_copypaste.create_buttons(["img2img", "inpaint", "extras"]) + if transfer: + if not shared.native: + buttons = generation_parameters_copypaste.create_buttons(["img2img", "inpaint", "extras"]) + else: + buttons = generation_parameters_copypaste.create_buttons(["txt2img", "img2img", "control", "extras", "caption"]) else: - buttons = generation_parameters_copypaste.create_buttons(["txt2img", "img2img", "control", "extras", "caption"]) + buttons = None download_files = gr.File(None, file_count="multiple", interactive=False, show_label=False, visible=False, elem_id=f'download_files_{tabname}') with gr.Group(): @@ -309,17 +312,18 @@ def create_output_panel(tabname, preview=True, prompt=None, height=None): else: paste_field_names = [] debug(f'Paste field: tab={tabname} fields={paste_field_names}') - for paste_tabname, paste_button in buttons.items(): - debug(f'Create output panel: source={tabname} target={paste_tabname} button={paste_button}') - bindings = generation_parameters_copypaste.ParamBinding( - paste_button=paste_button, - tabname=paste_tabname, - source_tabname=tabname, - source_image_component=result_gallery, - paste_field_names=paste_field_names, - source_text_component=prompt or generation_info - ) - generation_parameters_copypaste.register_paste_params_button(bindings) + if buttons is not None: + for paste_tabname, paste_button in buttons.items(): + debug(f'Create output panel: source={tabname} target={paste_tabname} button={paste_button}') + bindings = generation_parameters_copypaste.ParamBinding( + paste_button=paste_button, + tabname=paste_tabname, + source_tabname=tabname, + source_image_component=result_gallery, + paste_field_names=paste_field_names, + source_text_component=prompt or generation_info + ) + generation_parameters_copypaste.register_paste_params_button(bindings) return result_gallery, generation_info, html_info, html_info_formatted, html_log diff --git a/modules/ui_control.py b/modules/ui_control.py index 6aa6d0002..b10f2838e 100644 --- a/modules/ui_control.py +++ b/modules/ui_control.py @@ -645,6 +645,7 @@ def create_ui(_blocks: gr.Blocks=None): (mask_controls[5], "Mask dilate"), (mask_controls[6], "Mask auto"), # advanced + (cfg_scale, "Guidance scale"), (cfg_scale, "CFG scale"), (cfg_end, "CFG end"), (clip_skip, "Clip skip"), diff --git a/modules/ui_control_helpers.py b/modules/ui_control_helpers.py index 2b10e21d3..de8983be2 100644 --- a/modules/ui_control_helpers.py +++ b/modules/ui_control_helpers.py @@ -1,7 +1,7 @@ import os import gradio as gr from PIL import Image -from modules import shared, scripts, masking # pylint: disable=ungrouped-imports +from modules import shared, scripts, masking, video # pylint: disable=ungrouped-imports gr_height = None @@ -82,31 +82,9 @@ def display_units(num_units): return (num_units * [gr.update(visible=True)]) + ((max_units - num_units) * [gr.update(visible=False)]) -def get_video_params(filepath: str, capture: bool = False): - import cv2 - from modules.control.util import decode_fourcc - video = cv2.VideoCapture(filepath) - if not video.isOpened(): - msg = f'Control: video open failed: path="{filepath}"' - shared.log.error(msg) - raise RuntimeError(msg) - frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT)) - fps = video.get(cv2.CAP_PROP_FPS) - duration = float(frames) / fps - w, h = int(video.get(cv2.CAP_PROP_FRAME_WIDTH)), int(video.get(cv2.CAP_PROP_FRAME_HEIGHT)) - codec = decode_fourcc(video.get(cv2.CAP_PROP_FOURCC)) - frame = None - if capture: - _status, frame = video.read() - frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) - frame = Image.fromarray(frame) - video.release() - return frames, fps, duration, w, h, codec, frame - - def get_video(filepath: str): try: - frames, fps, duration, w, h, codec, _cap = get_video_params(filepath) + frames, fps, duration, w, h, codec, _cap = video.get_video_params(filepath) shared.log.debug(f'Control: input video: path={filepath} frames={frames} fps={fps} size={w}x{h} codec={codec}') msg = f'Control input | Video | Size {w}x{h} | Frames {frames} | FPS {fps:.2f} | Duration {duration:.2f} | Codec {codec}' return msg diff --git a/modules/ui_extensions.py b/modules/ui_extensions.py index f2e36004b..435ff95ae 100644 --- a/modules/ui_extensions.py +++ b/modules/ui_extensions.py @@ -5,7 +5,8 @@ import errno import html from datetime import datetime, timedelta import gradio as gr -from modules import extensions, shared, paths, errors, ui_symbols +from modules import extensions, shared, paths, errors, ui_symbols, call_queue + debug = shared.log.debug if os.environ.get('SD_EXT_DEBUG', None) is not None else lambda *args, **kwargs: None extensions_index = "https://vladmandic.github.io/sd-data/pages/extensions.json" @@ -98,7 +99,7 @@ def apply_changes(disable_list, update_list, disable_all): def check_updates(_id_task, disable_list, search_text, sort_column): if shared.cmd_opts.disable_extension_access: shared.log.error('Extension: apply changes disallowed because public access is enabled and insecure is not specified') - return + return create_html(search_text, sort_column) disabled = json.loads(disable_list) assert type(disabled) == list, f"wrong disable_list data for apply_and_restart: {disable_list}" exts = [ext for ext in extensions.extensions if ext.remote is not None and ext.name not in disabled] @@ -436,87 +437,84 @@ def create_html(search_text, sort_column): def create_ui(): - import modules.ui - with gr.Blocks(analytics_enabled=False) as ui: - extensions_disable_all = gr.Radio(label="Disable all extensions", choices=["none", "user", "all"], value=shared.opts.disable_all_extensions, elem_id="extensions_disable_all", visible=False) - extensions_disabled_list = gr.Text(elem_id="extensions_disabled_list", visible=False, container=False) - extensions_update_list = gr.Text(elem_id="extensions_update_list", visible=False, container=False) - with gr.Tabs(elem_id="tabs_extensions"): - with gr.TabItem("Manage extensions", id="manage"): - with gr.Row(elem_id="extensions_installed_top"): - extension_to_install = gr.Text(elem_id="extension_to_install", visible=False) - install_extension_button = gr.Button(elem_id="install_extension_button", visible=False) - uninstall_extension_button = gr.Button(elem_id="uninstall_extension_button", visible=False) - update_extension_button = gr.Button(elem_id="update_extension_button", visible=False) - with gr.Column(scale=4): - search_text = gr.Text(label="Search") - with gr.Column(scale=1): - sort_column = gr.Dropdown(value="default", label="Sort by", choices=list(sort_ordering.keys()), multiselect=False) - with gr.Column(scale=1): - refresh_extensions_button = gr.Button(value="Refresh extension list", variant="primary") - check = gr.Button(value="Update all installed", variant="primary") - apply = gr.Button(value="Apply changes", variant="primary") - list_extensions() - gr.HTML(''' -

Extension list

- - Refesh extension list to download latest list with status
- - Check status of an extension by looking at status icon before installing it
- - After any operation such as install/uninstall or enable/disable, please restart the server
- ''') - gr.HTML('') - info = gr.HTML('') - extensions_table = gr.HTML(create_html(search_text.value, sort_column.value)) - check.click( - fn=modules.ui.wrap_gradio_call(check_updates, extra_outputs=[gr.update()]), - _js="extensions_check", - inputs=[info, extensions_disabled_list, search_text, sort_column], - outputs=[extensions_table, info], - ) - apply.click( - fn=apply_changes, - _js="extensions_apply", - inputs=[extensions_disabled_list, extensions_update_list, extensions_disable_all], - outputs=[], - ) - refresh_extensions_button.click( - fn=modules.ui.wrap_gradio_call(refresh_extensions_list, extra_outputs=[gr.update(), gr.update()]), - inputs=[search_text, sort_column], - outputs=[extensions_table, info], - ) - install_extension_button.click( - fn=modules.ui.wrap_gradio_call(install_extension, extra_outputs=[gr.update(), gr.update(), gr.update()]), - inputs=[extension_to_install, search_text, sort_column], - outputs=[extensions_table, info], - ) - uninstall_extension_button.click( - fn=modules.ui.wrap_gradio_call(uninstall_extension, extra_outputs=[gr.update(), gr.update(), gr.update()]), - inputs=[extension_to_install, search_text, sort_column], - outputs=[extensions_table, info], - ) - update_extension_button.click( - fn=modules.ui.wrap_gradio_call(update_extension, extra_outputs=[gr.update(), gr.update(), gr.update()]), - inputs=[extension_to_install, search_text, sort_column], - outputs=[extensions_table, info], - ) - search_text.change( - fn=modules.ui.wrap_gradio_call(search_extensions, extra_outputs=[gr.update(), gr.update()]), - inputs=[search_text, sort_column], - outputs=[extensions_table, info], - ) - sort_column.change( - fn=modules.ui.wrap_gradio_call(search_extensions, extra_outputs=[gr.update(), gr.update()]), - inputs=[search_text, sort_column], - outputs=[extensions_table, info], - ) - with gr.TabItem("Manual install", id="install_from_url"): - install_url = gr.Text(label="Extension GIT repository URL") - install_branch = gr.Text(label="Specific branch name", placeholder="Leave empty for default main branch") - install_dirname = gr.Text(label="Local directory name", placeholder="Leave empty for auto") - install_button = gr.Button(value="Install", variant="primary") - info = gr.HTML(elem_id="extension_info") - install_button.click( - fn=modules.ui.wrap_gradio_call(install_extension_from_url, extra_outputs=[gr.update()]), - inputs=[install_dirname, install_url, install_branch, search_text, sort_column], - outputs=[extensions_table, info], - ) - return ui + extensions_disable_all = gr.Radio(label="Disable all extensions", choices=["none", "user", "all"], value=shared.opts.disable_all_extensions, elem_id="extensions_disable_all", visible=False) + extensions_disabled_list = gr.Text(elem_id="extensions_disabled_list", visible=False, container=False) + extensions_update_list = gr.Text(elem_id="extensions_update_list", visible=False, container=False) + with gr.Tabs(elem_id="tabs_extensions"): + with gr.TabItem("Manage extensions", id="manage"): + with gr.Row(elem_id="extensions_installed_top"): + extension_to_install = gr.Text(elem_id="extension_to_install", visible=False) + install_extension_button = gr.Button(elem_id="install_extension_button", visible=False) + uninstall_extension_button = gr.Button(elem_id="uninstall_extension_button", visible=False) + update_extension_button = gr.Button(elem_id="update_extension_button", visible=False) + with gr.Column(scale=4): + search_text = gr.Text(label="Search") + with gr.Column(scale=1): + sort_column = gr.Dropdown(value="default", label="Sort by", choices=list(sort_ordering.keys()), multiselect=False) + with gr.Column(scale=1): + refresh_extensions_button = gr.Button(value="Refresh extension list", variant="primary") + check = gr.Button(value="Update all installed", variant="primary") + apply = gr.Button(value="Apply changes", variant="primary") + list_extensions() + gr.HTML(''' +

Extension list

+ - Refesh extension list to download latest list with status
+ - Check status of an extension by looking at status icon before installing it
+ - After any operation such as install/uninstall or enable/disable, please restart the server
+ ''') + gr.HTML('') + info = gr.HTML('') + extensions_table = gr.HTML(create_html(search_text.value, sort_column.value)) + check.click( + fn=call_queue.wrap_gradio_call(check_updates, extra_outputs=[gr.update()]), + _js="extensions_check", + inputs=[info, extensions_disabled_list, search_text, sort_column], + outputs=[extensions_table, info], + ) + apply.click( + fn=apply_changes, + _js="extensions_apply", + inputs=[extensions_disabled_list, extensions_update_list, extensions_disable_all], + outputs=[], + ) + refresh_extensions_button.click( + fn=call_queue.wrap_gradio_call(refresh_extensions_list, extra_outputs=[gr.update(), gr.update()]), + inputs=[search_text, sort_column], + outputs=[extensions_table, info], + ) + install_extension_button.click( + fn=call_queue.wrap_gradio_call(install_extension, extra_outputs=[gr.update(), gr.update(), gr.update()]), + inputs=[extension_to_install, search_text, sort_column], + outputs=[extensions_table, info], + ) + uninstall_extension_button.click( + fn=call_queue.wrap_gradio_call(uninstall_extension, extra_outputs=[gr.update(), gr.update(), gr.update()]), + inputs=[extension_to_install, search_text, sort_column], + outputs=[extensions_table, info], + ) + update_extension_button.click( + fn=call_queue.wrap_gradio_call(update_extension, extra_outputs=[gr.update(), gr.update(), gr.update()]), + inputs=[extension_to_install, search_text, sort_column], + outputs=[extensions_table, info], + ) + search_text.change( + fn=call_queue.wrap_gradio_call(search_extensions, extra_outputs=[gr.update(), gr.update()]), + inputs=[search_text, sort_column], + outputs=[extensions_table, info], + ) + sort_column.change( + fn=call_queue.wrap_gradio_call(search_extensions, extra_outputs=[gr.update(), gr.update()]), + inputs=[search_text, sort_column], + outputs=[extensions_table, info], + ) + with gr.TabItem("Manual install", id="install_from_url"): + install_url = gr.Text(label="Extension GIT repository URL") + install_branch = gr.Text(label="Specific branch name", placeholder="Leave empty for default main branch") + install_dirname = gr.Text(label="Local directory name", placeholder="Leave empty for auto") + install_button = gr.Button(value="Install", variant="primary") + info = gr.HTML(elem_id="extension_info") + install_button.click( + fn=call_queue.wrap_gradio_call(install_extension_from_url, extra_outputs=[gr.update()]), + inputs=[install_dirname, install_url, install_branch, search_text, sort_column], + outputs=[extensions_table, info], + ) diff --git a/modules/ui_extra_networks.py b/modules/ui_extra_networks.py index f96974d28..e16a26bbf 100644 --- a/modules/ui_extra_networks.py +++ b/modules/ui_extra_networks.py @@ -56,6 +56,9 @@ preview_map = None def init_api(): def fetch_file(filename: str = ""): + global allowed_dirs # pylint: disable=global-statement + if len(allowed_dirs) == 0: + allowed_dirs = shared.demo.allowed_paths if not os.path.exists(filename): return JSONResponse({ "error": f"file {filename}: not found" }, status_code=404) if filename.startswith('html/') or filename.startswith('models/'): @@ -175,9 +178,9 @@ class ExtraNetworksPage: img = None try: img = Image.open(f) - except Exception: + except Exception as e: img = None - shared.log.warning(f'Extra network removing invalid image: {f}') + shared.log.warning(f'Network removing invalid: image={f} {e}') try: if img is None: img = None @@ -189,9 +192,9 @@ class ExtraNetworksPage: img.close() created += 1 except Exception as e: - shared.log.warning(f'Extra network error creating thumbnail: {f} {e}') + shared.log.warning(f'Network create thumbnail={f} {e}') if created > 0: - shared.log.info(f"Network thumbnails: {self.name} created={created}") + shared.log.info(f'Network thumbnails: {self.name} created={created}') self.missing_thumbs.clear() def create_items(self, tabname): @@ -221,7 +224,7 @@ class ExtraNetworksPage: return self.patch(self.html, tabname) self_name_id = self.name.replace(" ", "_") if skip: - return f"
Extra network page not ready
Click refresh to try again
" + return f"
Network page not ready
Click refresh to try again
" subdirs = {} allowed_folders = [os.path.abspath(x) for x in self.allowed_directories_for_previews() if os.path.exists(x)] for parentdir, dirs in {d: files_cache.walk(d, cached=True, recurse=files_cache.not_hidden) for d in allowed_folders}.items(): @@ -239,7 +242,7 @@ class ExtraNetworksPage: if not subdir: continue subdirs[subdir] = 1 - debug(f"Networks: page='{self.name}' subfolders={list(subdirs)}") + debug(f'Networks: page="{self.name}" subfolders={list(subdirs)}') subdirs = OrderedDict(sorted(subdirs.items())) if self.name == 'model' and shared.opts.extra_network_reference_enable: subdirs['Local'] = 1 @@ -289,7 +292,7 @@ class ExtraNetworksPage: self.html += ''.join(htmls) self.page_time = time.time() self.html = f"
{subdirs_html}
{self.html}
" - shared.log.debug(f"Networks: type='{self.name}' items={len(self.items)} subfolders={len(subdirs)} tab={tabname} folders={self.allowed_directories_for_previews()} list={self.list_time:.2f} thumb={self.preview_time:.2f} desc={self.desc_time:.2f} info={self.info_time:.2f} workers={shared.max_workers}") + shared.log.debug(f'Networks: type="{self.name}" items={len(self.items)} subfolders={len(subdirs)} tab={tabname} folders={self.allowed_directories_for_previews()} list={self.list_time:.2f} thumb={self.preview_time:.2f} desc={self.desc_time:.2f} info={self.info_time:.2f} workers={shared.max_workers}') if len(self.missing_thumbs) > 0: threading.Thread(target=self.create_thumb).start() return self.patch(self.html, tabname) @@ -473,6 +476,8 @@ def register_page(page: ExtraNetworksPage): def register_pages(): debug('EN register-pages') + shared.extra_networks.clear() + allowed_dirs.clear() from modules.ui_extra_networks_checkpoints import ExtraNetworksPageCheckpoints from modules.ui_extra_networks_vae import ExtraNetworksPageVAEs from modules.ui_extra_networks_styles import ExtraNetworksPageStyles @@ -677,13 +682,13 @@ def create_ui(container, button_parent, tabname, skip_indexing = False): if ui.gallery is not None: images = list(ui.gallery.temp_files) # gallery cannot be used as input component so looking at most recently registered temp files if len(images) < 1: - shared.log.warning(f'Extra network no image: item={ui.last_item.name}') + shared.log.warning(f'Network no image: item="{ui.last_item.name}"') return 'html/card-no-preview.png' try: images.sort(key=lambda f: os.path.getmtime(f), reverse=True) image = Image.open(images[0]) except Exception as e: - shared.log.error(f'Extra network error opening image: item={ui.last_item.name} {e}') + shared.log.error(f'Network error opening image: item="{ui.last_item.name}" {e}') return 'html/card-no-preview.png' fn_delete_img(image) if image.width > 512 or image.height > 512: @@ -691,9 +696,9 @@ def create_ui(container, button_parent, tabname, skip_indexing = False): image.thumbnail((512, 512), Image.Resampling.HAMMING) try: image.save(ui.last_item.local_preview, quality=50) - shared.log.debug(f'Extra network save image: item={ui.last_item.name} filename="{ui.last_item.local_preview}"') + shared.log.debug(f'Networks save image: item="{ui.last_item.name}" filename="{ui.last_item.local_preview}"') except Exception as e: - shared.log.error(f'Extra network save image: item={ui.last_item.name} filename="{ui.last_item.local_preview}" {e}') + shared.log.error(f'Network save image: item="{ui.last_item.name}" filename="{ui.last_item.local_preview}" {e}') return image def fn_delete_img(_image): @@ -702,7 +707,7 @@ def create_ui(container, button_parent, tabname, skip_indexing = False): for file in [f'{fn}{mid}{ext}' for ext in preview_extensions for mid in ['.thumb.', '.preview.', '.']]: if os.path.exists(file): os.remove(file) - shared.log.debug(f'Extra network delete image: item={ui.last_item.name} filename="{file}"') + shared.log.debug(f'Network delete image: item="{ui.last_item.name}" filename="{file}"') return 'html/card-no-preview.png' def fn_save_desc(desc): @@ -714,7 +719,7 @@ def create_ui(container, button_parent, tabname, skip_indexing = False): fn = os.path.splitext(ui.last_item.filename)[0] + '.txt' with open(fn, 'w', encoding='utf-8') as f: f.write(desc) - shared.log.debug(f'Extra network save desc: item={ui.last_item.name} filename="{fn}"') + shared.log.debug(f'Network save desc: item="{ui.last_item.name}" filename="{fn}"') return desc def fn_delete_desc(desc): @@ -722,7 +727,7 @@ def create_ui(container, button_parent, tabname, skip_indexing = False): return desc fn = os.path.splitext(ui.last_item.filename)[0] + '.txt' if os.path.exists(fn): - shared.log.debug(f'Extra network delete desc: item={ui.last_item.name} filename="{fn}"') + shared.log.debug(f'Network delete desc: item="{ui.last_item.name}" filename="{fn}"') os.remove(fn) return '' return desc @@ -730,7 +735,7 @@ def create_ui(container, button_parent, tabname, skip_indexing = False): def fn_save_info(info): fn = os.path.splitext(ui.last_item.filename)[0] + '.json' shared.writefile(info, fn, silent=True) - shared.log.debug(f'Extra network save info: item={ui.last_item.name} filename="{fn}"') + shared.log.debug(f'Network save info: item="{ui.last_item.name}" filename="{fn}"') return info def fn_delete_info(info): @@ -738,14 +743,14 @@ def create_ui(container, button_parent, tabname, skip_indexing = False): return info fn = os.path.splitext(ui.last_item.filename)[0] + '.json' if os.path.exists(fn): - shared.log.debug(f'Extra network delete info: item={ui.last_item.name} filename="{fn}"') + shared.log.debug(f'Network delete info: item="{ui.last_item.name}" filename="{fn}"') os.remove(fn) return '' return info def fn_save_style(info, description, prompt, negative, extra, wildcards): if not isinstance(info, dict) or isinstance(info, list): - shared.log.warning(f'Extra network save style skip: item={ui.last_item.name} not a dict: {type(info)}') + shared.log.warning(f'Network save style skip: item="{ui.last_item.name}" not a dict: {type(info)}') return info if ui.last_item is None: return info @@ -753,7 +758,7 @@ def create_ui(container, button_parent, tabname, skip_indexing = False): if hasattr(ui.last_item, 'type') and ui.last_item.type == 'Style': info.update(**{ 'description': description, 'prompt': prompt, 'negative': negative, 'extra': extra, 'wildcards': wildcards }) shared.writefile(info, fn, silent=True) - shared.log.debug(f'Extra network save style: item={ui.last_item.name} filename="{fn}"') + shared.log.debug(f'Network save style: item="{ui.last_item.name}" filename="{fn}"') return info def fn_delete_style(info): @@ -761,7 +766,7 @@ def create_ui(container, button_parent, tabname, skip_indexing = False): return info fn = os.path.splitext(ui.last_item.filename)[0] + '.json' if os.path.exists(fn): - shared.log.debug(f'Extra network delete style: item={ui.last_item.name} filename="{fn}"') + shared.log.debug(f'Network delete style: item="{ui.last_item.name}" filename="{fn}"') os.remove(fn) return {} return info @@ -785,6 +790,10 @@ def create_ui(container, button_parent, tabname, skip_indexing = False): if 'modelVersions' in fullinfo: # sanitize massive objects fullinfo['modelVersions'] = [] info = fullinfo + if isinstance(info, list): + item.filename = None + shared.log.warning('Network: show details not supported for compound item') + info = None """ if prompt is not None: item.prompt = prompt @@ -812,7 +821,7 @@ def create_ui(container, button_parent, tabname, skip_indexing = False): model = '' style = '' note = '' - if not os.path.exists(item.filename): + if item.filename is not None and not os.path.exists(item.filename): note = f'
Target filename: {item.filename}' if page.title == 'Model': merge = len(list(meta.get('sd_merge_models', {}))) @@ -904,7 +913,7 @@ def create_ui(container, button_parent, tabname, skip_indexing = False): page.refresh_time = 0 page.refresh() page.create_page(ui.tabname) - shared.log.debug(f"Networks: refresh page='{page.title}' items={len(page.items)} tab={ui.tabname}") + shared.log.debug(f'Networks: refresh page="{page.title}" items={len(page.items)} tab={ui.tabname}') pages.append(page.html) ui.search.update(title) return pages @@ -918,7 +927,7 @@ def create_ui(container, button_parent, tabname, skip_indexing = False): page.card = card_full if page.view == 'gallery' else card_list page.html = '' page.create_page(ui.tabname) - shared.log.debug(f"Networks: refresh page='{page.title}' items={len(page.items)} tab={ui.tabname} view={page.view}") + shared.log.debug(f'Networks: refresh page="{page.title}" items={len(page.items)} tab={ui.tabname} view={page.view}') pages.append(page.html) ui.search.update(title) return pages @@ -973,9 +982,9 @@ def create_ui(container, button_parent, tabname, skip_indexing = False): } shared.writefile(item, fn, silent=True) if len(prompt) > 0: - shared.log.debug(f"Network quick save style: item={name} filename='{fn}' unparsed={shared.opts.extra_networks_unparsed}") + shared.log.debug(f'Network quick save style: item="{name}" filename="{fn}" unparsed={shared.opts.extra_networks_unparsed}') else: - shared.log.warning(f"Network quick save model: item={name} filename='{fn}' prompt is empty") + shared.log.warning(f'Network quick save model: item="{name}" filename="{fn}" prompt is empty') def ui_sort_cards(sort_order): if shared.opts.extra_networks_sort != sort_order: diff --git a/modules/ui_extra_networks_lora.py b/modules/ui_extra_networks_lora.py index 9dd1b3573..194f16b41 100644 --- a/modules/ui_extra_networks_lora.py +++ b/modules/ui_extra_networks_lora.py @@ -1,8 +1,8 @@ import os import json import concurrent -import modules.lora.networks as networks from modules import shared, ui_extra_networks +from modules.lora import lora_load debug = os.environ.get('SD_LORA_DEBUG', None) is not None @@ -14,7 +14,7 @@ class ExtraNetworksPageLora(ui_extra_networks.ExtraNetworksPage): self.list_time = 0 def refresh(self): - networks.list_available_networks() + lora_load.list_available_networks() @staticmethod def get_tags(l, info): @@ -78,9 +78,9 @@ class ExtraNetworksPageLora(ui_extra_networks.ExtraNetworksPage): return clean_tags def create_item(self, name): - l = networks.available_networks.get(name) + l = lora_load.available_networks.get(name) if l is None: - shared.log.warning(f'Networks: type=lora registered={len(list(networks.available_networks))} file="{name}" not registered') + shared.log.warning(f'Networks: type=lora registered={len(list(lora_load.available_networks))} file="{name}" not registered') return None try: # path, _ext = os.path.splitext(l.filename) @@ -111,7 +111,7 @@ class ExtraNetworksPageLora(ui_extra_networks.ExtraNetworksPage): def list_items(self): items = [] with concurrent.futures.ThreadPoolExecutor(max_workers=shared.max_workers) as executor: - future_items = {executor.submit(self.create_item, net): net for net in networks.available_networks} + future_items = {executor.submit(self.create_item, net): net for net in lora_load.available_networks} for future in concurrent.futures.as_completed(future_items): item = future.result() if item is not None: diff --git a/modules/ui_gallery.py b/modules/ui_gallery.py index a1f317caa..40bcace03 100644 --- a/modules/ui_gallery.py +++ b/modules/ui_gallery.py @@ -3,7 +3,7 @@ from datetime import datetime from urllib.parse import unquote import gradio as gr from PIL import Image -from modules import shared, ui_symbols, ui_common, images, ui_control_helpers +from modules import shared, ui_symbols, ui_common, images, video from modules.ui_components import ToolButton def read_media(fn): @@ -13,7 +13,7 @@ def read_media(fn): return [[], None, '', '', f'Media not found: {fn}'] stat = os.stat(fn) if fn.lower().endswith('.mp4'): - frames, fps, duration, w, h, codec, _frame = ui_control_helpers.get_video_params(fn) + frames, fps, duration, w, h, codec, _frame = video.get_video_params(fn) geninfo = '' log = f'''

Video {w} x {h} diff --git a/modules/ui_img2img.py b/modules/ui_img2img.py index e6de84f6e..294065498 100644 --- a/modules/ui_img2img.py +++ b/modules/ui_img2img.py @@ -255,6 +255,7 @@ def create_ui(): (subseed, "Variation seed"), (subseed_strength, "Variation strength"), # advanced + (cfg_scale, "Guidance scale"), (cfg_scale, "CFG scale"), (cfg_end, "CFG end"), (image_cfg_scale, "Image CFG scale"), diff --git a/modules/ui_loadsave.py b/modules/ui_loadsave.py index 1135134d2..594b44c1d 100644 --- a/modules/ui_loadsave.py +++ b/modules/ui_loadsave.py @@ -13,7 +13,6 @@ class UiLoadsave: def __init__(self, filename): self.filename = filename self.component_mapping = {} - self.finalized_ui = False self.ui_defaults_view = None # button self.ui_defaults_apply = None # button self.ui_defaults_review = None # button @@ -24,8 +23,6 @@ class UiLoadsave: self.ui_settings = self.read_from_file() def add_component(self, path, x): - """adds component to the registry of tracked components""" - assert not self.finalized_ui def apply_field(obj, field, condition=None, init_field=None): key = f"{path}/{field}" @@ -253,7 +250,6 @@ class UiLoadsave: return "Restored system defaults for user interface" def create_ui(self): - """creates ui elements for editing defaults UI, without adding any logic to them""" with gr.Row(elem_id="config_row"): self.ui_defaults_apply = gr.Button(value='Set UI defaults', elem_id="ui_defaults_apply", variant="primary") self.ui_defaults_submenu = gr.Button(value='Set UI menu states', elem_id="ui_submenu_apply", variant="primary") @@ -262,9 +258,6 @@ class UiLoadsave: self.ui_defaults_review = gr.HTML("", elem_id="ui_defaults_review") def setup_ui(self): - """adds logic to elements created with create_ui; all add_block class must be made before this""" - assert not self.finalized_ui - self.finalized_ui = True self.ui_defaults_view.click(fn=self.ui_view, inputs=list(self.component_mapping.values()), outputs=[self.ui_defaults_review]) self.ui_defaults_apply.click(fn=self.ui_apply, inputs=list(self.component_mapping.values()), outputs=[self.ui_defaults_review]) self.ui_defaults_restore.click(fn=self.ui_restore, inputs=[], outputs=[self.ui_defaults_review]) diff --git a/modules/ui_models.py b/modules/ui_models.py index 4f6355a9a..59d8a9196 100644 --- a/modules/ui_models.py +++ b/modules/ui_models.py @@ -573,6 +573,8 @@ def create_ui(): def atomic_civit_search_metadata(item, res, rehash): from modules.modelloader import download_civit_preview, download_civit_meta + if item is None: + return meta = os.path.splitext(item['filename'])[0] + '.json' has_meta = os.path.isfile(meta) and os.stat(meta).st_size > 0 if ('card-no-preview.png' in item['preview'] or not has_meta) and os.path.isfile(item['filename']): @@ -629,6 +631,7 @@ def create_ui(): with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor: for fn in candidates: executor.submit(atomic_civit_search_metadata, fn, res, rehash) + atomic_civit_search_metadata(None, res, rehash) t1 = time.time() log.debug(f'CivitAI search metadata: items={i} time={t1-t0:.2f}') txt = '
'.join([r for r in res if len(r) > 0]) diff --git a/modules/ui_sections.py b/modules/ui_sections.py index 0d0e0eb8d..2785c40b6 100644 --- a/modules/ui_sections.py +++ b/modules/ui_sections.py @@ -4,7 +4,7 @@ from modules.ui_components import ToolButton from modules.interrogate import interrogate -def create_toprow(is_img2img: bool = False, id_part: str = None): +def create_toprow(is_img2img: bool = False, id_part: str = None, negative_visible: bool = True, reprocess_visible: bool = True): def apply_styles(prompt, prompt_neg, styles): prompt = shared.prompt_styles.apply_styles_to_prompt(prompt, styles, wildcards=not shared.opts.extra_networks_apply_unparsed) prompt_neg = shared.prompt_styles.apply_negative_styles_to_prompt(prompt_neg, styles, wildcards=not shared.opts.extra_networks_apply_unparsed) @@ -21,19 +21,20 @@ def create_toprow(is_img2img: bool = False, id_part: str = None): with gr.Row(): with gr.Column(scale=80): with gr.Row(elem_id=f"{id_part}_prompt_row"): - prompt = gr.Textbox(elem_id=f"{id_part}_prompt", label="Prompt", show_label=False, lines=3, placeholder="Prompt", elem_classes=["prompt"]) + prompt = gr.Textbox(elem_id=f"{id_part}_prompt", label="Prompt", show_label=False, lines=3 if negative_visible else 5, placeholder="Prompt", elem_classes=["prompt"]) with gr.Row(): with gr.Column(scale=80): with gr.Row(elem_id=f"{id_part}_negative_row"): - negative_prompt = gr.Textbox(elem_id=f"{id_part}_neg_prompt", label="Negative prompt", show_label=False, lines=3, placeholder="Negative prompt", elem_classes=["prompt"]) + negative_prompt = gr.Textbox(elem_id=f"{id_part}_neg_prompt", label="Negative prompt", show_label=False, lines=3, placeholder="Negative prompt", elem_classes=["prompt"], visible=negative_visible) with gr.Column(scale=1, elem_id=f"{id_part}_actions_column"): with gr.Row(elem_id=f"{id_part}_generate_box"): reprocess = [] submit = gr.Button('Generate', elem_id=f"{id_part}_generate", variant='primary') - reprocess.append(gr.Button('Reprocess', elem_id=f"{id_part}_reprocess", variant='primary', visible=True)) - reprocess.append(gr.Button('Reprocess decode', elem_id=f"{id_part}_reprocess_decode", variant='primary', visible=False)) - reprocess.append(gr.Button('Reprocess refine', elem_id=f"{id_part}_reprocess_refine", variant='primary', visible=False)) - reprocess.append(gr.Button('Reprocess face', elem_id=f"{id_part}_reprocess_detail", variant='primary', visible=False)) + if reprocess_visible: + reprocess.append(gr.Button('Reprocess', elem_id=f"{id_part}_reprocess", variant='primary', visible=True)) + reprocess.append(gr.Button('Reprocess decode', elem_id=f"{id_part}_reprocess_decode", variant='primary', visible=False)) + reprocess.append(gr.Button('Reprocess refine', elem_id=f"{id_part}_reprocess_refine", variant='primary', visible=False)) + reprocess.append(gr.Button('Reprocess face', elem_id=f"{id_part}_reprocess_detail", variant='primary', visible=False)) with gr.Row(elem_id=f"{id_part}_generate_line2"): interrupt = gr.Button('Stop', elem_id=f"{id_part}_interrupt") interrupt.click(fn=lambda: shared.state.interrupt(), _js="requestInterrupt", inputs=[], outputs=[]) @@ -79,9 +80,9 @@ def ar_change(ar, width, height): return gr.update(), gr.update() -def create_resolution_inputs(tab): - width = gr.Slider(minimum=64, maximum=4096, step=8, label="Width", value=1024, elem_id=f"{tab}_width") - height = gr.Slider(minimum=64, maximum=4096, step=8, label="Height", value=1024, elem_id=f"{tab}_height") +def create_resolution_inputs(tab, default_width=1024, default_height=1024): + width = gr.Slider(minimum=64, maximum=4096, step=8, label="Width", value=default_width, elem_id=f"{tab}_width") + height = gr.Slider(minimum=64, maximum=4096, step=8, label="Height", value=default_height, elem_id=f"{tab}_height") ar_list = ['AR'] + [x.strip() for x in shared.opts.aspect_ratios.split(',') if x.strip() != ''] ar_dropdown = gr.Dropdown(show_label=False, interactive=True, choices=ar_list, value=ar_list[0], elem_id=f"{tab}_ar", elem_classes=["ar-dropdown"]) for c in [ar_dropdown, width, height]: @@ -120,18 +121,18 @@ def create_batch_inputs(tab, accordion=True): return batch_count, batch_size -def create_seed_inputs(tab, reuse_visible=True): - with gr.Accordion(open=False, label="Seed", elem_id=f"{tab}_seed_group", elem_classes=["small-accordion"]): +def create_seed_inputs(tab, reuse_visible=True, accordion=True, subseed_visible=True, seed_resize_visible=False): + with gr.Accordion(open=False, label="Seed", elem_id=f"{tab}_seed_group", elem_classes=["small-accordion"]) if accordion else gr.Group(): with gr.Row(elem_id=f"{tab}_seed_row", variant="compact"): seed = gr.Number(label='Initial seed', value=-1, elem_id=f"{tab}_seed", container=True) random_seed = ToolButton(ui_symbols.random, elem_id=f"{tab}_random_seed", label='Random seed') reuse_seed = ToolButton(ui_symbols.reuse, elem_id=f"{tab}_reuse_seed", label='Reuse seed', visible=reuse_visible) - with gr.Row(elem_id=f"{tab}_subseed_row", variant="compact", visible=True): + with gr.Row(elem_id=f"{tab}_subseed_row", variant="compact", visible=subseed_visible): subseed = gr.Number(label='Variation', value=-1, elem_id=f"{tab}_subseed", container=True) random_subseed = ToolButton(ui_symbols.random, elem_id=f"{tab}_random_subseed") reuse_subseed = ToolButton(ui_symbols.reuse, elem_id=f"{tab}_reuse_subseed", visible=reuse_visible) subseed_strength = gr.Slider(label='Variation strength', value=0.0, minimum=0, maximum=1, step=0.01, elem_id=f"{tab}_subseed_strength") - with gr.Row(visible=False): + with gr.Row(visible=seed_resize_visible): seed_resize_from_w = gr.Slider(minimum=0, maximum=4096, step=8, label="Resize seed from width", value=0, elem_id=f"{tab}_seed_resize_from_w") seed_resize_from_h = gr.Slider(minimum=0, maximum=4096, step=8, label="Resize seed from height", value=0, elem_id=f"{tab}_seed_resize_from_h") random_seed.click(fn=lambda: -1, show_progress=False, inputs=[], outputs=[seed]) @@ -139,22 +140,22 @@ def create_seed_inputs(tab, reuse_visible=True): return seed, reuse_seed, subseed, reuse_subseed, subseed_strength, seed_resize_from_h, seed_resize_from_w -def create_video_inputs(tab:str): +def create_video_inputs(tab:str, show_always:bool=False): def video_type_change(video_type): return [ - gr.update(visible=video_type != 'None'), - gr.update(visible=video_type in ['GIF', 'PNG']), - gr.update(visible=video_type not in ['None', 'GIF', 'PNG']), - gr.update(visible=video_type not in ['None', 'GIF', 'PNG']), + gr.update(visible=video_type != 'None' or show_always), + gr.update(visible=video_type in ['GIF', 'PNG'] or show_always), + gr.update(visible=video_type not in ['None', 'GIF', 'PNG'] or show_always), + gr.update(visible=video_type not in ['None', 'GIF', 'PNG'] or show_always), ] with gr.Column(): video_codecs = ['None', 'GIF', 'PNG', 'MP4/MP4V', 'MP4/AVC1', 'MP4/JVT3', 'MKV/H264', 'AVI/DIVX', 'AVI/RGBA', 'MJPEG/MJPG', 'MPG/MPG1', 'AVR/AVR1'] - video_type = gr.Dropdown(label='Video type', choices=video_codecs, value='None', elem_id=f"{tab}_video_type") + video_type = gr.Dropdown(label='Save video', choices=video_codecs, value='None', elem_id=f"{tab}_video_type") with gr.Column(): - video_duration = gr.Slider(label='Duration', minimum=0.25, maximum=300, step=0.25, value=2, visible=False, elem_id=f"{tab}_video_duration") - video_loop = gr.Checkbox(label='Loop', value=True, visible=False, elem_id=f"{tab}_video_loop") - video_pad = gr.Slider(label='Pad frames', minimum=0, maximum=24, step=1, value=1, visible=False, elem_id=f"{tab}_video_pad") - video_interpolate = gr.Slider(label='Interpolate frames', minimum=0, maximum=24, step=1, value=0, visible=False, elem_id=f"{tab}_video_interpolate") + video_duration = gr.Slider(label='Duration', minimum=0.25, maximum=300, step=0.25, value=2, visible=show_always, elem_id=f"{tab}_video_duration") + video_loop = gr.Checkbox(label='Loop', value=True, visible=show_always, elem_id=f"{tab}_video_loop") + video_pad = gr.Slider(label='Pad frames', minimum=0, maximum=24, step=1, value=1, visible=show_always, elem_id=f"{tab}_video_pad") + video_interpolate = gr.Slider(label='Interpolate frames', minimum=0, maximum=24, step=1, value=0, visible=show_always, elem_id=f"{tab}_video_interpolate") video_type.change(fn=video_type_change, inputs=[video_type], outputs=[video_duration, video_loop, video_pad, video_interpolate]) return video_type, video_duration, video_loop, video_pad, video_interpolate @@ -218,7 +219,7 @@ def create_sampler_and_steps_selection(choices, tabname): sd_samplers.set_samplers() choices = [x for x in sd_samplers.samplers if not x.name == 'Same as primary'] with gr.Row(elem_classes=['flex-break']): - steps = gr.Slider(minimum=1, maximum=99, step=1, label="Steps", elem_id=f"{tabname}_steps", value=20) + steps = gr.Slider(minimum=1, maximum=100, step=1, label="Steps", elem_id=f"{tabname}_steps", value=20) sampler_index = gr.Dropdown(label='Sampling method', elem_id=f"{tabname}_sampling", choices=[x.name for x in choices], value='Default', type="index") return steps, sampler_index @@ -273,6 +274,13 @@ def create_sampler_options(tabname): shared.opts.schedulers_shift = sampler_shift shared.opts.save(shared.config_filename, silent=True) + def set_sigma_ajust(val, start, end): + shared.log.debug(f'Sampler set options: sigma={val} min={start} max={end}') + shared.opts.schedulers_sigma_adjust = val + shared.opts.schedulers_sigma_adjust_min = start + shared.opts.schedulers_sigma_adjust_max = end + shared.opts.save(shared.config_filename, silent=True) + # 'linear', 'scaled_linear', 'squaredcos_cap_v2' def set_sampler_preset(preset): if preset == 'AYS SD15': @@ -304,6 +312,10 @@ def create_sampler_options(tabname): with gr.Row(elem_classes=['flex-break']): sampler_presets = gr.Dropdown(label='Timesteps presets', elem_id=f"{tabname}_sampler_presets", choices=['None', 'AYS SD15', 'AYS SDXL'], value='None', type='value') sampler_timesteps = gr.Textbox(label='Timesteps override', elem_id=f"{tabname}_sampler_timesteps", value=shared.opts.schedulers_timesteps) + with gr.Row(elem_classes=['flex-break']): + sampler_sigma_adjust_val = gr.Slider(minimum=0.5, maximum=1.5, step=0.01, label='Sigma adjust', value=shared.opts.schedulers_sigma_adjust, elem_id=f"{tabname}_sampler_sigma_adjust") + sampler_sigma_adjust_min = gr.Slider(minimum=0.0, maximum=1.0, step=0.01, label='Adjust start', value=shared.opts.schedulers_sigma_adjust_min, elem_id=f"{tabname}_sampler_sigma_adjust_min") + sampler_sigma_adjust_max = gr.Slider(minimum=0.0, maximum=1.0, step=0.01, label='Adjust end', value=shared.opts.schedulers_sigma_adjust_max, elem_id=f"{tabname}_sampler_sigma_adjust_max") with gr.Row(elem_classes=['flex-break']): sampler_order = gr.Slider(minimum=0, maximum=5, step=1, label="Sampler order", value=shared.opts.schedulers_solver_order, elem_id=f"{tabname}_sampler_order") sampler_shift = gr.Slider(minimum=0, maximum=10, step=0.1, label="Flow shift", value=shared.opts.schedulers_shift, elem_id=f"{tabname}_sampler_shift") @@ -325,6 +337,9 @@ def create_sampler_options(tabname): sampler_order.change(fn=set_sampler_order, inputs=[sampler_order], outputs=[]) sampler_shift.change(fn=set_sampler_shift, inputs=[sampler_shift], outputs=[]) sampler_options.change(fn=set_sampler_options, inputs=[sampler_options], outputs=[]) + sampler_sigma_adjust_val.change(fn=set_sigma_ajust, inputs=[sampler_sigma_adjust_val, sampler_sigma_adjust_min, sampler_sigma_adjust_max], outputs=[]) + sampler_sigma_adjust_min.change(fn=set_sigma_ajust, inputs=[sampler_sigma_adjust_val, sampler_sigma_adjust_min, sampler_sigma_adjust_max], outputs=[]) + sampler_sigma_adjust_max.change(fn=set_sigma_ajust, inputs=[sampler_sigma_adjust_val, sampler_sigma_adjust_min, sampler_sigma_adjust_max], outputs=[]) def create_hires_inputs(tab): diff --git a/modules/ui_settings.py b/modules/ui_settings.py new file mode 100644 index 000000000..bda278a81 --- /dev/null +++ b/modules/ui_settings.py @@ -0,0 +1,375 @@ +import os +import gradio as gr +from modules import timer, shared, paths, theme, sd_models, modelloader, ui_common, ui_loadsave, generation_parameters_copypaste, call_queue, script_callbacks + + +text_settings = None # holds json of entire shared.opts +ui_system_tabs = None # required for system-info +dummy_component = gr.Textbox(visible=False, value='dummy') +loadsave = ui_loadsave.UiLoadsave(shared.cmd_opts.ui_config) +quicksettings_names = {x: i for i, x in enumerate(shared.opts.quicksettings_list) if x != 'quicksettings'} +quicksettings_list = [] +hidden_list = [] +components = [] + + +def apply_setting(key, value): + if value is None: + return gr.update() + if shared.cmd_opts.freeze: + return gr.update() + if key == 'sd_backend': + return gr.update() + if shared.opts.disable_weights_auto_swap and key in ['sd_model_checkpoint', 'sd_model_refiner', 'sd_model_dict', 'sd_vae', 'sd_unet', 'sd_text_encoder']: + return gr.update() + if key == "sd_model_checkpoint": + ckpt_info = sd_models.get_closet_checkpoint_match(value) + if ckpt_info is not None: + value = ckpt_info.title + else: + return gr.update() + comp_args = shared.opts.data_labels[key].component_args + if comp_args and isinstance(comp_args, dict) and comp_args.get('visible') is False: + return gr.update() + valtype = type(shared.opts.data_labels[key].default) + oldval = shared.opts.data.get(key, None) + shared.opts.data[key] = valtype(value) if valtype != type(None) else value + if oldval != value and shared.opts.data_labels[key].onchange is not None: + shared.opts.data_labels[key].onchange() + shared.opts.save(shared.config_filename) + return getattr(shared.opts, key) + + +def get_value_for_setting(key): + value = getattr(shared.opts, key) + info = shared.opts.data_labels[key] + args = info.component_args() if callable(info.component_args) else info.component_args or {} + args = {k: v for k, v in args.items() if k not in {'precision', 'multiselect', 'visible'}} + return gr.update(value=value, **args) + + +def ordered_ui_categories(): + return ['dimensions', 'sampler', 'seed', 'denoising', 'cfg', 'checkboxes', 'accordions', 'override_settings', 'scripts'] # a1111 compatibility item, not implemented + + +def create_setting_component(key, is_quicksettings=False): + def fun(): + return shared.opts.data[key] if key in shared.opts.data else shared.opts.data_labels[key].default + + info = shared.opts.data_labels[key] + t = type(info.default) + args = (info.component_args() if callable(info.component_args) else info.component_args) or {} + if info.component is not None: + comp = info.component + elif t == str: + comp = gr.Textbox + elif t == int: + comp = gr.Number + elif t == bool: + comp = gr.Checkbox + else: + raise ValueError(f'bad options item type: {t} for key {key}') + elem_id = f"setting_{key}" + dirty_indicator = None + + if not is_quicksettings: + dirtyable_setting = gr.Group(elem_classes="dirtyable", visible=args.get("visible", True)) + dirtyable_setting.__enter__() + dirty_indicator = gr.Button("", elem_classes="modification-indicator", elem_id=f"modification_indicator_{key}") + + if info.refresh is not None: + if is_quicksettings: + res = comp(label=info.label, value=fun(), elem_id=elem_id, **args) + ui_common.create_refresh_button(res, info.refresh, info.component_args, f"refresh_{key}") + else: + with gr.Row(): + res = comp(label=info.label, value=fun(), elem_id=elem_id, **args) + ui_common.create_refresh_button(res, info.refresh, info.component_args, f"refresh_{key}") + elif info.folder is not None: + with gr.Row(): + res = comp(label=info.label, value=fun(), elem_id=elem_id, elem_classes="folder-selector", **args) + # ui_common.create_browse_button(res, f"folder_{key}") + else: + try: + res = comp(label=info.label, value=fun(), elem_id=elem_id, **args) + except Exception as e: + shared.log.error(f'Error creating setting: {key} {e}') + res = None + + if res is not None and not is_quicksettings: + res.change(fn=None, inputs=res, _js=f'(val) => markIfModified("{key}", val)') + if dirty_indicator is not None: + dirty_indicator.click(fn=lambda: shared.opts.get_default(key), outputs=[res], show_progress=False) + dirtyable_setting.__exit__() + + return res + +def create_dirty_indicator(key, keys_to_reset, **kwargs): + def get_default_values(): + values = [shared.opts.get_default(key) for key in keys_to_reset] + shared.log.debug(f'Settings restore: section={key} keys={keys_to_reset} values={values}') + return values + + elements_to_reset = [shared.settings_components[_key] for _key in keys_to_reset if shared.settings_components[_key] is not None] + indicator = gr.Button('', elem_classes="modification-indicator", elem_id=f"modification_indicator_{key}", **kwargs) + indicator.click(fn=get_default_values, outputs=elements_to_reset, show_progress=True) + return indicator + + +def run_settings(*args): + changed = [] + for key, value, comp in zip(shared.opts.data_labels.keys(), args, components): + if comp == dummy_component or value=='dummy': # or getattr(comp, 'visible', True) is False or key in hidden_list: + continue + if not shared.opts.same_type(value, shared.opts.data_labels[key].default): + shared.log.error(f'Setting bad value: {key}={value} expecting={type(shared.opts.data_labels[key].default).__name__}') + continue + if shared.opts.set(key, value): + changed.append(key) + if shared.opts.cuda_compile_backend == "olive-ai": + from modules.onnx_impl import install_olive + install_olive() + if shared.cmd_opts.use_directml: + from modules.dml import directml_override_opts + directml_override_opts() + if shared.cmd_opts.use_openvino: + if "Model" not in shared.opts.cuda_compile: + shared.log.warning("OpenVINO: Enabling Torch Compile Model") + shared.opts.cuda_compile.append("Model") + if shared.opts.cuda_compile_backend != "openvino_fx": + shared.log.warning("OpenVINO: Setting Torch Compiler backend to OpenVINO FX") + shared.opts.cuda_compile_backend = "openvino_fx" + if shared.opts.sd_backend != "diffusers": + shared.log.warning("OpenVINO: Setting backend to Diffusers") + shared.opts.sd_backend = "diffusers" + try: + if len(changed) > 0: + shared.opts.save(shared.config_filename) + shared.log.info(f'Settings: changed={len(changed)} {changed}') + except RuntimeError: + shared.log.error(f'Settings failed: change={len(changed)} {changed}') + return shared.opts.dumpjson(), f'{len(changed)} Settings changed without save: {", ".join(changed)}' + return shared.opts.dumpjson(), f'{len(changed)} Settings changed{": " if len(changed) > 0 else ""}{", ".join(changed)}' + +def run_settings_single(value, key, progress=False): + if not shared.opts.same_type(value, shared.opts.data_labels[key].default): + return gr.update(visible=True), shared.opts.dumpjson() + if not shared.opts.set(key, value): + return gr.update(value=getattr(shared.opts, key)), shared.opts.dumpjson() + if key == "cuda_compile_backend" and value == "olive-ai": + from modules.onnx_impl import install_olive + install_olive() + if shared.cmd_opts.use_directml: + from modules.dml import directml_override_opts + directml_override_opts() + shared.opts.save(shared.config_filename) + shared.log.debug(f'Setting changed: {key}={value} progress={progress}') + return get_value_for_setting(key), shared.opts.dumpjson() + + +def create_ui(): + global text_settings # pylint: disable=global-statement + text_settings = gr.Textbox(elem_id="settings_json", elem_classes=["settings_json"], value=lambda: shared.opts.dumpjson(), visible=False) + with gr.Row(elem_id="system_row"): + restart_submit = gr.Button(value="Restart server", variant='primary', elem_id="restart_submit") + shutdown_submit = gr.Button(value="Shutdown server", variant='primary', elem_id="shutdown_submit") + unload_sd_model = gr.Button(value='Unload model', variant='primary', elem_id="sett_unload_sd_model") + reload_sd_model = gr.Button(value='Reload model', variant='primary', elem_id="sett_reload_sd_model") + enable_profiling = gr.Button(value='Start profiling', variant='primary', elem_id="enable_profiling") + + with gr.Tabs(elem_id="system") as system_tabs: + global ui_system_tabs # pylint: disable=global-statement + ui_system_tabs = system_tabs + with gr.TabItem("Settings", id="system_settings", elem_id="tab_settings"): + with gr.Row(elem_id="settings_row"): + settings_submit = gr.Button(value="Apply settings", variant='primary', elem_id="settings_submit") + preview_theme = gr.Button(value="Preview theme", variant='primary', elem_id="settings_preview_theme") + defaults_submit = gr.Button(value="Restore defaults", variant='primary', elem_id="defaults_submit") + with gr.Row(): + _settings_search = gr.Text(label="Search", elem_id="settings_search") + + result = gr.HTML(elem_id="settings_result") + script_callbacks.ui_settings_callback() # let extensions create settings + sections = [] + for item in shared.opts.data_labels.values(): # get unique sections from all items + if len(item.section) == 2: + section_id, section_text = item.section + elif len(item.section) == 3: # compatibility item with a1111 extensions + _category, section_id, section_text = item.section + item.section = section_id, section_text + else: + section_id = None + item.section = None, 'Hidden' + if (section_id, section_text) not in sections: + sections.append((section_id, section_text)) + + shared.log.debug(f'UI settings: sections={len(sections)} settings={len(list(shared.opts.data_labels))}') + with gr.Tabs(elem_id="settings"): + quicksettings_list.clear() + for (section_id, section_text) in sections: + items = [item for item in shared.opts.data_labels.items() if item[1].section[0] == section_id] # find all items in this section + hidden = section_id is None or 'hidden' in section_id.lower() or 'hidden' in section_text.lower() + # shared.log.trace(f'Settings: section="{section_id}" title="{section_text}" items={len(items)} hidden={hidden}') + if hidden: + for (key, _item) in items: + hidden_list.append(key) + components.append(dummy_component) + else: + with gr.TabItem(elem_id=f"settings_section_tab_{section_id}", label=section_text): + current_items = [] + for (key, item) in items: + if key in quicksettings_names: + quicksettings_list.append((key, item)) + components.append(dummy_component) + else: + with gr.Row(elem_id=f"settings_section_row_{section_id}"): # only so we can add dirty indicator at the start of the row + component = create_setting_component(key) + shared.settings_components[key] = component + current_items.append(key) + components.append(component) + create_dirty_indicator(section_id, current_items) + + with gr.TabItem("Show all pages", elem_id="settings_show_all_pages"): + create_dirty_indicator("show_all_pages", []) + request_notifications = gr.Button(value='Request browser notifications', elem_id="request_notifications", visible=False) + + + with gr.TabItem("Update", id="system_update", elem_id="tab_update"): + from modules import update + update.create_ui() + + with gr.TabItem("User interface", id="system_config", elem_id="tab_config"): + loadsave.create_ui() + create_dirty_indicator("tab_defaults", [], interactive=False) + + with gr.TabItem("ONNX", id="onnx_config", elem_id="tab_onnx"): + from modules.onnx_impl import ui as ui_onnx + ui_onnx.create_ui() + + def unload_sd_weights(): + sd_models.unload_model_weights(op='model') + sd_models.unload_model_weights(op='refiner') + + def reload_sd_weights(): + sd_models.reload_model_weights(force=True) + + def switch_profiling(): + shared.cmd_opts.profile = not shared.cmd_opts.profile + shared.log.warning(f'Profiling: {shared.cmd_opts.profile}') + return 'Stop profiling' if shared.cmd_opts.profile else 'Start profiling' + + unload_sd_model.click(fn=unload_sd_weights, inputs=[], outputs=[]) + reload_sd_model.click(fn=reload_sd_weights, inputs=[], outputs=[]) + enable_profiling.click(fn=switch_profiling, inputs=[], outputs=[enable_profiling]) + request_notifications.click(fn=lambda: None, inputs=[], outputs=[], _js='function(){}') + preview_theme.click(fn=None, _js='previewTheme', inputs=[], outputs=[]) + settings_submit.click( + fn=call_queue.wrap_gradio_call(run_settings, extra_outputs=[gr.update()]), + inputs=components, + outputs=[text_settings, result], + ) + defaults_submit.click(fn=lambda: shared.restore_defaults(restart=True), _js="restartReload") + restart_submit.click(fn=lambda: shared.restart_server(restart=True), _js="restartReload") + shutdown_submit.click(fn=lambda: shared.restart_server(restart=False), _js="restartReload") + + +def create_quicksettings(interfaces): + shared.tab_names = [] + for _interface, label, _ifid in interfaces: + shared.tab_names.append(label) + + with gr.Blocks(theme=theme.gradio_theme, analytics_enabled=False, title="SD.Next") as ui_app: + with gr.Row(elem_id="quicksettings", variant="compact"): + for k, _item in sorted(quicksettings_list, key=lambda x: quicksettings_names.get(x[1], x[0])): + component = create_setting_component(k, is_quicksettings=True) + shared.settings_components[k] = component + + generation_parameters_copypaste.connect_paste_params_buttons() + + with gr.Tabs(elem_id="tabs") as tabs: + for interface, label, ifid in interfaces: + if interface is None: + continue + with gr.TabItem(label, id=ifid, elem_id=f"tab_{ifid}"): + interface.render() + for interface, _label, ifid in interfaces: + if interface is None: + continue + if ifid in ["extensions", "system"]: + continue + loadsave.add_block(interface, ifid) + loadsave.add_component(f"webui/Tabs@{tabs.elem_id}", tabs) + loadsave.setup_ui() + + if shared.opts.notification_audio_enable and os.path.exists(os.path.join(paths.script_path, shared.opts.notification_audio_path)): + gr.Audio(interactive=False, value=os.path.join(paths.script_path, shared.opts.notification_audio_path), elem_id="audio_notification", visible=False) + + for k, _item in quicksettings_list: + component = shared.settings_components[k] + info = shared.opts.data_labels[k] + if isinstance(component, gr.components.Textbox): + change_handlers = [component.blur, component.submit] + else: + change_handlers = [component.release if hasattr(component, 'release') else component.change] + for change_handler in change_handlers: + change_handler( + fn=lambda value, k=k, progress=info.refresh is not None: run_settings_single(value, key=k, progress=progress), + inputs=[component], + outputs=[component, text_settings], + show_progress=info.refresh is not None, + ) + + button_set_checkpoint = gr.Button('Change model', elem_id='change_checkpoint', visible=False) + button_set_checkpoint.click( + fn=lambda value, _: run_settings_single(value, key='sd_model_checkpoint'), + _js="function(v){ var res = desiredCheckpointName; desiredCheckpointName = ''; return [res || v, null]; }", + inputs=[shared.settings_components['sd_model_checkpoint'], dummy_component], + outputs=[shared.settings_components['sd_model_checkpoint'], text_settings], + ) + button_set_refiner = gr.Button('Change refiner', elem_id='change_refiner', visible=False) + button_set_refiner.click( + fn=lambda value, _: run_settings_single(value, key='sd_model_checkpoint'), + _js="function(v){ var res = desiredCheckpointName; desiredCheckpointName = ''; return [res || v, null]; }", + inputs=[shared.settings_components['sd_model_refiner'], dummy_component], + outputs=[shared.settings_components['sd_model_refiner'], text_settings], + ) + button_set_vae = gr.Button('Change VAE', elem_id='change_vae', visible=False) + button_set_vae.click( + fn=lambda value, _: run_settings_single(value, key='sd_vae'), + _js="function(v){ var res = desiredVAEName; desiredVAEName = ''; return [res || v, null]; }", + inputs=[shared.settings_components['sd_vae'], dummy_component], + outputs=[shared.settings_components['sd_vae'], text_settings], + ) + + def reference_submit(model): + if '@' not in model: # diffusers + loaded = modelloader.load_reference(model) + return model if loaded else shared.opts.sd_model_checkpoint + else: # civitai + model, url = model.split('@') + loaded = modelloader.load_civitai(model, url) + return loaded if loaded is not None else shared.opts.sd_model_checkpoint + + button_set_reference = gr.Button('Change reference', elem_id='change_reference', visible=False) + button_set_reference.click( + fn=reference_submit, + _js="function(v){ return desiredCheckpointName; }", + inputs=[shared.settings_components['sd_model_checkpoint']], + outputs=[shared.settings_components['sd_model_checkpoint']], + ) + component_keys = [k for k in shared.opts.data_labels.keys() if k in shared.settings_components] + + def get_settings_values(): + return [get_value_for_setting(key) for key in component_keys] + + ui_app.load( + fn=get_settings_values, + inputs=[], + outputs=[shared.settings_components[k] for k in component_keys if shared.settings_components[k] is not None], + queue=False, + ) + + timer.startup.record("ui-defaults") + loadsave.dump_defaults() + ui_app.ui_loadsave = loadsave + return ui_app diff --git a/modules/ui_symbols.py b/modules/ui_symbols.py index c92ddb8f6..ef426e7e7 100644 --- a/modules/ui_symbols.py +++ b/modules/ui_symbols.py @@ -20,6 +20,7 @@ reuse = '♻️' info = 'ℹ' # noqa reset = '🔄' upload = '⬆️' +loading = '↺' reuse = '⬅️' search = '🔍' preview = '🖼️' diff --git a/modules/ui_txt2img.py b/modules/ui_txt2img.py index 30f9ab5ad..63d47920d 100644 --- a/modules/ui_txt2img.py +++ b/modules/ui_txt2img.py @@ -1,6 +1,6 @@ import gradio as gr from modules.call_queue import wrap_gradio_gpu_call, wrap_queued_call -from modules import timer, shared, ui_common, ui_sections, generation_parameters_copypaste, processing, processing_vae, devices +from modules import timer, shared, ui_common, ui_sections, generation_parameters_copypaste, processing, processing_vae, devices, images from modules.ui_components import ToolButton # pylint: disable=unused-import @@ -23,7 +23,7 @@ def create_ui(): txt2img_prompt, txt2img_prompt_styles, txt2img_negative_prompt, txt2img_submit, txt2img_reprocess, txt2img_paste, txt2img_extra_networks_button, txt2img_token_counter, txt2img_token_button, txt2img_negative_token_counter, txt2img_negative_token_button = ui_sections.create_toprow(is_img2img=False, id_part="txt2img") txt_prompt_img = gr.File(label="", elem_id="txt2img_prompt_image", file_count="single", type="binary", visible=False) - txt_prompt_img.change(fn=modules.images.image_data, inputs=[txt_prompt_img], outputs=[txt2img_prompt, txt_prompt_img]) + txt_prompt_img.change(fn=images.image_data, inputs=[txt_prompt_img], outputs=[txt2img_prompt, txt_prompt_img]) with gr.Row(variant='compact', elem_id="txt2img_extra_networks", elem_classes=["extra_networks_root"], visible=False) as extra_networks_ui: from modules import ui_extra_networks @@ -116,6 +116,7 @@ def create_ui(): (subseed, "Variation seed"), (subseed_strength, "Variation strength"), # advanced + (cfg_scale, "Guidance scale"), (cfg_scale, "CFG scale"), (cfg_end, "CFG end"), (clip_skip, "Clip skip"), diff --git a/modules/ui_video.py b/modules/ui_video.py new file mode 100644 index 000000000..b53810d4b --- /dev/null +++ b/modules/ui_video.py @@ -0,0 +1,185 @@ +import os +import gradio as gr +from modules import shared, sd_models, timer, images, ui_common, ui_sections, ui_symbols, call_queue, generation_parameters_copypaste +from modules.ui_components import ToolButton +from modules.video_models import models_def, video_utils + + +debug = shared.log.trace if os.environ.get('SD_VIDEO_DEBUG', None) is not None else lambda *args, **kwargs: None + + +def engine_change(engine): + debug(f'Video change: engine="{engine}"') + found = [model.name for model in models_def.models.get(engine, [])] + return gr.update(choices=found, value=found[0] if len(found) > 0 else None) + + +def get_selected(engine, model): + found = [model.name for model in models_def.models.get(engine, [])] + if len(models_def.models[engine]) > 0 and len(found) > 0: + selected = [m for m in models_def.models[engine] if m.name == model][0] + return selected + return None + + +def model_change(engine, model): + debug(f'Video change: engine="{engine}" model="{model}"') + found = [model.name for model in models_def.models.get(engine, [])] + selected = [m for m in models_def.models[engine] if m.name == model][0] if len(found) > 0 else None + return video_utils.get_url(selected.url if selected else None) + + +def model_load(engine, model): + debug(f'Video load: engine="{engine}" model="{model}"') + selected = get_selected(engine, model) + yield f'Video model loading: {selected.name}' + if selected: + if 'None' in selected.name: + sd_models.unload_model_weights() + msg = 'Video model unloaded' + else: + from modules.video_models import video_load + msg = video_load.load_model(selected) + else: + sd_models.unload_model_weights() + msg = 'Video model unloaded' + yield msg + return msg + + +def run_video(*args): + engine, model = args[2], args[3] + debug(f'Video run: engine="{engine}" model="{model}"') + selected = get_selected(engine, model) + if not selected or engine is None or model is None or engine == 'None' or model == 'None': + return video_utils.queue_err('model not selected') + debug(f'Video run: {str(selected)}') + from modules.video_models import video_run + if selected and 'Hunyuan' in selected.name: + return video_run.generate(*args) + elif selected and 'LTX' in selected.name: + return video_run.generate(*args) + elif selected and 'Mochi' in selected.name: + return video_run.generate(*args) + elif selected and 'Cog' in selected.name: + return video_run.generate(*args) + elif selected and 'Allegro' in selected.name: + return video_run.generate(*args) + elif selected and 'WAN' in selected.name: + return video_run.generate(*args) + elif selected and 'Latte' in selected.name: + return video_run.generate(*args) + return video_utils.queue_err(f'model not found: engine="{engine}" model="{model}"') + + +def create_ui(): + shared.log.debug('UI initialize: video') + with gr.Blocks(analytics_enabled=False) as _video_interface: + prompt, styles, negative, generate, _reprocess, paste, networks_button, _token_counter, _token_button, _token_counter_negative, _token_button_negative = ui_sections.create_toprow(is_img2img=False, id_part="video", negative_visible=True, reprocess_visible=False) + prompt_image = gr.File(label="", elem_id="video_prompt_image", file_count="single", type="binary", visible=False) + prompt_image.change(fn=images.image_data, inputs=[prompt_image], outputs=[prompt, prompt_image]) + + with gr.Row(variant='compact', elem_id="video_extra_networks", elem_classes=["extra_networks_root"], visible=False) as extra_networks_ui: + from modules import ui_extra_networks + extra_networks_ui = ui_extra_networks.create_ui(extra_networks_ui, networks_button, 'video', skip_indexing=shared.opts.extra_network_skip_indexing) + timer.startup.record('ui-networks') + + with gr.Row(elem_id="video_interface", equal_height=False): + with gr.Column(variant='compact', elem_id="video_settings", scale=1): + + with gr.Row(): + engine = gr.Dropdown(label='Engine', choices=list(models_def.models), value='None', elem_id="video_engine") + model = gr.Dropdown(label='Model', choices=[''], value=None, elem_id="video_model") + btn_load = ToolButton(ui_symbols.loading, elem_id="video_model_load", label='Load model') + with gr.Row(): + url = gr.HTML(label='Model URL', elem_id='video_model_url', value='

') + with gr.Accordion(open=True, label="Size", elem_id='video_size_accordion'): + with gr.Row(): + width, height = ui_sections.create_resolution_inputs('video', default_width=832, default_height=480) + with gr.Row(): + frames = gr.Slider(label='Frames', minimum=1, maximum=1024, step=1, value=15, elem_id="video_frames") + seed = gr.Number(label='Initial seed', value=-1, elem_id="video_seed", container=True) + random_seed = ToolButton(ui_symbols.random, elem_id="video_random_seed", label='Random seed') + reuse_seed = ToolButton(ui_symbols.reuse, elem_id="video_reuse_seed", label='Reuse seed') + with gr.Accordion(open=True, label="Parameters", elem_id='video_parameters_accordion'): + steps, sampler_index = ui_sections.create_sampler_and_steps_selection(None, "video") + with gr.Row(): + sampler_shift = gr.Slider(label='Sampler shift', minimum=-1.0, maximum=20.0, step=0.1, value=-1.0, elem_id="video_scheduler_shift") + dynamic_shift = gr.Checkbox(label='Dynamic shift', value=False, elem_id="video_dynamic_shift") + with gr.Row(): + guidance_scale = gr.Slider(label='Guidance scale', minimum=-1.0, maximum=14.0, step=0.1, value=-1.0, elem_id="video_guidance_scale") + guidance_true = gr.Slider(label='True guidance', minimum=-1.0, maximum=14.0, step=0.1, value=-1.0, elem_id="video_guidance_true") + with gr.Accordion(open=True, label="Decode", elem_id='video_decode_accordion'): + with gr.Row(): + vae_type = gr.Dropdown(label='VAE decode', choices=['Default', 'Tiny', 'Remote'], value='Default', elem_id="video_vae_type") + vae_tile_frames = gr.Slider(label='Tile frames', minimum=1, maximum=64, step=1, value=16, elem_id="video_vae_tile_frames") + with gr.Accordion(open=False, label="Init image", elem_id='video_init_accordion'): + gr.HTML("
  Init image") + init_image = gr.Image(elem_id="video_image", show_label=False, type="pil", image_mode="RGB", height=512) + init_strength = gr.Slider(label='Init strength', minimum=0.0, maximum=1.0, step=0.01, value=0.5, elem_id="video_denoising_strength") + with gr.Accordion(open=True, label="Output", elem_id='video_output_accordion'): + with gr.Row(): + save_frames = gr.Checkbox(label='Save image frames', value=False, elem_id="video_save_frames") + with gr.Row(): + video_type, video_duration, video_loop, video_pad, video_interpolate = ui_sections.create_video_inputs(tab='video', show_always=True) + override_settings = ui_common.create_override_inputs('video') + + # output panel with gallery and video tabs + with gr.Column(elem_id='video-output-column', scale=2) as _column_output: + with gr.Tabs(elem_classes=['video-output-tabs'], elem_id='video-output-tabs'): + with gr.Tab('Frames', id='out-gallery'): + gallery, gen_info, html_info, _html_info_formatted, html_log = ui_common.create_output_panel("video", prompt=prompt, preview=False, transfer=False, scale=2) + with gr.Tab('Video', id='out-video'): + video = gr.Video(label="Output", show_label=False, elem_id='control_output_video', elem_classes=['control-image'], height=512, autoplay=False) + + # connect reuse seed button + ui_common.connect_reuse_seed(seed, reuse_seed, gen_info, is_subseed=False) + random_seed.click(fn=lambda: -1, show_progress=False, inputs=[], outputs=[seed]) + # handle engine and model change + engine.change(fn=engine_change, inputs=[engine], outputs=[model]) + model.change(fn=model_change, inputs=[engine, model], outputs=[url]) + btn_load.click(fn=model_load, inputs=[engine, model], outputs=[html_log]) + # setup extra networks + ui_extra_networks.setup_ui(extra_networks_ui, gallery) + + # handle restore fields + paste_fields = [ + (prompt, "Prompt"), + # main + (width, "Size-1"), + (height, "Size-2"), + (frames, "Frames"), + ] + generation_parameters_copypaste.add_paste_fields("video", None, paste_fields, override_settings) + bindings = generation_parameters_copypaste.ParamBinding(paste_button=paste, tabname="video", source_text_component=prompt, source_image_component=None) + generation_parameters_copypaste.register_paste_params_button(bindings) + # hidden fields + task_id = gr.Textbox(visible=False, value='') + ui_state = gr.Textbox(visible=False, value='') + # generate args + video_args = [ + task_id, ui_state, + engine, model, + prompt, negative, styles, + width, height, + frames, + steps, sampler_index, + sampler_shift, dynamic_shift, + seed, + guidance_scale, guidance_true, + init_image, init_strength, + vae_type, vae_tile_frames, + save_frames, + video_type, video_duration, video_loop, video_pad, video_interpolate, + override_settings, + ] + # generate function + video_dict = dict( + fn=call_queue.wrap_gradio_gpu_call(run_video, extra_outputs=[None, '', ''], name='Video'), + _js="submit_video", + inputs=video_args, + outputs=[gallery, video, gen_info, html_info, html_log], + show_progress=False, + ) + prompt.submit(**video_dict) + generate.click(**video_dict) diff --git a/modules/upscaler_simple.py b/modules/upscaler_simple.py index a28d540f5..bb342d5c8 100644 --- a/modules/upscaler_simple.py +++ b/modules/upscaler_simple.py @@ -98,21 +98,29 @@ class UpscalerAsymmetricVAE(Upscaler): super().__init__(False) self.name = "Asymmetric VAE" self.vae = None + self.selected = None self.scalers = [ - UpscalerData("Asymmetric VAE", None, self), + UpscalerData("Asymmetric VAE v1", None, self), + UpscalerData("Asymmetric VAE v2", None, self), ] def do_upscale(self, img: Image, selected_model=None): + if selected_model is None: + return img import torchvision.transforms.functional as F import diffusers from modules import shared, devices - - if self.vae is None: - self.vae = diffusers.AsymmetricAutoencoderKL.from_pretrained("Heasterian/AsymmetricAutoencoderKLUpscaler", cache_dir=shared.opts.hfcache_dir) + if self.vae is None or selected_model != self.selected: + if 'v1' in selected_model: + repo_id = 'Heasterian/AsymmetricAutoencoderKLUpscaler' + else: + repo_id = 'Heasterian/AsymmetricAutoencoderKLUpscaler_v2' + self.vae = diffusers.AsymmetricAutoencoderKL.from_pretrained(repo_id, cache_dir=shared.opts.hfcache_dir) + shared.log.debug(f'Upscaler load: vae="{repo_id}"') self.vae.requires_grad_(False) self.vae = self.vae.to(device=devices.device, dtype=devices.dtype) self.vae.eval() - img = img.resize((8 * (img.width // 8), 8 * (img.height // 8)), resample=Image.Resampling.BILINEAR).convert('RGB') + img = img.resize((8 * (img.width // 8), 8 * (img.height // 8)), resample=Image.Resampling.LANCZOS).convert('RGB') tensor = (F.pil_to_tensor(img).unsqueeze(0) / 255.0).to(device=devices.device, dtype=devices.dtype) self.vae = self.vae.to(device=devices.device) tensor = self.vae(tensor).sample @@ -141,3 +149,54 @@ class UpscalerDCC(Upscaler): upscaled = (255.0 * upscaled).astype(np.uint8) upscaled = Image.fromarray(upscaled) return upscaled + + +class UpscalerVIPS(Upscaler): + def __init__(self, dirname=None): # pylint: disable=unused-argument + super().__init__(False) + self.name = "VIPS" + self.scalers = [ + UpscalerData("VIPS Lanczos 2", None, self), + UpscalerData("VIPS Lanczos 3", None, self), + UpscalerData("VIPS Mitchell", None, self), + UpscalerData("VIPS MagicKernelSharp 2013", None, self), + UpscalerData("VIPS MagicKernelSharp 2021", None, self), + ] + + def do_upscale(self, img: Image, selected_model=None): + if selected_model is None: + return img + from installer import install + from modules.shared import log + install('pyvips') + try: + import pyvips + except Exception as e: + log.error(f"Upscaler: vips {e}") + return img + vips_image = pyvips.Image.new_from_array(img) + # import numpy as np + # np_image = np.array(img) + # h, w, c = np_image.shape + # np_linear = np_image.reshape(w * h * c) + # vips_image = pyvips.Image.new_from_memory(np_linear.data, w, h, c, 'uchar') + try: + if selected_model is None: + return img + elif selected_model == "VIPS Lanczos 2": + vips_image = vips_image.resize(2, kernel='lanczos2') + elif selected_model == "VIPS Lanczos 3": + vips_image = vips_image.resize(2, kernel='lanczos3') + elif selected_model == "VIPS Mitchell": + vips_image = vips_image.resize(2, kernel='mitchell') + elif selected_model == "VIPS MagicKernelSharp 2013": + vips_image = vips_image.resize(2, kernel='mks2013') + elif selected_model == "VIPS MagicKernelSharp 2021": + vips_image = vips_image.resize(2, kernel='mks2021') + else: + return img + except Exception as e: + log.error(f"Upscaler: vips {e}") + return img + upscaled = Image.fromarray(vips_image.numpy()) + return upscaled diff --git a/modules/video.py b/modules/video.py index d9e40a27f..bd26aac04 100644 --- a/modules/video.py +++ b/modules/video.py @@ -1,6 +1,7 @@ import os import threading import numpy as np +from PIL import Image from modules import shared, errors from modules.images_namegen import FilenameGenerator # pylint: disable=unused-import @@ -83,3 +84,25 @@ def save_video(p, images, filename = None, video_type: str = 'none', duration: f else: save_video_atomic(images, filename, video_type, duration, loop, interpolate, scale, pad, change) return filename + + +def get_video_params(filepath: str, capture: bool = False): + import cv2 + from modules.control.util import decode_fourcc + video = cv2.VideoCapture(filepath) + if not video.isOpened(): + msg = f'Video open failed: path="{filepath}"' + shared.log.error(msg) + raise RuntimeError(msg) + frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT)) + fps = round(video.get(cv2.CAP_PROP_FPS), 2) + duration = round(float(frames) / fps, 2) + w, h = int(video.get(cv2.CAP_PROP_FRAME_WIDTH)), int(video.get(cv2.CAP_PROP_FRAME_HEIGHT)) + codec = decode_fourcc(video.get(cv2.CAP_PROP_FOURCC)) + frame = None + if capture: + _status, frame = video.read() + frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) + frame = Image.fromarray(frame) + video.release() + return frames, fps, duration, w, h, codec, frame diff --git a/modules/video_models/models_def.py b/modules/video_models/models_def.py new file mode 100644 index 000000000..69287bb65 --- /dev/null +++ b/modules/video_models/models_def.py @@ -0,0 +1,198 @@ +from dataclasses import dataclass +import diffusers +import transformers + + +@dataclass +class Model(): + name: str + url: str = '' + repo: str = None + repo_cls: classmethod = None + dit: str = None + dit_cls: classmethod = None + dit_folder: str = 'transformer' + te: str = None + te_cls: classmethod = None + te_folder: str = 'text_encoder' + te_hijack: bool = True + image_hijack: bool = True + vae_hijack: bool = True + vae_remote: bool = False + + def __str__(self): + return f'name="{self.name}" url="{self.url}" repo="{self.repo}" repo_cls="{self.repo_cls}" dit="{self.dit}" dit_cls="{self.dit_cls}" dit_folder="{self.dit_folder}" te="{self.te}" te_cls="{self.te_cls}" te_folder="{self.te_folder}" te_hijack={self.te_hijack} vae_hijack={self.vae_hijack} vae_remote={self.vae_remote}' + + +models = { + 'None': [], + 'Hunyuan Video': [ + Model(name='None'), + Model(name='Hunyuan Video T2V', + url='https://huggingface.co/tencent/HunyuanVideo', + vae_remote=True, + repo='hunyuanvideo-community/HunyuanVideo', + repo_cls=diffusers.HunyuanVideoPipeline, + te_cls=transformers.LlamaModel, + dit_cls=diffusers.HunyuanVideoTransformer3DModel), + Model(name='Hunyuan Video I2V', # https://github.com/huggingface/diffusers/pull/10983 + url='https://huggingface.co/tencent/HunyuanVideo-I2V', + vae_remote=True, + repo='hunyuanvideo-community/HunyuanVideo-I2V', + repo_cls=diffusers.HunyuanVideoImageToVideoPipeline, + te_cls=transformers.LlavaForConditionalGeneration, + dit_cls=diffusers.HunyuanVideoTransformer3DModel), + Model(name='SkyReels Hunyuan T2V', # https://github.com/huggingface/diffusers/pull/10837 + url='https://huggingface.co/Skywork/SkyReels-V1-Hunyuan-T2V', + vae_remote=True, + repo='hunyuanvideo-community/HunyuanVideo', + repo_cls=diffusers.HunyuanVideoPipeline, + te_cls=transformers.LlamaModel, + dit='Skywork/SkyReels-V1-Hunyuan-T2V', + dit_folder=None, + dit_cls=diffusers.HunyuanVideoTransformer3DModel), + Model(name='SkyReels Hunyuan I2V', # https://github.com/huggingface/diffusers/pull/10837 + url='https://huggingface.co/Skywork/SkyReels-V1-Hunyuan-I2V', + vae_remote=True, + repo='hunyuanvideo-community/HunyuanVideo', + repo_cls=diffusers.HunyuanSkyreelsImageToVideoPipeline, + te_cls=transformers.LlamaModel, + dit='Skywork/SkyReels-V1-Hunyuan-I2V', + dit_folder=None, + dit_cls=diffusers.HunyuanVideoTransformer3DModel), + Model(name='Fast Hunyuan T2V', # https://github.com/hao-ai-lab/FastVideo/blob/8a77cf22c9b9e7f931f42bc4b35d21fd91d24e45/fastvideo/models/hunyuan/inference.py#L213 + url='https://huggingface.co/FastVideo/FastHunyuan', + vae_remote=True, + repo='hunyuanvideo-community/HunyuanVideo', + repo_cls=diffusers.HunyuanVideoPipeline, + te_cls=transformers.LlamaModel, + dit='FastVideo/FastHunyuan-diffusers', + dit_cls=diffusers.HunyuanVideoTransformer3DModel), + ], + 'LTX Video': [ + Model(name='None'), + 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, + te_cls=transformers.T5EncoderModel, + dit_cls=diffusers.LTXVideoTransformer3DModel), + Model(name='LTXVideo 0.9.5 I2V', + url='https://huggingface.co/Lightricks/LTX-Video-0.9.5', + repo='Lightricks/LTX-Video-0.9.5', + repo_cls=diffusers.LTXConditionPipeline, + te_cls=transformers.T5EncoderModel, + dit_cls=diffusers.LTXVideoTransformer3DModel), + Model(name='LTXVideo 0.9.1 T2V', + url='https://huggingface.co/a-r-r-o-w/LTX-Video-0.9.1-diffusers', + repo='a-r-r-o-w/LTX-Video-0.9.1-diffusers', + repo_cls=diffusers.LTXPipeline, + te_cls=transformers.T5EncoderModel, + dit_cls=diffusers.LTXVideoTransformer3DModel), + Model(name='LTXVideo 0.9.1 I2V', + url='https://huggingface.co/a-r-r-o-w/LTX-Video-0.9.1-diffusers', + repo='a-r-r-o-w/LTX-Video-0.9.1-diffusers', + repo_cls=diffusers.LTXImageToVideoPipeline, + te_cls=transformers.T5EncoderModel, + dit_cls=diffusers.LTXVideoTransformer3DModel), + Model(name='LTXVideo 0.9.0 T2V', + url='https://huggingface.co/a-r-r-o-w/LTX-Video-diffusers', + repo='a-r-r-o-w/LTX-Video-diffusers', + repo_cls=diffusers.LTXPipeline, + te_cls=transformers.T5EncoderModel, + dit_cls=diffusers.LTXVideoTransformer3DModel), + Model(name='LTXVideo 0.9.0 I2V', + url='https://huggingface.co/a-r-r-o-w/LTX-Video-diffusers', + repo='a-r-r-o-w/LTX-Video-diffusers', + repo_cls=diffusers.LTXImageToVideoPipeline, + te_cls=transformers.T5EncoderModel, + dit_cls=diffusers.LTXVideoTransformer3DModel), + ], + 'WAN Video': [ + Model(name='None'), + Model(name='WAN 2.1 1.3B T2V', + url='https://huggingface.co/Wan-AI/Wan2.1-T2V-1.3B-Diffusers', + repo='Wan-AI/Wan2.1-T2V-1.3B-Diffusers', + repo_cls=diffusers.WanPipeline, + te_cls=transformers.T5EncoderModel, + dit_cls=diffusers.WanTransformer3DModel), + Model(name='WAN 2.1 14B T2V', + url='https://huggingface.co/Wan-AI/Wan2.1-T2V-14B-Diffusers', + repo='Wan-AI/Wan2.1-T2V-14B-Diffusers', + repo_cls=diffusers.WanPipeline, + te_cls=transformers.T5EncoderModel, + dit_cls=diffusers.WanTransformer3DModel), + Model(name='WAN 2.1 14B I2V 480p', + url='https://huggingface.co/Wan-AI/Wan2.1-I2V-14B-480P-Diffusers', + repo='Wan-AI/Wan2.1-I2V-14B-480P-Diffusers', + repo_cls=diffusers.WanImageToVideoPipeline, + te_cls=transformers.T5EncoderModel, + dit_cls=diffusers.WanTransformer3DModel), + Model(name='WAN 2.1 14B I2V 720p', + url='https://huggingface.co/Wan-AI/Wan2.1-I2V-14B-720P-Diffusers', + repo='Wan-AI/Wan2.1-I2V-14B-720P-Diffusers', + repo_cls=diffusers.WanImageToVideoPipeline, + te_cls=transformers.T5EncoderModel, + dit_cls=diffusers.WanTransformer3DModel), + ], + 'Mochi Video': [ + Model(name='None'), + Model(name='Mochi 1 T2V', + url='https://huggingface.co/genmo/mochi-1-preview', + repo='genmo/mochi-1-preview', + repo_cls=diffusers.MochiPipeline, + te_cls=transformers.T5EncoderModel, + dit_cls=diffusers.MochiTransformer3DModel), + ], + 'Latte Video': [ + Model(name='None'), + Model(name='Latte 1 T2V', + url='https://huggingface.co/maxin-cn/Latte-1', + repo='maxin-cn/Latte-1', + repo_cls=diffusers.LattePipeline, + te_cls=transformers.T5EncoderModel, + dit_cls=diffusers.LatteTransformer3DModel), + ], + 'Allegro Video': [ + Model(name='None'), + Model(name='Allegro T2V', + url='https://huggingface.co/rhymes-ai/Allegro', + repo='rhymes-ai/Allegro', + repo_cls=diffusers.AllegroPipeline, + te_cls=transformers.T5EncoderModel, + dit_cls=diffusers.AllegroTransformer3DModel), + ], + 'Cog Video': [ + Model(name='None'), + Model(name='CogVideoX 1.0 2B T2V', + url='https://huggingface.co/THUDM/CogVideoX-2b', + repo='THUDM/CogVideoX-2b', + repo_cls=diffusers.CogVideoXPipeline, + te_cls=transformers.T5EncoderModel, + dit_cls=diffusers.CogVideoXTransformer3DModel), + Model(name='CogVideoX 1.0 5B T2V', + url='https://huggingface.co/THUDM/CogVideoX-5b', + repo='THUDM/CogVideoX-5b', + repo_cls=diffusers.CogVideoXPipeline, + te_cls=transformers.T5EncoderModel, + dit_cls=diffusers.CogVideoXTransformer3DModel), + Model(name='CogVideoX 1.0 5B I2V', + url='https://huggingface.co/THUDM/CogVideoX-5b-I2V', + repo='THUDM/CogVideoX-5b-I2V', + repo_cls=diffusers.CogVideoXImageToVideoPipeline, + te_cls=transformers.T5EncoderModel, + dit_cls=diffusers.CogVideoXTransformer3DModel), + Model(name='CogVideoX 1.5 5B T2V', + url='https://huggingface.co/THUDM/CogVideoX1.5-5B', + repo='THUDM/CogVideoX1.5-5B', + repo_cls=diffusers.CogVideoXPipeline, + te_cls=transformers.T5EncoderModel, + dit_cls=diffusers.CogVideoXTransformer3DModel), + Model(name='CogVideoX 1.5 5B I2V', + url='https://huggingface.co/THUDM/CogVideoX1.5-5B-I2V', + repo='THUDM/CogVideoX1.5-5B-I2V', + repo_cls=diffusers.CogVideoXImageToVideoPipeline, + te_cls=transformers.T5EncoderModel, + dit_cls=diffusers.CogVideoXTransformer3DModel), + ], +} diff --git a/modules/video_models/video_load.py b/modules/video_models/video_load.py new file mode 100644 index 000000000..13d0dc935 --- /dev/null +++ b/modules/video_models/video_load.py @@ -0,0 +1,90 @@ +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 + + +loaded_model = None +debug = shared.log.trace if os.environ.get('SD_VIDEO_DEBUG', None) is not None else lambda *args, **kwargs: None + + +def load_model(selected: models_def.Model): + if selected is None: + return '' + global loaded_model # pylint: disable=global-statement + if loaded_model == selected.name: + return '' + sd_models.unload_model_weights() + t0 = time.time() + + # 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)}') + text_encoder = selected.te_cls.from_pretrained( + pretrained_model_name_or_path=selected.te or selected.repo, + subfolder=selected.te_folder, + cache_dir=shared.opts.hfcache_dir, + torch_dtype=devices.dtype, + **quant_args + ) + except Exception as e: + shared.log.error(f'video load: module=te cls={selected.te_cls.__name__} {e}') + errors.display(e, 'video') + text_encoder = None + + # 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)}') + transformer = selected.dit_cls.from_pretrained( + pretrained_model_name_or_path=selected.dit or selected.repo, + subfolder=selected.dit_folder, + torch_dtype=devices.dtype, + cache_dir=shared.opts.hfcache_dir, + **quant_args + ) + except Exception as e: + shared.log.error(f'video load: module=transformer cls={selected.dit_cls.__name__} {e}') + errors.display(e, 'video') + transformer = None + + # overrides + kwargs = video_overrides.load_override(selected) + + # model + try: + debug(f'Video load: module=pipe repo="{selected.repo}" cls={selected.repo_cls.__name__}') + shared.sd_model = selected.repo_cls.from_pretrained( + pretrained_model_name_or_path=selected.repo, + transformer=transformer, + text_encoder=text_encoder, + cache_dir=shared.opts.hfcache_dir, + torch_dtype=devices.dtype, + **kwargs, + ) + except Exception as e: + shared.log.error(f'video load: module=pipe repo="{selected.repo}" cls={selected.repo_cls.__name__} {e}') + errors.display(e, 'video') + + t1 = time.time() + shared.sd_model.sd_checkpoint_info = sd_checkpoint.CheckpointInfo(selected.repo) + shared.sd_model.sd_model_hash = None + sd_models.set_diffuser_options(shared.sd_model) + if selected.vae_hijack and hasattr(shared.sd_model.vae, 'decode'): + shared.sd_model.vae.orig_decode = shared.sd_model.vae.decode + shared.sd_model.vae.decode = video_vae.hijack_vae_decode + shared.sd_model.vae.orig_encode = shared.sd_model.vae.encode + shared.sd_model.vae.encode = video_vae.hijack_vae_encode + if selected.te_hijack and hasattr(shared.sd_model, 'encode_prompt'): + shared.sd_model.orig_encode_prompt = shared.sd_model.encode_prompt + shared.sd_model.encode_prompt = video_utils.hijack_encode_prompt + if selected.image_hijack and hasattr(shared.sd_model, 'encode_image'): + shared.sd_model.orig_encode_image = shared.sd_model.encode_image + shared.sd_model.encode_image = video_utils.hijack_encode_image + if hasattr(shared.sd_model.vae, 'enable_slicing'): + shared.sd_model.vae.enable_slicing() + loaded_model = selected.name + msg = f'Video load: cls={shared.sd_model.__class__.__name__} model="{selected.name}" time={t1-t0:.2f}' + shared.log.info(msg) + return msg diff --git a/modules/video_models/video_overrides.py b/modules/video_models/video_overrides.py new file mode 100644 index 000000000..4168a1441 --- /dev/null +++ b/modules/video_models/video_overrides.py @@ -0,0 +1,45 @@ +import os +import torch +import diffusers +from modules import shared, processing +from modules.video_models.models_def import Model + + +debug = shared.log.trace if os.environ.get('SD_VIDEO_DEBUG', None) is not None else lambda *args, **kwargs: None + + +def load_override(selected: Model): + kwargs = {} + if selected.name == 'Allegro T2V': + kwargs['vae'] = diffusers.AutoencoderKLAllegro.from_pretrained(selected.repo, subfolder="vae", torch_dtype=torch.float32, cache_dir=shared.opts.hfcache_dir) + if selected.name == 'LTXVideo 0.9.5 I2V': + kwargs['vae'] = diffusers.AutoencoderKLLTXVideo.from_pretrained(selected.repo, subfolder="vae", torch_dtype=torch.float32, cache_dir=shared.opts.hfcache_dir) + if selected.name == 'WAN 2.1 14B I2V 480p' or selected.name == 'WAN 2.1 14B I2V 720p': + kwargs['vae'] = diffusers.AutoencoderKLWan.from_pretrained(selected.repo, subfolder="vae", torch_dtype=torch.float32, cache_dir=shared.opts.hfcache_dir) + debug(f'Video overrides: model="{selected.name}" kwargs={list(kwargs)}') + return kwargs + + +def set_overrides(p: processing.StableDiffusionProcessingVideo, selected: Model): + cls = shared.sd_model.__class__.__name__ + # Allegro + if selected.name == 'Allegro T2V': + shared.sd_model.vae.enable_tiling() + # Latte + if selected.name == 'Latte 1 T2V': + p.task_args['enable_temporal_attentions'] = True + p.task_args['video_length'] = 16 * (max(p.frames // 16, 1)) + # LTX + if cls == 'LTXImageToVideoPipeline' or cls == 'LTXConditionPipeline': + p.task_args['generator'] = None + if cls == 'LTXConditionPipeline': + p.task_args['strength'] = p.denoising_strength + # WAN + if 'Wan' in cls: + p.task_args['width'] = 16 * (p.width // 16) + p.task_args['height'] = 16 * (p.height // 16) + p.frames = 4 * (max(p.frames // 4, 1)) + 1 + # LTX + if 'LTX' in cls: + p.task_args['width'] = 32 * (p.width // 32) + p.task_args['height'] = 32 * (p.height // 32) diff --git a/modules/video_models/video_run.py b/modules/video_models/video_run.py new file mode 100644 index 000000000..2b4fbcbea --- /dev/null +++ b/modules/video_models/video_run.py @@ -0,0 +1,105 @@ +import os +import time +from modules import shared, errors, sd_models, processing, devices, images, ui_common +from modules.video_models import models_def, video_utils, video_load, video_vae, video_overrides + + +debug = shared.log.trace if os.environ.get('SD_VIDEO_DEBUG', None) is not None else lambda *args, **kwargs: None + + +def generate(*args, **kwargs): + task_id, ui_state, engine, model, prompt, negative, styles, width, height, frames, steps, sampler_index, sampler_shift, dynamic_shift, seed, guidance_scale, guidance_true, init_image, init_strength, vae_type, vae_tile_frames, save_frames, video_type, video_duration, video_loop, video_pad, video_interpolate, override_settings = args + if engine is None or model is None or engine == 'None' or model == 'None': + return video_utils.queue_err('model not selected') + found = [model.name for model in models_def.models.get(engine, [])] + selected: models_def.Model = [m for m in models_def.models[engine] if m.name == model][0] if len(found) > 0 else None + if not shared.sd_loaded: + debug('Video: model not yet loaded') + video_load.load_model(selected) + if selected.name != video_load.loaded_model: + debug('Video: force reload') + video_load.load_model(selected) + if not shared.sd_loaded: + debug('Video: model still not loaded') + return video_utils.queue_err('model not loaded') + debug(f'Video generate: task={task_id} args={args} kwargs={kwargs}') + + p = processing.StableDiffusionProcessingVideo( + sd_model=shared.sd_model, + prompt=prompt, + negative_prompt=negative, + styles=styles, + seed=int(seed), + sampler_name = processing.get_sampler_name(sampler_index), + sampler_shift=float(sampler_shift), + steps=int(steps), + width=16 * int(width // 16), + height=16 * int(height // 16), + frames=int(frames), + denoising_strength=float(init_strength), + init_image=init_image, + cfg_scale=float(guidance_scale), + pag_scale=float(guidance_true), + vae_type=vae_type, + vae_tile_frames=int(vae_tile_frames), + override_settings=override_settings, + ) + if p.vae_type == 'Remote' and not selected.vae_remote: + shared.log.warning(f'Video: model={selected.name} remote vae not supported') + p.vae_type = 'Default' + p.scripts = None + p.script_args = None + p.state = ui_state + p.do_not_save_grid = True + p.do_not_save_samples = not save_frames + if 'I2V' in model: + 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') + + # cleanup memory + shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model) + devices.torch_gc(force=True) + + # set args + processing.fix_seed(p) + video_vae.set_vae_params(p) + video_utils.set_prompt(p) + p.task_args['num_inference_steps'] = p.steps + p.task_args['width'] = p.width + p.task_args['height'] = p.height + p.task_args['output_type'] = 'latent' if (p.vae_type == 'Remote') else 'pil' + p.ops.append('video') + orig_dynamic_shift = shared.opts.schedulers_dynamic_shift + orig_sampler_shift = shared.opts.schedulers_shift + shared.opts.data['schedulers_dynamic_shift'] = dynamic_shift + shared.opts.data['schedulers_shift'] = sampler_shift + video_overrides.set_overrides(p, selected) + debug(f'Video: task_args={p.task_args}') + + # run processing + shared.state.disable_preview = True + shared.log.debug(f'Video: cls={shared.sd_model.__class__.__name__} width={p.width} height={p.height} frames={p.frames} steps={p.steps}') + err = None + t0 = time.time() + processed = None + try: + processed = processing.process_images(p) + except Exception as e: + err = str(e) + errors.display(e, 'video') + t1 = time.time() + shared.state.disable_preview = False + shared.opts.data['schedulers_dynamic_shift'] = orig_dynamic_shift + shared.opts.data['schedulers_shift'] = orig_sampler_shift + p.close() + + # done + if err: + return video_utils.queue_err(err) + if processed is None or len(processed.images) == 0: + return video_utils.queue_err('processing failed') + shared.log.info(f'Video: name="{selected.name}" cls={shared.sd_model.__class__.__name__} frames={len(processed.images)} time={t1-t0:.2f}') + video_file = images.save_video(p, filename=None, images=processed.images, video_type=video_type, duration=video_duration, loop=video_loop, pad=video_pad, interpolate=video_interpolate) + generation_info_js = processed.js() if processed is not None else '' + return processed.images, video_file, generation_info_js, processed.info, ui_common.plaintext_to_html(processed.comments) diff --git a/modules/video_models/video_utils.py b/modules/video_models/video_utils.py new file mode 100644 index 000000000..d0a80d42a --- /dev/null +++ b/modules/video_models/video_utils.py @@ -0,0 +1,62 @@ +import os +import time +from modules import shared, sd_models, timer, errors, devices + + +debug = shared.log.trace if os.environ.get('SD_VIDEO_DEBUG', None) is not None else lambda *args, **kwargs: None + + +def queue_err(msg): + shared.log.error(f'Video: {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 '

' + + +def set_prompt(p): + p.prompt = shared.prompt_styles.apply_styles_to_prompt(p.prompt, p.styles) + p.negative_prompt = shared.prompt_styles.apply_negative_styles_to_prompt(p.negative_prompt, p.styles) + shared.prompt_styles.apply_styles_to_extra(p) + p.styles = [] + p.task_args['prompt'] = p.prompt + p.task_args['negative_prompt'] = p.negative_prompt + + +def hijack_encode_prompt(*args, **kwargs): + t0 = time.time() + try: + sd_models.move_model(shared.sd_model.text_encoder, devices.device) + res = shared.sd_model.orig_encode_prompt(*args, **kwargs) + except Exception as e: + shared.log.error(f'Video encode prompt: {e}') + errors.display(e, 'Video encode prompt') + res = None + t1 = time.time() + timer.process.add('te', t1-t0) + debug(f'Video encode prompt: 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 hijack_encode_image(*args, **kwargs): + t0 = time.time() + try: + sd_models.move_model(shared.sd_model.image_encoder, devices.device) + res = shared.sd_model.orig_encode_image(*args, **kwargs) + except Exception as e: + shared.log.error(f'Video encode image: {e}') + errors.display(e, 'Video encode image') + res = None + t1 = time.time() + timer.process.add('te', t1-t0) + debug(f'Video encode image: te={shared.sd_model.image_encoder.__class__.__name__} time={t1-t0:.2f}') + shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model) + return res diff --git a/modules/video_models/video_vae.py b/modules/video_models/video_vae.py new file mode 100644 index 000000000..0dab30c72 --- /dev/null +++ b/modules/video_models/video_vae.py @@ -0,0 +1,97 @@ +import os +import time +import torch +from modules import shared, sd_models, devices, timer, errors + + +debug = shared.log.trace if os.environ.get('SD_VIDEO_DEBUG', None) is not None else lambda *args, **kwargs: None +vae_type = None + + +def set_vae_params(p): + global vae_type # pylint: disable=global-statement + vae_type = p.vae_type + if p.vae_tile_frames > p.frames: + if hasattr(shared.sd_model.vae, 'tile_sample_min_num_frames'): + shared.sd_model.vae.tile_sample_min_num_frames = p.vae_tile_frames + if hasattr(shared.sd_model.vae, 'use_framewise_decoding'): + shared.sd_model.vae.use_framewise_decoding = True + if hasattr(shared.sd_model.vae, 'enable_tiling'): + shared.sd_model.vae.enable_tiling() + else: + if hasattr(shared.sd_model.vae, 'use_framewise_decoding'): + shared.sd_model.vae.use_framewise_decoding = False + if hasattr(shared.sd_model.vae, 'disable_tiling'): + shared.sd_model.vae.disable_tiling() + + +def vae_decode_tiny(latents): + if 'Hunyuan' in shared.sd_model.__class__.__name__: + variant = 'TAE HunyuanVideo' + elif 'Mochi' in shared.sd_model.__class__.__name__: + variant = 'TAE MochiVideo' + elif 'WAN' in shared.sd_model.__class__.__name__: + variant = 'TAE WanVideo' + else: + shared.log.warning(f'Video VAE: type=Tiny cls={shared.sd_model.__class__.__name__} not supported') + return None + from modules import sd_vae_taesd + vae = sd_vae_taesd.get_model(variant) + if vae is None: + return None + debug(f'Video VAE: type=Tiny cls={vae.__class__.__name__} variant="{variant}" latents={latents.shape}') + vae = vae.to(device=devices.device, dtype=devices.dtype) + latents = latents.transpose(1, 2).to(device=devices.device, dtype=devices.dtype) + images = vae.decode_video(latents, parallel=False).transpose(1, 2).mul_(2).sub_(1) + images = images.transpose(1, 2).mul_(2).sub_(1) + return (images, None) + + +def hijack_vae_decode(*args, **kwargs): + t0 = time.time() + res = None + if vae_type == 'Tiny': + res = vae_decode_tiny(args[0]) + if vae_type == 'Remote': + pass + if res is None: + shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model, exclude=['vae']) + try: + sd_models.move_model(shared.sd_model.vae, devices.device) + if torch.is_tensor(args[0]): + latent = args[0] + latent = latent.to(device=devices.device, dtype=shared.sd_model.vae.dtype) # upcast to vae dtype + res = shared.sd_model.vae.orig_decode(latent, *args[1:], **kwargs) + else: + res = shared.sd_model.vae.orig_decode(*args, **kwargs) + except Exception as e: + shared.log.error(f'Video VAE decode: type={vae_type} {e}') + errors.display(e, 'Video VAE') + res = None + t1 = time.time() + timer.process.add('vae', t1-t0) + debug(f'Video VAE decode: type={vae_type} vae={shared.sd_model.vae.__class__.__name__} latents={args[0].shape} time={t1-t0:.2f}') + return res + + +def hijack_vae_encode(*args, **kwargs): + t0 = time.time() + res = None + if res is None: + shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model, exclude=['vae']) + try: + sd_models.move_model(shared.sd_model.vae, devices.device) + if torch.is_tensor(args[0]): + latent = args[0] + latent = latent.to(device=devices.device, dtype=shared.sd_model.vae.dtype) # upcast to vae dtype + res = shared.sd_model.vae.orig_encode(latent, *args[1:], **kwargs) + else: + res = shared.sd_model.vae.orig_encode(*args, **kwargs) + except Exception as e: + shared.log.error(f'Video VAE encode: type={vae_type} {e}') + errors.display(e, 'Video VAE') + res = None + t1 = time.time() + timer.process.add('vae', t1-t0) + debug(f'Video VAE encode: type={vae_type} vae={shared.sd_model.vae.__class__.__name__} latents={args[0].shape} time={t1-t0:.2f}') + return res diff --git a/modules/zluda.py b/modules/zluda.py index 0203a6398..431ab2c8c 100644 --- a/modules/zluda.py +++ b/modules/zluda.py @@ -3,7 +3,8 @@ from typing import Union import torch from torch._prims_common import DeviceLikeType import onnxruntime as ort -from modules import shared, devices +from modules import shared, devices, zluda_installer +from modules.zluda_installer import core, default_agent # pylint: disable=unused-import from modules.onnx_impl.execution_providers import available_execution_providers, ExecutionProvider @@ -32,7 +33,6 @@ def initialize_zluda(): from modules.zluda_hijacks import do_hijack do_hijack() - from modules import zluda_installer torch.backends.cudnn.enabled = zluda_installer.MIOpen_available if not zluda_installer.MIOpen_available: torch.backends.cuda.enable_cudnn_sdp(False) diff --git a/modules/zluda_hijacks.py b/modules/zluda_hijacks.py index 872127cf1..bbbec7a81 100644 --- a/modules/zluda_hijacks.py +++ b/modules/zluda_hijacks.py @@ -1,5 +1,21 @@ +from functools import wraps import torch -from modules import rocm +import torch._dynamo.device_interface +from modules import shared, zluda # pylint: disable=unused-import + + +MEM_BUS_WIDTH = { + "AMD Radeon RX 9070 XT": 256, + "AMD Radeon RX 9070": 256, + "AMD Radeon RX 9060 XT": 192, + "AMD Radeon RX 7900 XTX": 384, + "AMD Radeon RX 7900 XT": 320, + "AMD Radeon RX 7900 GRE": 256, + "AMD Radeon RX 7800 XT": 256, + "AMD Radeon RX 7700 XT": 192, + "AMD Radeon RX 7600 XT": 128, + "AMD Radeon RX 7600": 128, +} _topk = torch.topk @@ -9,6 +25,84 @@ def topk(input: torch.Tensor, *args, **kwargs): # pylint: disable=redefined-buil return torch.return_types.topk((values.to(device), indices.to(device),)) +class DeviceProperties: + PROPERTIES_OVERRIDE = {"regs_per_multiprocessor": 65535, "gcnArchName": "UNKNOWN ARCHITECTURE"} + internal: torch._C._CudaDeviceProperties + + def __init__(self, props: torch._C._CudaDeviceProperties): + self.internal = props + + def __getattr__(self, name): + if name in DeviceProperties.PROPERTIES_OVERRIDE: + return DeviceProperties.PROPERTIES_OVERRIDE[name] + return getattr(self.internal, name) + + +__get_device_properties = torch.cuda._get_device_properties # pylint: disable=protected-access +def torch_cuda__get_device_properties(device): + return DeviceProperties(__get_device_properties(device)) + + +_cuda_getCurrentRawStream = torch._C._cuda_getCurrentRawStream # pylint: disable=protected-access +def torch__C__cuda_getCurrentRawStream(device): + return zluda.core.to_hip_stream(_cuda_getCurrentRawStream(device)) + + def do_hijack(): - torch.version.hip = rocm.version torch.topk = topk + if zluda.default_agent is not None: + DeviceProperties.PROPERTIES_OVERRIDE["gcnArchName"] = zluda.default_agent.name + torch.cuda._get_device_properties = torch_cuda__get_device_properties # pylint: disable=protected-access + torch._C._cuda_getCurrentRawStream = torch__C__cuda_getCurrentRawStream # pylint: disable=protected-access + torch._dynamo.device_interface.CudaInterface.get_raw_stream = staticmethod(torch__C__cuda_getCurrentRawStream) # pylint: disable=protected-access + + # Triton + try: + import triton + _get_device_properties = triton.runtime.driver.active.utils.get_device_properties + def triton_runtime_driver_active_utils_get_device_properties(device): + props = _get_device_properties(device) + name = torch.cuda.get_device_name()[:-8] + if name in MEM_BUS_WIDTH: + props["mem_bus_width"] = MEM_BUS_WIDTH[name] + else: + props["mem_bus_width"] = 128 + shared.log.warning(f'[TRITON] defaulting mem_bus_width=128 for device "{name}".') + return props + triton.runtime.driver.active.utils.get_device_properties = triton_runtime_driver_active_utils_get_device_properties + + if 'Flash attention' in shared.opts.sdp_options: + from modules.flash_attn_triton_amd import interface_fa + sdpa_pre_flash_atten = torch.nn.functional.scaled_dot_product_attention + @wraps(sdpa_pre_flash_atten) + def sdpa_flash_atten(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attn_mask=None, dropout_p=0.0, is_causal=False, scale=None): + if query.shape[-1] <= 128 and attn_mask is None and query.dtype != torch.float32: + if scale is None: + scale = query.shape[-1] ** (-0.5) + head_size_og = query.size(3) + if head_size_og % 8 != 0: + query = torch.nn.functional.pad(query, [0, 8 - head_size_og % 8]) + key = torch.nn.functional.pad(key, [0, 8 - head_size_og % 8]) + value = torch.nn.functional.pad(value, [0, 8 - head_size_og % 8]) + out_padded, _, _, _ = interface_fa.fwd( + query.transpose(1, 2), + key.transpose(1, 2), + value.transpose(1, 2), + None, + None, + dropout_p, + scale, + is_causal, + -1, + -1, + 0.0, + False, + None, + ) + return out_padded[..., :head_size_og].transpose(1, 2) + else: + return sdpa_pre_flash_atten(query=query, key=key, value=value, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal, scale=scale) + torch.nn.functional.scaled_dot_product_attention = sdpa_flash_atten + shared.log.debug('Torch attention: type="triton flash attention"') + except Exception: + pass diff --git a/modules/zluda_installer.py b/modules/zluda_installer.py index dc3d85e61..54a0f9234 100644 --- a/modules/zluda_installer.py +++ b/modules/zluda_installer.py @@ -6,6 +6,7 @@ import shutil import zipfile import urllib.request from typing import Union +from installer import args, log from modules import rocm @@ -17,7 +18,6 @@ DLL_MAPPING = { 'nvrtc.dll': 'nvrtc64_112_0.dll', } HIPSDK_TARGETS = ['rocblas.dll', 'rocsolver.dll', 'hipfft.dll',] -ZLUDA_TARGETS = ('nvcuda.dll', 'nvml.dll',) hipBLASLt_available = False MIOpen_available = False @@ -26,8 +26,49 @@ path = os.path.abspath(os.environ.get('ZLUDA', '.zluda')) default_agent: Union[rocm.Agent, None] = None hipBLASLt_enabled = False -nightly = os.environ.get("ZLUDA_NIGHTLY", "0") == "1" -skip_arch_test = os.environ.get("ZLUDA_SKIP_ARCH_TEST", "0") == "1" + +class ZLUDAResult(ctypes.Structure): + _fields_ = [ + ('return_code', ctypes.c_int), + ('value', ctypes.c_ulonglong), + ] + + +class ZLUDALibrary: + internal: ctypes.WinDLL + + def __init__(self, internal: ctypes.WinDLL): + self.internal = internal + + +class Core(ZLUDALibrary): + def __init__(self, internal: ctypes.WinDLL): + internal.zluda_get_hip_object.restype = ZLUDAResult + internal.zluda_get_hip_object.argtypes = [ctypes.c_void_p, ctypes.c_int] + + try: + internal.zluda_get_nightly_flag.restype = ctypes.c_int + internal.zluda_get_nightly_flag.argtypes = [] + except AttributeError: + internal.zluda_get_nightly_flag = lambda: 0 + + super().__init__(internal) + + def to_hip_stream(self, zluda_object: ctypes.c_void_p): + return self.internal.zluda_get_hip_object(zluda_object, 1).value + + def get_nightly_flag(self) -> int: + return self.internal.zluda_get_nightly_flag() + + +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): @@ -36,19 +77,17 @@ def set_default_agent(agent: rocm.Agent): is_nightly = False try: - nvcuda = ctypes.windll.LoadLibrary(os.path.join(path, 'nvcuda.dll')) - nvcuda.zluda_get_nightly_flag.restype = ctypes.c_int - nvcuda.zluda_get_nightly_flag.argtypes = [] - is_nightly = nvcuda.zluda_get_nightly_flag() == 1 + 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 ((not os.path.exists(path) and nightly) or os.path.exists(os.path.join(path, 'cublasLt.dll'))) + 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 (skip_arch_test or agent.gfx_version in (0x908, 0x90a, 0x940, 0x941, 0x942, 0x1030, 0x1100, 0x1101, 0x1102, 0x1150,)) + 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 @@ -60,8 +99,11 @@ def install() -> None: return platform = "windows" - commit = os.environ.get("ZLUDA_HASH", "4d14bf95d4c500863e240a0b1fa82793d0da789b") - if nightly: + commit = os.environ.get("ZLUDA_HASH", "dba64c0966df2c71e82255e942c96e2e1cea3a2d") + if os.environ.get("ZLUDA_NIGHTLY", "0") == "1": + log.warning("Environment variable 'ZLUDA_NIGHTLY' will be removed. Please use command-line argument '--use-nightly' instead.") + args.use_nightly = True + if args.use_nightly: platform = "nightly-" + platform urllib.request.urlretrieve(f'https://github.com/lshqqytiger/ZLUDA/releases/download/rel.{commit}/ZLUDA-{platform}-rocm{rocm.version[0]}-amd64.zip', '_zluda') with zipfile.ZipFile('_zluda', 'r') as archive: @@ -110,13 +152,14 @@ def make_copy() -> None: 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 ZLUDA_TARGETS: - ctypes.windll.LoadLibrary(os.path.join(path, v)) for v in DLL_MAPPING.values(): ctypes.windll.LoadLibrary(os.path.join(path, v)) @@ -132,7 +175,8 @@ def load() -> None: ctypes.windll.LoadLibrary(os.path.join(path, 'cudnn64_9.dll')) def conceal(): - import torch # pylint: disable=unused-import + import torch + torch.version.hip = rocm.version platform = sys.platform sys.platform = "" from torch.utils import cpp_extension diff --git a/requirements.txt b/requirements.txt index f61890cbe..f1de3abd7 100644 --- a/requirements.txt +++ b/requirements.txt @@ -32,7 +32,7 @@ pi-heif # versioned rich==13.9.4 -safetensors==0.5.2 +safetensors==0.5.3 tensordict==0.1.2 peft==0.14.0 httpx==0.24.1 @@ -41,18 +41,18 @@ torchsde==0.2.6 antlr4-python3-runtime==4.9.3 requests==2.32.3 tqdm==4.67.1 -accelerate==1.3.0 +accelerate==1.6.0 opencv-contrib-python-headless==4.9.0.80 einops==0.4.1 gradio==3.43.2 -huggingface_hub==0.28.1 +huggingface_hub==0.30.1 numexpr==2.8.8 numpy==1.26.4 numba==0.59.1 protobuf==4.25.3 pytorch_lightning==1.9.4 -tokenizers==0.21.0 -transformers==4.48.3 +tokenizers==0.21.1 +transformers==4.50.3 urllib3==1.26.19 Pillow==10.4.0 timm==0.9.16 diff --git a/scripts/allegrovideo.py b/scripts/allegrovideo.py index 003211a60..cf35500fb 100644 --- a/scripts/allegrovideo.py +++ b/scripts/allegrovideo.py @@ -31,7 +31,7 @@ def hijack_encode_prompt(*args, **kwargs): class Script(scripts.Script): def title(self): - return 'Video: Allegro' + return 'Video: Allegro (Legacy)' def show(self, is_img2img): return not is_img2img if shared.native else False @@ -61,14 +61,7 @@ class Script(scripts.Script): if shared.sd_model.__class__ != diffusers.AllegroPipeline: sd_models.unload_model_weights() t0 = time.time() - quant_args = {} - quant_args = model_quant.create_bnb_config(quant_args) - if quant_args: - model_quant.load_bnb(f'Load model: type=Allegro quant={quant_args}') - if not quant_args: - quant_args = model_quant.create_ao_config(quant_args) - if quant_args: - model_quant.load_torchao(f'Load model: type=Allegro quant={quant_args}') + quant_args = model_quant.create_config() transformer = diffusers.AllegroTransformer3DModel.from_pretrained( repo_id, subfolder="transformer", diff --git a/scripts/animatediff.py b/scripts/animatediff.py index 34b7829c1..6fb77a45b 100644 --- a/scripts/animatediff.py +++ b/scripts/animatediff.py @@ -140,6 +140,8 @@ def set_scheduler(p, model, override: bool = False): def set_prompt(p): p.prompt = shared.prompt_styles.apply_styles_to_prompt(p.prompt, p.styles) p.negative_prompt = shared.prompt_styles.apply_negative_styles_to_prompt(p.negative_prompt, p.styles) + shared.prompt_styles.apply_styles_to_extra(p) + p.styles = [] prompts = p.prompt.split('\n') try: prompt = {} diff --git a/scripts/cogvideo.py b/scripts/cogvideo.py index 7184dd946..de3c7736c 100644 --- a/scripts/cogvideo.py +++ b/scripts/cogvideo.py @@ -22,7 +22,7 @@ debug = (os.environ.get('SD_LOAD_DEBUG', None) is not None) or (os.environ.get(' class Script(scripts.Script): def title(self): - return 'Video: CogVideoX' + return 'Video: CogVideoX (Legacy)' def show(self, is_img2img): return shared.native @@ -50,7 +50,7 @@ class Script(scripts.Script): return [model, sampler, frames, guidance, offload, override, video_type, duration, loop, pad, interpolate, image, video] def load(self, model): - if (shared.sd_model_type != 'cogvideox' or shared.sd_model.sd_model_checkpoint != model) and model != 'None': + if (shared.sd_model_type != 'cogvideo' or shared.sd_model.sd_model_checkpoint != model) and model != 'None': sd_models.unload_model_weights('model') shared.log.info(f'CogVideoX load: model="{model}"') try: @@ -64,7 +64,7 @@ class Script(scripts.Script): shared.log.error(f'Load CogVideoX: {e}') if debug: errors.display(e, 'CogVideoX') - if shared.sd_model_type == 'cogvideox' and model != 'None': + if shared.sd_model_type == 'cogvideo' and model != 'None': shared.sd_model.set_progress_bar_config(bar_format='Progress {rate_fmt}{postfix} {bar} {percentage:3.0f}% {n_fmt}/{total_fmt} {elapsed} {remaining} ' + '\x1b[38;5;71m', ncols=80, colour='#327fba') shared.log.debug(f'CogVideoX load: class="{shared.sd_model.__class__.__name__}"') if shared.sd_model is not None and model == 'None': @@ -74,7 +74,7 @@ class Script(scripts.Script): devices.torch_gc() def offload(self, offload): - if shared.sd_model_type != 'cogvideox': + if shared.sd_model_type != 'cogvideo': return if offload == 'none': sd_models.move_model(shared.sd_model, devices.device) @@ -131,7 +131,7 @@ class Script(scripts.Script): return img def generate(self, p: processing.StableDiffusionProcessing, model: str): - if shared.sd_model_type != 'cogvideox': + if shared.sd_model_type != 'cogvideo': return [] shared.log.info(f'CogVideoX: sampler={p.sampler} steps={p.steps} frames={p.frames} width={p.width} height={p.height} seed={p.seed} guidance={p.guidance}') if p.sampler == 'DDIM': diff --git a/scripts/consistory_ext.py b/scripts/consistory_ext.py index 7a8f21e3b..c02ca1e50 100644 --- a/scripts/consistory_ext.py +++ b/scripts/consistory_ext.py @@ -118,6 +118,8 @@ class Script(scripts.Script): settings = [p.strip() for p in prompts.split('\n') if p.strip() != ''] anchors = [f'{subject} {p}' for p in settings] prompt = shared.prompt_styles.apply_styles_to_prompt(p.prompt, p.styles) + shared.prompt_styles.apply_styles_to_extra(p) + p.styles = [] prompts = [p.strip() for p in prompt.split('\n') if p.strip() != ''] for i, prompt in enumerate(prompts): if subject not in prompt: diff --git a/scripts/flux_prompt_enhance.py b/scripts/flux_prompt_enhance.py new file mode 100644 index 000000000..0ab087e1b --- /dev/null +++ b/scripts/flux_prompt_enhance.py @@ -0,0 +1,102 @@ +# repo: https://huggingface.co/gokaygokay/Flux-Prompt-Enhance + +import time +import random +import threading +from transformers import AutoTokenizer, AutoModelForSeq2SeqLM +import gradio as gr +from modules import shared, scripts, devices, processing + + +repo_id = "gokaygokay/Flux-Prompt-Enhance" +num_return_sequences = 5 +load_lock = threading.Lock() + + +class Script(scripts.Script): + prompts = [['']] + tokenizer: AutoTokenizer = None + model: AutoModelForSeq2SeqLM = None + prefix: str = "enhance prompt: " + button: gr.Button = None + auto_apply: gr.Checkbox = None + max_length: gr.Slider = None + temperature: gr.Slider = None + repetition_penalty: gr.Slider = None + table: gr.DataFrame = None + prompt: gr.Textbox = None + + def title(self): + return 'Flux Prompt enhance (Legacy)' + + def show(self, is_img2img): + return shared.native + + def load(self): + with load_lock: + if self.tokenizer is None: + self.tokenizer = AutoTokenizer.from_pretrained('gokaygokay/Flux-Prompt-Enhance', cache_dir=shared.opts.hfcache_dir) + if self.model is None: + shared.log.info(f'Prompt enhance: model="{repo_id}"') + self.model = AutoModelForSeq2SeqLM.from_pretrained('gokaygokay/Flux-Prompt-Enhance', cache_dir=shared.opts.hfcache_dir).to(device=devices.cpu, dtype=devices.dtype) + + def enhance(self, prompt, auto_apply: bool = False, temperature: float = 0.7, repetition_penalty: float = 1.2, max_length: int = 128): + self.load() + t0 = time.time() + input_text = self.prefix + prompt + input_ids = self.tokenizer(input_text, return_tensors="pt").input_ids.to(devices.device) + self.model = self.model.to(devices.device) + kwargs = { + 'max_length': int(max_length), + 'num_return_sequences': int(num_return_sequences), + 'do_sample': True, + 'temperature': float(temperature), + 'repetition_penalty': float(repetition_penalty), + } + try: + outputs = self.model.generate(input_ids, **kwargs) + except Exception as e: + shared.log.error(f'Prompt enhance: error="{e}"') + return [['']] + self.model = self.model.to(devices.cpu) + prompts = self.tokenizer.batch_decode(outputs, skip_special_tokens=True) + prompts = [[p] for p in prompts] + t1 = time.time() + shared.log.info(f'Prompt enhance: temperature={temperature} repetition={repetition_penalty} length={max_length} sequences={num_return_sequences} apply={auto_apply} time={t1-t0:.2f}s') + return prompts + + def select(self, cell: gr.SelectData, _table): + prompt = cell.value if hasattr(cell, 'value') else cell + shared.log.info(f'Prompt enhance: prompt="{prompt}"') + return prompt + + def ui(self, _is_img2img): + with gr.Row(): + self.button = gr.Button(value='Enhance prompt') + self.auto_apply = gr.Checkbox(label='Auto apply', default=False) + with gr.Row(): + self.max_length = gr.Slider(label='Length', minimum=64, maximum=512, step=1, value=128) + self.temperature = gr.Slider(label='Temperature', minimum=0.1, maximum=2.0, step=0.05, value=0.7) + self.repetition_penalty = gr.Slider(label='Penalty', minimum=0.1, maximum=2.0, step=0.05, value=1.2) + with gr.Row(): + self.table = gr.DataFrame(self.prompts, label='', show_label=False, interactive=False, wrap=True, datatype="str", col_count=1, max_rows=num_return_sequences, headers=['Prompts']) + + if self.prompt is not None: + self.button.click(fn=self.enhance, inputs=[self.prompt, self.auto_apply, self.temperature, self.repetition_penalty, self.max_length], outputs=[self.table]) + self.table.select(fn=self.select, inputs=[self.table], outputs=[self.prompt]) + return [self.auto_apply, self.temperature, self.repetition_penalty, self.max_length] + + def run(self, p: processing.StableDiffusionProcessing, auto_apply, temperature, repetition_penalty, max_length): # pylint: disable=arguments-differ + if auto_apply: + p.prompt = shared.prompt_styles.apply_styles_to_prompt(p.prompt, p.styles) + p.negative_prompt = shared.prompt_styles.apply_negative_styles_to_prompt(p.negative_prompt, p.styles) + shared.prompt_styles.apply_styles_to_extra(p) + p.styles = [] + shared.log.debug(f'Prompt enhance: source="{p.prompt}"') + prompts = self.enhance(p.prompt, auto_apply, temperature, repetition_penalty, max_length) + p.prompt = random.choice(prompts)[0] + shared.log.debug(f'Prompt enhance: prompt="{p.prompt}"') + + def after_component(self, component, **kwargs): # searching for actual ui prompt components + if getattr(component, 'elem_id', '') in ['txt2img_prompt', 'img2img_prompt', 'control_prompt', 'video_prompt']: + self.prompt = component diff --git a/scripts/hunyuanvideo.py b/scripts/hunyuanvideo.py index d67fd7d31..c39cec688 100644 --- a/scripts/hunyuanvideo.py +++ b/scripts/hunyuanvideo.py @@ -60,7 +60,7 @@ def hijack_encode_prompt(*args, **kwargs): class Script(scripts.Script): def title(self): - return 'Video: Hunyuan Video' + return 'Video: Hunyuan Video (Legacy)' def show(self, is_img2img): return not is_img2img if shared.native else False @@ -91,14 +91,7 @@ class Script(scripts.Script): if shared.sd_model.__class__ != diffusers.HunyuanVideoPipeline or model != loaded_model: sd_models.unload_model_weights() t0 = time.time() - quant_args = {} - quant_args = model_quant.create_bnb_config(quant_args) - if quant_args: - model_quant.load_bnb(f'Load model: type=HunyuanVideo quant={quant_args}') - if not quant_args: - quant_args = model_quant.create_ao_config(quant_args) - if quant_args: - model_quant.load_torchao(f'Load model: type=HunyuanVideo quant={quant_args}') + quant_args = model_quant.create_config() transformer = diffusers.HunyuanVideoTransformer3DModel.from_pretrained( pretrained_model_name_or_path='tencent/HunyuanVideo', subfolder="transformer", diff --git a/scripts/infiniteyou_ext.py b/scripts/infiniteyou_ext.py new file mode 100644 index 000000000..59432cd40 --- /dev/null +++ b/scripts/infiniteyou_ext.py @@ -0,0 +1,121 @@ +# https://huggingface.co/ByteDance/InfiniteYou +# https://github.com/bytedance/InfiniteYou +# flux base model + 11.8gb controlnet module + 338mb image module + 428 insightface module + +import gradio as gr +from PIL import Image +from modules import scripts, processing, shared, sd_models, devices + + +prefix = 'InfiniteYou' +model_versions = ['aes_stage2', 'sim_stage1'] +orig_pipeline, orig_prompt_attention = None, None + + +def verify_insightface(): + from installer import installed, install, reload + if not installed('insightface', reload=False, quiet=True): + install('insightface==0.7.3', ignore=False) + install('albumentations==1.4.3', ignore=False, reinstall=True) + install('pydantic==1.10.21', ignore=False, reinstall=True, force=True) + reload('pydantic') + + +def load_infiniteyou(model: str): + from modules.infiniteyou import InfUFluxPipeline + shared.sd_model = InfUFluxPipeline( + pipe=shared.sd_model, + model_version=model, + ) + sd_models.copy_diffuser_options(shared.sd_model, orig_pipeline) + sd_models.set_diffuser_options(shared.sd_model) + + +class Script(scripts.Script): + def title(self): + return f'{prefix}: Flexible Photo Recrafting' + + def show(self, is_img2img): + return not is_img2img if shared.native else False + + # return signature is array of gradio components + def ui(self, _is_img2img): + with gr.Row(): + gr.HTML(f'  {prefix}: Flexible Photo Recrafting
') + with gr.Row(): + model = gr.Dropdown(label='IY model', choices=model_versions, value=model_versions[0]) + restore = gr.Checkbox(label='Restore pipeline on end', value=False) + with gr.Row(): + scale = gr.Slider(label='IY scale', value=1.0, minimum=0.0, maximum=2.0, step=0.05) + with gr.Row(): + start = gr.Slider(label='IY start', value=0.0, minimum=0.0, maximum=1.0, step=0.05) + end = gr.Slider(label='IY end', value=1.0, minimum=0.0, maximum=1.0, step=0.05) + with gr.Row(): + id_guidance = gr.Slider(label='Identity guidance', value=3.5, minimum=0.0, maximum=14.0, step=0.05) + with gr.Row(): + id_image = gr.Image(label='Identity image', type='pil') + with gr.Row(): + control_guidance = gr.Slider(label='Control guidance', value=1.0, minimum=0.0, maximum=14.0, step=0.05) + with gr.Row(): + control_image = gr.Image(label='Control image', type='pil') + return [model, id_image, control_image, scale, start, end, id_guidance, control_guidance, restore] + + def run(self, p: processing.StableDiffusionProcessing, + model: str = None, + id_image: Image.Image = None, + control_image: Image.Image = None, + scale: float = 1.0, + start: float = 0.0, + end: float = 1.0, + id_guidance: float = 3.5, + control_guidance: float = 1.0, + restore: bool = False, + ): # pylint: disable=arguments-differ, unused-argument + + if model is None or model not in model_versions: + return None + if id_image is None: + shared.log.error(f'{prefix}: no init_images') + return None + if shared.sd_model_type != 'f1': + shared.log.error(f'{prefix}: invalid model type: {shared.sd_model_type}') + return None + + global orig_pipeline, orig_prompt_attention # pylint: disable=global-statement + orig_pipeline = shared.sd_model + if shared.sd_model.__class__.__name__ != 'InfUFluxPipeline': + verify_insightface() + load_infiniteyou(model) + devices.torch_gc() + shared.log.info(f'{prefix}: cls={shared.sd_model.__class__.__name__} loaded') + + 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['seed'] = p.seed + p.task_args['negative_prompt'] = None + p.task_args['guidance_scale'] = id_guidance + p.task_args['controlnet_guidance_scale'] = control_guidance + p.extra_generation_params['IY model'] = model + p.extra_generation_params['IY guidance'] = f'{scale:.1f}/{start:.1f}/{end:.1f}' + orig_prompt_attention = shared.opts.prompt_attention + shared.opts.data['prompt_attention'] = 'fixed' + shared.log.debug(f'{prefix}: args={p.task_args}') + + processed = processing.process_images(p) + return processed + + def after(self, p: processing.StableDiffusionProcessing, processed: processing.Processed, *args, **kwargs): # pylint: disable=unused-argument + # restore pipeline + global orig_pipeline, orig_prompt_attention # pylint: disable=global-statement + restore = args[-1] + if orig_prompt_attention is not None: + shared.opts.data['prompt_attention'] = orig_prompt_attention + orig_prompt_attention = None + if restore and orig_pipeline is not None: + shared.log.info(f'{prefix}: restoring pipeline') + shared.sd_model = orig_pipeline + orig_pipeline = None diff --git a/scripts/legacy_allegrovideo.py b/scripts/legacy_allegrovideo.py new file mode 100644 index 000000000..cf35500fb --- /dev/null +++ b/scripts/legacy_allegrovideo.py @@ -0,0 +1,131 @@ +import time +import gradio as gr +import transformers +import diffusers +from modules import scripts, processing, shared, images, devices, sd_models, sd_checkpoint, model_quant, timer + + +repo_id = 'rhymes-ai/Allegro' + + +def hijack_decode(*args, **kwargs): + t0 = time.time() + vae: diffusers.AutoencoderKLAllegro = shared.sd_model.vae + shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model, exclude=['vae']) + res = shared.sd_model.vae.orig_decode(*args, **kwargs) + t1 = time.time() + timer.process.add('vae', t1-t0) + shared.log.debug(f'Video: vae={vae.__class__.__name__} time={t1-t0:.2f}') + return res + + +def hijack_encode_prompt(*args, **kwargs): + t0 = time.time() + res = shared.sd_model.vae.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}') + shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model) + return res + + +class Script(scripts.Script): + def title(self): + return 'Video: Allegro (Legacy)' + + def show(self, is_img2img): + return not is_img2img if shared.native else False + + # return signature is array of gradio components + def ui(self, is_img2img): + with gr.Row(): + gr.HTML('  Allegro Video
') + with gr.Row(): + num_frames = gr.Slider(label='Frames', minimum=4, maximum=88, step=1, value=22) + with gr.Row(): + override_scheduler = gr.Checkbox(label='Override scheduler', value=True) + with gr.Row(): + from modules.ui_sections import create_video_inputs + video_type, duration, gif_loop, mp4_pad, mp4_interpolate = create_video_inputs(tab='img2img' if is_img2img else 'txt2img') + return [num_frames, override_scheduler, video_type, duration, gif_loop, mp4_pad, mp4_interpolate] + + def run(self, p: processing.StableDiffusionProcessing, num_frames, override_scheduler, video_type, duration, gif_loop, mp4_pad, mp4_interpolate): # pylint: disable=arguments-differ, unused-argument + # set params + num_frames = int(num_frames) + p.width = 8 * int(p.width // 8) + p.height = 8 * int(p.height // 8) + p.do_not_save_grid = True + p.ops.append('video') + + # load model + if shared.sd_model.__class__ != diffusers.AllegroPipeline: + sd_models.unload_model_weights() + t0 = time.time() + quant_args = model_quant.create_config() + transformer = diffusers.AllegroTransformer3DModel.from_pretrained( + repo_id, + subfolder="transformer", + torch_dtype=devices.dtype, + cache_dir=shared.opts.hfcache_dir, + **quant_args + ) + shared.log.debug(f'Video: module={transformer.__class__.__name__}') + text_encoder = transformers.T5EncoderModel.from_pretrained( + repo_id, + subfolder="text_encoder", + cache_dir=shared.opts.hfcache_dir, + torch_dtype=devices.dtype, + **quant_args + ) + shared.log.debug(f'Video: module={text_encoder.__class__.__name__}') + shared.sd_model = diffusers.AllegroPipeline.from_pretrained( + repo_id, + # transformer=transformer, + # text_encoder=text_encoder, + cache_dir=shared.opts.hfcache_dir, + torch_dtype=devices.dtype, + **quant_args + ) + t1 = time.time() + shared.log.debug(f'Video: load cls={shared.sd_model.__class__.__name__} repo="{repo_id}" dtype={devices.dtype} time={t1-t0:.2f}') + sd_models.set_diffuser_options(shared.sd_model) + 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.vae.decode = hijack_decode + shared.sd_model.encode_prompt = hijack_encode_prompt + shared.sd_model.vae.enable_tiling() + # shared.sd_model.vae.enable_slicing() + + shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model) + devices.torch_gc(force=True) + + processing.fix_seed(p) + if override_scheduler: + p.sampler_name = 'Default' + p.steps = 100 + p.task_args['num_frames'] = num_frames + p.task_args['output_type'] = 'pil' + p.task_args['clean_caption'] = False + + p.all_prompts, p.all_negative_prompts = shared.prompt_styles.apply_styles_to_prompts([p.prompt], [p.negative_prompt], p.styles, [p.seed]) + p.task_args['prompt'] = p.all_prompts[0] + p.task_args['negative_prompt'] = p.all_negative_prompts[0] + + # w = shared.sd_model.transformer.config.sample_width * shared.sd_model.vae_scale_factor_spatial + # h = shared.sd_model.transformer.config.sample_height * shared.sd_model.vae_scale_factor_spatial + # n = shared.sd_model.transformer.config.sample_frames * shared.sd_model.vae_scale_factor_temporal + + # run processing + t0 = time.time() + shared.state.disable_preview = True + shared.log.debug(f'Video: cls={shared.sd_model.__class__.__name__} width={p.width} height={p.height} frames={num_frames}') + processed = processing.process_images(p) + shared.state.disable_preview = False + t1 = time.time() + if processed is not None and len(processed.images) > 0: + shared.log.info(f'Video: frames={len(processed.images)} time={t1-t0:.2f}') + if video_type != 'None': + images.save_video(p, filename=None, images=processed.images, video_type=video_type, duration=duration, loop=gif_loop, pad=mp4_pad, interpolate=mp4_interpolate) + return processed diff --git a/scripts/ltxvideo.py b/scripts/ltxvideo.py index 5599ceed4..697e64021 100644 --- a/scripts/ltxvideo.py +++ b/scripts/ltxvideo.py @@ -11,26 +11,20 @@ from modules.teacache.teacache_ltx import teacache_forward repos = { '0.9.0': 'a-r-r-o-w/LTX-Video-diffusers', '0.9.1': 'a-r-r-o-w/LTX-Video-0.9.1-diffusers', + '0.9.5': 'Lightricks/LTX-Video-0.9.5', 'custom': None, } def load_quants(kwargs, repo_id): - quant_args = {} - quant_args = model_quant.create_bnb_config(quant_args) - if quant_args: - model_quant.load_bnb(f'Load model: type=LTXVideo quant={quant_args}') - if not quant_args: - quant_args = model_quant.create_ao_config(quant_args) - if quant_args: - model_quant.load_torchao(f'Load model: type=LTXVideo quant={quant_args}') + quant_args = model_quant.create_config() if not quant_args: return kwargs model_quant.load_bnb(f'Load model: type=LTX quant={quant_args}') if 'transformer' not in kwargs and ('Model' in shared.opts.bnb_quantization or 'Model' in shared.opts.torchao_quantization): kwargs['transformer'] = diffusers.LTXVideoTransformer3DModel.from_pretrained(repo_id, subfolder="transformer", cache_dir=shared.opts.hfcache_dir, torch_dtype=devices.dtype, **quant_args) shared.log.debug(f'Quantization: module=transformer type=bnb dtype={shared.opts.bnb_quantization_type} storage={shared.opts.bnb_quantization_storage}') - if 'text_encoder' not in kwargs and ('Text Encoder' in shared.opts.bnb_quantization or 'Text Encoder' in shared.opts.torchao_quantization): + if 'text_encoder' not in kwargs and ('TE' in shared.opts.bnb_quantization or 'TE' in shared.opts.torchao_quantization): kwargs['text_encoder'] = transformers.T5EncoderModel.from_pretrained(repo_id, subfolder="text_encoder", cache_dir=shared.opts.hfcache_dir, torch_dtype=devices.dtype, **quant_args) shared.log.debug(f'Quantization: module=t5 type=bnb dtype={shared.opts.bnb_quantization_type} storage={shared.opts.bnb_quantization_storage}') return kwargs @@ -38,7 +32,6 @@ def load_quants(kwargs, repo_id): def hijack_decode(*args, **kwargs): t0 = time.time() - # vae: diffusers.AutoencoderKLHunyuanVideo = shared.sd_model.vae shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model, exclude=['vae']) res = shared.sd_model.vae.orig_decode(*args, **kwargs) t1 = time.time() @@ -59,7 +52,7 @@ def hijack_encode_prompt(*args, **kwargs): class Script(scripts.Script): def title(self): - return 'Video: LTX Video' + return 'Video: LTX Video (Legacy)' def show(self, is_img2img): return shared.native @@ -119,9 +112,7 @@ class Script(scripts.Script): repo_id = model_custom if shared.sd_model.__class__ != cls: sd_models.unload_model_weights() - kwargs = {} - kwargs = model_quant.create_bnb_config(kwargs) - kwargs = model_quant.create_ao_config(kwargs) + kwargs = model_quant.create_config() diffusers.LTXVideoTransformer3DModel.forward = teacache_forward if os.path.isfile(repo_id): shared.sd_model = cls.from_single_file( diff --git a/scripts/mixture_of_diffusers.py b/scripts/mixture_of_diffusers.py index 463948d36..58598ec66 100644 --- a/scripts/mixture_of_diffusers.py +++ b/scripts/mixture_of_diffusers.py @@ -91,6 +91,8 @@ class Script(scripts.Script): p.prompt = shared.prompt_styles.apply_styles_to_prompt(p.prompt, p.styles) p.negative_prompt = shared.prompt_styles.apply_negative_styles_to_prompt(p.negative_prompt, p.styles) + shared.prompt_styles.apply_styles_to_extra(p) + p.styles = [] p.prompts, guidance = self.get_prompts(x_tiles, y_tiles, prompts, p.prompt, p.cfg_scale) p.all_prompts = p.prompts p.task_args['prompts'] = p.prompts diff --git a/scripts/mochivideo.py b/scripts/mochivideo.py index cbc9dad20..1e7ba5541 100644 --- a/scripts/mochivideo.py +++ b/scripts/mochivideo.py @@ -10,7 +10,7 @@ repo_id = 'genmo/mochi-1-preview' class Script(scripts.Script): def title(self): - return 'Video: Mochi.1 Video' + return 'Video: Mochi.1 Video (Legacy)' def show(self, is_img2img): return not is_img2img if shared.native else False @@ -42,9 +42,7 @@ class Script(scripts.Script): cls = diffusers.MochiPipeline if shared.sd_model.__class__ != cls: sd_models.unload_model_weights() - kwargs = {} - kwargs = model_quant.create_bnb_config(kwargs) - kwargs = model_quant.create_ao_config(kwargs) + kwargs = model_quant.create_config() shared.sd_model = cls.from_pretrained( repo_id, cache_dir = shared.opts.hfcache_dir, diff --git a/scripts/postprocessing_upscale.py b/scripts/postprocessing_upscale.py index 104a0fb37..cffff8ed4 100644 --- a/scripts/postprocessing_upscale.py +++ b/scripts/postprocessing_upscale.py @@ -54,7 +54,7 @@ class ScriptPostprocessingUpscale(scripts_postprocessing.ScriptPostprocessing): info["Postprocess upscale to"] = f"{upscale_to_width}x{upscale_to_height}" else: info["Postprocess upscale by"] = upscale_by - image = upscaler.scaler.upscale(image, upscale_by, upscaler.data_path) + image = upscaler.scaler.upscale(image, upscale_by, upscaler.name) if upscale_mode == 1 and upscale_crop: cropped = Image.new("RGB", (upscale_to_width, upscale_to_height)) cropped.paste(image, box=(upscale_to_width // 2 - image.width // 2, upscale_to_height // 2 - image.height // 2)) diff --git a/scripts/prompt_enhance.py b/scripts/prompt_enhance.py index 0ad17bba4..bbc52b0be 100644 --- a/scripts/prompt_enhance.py +++ b/scripts/prompt_enhance.py @@ -1,99 +1,377 @@ -# repo: https://huggingface.co/gokaygokay/Flux-Prompt-Enhance - +from dataclasses import dataclass +import os +import re import time -import random -import threading -from transformers import AutoTokenizer, AutoModelForSeq2SeqLM import gradio as gr -from modules import shared, scripts, devices, processing +import transformers +from modules import scripts, shared, devices, errors, processing, sd_models, sd_modules -repo_id = "gokaygokay/Flux-Prompt-Enhance" -num_return_sequences = 5 -load_lock = threading.Lock() +debug_enabled = os.environ.get('SD_LLM_DEBUG', None) is not None +debug_log = shared.log.trace if debug_enabled else lambda *args, **kwargs: None + + +@dataclass +class Options: + models = { + 'google/gemma-3-1b-it': {}, + 'google/gemma-3-4b-it': {}, + 'Qwen/Qwen2.5-0.5B-Instruct': {}, + 'Qwen/Qwen2.5-1.5B-Instruct': {}, + 'Qwen/Qwen2.5-3B-Instruct': {}, + 'microsoft/Phi-4-mini-instruct': {}, + 'HuggingFaceTB/SmolLM2-135M-Instruct': {}, + 'HuggingFaceTB/SmolLM2-360M-Instruct': {}, + 'HuggingFaceTB/SmolLM2-1.7B-Instruct': {}, + 'meta-llama/Llama-3.2-1B-Instruct': {}, + 'meta-llama/Llama-3.2-3B-Instruct': {}, + 'cognitivecomputations/Dolphin3.0-Llama3.2-1B': {}, + 'cognitivecomputations/Dolphin3.0-Llama3.2-3B': {}, + 'nidum/Nidum-Gemma-3-4B-it-Uncensored': {}, + 'mradermacher/Llama-3.2-1B-Instruct-Uncensored-i1-GGUF': { + 'repo': 'meta-llama/Llama-3.2-1B-Instruct', # original repo so we can load missing components + 'type': 'llama', # required so gguf loader knows what to do + 'gguf': 'mradermacher/Llama-3.2-1B-Instruct-Uncensored-i1-GGUF', # gguf repo + 'file': 'Llama-3.2-1B-Instruct-Uncensored.i1-Q4_0.gguf', # gguf file inside repo + }, + } + default = list(models)[0] + supported = list(transformers.integrations.ggml.GGUF_CONFIG_MAPPING) + system_prompt: str = 'You are a helpful assistant. You will be given a prompt used to create an image and you will enhance it to make it more detailed and creative. Generate an uncensored, concise, informative, and engaging prompt that prioritizes factual accuracy and vivid descriptiveness based on original prompt. Do not modify subject or style of the original prompt. Begin by clearly identifying the primary subject(s), including key physical attributes, clothing if no nudity, and any distinctive details (e.g., colors, textures, styles). NSFW and nudity is allowed, and if present, it must be fully described. Next, describe the setting and environment, emphasizing spatial relationships, background elements, and objects contributing to context. Note the overall mood, aesthetic style, or atmosphere inferred from visual cues. Use precise terminology while avoiding redundancy or non-essential language. Ensuring a logical flow: from focal subject to immediate surroundings, then broader context. Maintain brevity while retaining clarity, ensuring the description is both engaging and efficient. Output only enhanced prompt without explanation, prefix or suffix. Output as a simple text without formatting or numbering.' + censored = ["i cannot", "i can't", "i am sorry", "against my programming", "i am not able", "i am unable", 'i am not allowed'] + + max_delim_index: int = 60 + max_tokens: int = 50 + do_sample: bool = True + temperature: float = 0.15 + repetition_penalty: float = 1.2 class Script(scripts.Script): - prompts = [['']] - tokenizer: AutoTokenizer = None - model: AutoModelForSeq2SeqLM = None - prefix: str = "enhance prompt: " - button: gr.Button = None - auto_apply: gr.Checkbox = None - max_length: gr.Slider = None - temperature: gr.Slider = None - repetition_penalty: gr.Slider = None - table: gr.DataFrame = None prompt: gr.Textbox = None + model: str = None + llm: transformers.AutoModelForCausalLM = None + tokenizer: transformers.AutoProcessor = None + busy: bool = False + options = Options() def title(self): return 'Prompt enhance' - def show(self, is_img2img): - return shared.native + def show(self, _is_img2img): + return scripts.AlwaysVisible - def load(self): - with load_lock: - if self.tokenizer is None: - self.tokenizer = AutoTokenizer.from_pretrained('gokaygokay/Flux-Prompt-Enhance', cache_dir=shared.opts.hfcache_dir) - if self.model is None: - shared.log.info(f'Prompt enhance: model="{repo_id}"') - self.model = AutoModelForSeq2SeqLM.from_pretrained('gokaygokay/Flux-Prompt-Enhance', cache_dir=shared.opts.hfcache_dir).to(device=devices.cpu, dtype=devices.dtype) + def load(self, name:str=None, model_repo:str=None, model_gguf:str=None, model_type:str=None, model_file:str=None): + name = name or self.options.default + if self.busy: + shared.log.debug('Prompt enhance: busy') + return + self.busy = True + if self.model is not None and self.model == name: + return + + from modules import modelloader, model_quant, ggml + modelloader.hf_login() + model_repo = model_repo or self.options.models.get(name, {}).get('repo', None) or name + model_gguf = model_gguf or self.options.models.get(name, {}).get('gguf', None) or model_repo + model_type = model_type or self.options.models.get(name, {}).get('type', None) + model_file = model_file or self.options.models.get(name, {}).get('file', None) + + gguf_args = {} + if model_type is not None and model_file is not None and len(model_type) > 2 and len(model_file) > 2: + debug_log(f'Prompt enhance: gguf supported={self.options.supported}') + if model_type not in self.options.supported: + shared.log.error(f'Prompt enhance: name="{name}" repo="{model_repo}" fn="{model_file}" type={model_type} gguf not supported') + shared.log.trace(f'Prompt enhance: gguf supported={self.options.supported}') + self.busy = False + return + ggml.install_gguf() + gguf_args['model_type'] = model_type + gguf_args['gguf_file'] = model_file + + quant_args = model_quant.create_config(module='LLM') if not gguf_args else {} - def enhance(self, prompt, auto_apply: bool = False, temperature: float = 0.7, repetition_penalty: float = 1.2, max_length: int = 128): - self.load() - t0 = time.time() - input_text = self.prefix + prompt - input_ids = self.tokenizer(input_text, return_tensors="pt").input_ids.to(devices.device) - self.model = self.model.to(devices.device) - kwargs = { - 'max_length': int(max_length), - 'num_return_sequences': int(num_return_sequences), - 'do_sample': True, - 'temperature': float(temperature), - 'repetition_penalty': float(repetition_penalty), - } try: - outputs = self.model.generate(input_ids, **kwargs) + t0 = time.time() + if self.llm is not None: + self.llm = None + shared.log.debug(f'Prompt enhance: name="{self.model}" unload') + self.model = None + self.llm = transformers.AutoModelForCausalLM.from_pretrained( + pretrained_model_name_or_path=model_repo if not gguf_args else model_gguf, + trust_remote_code=True, + torch_dtype=devices.dtype, + cache_dir=shared.opts.hfcache_dir, + _attn_implementation="eager", + **gguf_args, + **quant_args, + ) + self.llm.eval() + self.tokenizer = transformers.AutoTokenizer.from_pretrained( + pretrained_model_name_or_path=model_repo, + cache_dir=shared.opts.hfcache_dir, + ) + if debug_enabled: + modules = sd_modules.get_model_stats(self.llm) + sd_modules.get_model_stats(self.tokenizer) + for m in modules: + shared.log.trace(f'Prompt enhance: {m}') + self.model = name + t1 = time.time() + shared.log.info(f'Prompt enhance: cls={self.llm.__class__.__name__} name="{name}" repo="{model_repo}" fn="{model_file}" time={t1-t0:.2f} loaded') except Exception as e: - shared.log.error(f'Prompt enhance: error="{e}"') - return [['']] - self.model = self.model.to(devices.cpu) - prompts = self.tokenizer.batch_decode(outputs, skip_special_tokens=True) - prompts = [[p] for p in prompts] - t1 = time.time() - shared.log.info(f'Prompt enhance: temperature={temperature} repetition={repetition_penalty} length={max_length} sequences={num_return_sequences} apply={auto_apply} time={t1-t0:.2f}s') - return prompts + shared.log.error(f'Prompt enhance: load {e}') + errors.display(e, 'Prompt enhance') + devices.torch_gc() + self.busy = False - def select(self, cell: gr.SelectData, _table): - prompt = cell.value if hasattr(cell, 'value') else cell - shared.log.info(f'Prompt enhance: prompt="{prompt}"') - return prompt + def censored(self, response): + text = response.lower().replace("i'm", "i am") + return any(c.lower() in text for c in self.options.censored) + + def unload(self): + if self.llm is not None: + sd_models.move_model(self.llm, devices.cpu) + self.model = None + self.llm = None + self.tokenizer = None + devices.torch_gc() + shared.log.debug('Prompt enhance: model unloaded') + + def clean(self, response): + # remove special characters + response = response.replace('"', '').replace("'", "").replace('“', '').replace('”', '').replace('**', '') + # remove repeating characters + response = response.replace('\n\n', '\n').replace(' ', ' ').replace('...', '.') + + # remove comments between brackets + response = re.sub(r'<.*?>', '', response) + response = re.sub(r'\[.*?\]', '', response) + response = re.sub(r'\/.*?\/', '', response) + + # remove llm commentary + removed = '' + if response.startswith('Prompt'): + removed, response = response.split('Prompt', maxsplit=1) + if 0 <= response.find(':') < self.options.max_delim_index: + removed, response = response.split(':', maxsplit=1) + if 0 <= response.find('---') < self.options.max_delim_index: + response, removed = response.split('---', maxsplit=1) + if len(removed) > 0: + debug_log(f'Prompt enhance: max={self.options.max_delim_index} removed="{removed}"') + + # remove bullets and lists + lines = [re.sub(r'^(\s*[-*]|\s*\d+)\s+', '', line).strip() for line in response.splitlines()] + response = '\n'.join(lines) + + response = response.strip() + return response + + def post(self, response, prefix, suffix, networks): + response = response.strip() + prefix = prefix.strip() + suffix = suffix.strip() + if len(prefix) > 0: + response = f'{prefix} {response}' + if len(suffix) > 0: + response = f'{response} {suffix}' + if len(networks) > 0: + response = f'{response} {" ".join(networks)}' + return response + + def extract(self, prompt): + pattern = r'(<.*?>)' + matches = re.findall(pattern, prompt) + filtered = re.sub(pattern, '', prompt) + return filtered, matches + + def enhance(self, model: str=None, prompt:str=None, system:str=None, prefix:str=None, suffix:str=None, sample:bool=None, tokens:int=None, temperature:float=None, penalty:float=None): + model = model or self.options.default + prompt = prompt or self.prompt.value + prefix = prefix or '' + suffix = suffix or '' + system = system or self.options.system_prompt + tokens = tokens or self.options.max_tokens + penalty = penalty or self.options.repetition_penalty + temperature = temperature or self.options.temperature + sample = sample if sample is not None else self.options.do_sample + while self.busy: + time.sleep(0.1) + self.load(model) + if self.llm is None: + shared.log.error('Prompt enhance: model not loaded') + return prompt + prompt, networks = self.extract(prompt) + debug_log(f'Prompt enhance: networks={networks}') + chat_template = [ + { "role": "system", "content": system }, + { "role": "user", "content": prompt }, + ] + t0 = time.time() + self.busy = True + try: + inputs = self.tokenizer.apply_chat_template( + chat_template, + add_generation_prompt=True, + tokenize=True, + return_dict=True, + return_tensors="pt", + ).to(devices.device).to(devices.dtype) + input_len = inputs['input_ids'].shape[1] + except Exception as e: + shared.log.error(f'Prompt enhance tokenize: {e}') + errors.display(e, 'Prompt enhance') + self.busy = False + return prompt + try: + with devices.inference_context(): + sd_models.move_model(self.llm, devices.device) + outputs = self.llm.generate( + **inputs, + do_sample=sample, + temperature=float(temperature), + max_new_tokens=int(input_len + tokens), + repetition_penalty=float(penalty), + ) + if shared.opts.diffusers_offload_mode != 'none': + sd_models.move_model(self.llm, devices.cpu) + devices.torch_gc() + if debug_enabled: + raw_response = self.tokenizer.batch_decode(outputs, skip_special_tokens=True, clean_up_tokenization_spaces=True) + shared.log.trace(f'Prompt enhance: raw="{raw_response}"') + outputs_cropped = outputs[:, input_len:] + response = self.tokenizer.batch_decode( + outputs_cropped, + skip_special_tokens=True, + clean_up_tokenization_spaces=True, + ) + except Exception as e: + shared.log.error(f'Prompt enhance generate: {e}') + errors.display(e, 'Prompt enhance') + self.busy = False + response = f'Error: {str(e)}' + t1 = time.time() + + if isinstance(response, list): + response = response[0] + is_censored = self.censored(response) + if not is_censored: + response = self.clean(response) + response = self.post(response, prefix, suffix, networks) + shared.log.info(f'Prompt enhance: model="{model}" time={t1-t0:.2f} inputs={input_len} outputs={outputs.shape[-1]} prompt={len(prompt)} response={len(response)}') + if debug_enabled: + shared.log.trace(f'Prompt enhance: sample={sample} tokens={tokens} temperature={temperature} penalty={penalty}') + shared.log.trace(f'Prompt enhance: prompt="{prompt}"') + shared.log.trace(f'Prompt enhance: response="{response}"') + self.busy = False + if is_censored: + shared.log.warning(f'Prompt enhance: censored response="{response}"') + return prompt + return response + + def apply(self, prompt, apply_prompt, llm_model, prompt_system, prompt_prefix, prompt_suffix, max_tokens, do_sample, temperature, repetition_penalty): + response = self.enhance( + prompt=prompt, + prefix=prompt_prefix, + suffix=prompt_suffix, + model=llm_model, + system=prompt_system, + sample=do_sample, + tokens=max_tokens, + temperature=temperature, + penalty=repetition_penalty, + ) + if apply_prompt: + return [response, response] + return [response, gr.update()] + + def get_custom(self, name): + model_repo = self.options.models.get(name, {}).get('repo', None) or name + model_gguf = self.options.models.get(name, {}).get('gguf', None) + model_type = self.options.models.get(name, {}).get('type', None) + model_file = self.options.models.get(name, {}).get('file', None) + return [model_repo, model_gguf, model_type, model_file] def ui(self, _is_img2img): - with gr.Row(): - self.button = gr.Button(value='Enhance prompt') - self.auto_apply = gr.Checkbox(label='Auto apply', default=False) - with gr.Row(): - self.max_length = gr.Slider(label='Length', minimum=64, maximum=512, step=1, value=128) - self.temperature = gr.Slider(label='Temperature', minimum=0.1, maximum=2.0, step=0.05, value=0.7) - self.repetition_penalty = gr.Slider(label='Penalty', minimum=0.1, maximum=2.0, step=0.05, value=1.2) - with gr.Row(): - self.table = gr.DataFrame(self.prompts, label='', show_label=False, interactive=False, wrap=True, datatype="str", col_count=1, max_rows=num_return_sequences, headers=['Prompts']) - - if self.prompt is not None: - self.button.click(fn=self.enhance, inputs=[self.prompt, self.auto_apply, self.temperature, self.repetition_penalty, self.max_length], outputs=[self.table]) - self.table.select(fn=self.select, inputs=[self.table], outputs=[self.prompt]) - return [self.auto_apply, self.temperature, self.repetition_penalty, self.max_length] - - def run(self, p: processing.StableDiffusionProcessing, auto_apply, temperature, repetition_penalty, max_length): # pylint: disable=arguments-differ - if auto_apply: - p.prompt = shared.prompt_styles.apply_styles_to_prompt(p.prompt, p.styles) - shared.log.debug(f'Prompt enhance: source="{p.prompt}"') - prompts = self.enhance(p.prompt, auto_apply, temperature, repetition_penalty, max_length) - p.prompt = random.choice(prompts)[0] - shared.log.debug(f'Prompt enhance: prompt="{p.prompt}"') + with gr.Accordion('Prompt enhance', open=False, elem_id='prompt_enhance'): + with gr.Row(): + apply_btn = gr.Button(value='Enhance now', elem_id='prompt_enhance_apply', variant='primary') + with gr.Row(): + apply_prompt = gr.Checkbox(label='Apply to prompt', value=False) + apply_auto = gr.Checkbox(label='Auto enhance', value=False) + gr.HTML('
') + with gr.Group(): + with gr.Row(): + llm_model = gr.Dropdown(label='LLM model', choices=list(self.options.models), value=self.options.default, interactive=True, allow_custom_value=True, elem_id='prompt_enhance_model') + with gr.Row(): + load_btn = gr.Button(value='Load model', elem_id='prompt_enhance_load', variant='secondary') + load_btn.click(fn=self.load, inputs=[llm_model], outputs=[]) + unload_btn = gr.Button(value='Unload model', elem_id='prompt_enhance_unload', variant='secondary') + unload_btn.click(fn=self.unload, inputs=[], outputs=[]) + with gr.Accordion('Custom model', open=False, elem_id='prompt_enhance_custom'): + with gr.Row(): + model_repo = gr.Textbox(label='Model repo', value=None, interactive=True, elem_id='prompt_enhance_model_repo', placeholder='Original model repo on huggingface') + with gr.Row(): + model_gguf = gr.Textbox(label='Model gguf', value=None, interactive=True, elem_id='prompt_enhance_model_gguf', placeholder='Optional GGUF model repo on huggingface') + with gr.Row(): + model_type = gr.Textbox(label='Model type', value=None, interactive=True, elem_id='prompt_enhance_model_type', placeholder='Optional GGUF model type') + with gr.Row(): + model_file = gr.Textbox(label='Model file', value=None, interactive=True, elem_id='prompt_enhance_model_file', placeholder='Optional GGUF model file inside GGUF model repo') + with gr.Row(): + custom_btn = gr.Button(value='Load custom model', elem_id='prompt_enhance_custom_load', variant='secondary') + custom_btn.click(fn=self.load, inputs=[model_file, model_repo, model_gguf, model_type, model_file], outputs=[]) + llm_model.change(fn=self.get_custom, inputs=[llm_model], outputs=[model_repo, model_gguf, model_type, model_file]) + gr.HTML('
') + with gr.Accordion('Options', open=False, elem_id='prompt_enhance_options'): + with gr.Row(): + max_tokens = gr.Slider(label='Max tokens', value=self.options.max_tokens, minimum=10, maximum=1024, step=1, interactive=True) + do_sample = gr.Checkbox(label='Do sample', value=self.options.do_sample, interactive=True) + with gr.Row(): + temperature = gr.Slider(label='Temperature', value=self.options.temperature, minimum=0.0, maximum=1.0, step=0.01, interactive=True) + repetition_penalty = gr.Slider(label='Repetition penalty', value=self.options.repetition_penalty, minimum=0.0, maximum=2.0, step=0.01, interactive=True) + gr.HTML('
') + with gr.Accordion('Input', open=False, elem_id='prompt_enhance_system_prompt'): + with gr.Row(): + prompt_prefix = gr.Textbox(label='Prompt prefix', value='', placeholder='Optional prompt prefix', interactive=True, lines=2, elem_id='prompt_enhance_prefix') + with gr.Row(): + prompt_suffix = gr.Textbox(label='Prompt suffix', value='', placeholder='Optional prompt suffix', interactive=True, lines=2, elem_id='prompt_enhance_suffix') + with gr.Row(): + prompt_system = gr.Textbox(label='System prompt', value=self.options.system_prompt, interactive=True, lines=4, elem_id='prompt_enhance_system') + with gr.Accordion('Output', open=True, elem_id='prompt_enhance_system_prompt'): + with gr.Row(): + prompt_output = gr.Textbox(label='Enhanced prompt', value='', interactive=True, lines=4) + with gr.Row(): + clear_btn = gr.Button(value='Clear', elem_id='prompt_enhance_clear', variant='secondary') + clear_btn.click(fn=lambda: '', inputs=[], outputs=[prompt_output]) + copy_btn = gr.Button(value='Set prompt', elem_id='prompt_enhance_copy', variant='secondary') + copy_btn.click(fn=lambda x: x, inputs=[prompt_output], outputs=[self.prompt]) + apply_btn.click(fn=self.apply, inputs=[self.prompt, apply_prompt, llm_model, prompt_system, prompt_prefix, prompt_suffix, max_tokens, do_sample, temperature, repetition_penalty], outputs=[prompt_output, self.prompt]) + return [apply_auto, llm_model, prompt_system, prompt_prefix, prompt_suffix, max_tokens, do_sample, temperature, repetition_penalty] def after_component(self, component, **kwargs): # searching for actual ui prompt components - if getattr(component, 'elem_id', '') in ['txt2img_prompt', 'img2img_prompt', 'control_prompt']: + if getattr(component, 'elem_id', '') in ['txt2img_prompt', 'img2img_prompt', 'control_prompt', 'video_prompt']: self.prompt = component + + def before_process(self, p: processing.StableDiffusionProcessing, *args, **kwargs): # pylint: disable=unused-argument + apply_auto, llm_model, prompt_system, prompt_prefix, prompt_suffix, max_tokens, do_sample, temperature, repetition_penalty = args + if not apply_auto and not p.enhance_prompt: + return + if shared.state.skipped or shared.state.interrupted: + return + p.prompt = shared.prompt_styles.apply_styles_to_prompt(p.prompt, p.styles) + p.negative_prompt = shared.prompt_styles.apply_negative_styles_to_prompt(p.negative_prompt, p.styles) + shared.prompt_styles.apply_styles_to_extra(p) + p.styles = [] + shared.state.begin('LLM') + p.prompt = self.enhance( + prompt=p.prompt, + prefix=prompt_prefix, + suffix=prompt_suffix, + model=llm_model, + system=prompt_system, + sample=do_sample, + tokens=max_tokens, + temperature=temperature, + penalty=repetition_penalty, + ) + p.extra_generation_params['LLM'] = llm_model + shared.state.end() diff --git a/scripts/pulid_ext.py b/scripts/pulid_ext.py index 40e726549..dbc62715d 100644 --- a/scripts/pulid_ext.py +++ b/scripts/pulid_ext.py @@ -225,6 +225,8 @@ class Script(scripts.Script): processing.fix_seed(p) p.prompt = shared.prompt_styles.apply_styles_to_prompt(p.prompt, p.styles) p.negative_prompt = shared.prompt_styles.apply_negative_styles_to_prompt(p.negative_prompt, p.styles) + shared.prompt_styles.apply_styles_to_extra(p) + p.styles = [] with devices.inference_context(): output = shared.sd_model( prompt=p.prompt, diff --git a/scripts/x_adapter.py b/scripts/x_adapter.py index c67eca18b..08874aac9 100644 --- a/scripts/x_adapter.py +++ b/scripts/x_adapter.py @@ -110,6 +110,8 @@ class Script(scripts.Script): shared.opts.data['prompt_attention'] = 'fixed' prompt = shared.prompt_styles.apply_styles_to_prompt(p.prompt, p.styles) negative = shared.prompt_styles.apply_negative_styles_to_prompt(p.negative_prompt, p.styles) + shared.prompt_styles.apply_styles_to_extra(p) + p.styles = [] p.task_args['prompt'] = prompt p.task_args['negative_prompt'] = negative p.task_args['prompt_sd1_5'] = prompt diff --git a/scripts/xyz_grid_classes.py b/scripts/xyz_grid_classes.py index cd4df56e8..7c7cd1b48 100644 --- a/scripts/xyz_grid_classes.py +++ b/scripts/xyz_grid_classes.py @@ -25,26 +25,56 @@ class AxisOptionTxt2Img(AxisOption): class SharedSettingsStackHelper(object): - vae = None - schedulers_solver_order = None - tome_ratio = None - todo_ratio = None sd_model_checkpoint = None sd_model_refiner = None sd_model_dict = None sd_vae = None sd_unet = None sd_text_encoder = None + prompt_attention = None + freeu_b1 = None + freeu_b2 = None + freeu_s1 = None + freeu_s2 = None + schedulers_sigma_adjust = None + schedulers_beta_schedule = None + schedulers_beta_start = None + schedulers_beta_end = None + schedulers_shift = None + schedulers_sigma = None + schedulers_timestep_spacing = None + schedulers_timesteps_range = None + schedulers_beta_schedule = None + schedulers_beta_start = None + schedulers_beta_end = None + schedulers_shift = None + scheduler_eta = None + schedulers_solver_order = None + eta_noise_seed_delta = None + tome_ratio = None + todo_ratio = None extra_networks_default_multiplier = None disable_weights_auto_swap = None - prompt_attention = None def __enter__(self): - #Save overridden settings so they can be restored later. - self.vae = shared.opts.sd_vae + # Save overridden settings so they can be restored later + self.prompt_attention = shared.opts.prompt_attention + self.schedulers_sigma_adjust = shared.opts.schedulers_sigma_adjust + self.schedulers_timestep_spacing = shared.opts.schedulers_timestep_spacing + self.schedulers_timesteps_range = shared.opts.schedulers_timesteps_range self.schedulers_solver_order = shared.opts.schedulers_solver_order + self.schedulers_beta_schedule = shared.opts.schedulers_beta_schedule + self.schedulers_beta_start = shared.opts.schedulers_beta_start + self.schedulers_beta_end = shared.opts.schedulers_beta_end + self.schedulers_shift = shared.opts.schedulers_shift + self.scheduler_eta = shared.opts.scheduler_eta + self.eta_noise_seed_delta = shared.opts.eta_noise_seed_delta self.tome_ratio = shared.opts.tome_ratio self.todo_ratio = shared.opts.todo_ratio + self.freeu_b1 = shared.opts.freeu_b1 + self.freeu_b2 = shared.opts.freeu_b2 + self.freeu_s1 = shared.opts.freeu_s1 + self.freeu_s2 = shared.opts.freeu_s2 self.sd_model_checkpoint = shared.opts.sd_model_checkpoint self.sd_model_refiner = shared.opts.sd_model_refiner self.sd_model_dict = shared.opts.sd_model_dict @@ -53,18 +83,30 @@ class SharedSettingsStackHelper(object): self.sd_text_encoder = shared.opts.sd_text_encoder self.extra_networks_default_multiplier = shared.opts.extra_networks_default_multiplier self.disable_weights_auto_swap = shared.opts.disable_weights_auto_swap - self.prompt_attention = shared.opts.prompt_attention shared.opts.data["disable_weights_auto_swap"] = False def __exit__(self, exc_type, exc_value, tb): - #Restore overriden settings after plot generation. + # Restore overriden settings after plot generation shared.opts.data["disable_weights_auto_swap"] = self.disable_weights_auto_swap - shared.opts.data["sd_vae"] = self.vae - shared.opts.data["schedulers_solver_order"] = self.schedulers_solver_order - shared.opts.data["tome_ratio"] = self.tome_ratio - shared.opts.data["todo_ratio"] = self.todo_ratio shared.opts.data["extra_networks_default_multiplier"] = self.extra_networks_default_multiplier shared.opts.data["prompt_attention"] = self.prompt_attention + shared.opts.data["schedulers_solver_order"] = self.schedulers_solver_order + shared.opts.data["schedulers_sigma_adjust"] = self.schedulers_sigma_adjust + shared.opts.data["schedulers_timestep_spacing"] = self.schedulers_timestep_spacing + shared.opts.data["schedulers_timesteps_range"] = self.schedulers_timesteps_range + shared.opts.data["schedulers_beta_schedule"] = self.schedulers_beta_schedule + shared.opts.data["schedulers_beta_start"] = self.schedulers_beta_start + shared.opts.data["schedulers_beta_end"] = self.schedulers_beta_end + shared.opts.data["schedulers_shift"] = self.schedulers_shift + shared.opts.data["scheduler_eta"] = self.scheduler_eta + shared.opts.data["eta_noise_seed_delta"] = self.eta_noise_seed_delta + shared.opts.data["freeu_b1"] = self.freeu_b1 + shared.opts.data["freeu_b2"] = self.freeu_b2 + shared.opts.data["freeu_s1"] = self.freeu_s1 + shared.opts.data["freeu_s2"] = self.freeu_s2 + shared.opts.data["tome_ratio"] = self.tome_ratio + shared.opts.data["todo_ratio"] = self.todo_ratio + if self.sd_model_checkpoint != shared.opts.sd_model_checkpoint: shared.opts.data["sd_model_checkpoint"] = self.sd_model_checkpoint sd_models.reload_model_weights(op='model') @@ -119,6 +161,7 @@ axis_options = [ AxisOptionTxt2Img("[Sampler] Name", str, apply_sampler, fmt=format_value_add_label, confirm=confirm_samplers, choices=lambda: [x.name for x in sd_samplers.samplers]), AxisOptionImg2Img("[Sampler] Name", str, apply_sampler, fmt=format_value_add_label, confirm=confirm_samplers, choices=lambda: [x.name for x in sd_samplers.samplers_for_img2img]), AxisOption("[Sampler] Sigma method", str, apply_setting("schedulers_sigma"), choices=lambda: ['default', 'karras', 'betas', 'exponential', 'lambdas']), + AxisOption("[Sampler] Sigma adjust", float, apply_setting("schedulers_sigma_adjust")), AxisOption("[Sampler] Timestep spacing", str, apply_setting("schedulers_timestep_spacing"), choices=lambda: ['default', 'linspace', 'leading', 'trailing']), AxisOption("[Sampler] Timestep range", int, apply_setting("schedulers_timesteps_range")), AxisOption("[Sampler] Solver order", int, apply_setting("schedulers_solver_order")), diff --git a/scripts/xyz_grid_draw.py b/scripts/xyz_grid_draw.py index bac96bb8b..d654a2c7f 100644 --- a/scripts/xyz_grid_draw.py +++ b/scripts/xyz_grid_draw.py @@ -128,6 +128,6 @@ def draw_xyz_grid(p, xs, ys, zs, x_labels, y_labels, z_labels, cell, draw_legend processed_result.infotexts.insert(0, processed_result.infotexts[0]) t2 = time.time() - shared.log.info(f'XYZ grid complete: images={list_size} results={len(processed_result.images)}size={grid.size if grid is not None else None} time={t1-t0:.2f} save={t2-t1:.2f}') + shared.log.info(f'XYZ grid complete: images={list_size} results={len(processed_result.images)} size={grid.size if grid is not None else None} time={t1-t0:.2f} save={t2-t1:.2f}') p.skip_processing = True return processed_result diff --git a/scripts/xyz_grid_shared.py b/scripts/xyz_grid_shared.py index a3f7a3da9..f624458d6 100644 --- a/scripts/xyz_grid_shared.py +++ b/scripts/xyz_grid_shared.py @@ -211,7 +211,7 @@ def apply_vae(p, x, xs): def list_lora(): import sys - lora = [v for k, v in sys.modules.items() if k == 'networks' or k == 'modules.lora.networks'][0] + lora = [v for k, v in sys.modules.items() if k == 'networks' or k == 'modules.lora.lora_load'][0] loras = [v.fullname for v in lora.available_networks.values()] return ['None'] + sorted(loras) diff --git a/webui.py b/webui.py index ecfafe3f1..1996c41b5 100644 --- a/webui.py +++ b/webui.py @@ -37,6 +37,7 @@ import modules.hypernetworks.hypernetwork import modules.script_callbacks import modules.api.middleware + if not modules.loader.initialized: timer.startup.record("libraries") import modules.sd_hijack # runs conditional load of ldm if not shared.native @@ -88,8 +89,8 @@ def initialize(): timer.startup.record("models") if not shared.opts.lora_legacy: - import modules.lora.networks as lora_networks - lora_networks.list_available_networks() + from modules.lora import lora_load + lora_load.list_available_networks() timer.startup.record("lora") shared.prompt_styles.reload() @@ -292,6 +293,8 @@ def start_ui(): allowed_paths = [os.path.dirname(__file__)] if shared.cmd_opts.data_dir is not None and os.path.isdir(shared.cmd_opts.data_dir): allowed_paths.append(shared.cmd_opts.data_dir) + if shared.cmd_opts.models_dir is not None and os.path.isdir(shared.cmd_opts.models_dir): + allowed_paths.append(shared.cmd_opts.models_dir) if shared.cmd_opts.allowed_paths is not None: allowed_paths += [p for p in shared.cmd_opts.allowed_paths if os.path.isdir(p)] shared.log.debug(f'Root paths: {allowed_paths}') diff --git a/wiki b/wiki index 910dc3083..9408b299f 160000 --- a/wiki +++ b/wiki @@ -1 +1 @@ -Subproject commit 910dc3083c087eb806f548f05b6cf86ce4666268 +Subproject commit 9408b299fffbb8efaccae968b2ac64d9216326fb