Public/automatic
1
0
Fork 0
mirror of https://github.com/vladmandic/automatic synced 2026-09-19 09:14:35 +02:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #3845 from vladmandic/dev

Browse Source 
dev merge
This commit is contained in:
Vladimir Mandic
2025-04-03 11:31:36 -04:00
committed by GitHub
parent 6633e8e5e7 a5f796a166
commit f33b50bb25
187 changed files with 9599 additions and 2044 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.pylintrc
+2
View File
@@ -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,
.ruff.toml
+6
View File
@@ -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 = []
CHANGELOG.md
+178 -11
View File
@@ -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
TODO.md
+28 -19
View File
@@ -4,28 +4,37 @@ Main ToDo list can be found at [GitHub projects](https://github.com/users/vladma
## Current
- <https://github.com/KohakuBlueleaf/z-tipo-extension/pull/73>
### Issues/Limitations
N/A
## Future Candidates
- Redesign postprocessing
- Flux NF4 loader: <https://github.com/huggingface/diffusers/issues/9996>
- IPAdapter negative: <https://github.com/huggingface/diffusers/discussions/7167>
- Control API enhance scripts compatibility
- CogView4
- Flux: NF4 loader: <https://github.com/huggingface/diffusers/issues/9996>
- IPAdapter: negative guidance: <https://github.com/huggingface/diffusers/discussions/7167>
- Control: API enhance scripts compatibility
- Video: add generate context menu
- Video: API support
- Video: STG: <https://github.com/huggingface/diffusers/blob/main/examples/community/README.md#spatiotemporal-skip-guidance>
- 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
cli/api-grid.py
Executable
+198
View File
@@ -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:{v}> ' # 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)
cli/api-progress.py
+13 -6
View File
@@ -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)
cli/api-txt2img.py
+3
View File
@@ -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__":
extensions-builtin/Lora/extra_networks_lora.py
+2 -2
View File
@@ -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()
extensions-builtin/Lora/networks.py
+18 -18
View File
@@ -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)
extensions-builtin/sd-extension-system-info
+1 -1
Submodule extensions-builtin/sd-extension-system-info updated: 8c7edb3be1...ce373b9c27
extensions-builtin/sdnext-modernui
+1 -1
Submodule extensions-builtin/sdnext-modernui updated: 7fc52bb976...770db00768
html/licenses.html
+34
View File
@@ -637,6 +637,40 @@ SOFTWARE.
limitations under the License.
</pre>
<h2><a href="https://github.com/Dao-AILab/flash-attention/blob/main/LICENSE">Flash Attention</a></h2>
<small>Fast and memory-efficient exact attention</small>
<pre>
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.
</pre>
<h2><a href="https://github.com/explosion/curated-transformers/blob/main/LICENSE">Curated transformers</a></h2>
<small>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</small>
<pre>
html/reference.json
+29 -5
View File
@@ -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
},
installer.py
+10 -8
View File
@@ -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")
javascript/amethyst-nightfall.css
-2
View File
@@ -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; }
javascript/base.css
+5 -5
View File
@@ -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; }
javascript/black-gray.css
-5
View File
@@ -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; }
javascript/black-orange.css
-2
View File
@@ -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; }
javascript/black-teal-reimagined.css
+2
View File
@@ -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;
javascript/black-teal.css
+9 -15
View File
@@ -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; }
javascript/emerald-paradise.css
-2
View File
@@ -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; }
javascript/extraNetworks.js
+13 -1
View File
@@ -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`);
}
javascript/imageParams.js
+1
View File
@@ -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"]');
javascript/invoked.css
-2
View File
@@ -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; }
javascript/light-teal.css
-4
View File
@@ -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; }
javascript/midnight-barbie.css
-2
View File
@@ -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; }
javascript/orchid-dreams.css
-2
View File
@@ -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; }
javascript/progressBar.js
+5 -1
View File
@@ -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 = '';
javascript/promptChecker.js
+2
View File
@@ -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');
}
javascript/script.js
+1
View File
@@ -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) {
javascript/sdnext.css
+43 -34
View File
@@ -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; }
javascript/settings.js
+10 -6
View File
@@ -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'] });
});
javascript/startup.js
+8 -1
View File
@@ -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);
javascript/timeless-beige.css
-2
View File
@@ -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; }
javascript/ui.js
+17
View File
@@ -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');
launch.py
+17 -6
View File
@@ -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...')
models/Reference/Alpha-VLLM--Lumina-Image-2.0.jpg
Executable → Regular
View File
Side by Side Swipe Overlay

Before

Width:  |  Height:  |  Size: 37 KiB

After

Width:  |  Height:  |  Size: 37 KiB

models/Reference/THUDM--CogView4-6B.jpg
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 39 KiB

modules/api/endpoints.py
+5 -5
View File
@@ -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
modules/api/gallery.py
+1 -1
View File
@@ -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)
modules/api/server.py
+2 -2
View File
@@ -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():
modules/call_queue.py
+6 -2
View File
@@ -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 ''
modules/cmd_args.py
+1
View File
@@ -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')
modules/devices.py
+32 -7
View File
@@ -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
modules/extra_networks.py
+4 -3
View File
@@ -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
modules/face/faceswap.py
+10 -4
View File
@@ -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)
modules/face/insightface.py
+2 -1
View File
@@ -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)
modules/flash_attn_triton_amd/__init__.py
View File
modules/flash_attn_triton_amd/bwd_prefill.py
+606
View File
@@ -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
modules/flash_attn_triton_amd/bwd_ref.py
+271
View File
@@ -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
modules/flash_attn_triton_amd/fwd_decode.py
+700
View File
@@ -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
modules/flash_attn_triton_amd/fwd_prefill.py
+634
View File
@@ -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
modules/flash_attn_triton_amd/fwd_ref.py
+258
View File
@@ -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)
modules/flash_attn_triton_amd/interface_fa.py
+394
View File
@@ -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
modules/flash_attn_triton_amd/utils.py
+280
View File
@@ -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")
modules/generation_parameters_copypaste.py
+3
View File
@@ -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:
modules/img2img.py
+2 -1
View File
@@ -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')
modules/infiniteyou/__init__.py
+2
View File
@@ -0,0 +1,2 @@
from .pipeline_flux_infusenet import FluxInfuseNetPipeline
from .pipeline_infu_flux import InfUFluxPipeline
modules/infiniteyou/pipeline_flux_infusenet.py
+612
View File
@@ -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)
modules/infiniteyou/pipeline_infu_flux.py
+325
View File
@@ -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
modules/infiniteyou/resampler.py
+121
View File
@@ -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)
modules/infotext.py
+4 -4
View File
@@ -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:
modules/intel/ipex/attention.py
+6 -6
View File
@@ -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
modules/intel/ipex/hijacks.py
+1 -1
View File
@@ -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
modules/intel/openvino/__init__.py
+2 -2
View File
@@ -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 = ""
modules/interrogate/deepbooru.py
+3 -3
View File
@@ -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):
modules/interrogate/deepseek.py
+2 -2
View File
@@ -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,
modules/interrogate/interrogate.py
+1 -1
View File
@@ -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:
modules/interrogate/openclip.py
+8 -8
View File
@@ -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()
modules/interrogate/vqa.py
+152 -35
View File
@@ -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 = [
'<CAPTION>',
'<DETAILED_CAPTION>',
@@ -70,6 +75,7 @@ def b64(image):
def clean(response, question):
strip = ['---', '\r', '\t', '**', '"', '“', '”', 'Assistant:', 'Caption:', '<|im_end|>', '<pad>']
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 = '<MORE_DETAILED_CAPTION>'
input_dict = {
'image': image,
'text': f'<image>{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)
modules/lora/extra_networks_lora.py
+46 -35
View File
@@ -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()
modules/lora/lora_apply.py
+203
View File
@@ -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
modules/lora/lora_common.py
+21
View File
@@ -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
modules/lora/lora_extract.py
+3 -3
View File
@@ -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:
modules/lora/lora_load.py
+283
View File
@@ -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
modules/lora/network_overrides.py → modules/lora/lora_overrides.py
+7 -1
View File
@@ -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
modules/lora/network.py
+9 -2
View File
@@ -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
modules/lora/networks.py
+78 -571
View File
@@ -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")
modules/memstats.py
+16 -6
View File
@@ -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
modules/model_auraflow.py
+1 -1
View File
@@ -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
modules/model_cogview.py
+95
View File
@@ -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
modules/model_flux.py
+25 -79
View File
@@ -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
modules/model_flux_nf4.py
+1 -1
View File
@@ -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",
)
modules/model_kolors.py
+1 -1
View File
@@ -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
modules/model_lumina.py
+5 -6
View File
@@ -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
modules/model_meissonic.py
+1 -1
View File
@@ -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
modules/model_omnigen.py
+1 -8
View File
@@ -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
modules/model_pixart.py
+1 -1
View File
@@ -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
modules/model_quant.py
+91 -36
View File
@@ -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
modules/model_sana.py
+16 -21
View File
@@ -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
modules/model_sd3.py
+11 -21
View File
@@ -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
modules/model_stablecascade.py
+3 -2
View File
@@ -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
modules/model_tools.py
+2 -4
View File
@@ -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}')
modules/modeldata.py
+16 -5
View File
@@ -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
modules/modelloader.py
+37 -5
View File
@@ -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()
modules/omnigen/utils.py
+1 -4
View File
@@ -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:
modules/onnx_impl/execution_providers.py
+5 -1
View File
@@ -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 += '</pre><br>'
res += 'Server restart required'
log.info("Server restart required")
importlib.reload(ort)
try:
importlib.reload(ort)
except Exception:
pass
return res
modules/onnx_impl/pipelines/__init__.py
+1 -1
View File
@@ -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:
modules/pixelsmith/pixelsmith_pipeline.py
-2
View File
@@ -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
)
modules/postprocess/sdupscaler_model.py
+7 -3
View File
@@ -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]

Some files were not shown because too many files have changed in this diff Show More