diff --git a/.eslintrc.json b/.eslintrc.json index 7a730c1c2..32283e620 100644 --- a/.eslintrc.json +++ b/.eslintrc.json @@ -87,6 +87,7 @@ "get_tab_index": "readonly", "create_submit_args": "readonly", "restartReload": "readonly", + "markSelectedCards": "readonly", "updateInput": "readonly", "toggleCompact": "readonly", "setFontSize": "readonly", @@ -107,7 +108,8 @@ "getExif": "readonly", "jobStatusEl": "readonly", "removeSplash": "readonly", - "initNVML": "readonly", + "initGPU": "readonly", + "startGPU": "readonly", "disableNVML": "readonly", "idbGet": "readonly", "idbPut": "readonly", diff --git a/.pylintrc b/.pylintrc index 1929c26fb..0f77d8aa4 100644 --- a/.pylintrc +++ b/.pylintrc @@ -24,7 +24,6 @@ ignore-paths=/usr/lib/.*$, modules/intel, modules/intel/ipex, modules/framepack/pipeline, - modules/ldsr, modules/onnx_impl, modules/pag, modules/postprocess/aurasr_arch.py, diff --git a/CHANGELOG.md b/CHANGELOG.md index c53e8242d..2fbb69c26 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,49 +1,172 @@ # Change Log for SD.Next -## Update for 2025-08-01 +## Update for 2025-08-15 + +### Highlights for 2025-08-15 + +New release two weeks after the last one and its a big one with over 150 commits! +- Several new models: [Qwen-Image](https://qwenlm.github.io/blog/qwen-image/) (plus *Lightning* variant) and [FLUX.1-Krea-Dev](https://www.krea.ai/blog/flux-krea-open-source-release) +- Several updated models: [Chroma](https://huggingface.co/lodestones/Chroma), [SkyReels-V2](https://huggingface.co/Skywork/SkyReels-V2-DF-14B-720P-Diffusers), [Wan-VACE](https://huggingface.co/Wan-AI/Wan2.1-VACE-14B-diffusers), [HunyuanDiT](https://huggingface.co/Tencent-Hunyuan/HunyuanDiT-v1.2-Diffusers-Distilled) +- Plus continuing with major **UI** work with new embedded **Docs/Wiki** search, redesigned real-time **hints**, **wildcards** UI selector, built-in **GPU monitor**, **CivitAI** integration and more! +- On the compute side, new profiles for high-vram GPUs, offloading improvements, parallel-load for large models, support for new `torch` release and improved quality when using low-bit quantization! +- [SD.Next Model Samples Gallery](https://vladmandic.github.io/sd-samples/compare.html): pre-generated image gallery with 60 models (45 base and 15 finetunes) and 40 different styles resulting in 2,400 high resolution images! + gallery additionally includes model details such as typical load and inference times as well as sizes and types of each model component (*e.g. unet, transformer, text-encoder, vae*) +- And (*as always*) many bugfixes and improvements to existing features! + +![sd-samples](https://github.com/user-attachments/assets/3efc8603-0766-4e4e-a4cb-d8c9b13d1e1d) + +[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) + +*Note*: Change-in-behavior - locations of downloaded HuggingFace models and components are changed to allow for de-duplication of common modules and switched from using system default cache folder to `models/huggingface` +SD.Next will warn on startup on unused cache entries that can be removed. Also, to take advantage of de-duplication, you'll need to delete models from your `models/Diffusers` folder and let SD.Next re-download them! + +### Details for 2025-08-15 - **Models** + - [Qwen-Image](https://qwenlm.github.io/blog/qwen-image/) + new image foundational model with *20B* params DiT and using *Qwen2.5-VL-7B* as the text-encoder! + available via *networks -> models -> reference* + *note*: this model is almost 2x the size of Flux, quantization and offloading are highly recommended! + *recommended* params: *steps=50, attention-guidance=4* + also available is pre-packaged [Qwen-Lightning](https://huggingface.co/vladmandic/Qwen-Lightning) + which is an unofficial merge of [Qwen-Image](https://qwenlm.github.io/blog/qwen-image/) with [Qwen-Lightning-LoRA](https://github.com/ModelTC/Qwen-Image-Lightning/) to improve quality and allow for generating in 8-steps! - [FLUX.1-Krea-Dev](https://www.krea.ai/blog/flux-krea-open-source-release) new 12B base model compatible with FLUX.1-Dev from *Black Forest Labs* with opinionated aesthetics and aesthetic preferences in mind - simply select in *networks -> models -> reference* + available via *networks -> models -> reference* - [Chroma](https://huggingface.co/lodestones/Chroma) great model based on FLUX.1 and then redesigned and retrained by *lodestones* - update with latest **v48**, **v48 Detail Calibrated** and **v46 Flash** variants - simply select in *networks -> models -> reference* + update with latest **HD**, **HD Flash** and **HD Annealed** variants which are based on *v50* release + available via *networks -> models -> reference* + - [SkyReels-V2](https://huggingface.co/Skywork/SkyReels-V2-DF-14B-720P-Diffusers) + SkyReels-V2 is a genarative video model based on Wan-2.1 but with heavily modified execution to allow for infinite-length video generation + supported variants are: + - diffusion-forcing: *T2I DF 1.3B* for 540p videos, *T2I DF 14B* for 720p videos, *I2I DF 14B* for 720p videos + - standard: *T2I 14B* for 720p videos and *I2I 14B* for 720p videos + - [Wan-VACE](https://huggingface.co/Wan-AI/Wan2.1-VACE-14B-diffusers) + basic support for *Wan 2.1 VACE 1.3B* and *14B* variants + optimized support with granular guidance control will follow soon + - [HunyuanDiT-Distilled](https://huggingface.co/Tencent-Hunyuan/HunyuanDiT-v1.2-Diffusers-Distilled) + variant of HunyuanDiT with reduced steps and improved performance +**Torch** + - Set default to `torch==2.8.0` for *CUDA, ROCm and OpenVINO* + - Add support for `torch==2.9.0-nightly` - **UI** - new embedded docs/wiki search! **Docs** search: fully-local and works in real-time on all document pages **Wiki** search: uses github api to search online wiki pages + - updated real-time hints, thanks @CalamitousFelicitousness + - add **Wilcards** UI + in networks display + - every heading element is collapsible! - quicksettings reset button to restore all quicksettings to default values because things do sometimes get wrong... - - updated real-time hints, thanks @CalamitousFelicitousness - - modernui checkbox/radio styling -- **Offloading** + - configurable image fit in all image views + - rewritten **CivitAI downloader** + in *models -> civitai* + *hint*: you can enter model id in a search bar to pull information on specific model directly + *hint*: you can download individual versions or batch-download all-at-once! + - redesigned **GPU monitor** + - standard-ui: *system -> gpu monitor* + - modern-ui: *aside -> console -> gpu monitor* + - supported for *nVidia CUDA* and *AMD ROCm* platforms + - configurable interval in *settings -> user interface* + - updated *models* tab + - updated *models -> current* tab + - updated *models -> list models* tab + - updated *models -> metadata* tab + - updated *extensions* tab + - redesigned *settings -> user interface* + - gallery bypass browser cache for thumbnails + - gallery safer delete operation + - networks display indicator for currently active items + applies to: *styles, loras* + - apply privacy blur to hf and civitai tokens + - image download will now use actual image filename + - increase default and maximum ui request timeout to 2min/5min + - *hint*: card layout + card layout is used by networks, gallery, civitai search, etc. + you can change card size in *settings -> user interface* +- **Offloading** - changed **default** values for offloading based on detected gpu memory see [offloading docs](https://vladmandic.github.io/sdnext-docs/Offload/) for details - new feature to specify which modules to offload always or never in *settings -> model offloading -> offload always/never* - new `highvram` profile provides significant performance boost on gpus with more than 24gb - - new `offload during pre-forward` option - in *settings -> model offloading* + - new `offload during pre-forward` option + in *settings -> model offloading* switches from explicit offloading to implicit offloading on module execution change - new `diffusers_offload_nonblocking` exerimental setting instructs torch to use non-blocking move operations when possible - **Features** + - new `T5: Use shared instance of text encoder` option + in *settings -> text encoder* + since a lot of new models use T5 text encoder, this option allows to share + the same instance across all models without duplicate downloads + *note* this will not reduce size of your already downloaded models, but will reduce size of future downloads - **Wan** select which stage to run: *first/second/both* with configurable *boundary ration* when running both stages in settings -> model options - prompt parser allow explict `BOS` and `EOS` tokens in prompt - **Nunchaku** support for *FLUX.1-Fill* and *FLUX.1-Depth* models + - update requirements/packages + - use model vae scale-factor for image width/heigt calculations + - **SDNQ** add `modules_dtype_dict` to quantize *Qwen Image* with mixed dtype + - **prompt enhance** + add `allura-org/Gemma-3-Glitter-4B`, `Qwen/Qwen3-4B-Instruct-2507`, `Qwen/Qwen2.5-VL-3B-Instruct` model support + improve system prompt + - **schedulers** add **Flash FlowMatch** + - **model loader** add parallel loader option + enabled by default, selectable in *settings -> model loading* + - **filename namegen** use exact sequence number instead of next available + this allows for more predictable and consistent filename generation + - **network delete** new feature that allows to delete network from disk + in *networks -> show details -> delete* + this will also delete description, metadata and previews associated with the network + only applicable to safetensors networks, not downloaded diffuser models +- **Wiki** + - Models page updated with links to original model repos and model licenses, thanks @alerikaisattera + - Updated Model-Support with newly supported models + - Updated Offload, Prompting, API pages +- **API** + - add `/sdapi/v1/checkpoint` POST endpoint to simply load a model + - add `/sdapi/v1/modules` GET endpoint to get info on model components/modules + - all generate endpoints now support `sd_model_checkpoint` parameter + this allows to specify which model to use for generation without needing to use additional endpoints +- **Refactor** + - change default huggingface cache folder from system default to `models/huggingface` + sd.next will warn on startup on unused cache entries + - new unified pipeline component loader in `pipelines/generic` + - remove **LDSR** + - remove `api-only` cli option +- **Docker** + - update cuda base image: `pytorch/pytorch:2.8.0-cuda12.8-cudnn9-runtime` + - update official builds: - **Fixes** - refactor legacy processing loop - - fix Wan 2.2-5B I2V workflow + - fix settings components mismatch + - fix *Wan 2.2-5B I2V* workflow + - fix *Wan* T2I workflow - fix OpenVINO + - fix video model vs pipeline mismatch + - fix video generic save frames - fix inpaint image metadata - fix processing image save loop - fix progress bar with refine/detailer - fix api progress reporting endpoint - - fix openvino backend failing to compile - - fix nunchaku fallback on unsupported model + - fix `openvino` backend failing to compile + - fix `zluda` with hip-sdk==6.4 + - fix `nunchaku` fallback on unsupported model + - fix `nunchaku` windows download links + - fix *Flux.1-Kontext-Dev* with variable resolution + - use `utf_16_be` as primary metadata decoding + - fix `sd35` width/height alignment + - fix `nudenet` api + - fix global state tracking + - fix ui tab detection for networks + - fix ui checkbox/radio styling for non-default themes + - fix loading custom transformers and t5 safetensors tunes + - add mtime to reference models + - patch torch version so 3rd party libraries can use expected format + - unified stat size/mtime calls - reapply offloading on ipadapter load - api set default script-name - avoid forced gc and rely on thresholds @@ -59,7 +182,7 @@ Feature highlights include: - [ModernUI](https://github.com/user-attachments/assets/6f156154-0b0a-4be2-94f0-979e9f679501) has quite some redesign which should make it more user friendly and easier to navigate plus several new UI themes If you're still using **StandardUI**, give [ModernUI](https://vladmandic.github.io/sdnext-docs/Themes/) a try! - New models such as [WanAI 2.2](https://wan.video/) in 5B and A14B variants for both *text-to-video* and *image-to-video* workflows as well as *text-to-image* workflow! - and also [FreePix F-Lite](https://huggingface.co/Freepik/F-Lite), [Bria 3.2](https://huggingface.co/briaai/BRIA-3.2) and [bigASP 2.5](https://civitai.com/models/1789765?modelVersionId=2025412) + and also [FreePik F-Lite](https://huggingface.co/Freepik/F-Lite), [Bria 3.2](https://huggingface.co/briaai/BRIA-3.2) and [bigASP 2.5](https://civitai.com/models/1789765?modelVersionId=2025412) - Redesigned [Video](https://vladmandic.github.io/sdnext-docs/Video) interface with support for general video models plus optimized [FramePack](https://vladmandic.github.io/sdnext-docs/FramePack) and [LTXVideo](https://vladmandic.github.io/sdnext-docs/LTX) support - Fully integrated nudity detection and optional censorship with [NudeNet](https://vladmandic.github.io/sdnext-docs/NudeNet) - New background replacement and relightning methods using **Latent Bridge Matching** and new **PixelArt** processing filter @@ -104,7 +227,7 @@ For details, see [ChangeLog](https://github.com/vladmandic/automatic/blob/master can also load unet/transformer fine-tunes in safetensors format using UNET loader simply select in *networks -> models -> reference* *note* 1.3B model is a bit too small for good results and 14B is very large at 78GB even without second-stage so aggressive quantization and offloading are recommended - - [FreePix F-Lite](https://huggingface.co/Freepik/F-Lite) in *7B, 10B and Texture* variants + - [FreePik F-Lite](https://huggingface.co/Freepik/F-Lite) in *7B, 10B and Texture* variants F-Lite is a 7B/10B model trained exclusively on copyright-safe and SFW content, trained on internal dataset comprising approximately 80 million copyright-safe images available via *networks -> models -> reference* - [Bria 3.2](https://huggingface.co/briaai/BRIA-3.2) @@ -338,6 +461,7 @@ And (as always) many bugfixes and improvements to existing features! *note*: you need to enable quantization and choose what it applies on, then xyz grid can change quantization mode *note*: you can also enable 'add time info' to compare performance of different quantization modes - **API** + - Add `/sdapi/v1/network?page=&item=` endpoint that returns full network info - Add `/sdapi/v1/lora?lora=` endpoint that returns full lora info and metadata - Add `/sdapi/v1/controlnets?model_type=` endpoints that returns list of available controlnets for specific model type - Set default sampler to `Default` diff --git a/TODO.md b/TODO.md index 0bbf4d8ae..389be8963 100644 --- a/TODO.md +++ b/TODO.md @@ -2,14 +2,13 @@ Main ToDo list can be found at [GitHub projects](https://github.com/users/vladmandic/projects) -## Current - -- Gallery: force refresh on delete - ## Future Candidates +- Remote TE +- Unified `CLIPTextModelWithProjection` loader - [Modular pipelines and guiders](https://github.com/huggingface/diffusers/issues/11915) - Refactor: Sampler options +- Refactor: [GGUF](https://huggingface.co/docs/diffusers/main/en/quantization/gguf) - Feature: Diffusers [group offloading](https://github.com/vladmandic/sdnext/issues/4049) - Feature: Common repo for `T5` and `CLiP` - Feature: LoRA add OMI format support for SD35/FLUX.1 @@ -19,7 +18,6 @@ Main ToDo list can be found at [GitHub projects](https://github.com/users/vladma - Video: LTX PromptEnhance - Video: LTX Conditioning preprocess - [WanAI-2.1 VACE](https://huggingface.co/Wan-AI/Wan2.1-VACE-14B)(https://github.com/huggingface/diffusers/pull/11582) -- [SkyReels-v2](https://github.com/SkyworkAI/SkyReels-V2)(https://github.com/huggingface/diffusers/pull/11518) - [Cosmos-Predict2-Video](https://huggingface.co/nvidia/Cosmos-Predict2-2B-Video2World)(https://github.com/huggingface/diffusers/pull/11695) ### Blocked items @@ -28,23 +26,13 @@ Main ToDo list can be found at [GitHub projects](https://github.com/users/vladma - see - blocked by `insightface` -## ModernUI - -- Extensions tab: - - full CSS redesign -- Models tab: - - Validate subtab: replace table with custom html - - Update subtab: replace table with custom html - - CivitAI subtab: redesign downloader - ### Under Consideration - [IPAdapter negative guidance](https://github.com/huggingface/diffusers/discussions/7167) - [IPAdapter composition](https://huggingface.co/ostris/ip-composition-adapter) - [STG](https://github.com/huggingface/diffusers/blob/main/examples/community/README.md#spatiotemporal-skip-guidance) - [SmoothCache](https://github.com/huggingface/diffusers/issues/11135) -- [MagCache](https://github.com/lllyasviel/FramePack/pull/673/files) -- [HiDream GGUF](https://github.com/huggingface/diffusers/pull/11550) +- [MagCache](https://github.com/lllyasviel/FramePack/pull/673/files) - [Nunchaku PulID](https://github.com/mit-han-lab/nunchaku/pull/274) - [Dream0 guidance](https://huggingface.co/ByteDance/DreamO) - [SUPIR upscaler](https://github.com/Fanghua-Yu/SUPIR) @@ -86,7 +74,6 @@ Main ToDo list can be found at [GitHub projects](https://github.com/users/vladma - control: support scripts via api - fc: autodetect distilled based on model - fc: autodetect tensor format based on model -- flux: loader for civitai nf4 models - hypertile: vae breaks when using non-standard sizes - install: enable ROCm for windows when available - loader: load receipe @@ -94,11 +81,10 @@ Main ToDo list can be found at [GitHub projects](https://github.com/users/vladma - lora: add other quantization types - lora: add t5 key support for sd35/f1 - lora: maybe force imediate quantization -- lora: support pre-quantized flux -- model load: cogview4: balanced offload does not work for GlmModel -- model load: add ChromaFillPipeline, ChromaControlPipeline, ChromaImg2ImgPipeline etc when available - model load: force-reloading entire model as loading transformers only leads to massive memory usage +- model load: group offload - model load: implement model in-memory caching - modernui: monkey-patch for missing tabs.select event +- modules/lora/lora_extract.py:188:9: W0511: TODO: lora: support pre-quantized flux - processing: remove duplicate mask params - resize image: enable full VAE mode for resize-latent diff --git a/cli/api-checkpoint.py b/cli/api-checkpoint.py index 61f4e4370..ff939f64e 100755 --- a/cli/api-checkpoint.py +++ b/cli/api-checkpoint.py @@ -35,3 +35,5 @@ def get(endpoint: str, dct: dict = None): if __name__ == "__main__": model = get('/sdapi/v1/checkpoint') log.info(f'api-checkpoint: {model}') + model = get('/sdapi/v1/modules') + log.info(f'api-modules: {model}') diff --git a/cli/api-control.py b/cli/api-control.py index 79667a2e9..925c77599 100755 --- a/cli/api-control.py +++ b/cli/api-control.py @@ -108,6 +108,7 @@ def generate(args): # pylint: disable=redefined-outer-name if args.mask is not None: options['mask'] = encode(args.mask) + data = post('/sdapi/v1/control', options) t1 = time.time() if 'info' in data: diff --git a/cli/civitai-search.py b/cli/civitai-search.py new file mode 100755 index 000000000..3d91e8711 --- /dev/null +++ b/cli/civitai-search.py @@ -0,0 +1,217 @@ +#!/usr/bin/env python +from dataclasses import dataclass +import os +import sys +import json +import time +import logging + + +full_dct = False +full_html = False +debug = False +logging.basicConfig(level = logging.INFO, format = '%(asctime)s %(levelname)s: %(message)s') +log = logging.getLogger(__name__) + + +@dataclass +class ModelImage(): + def __init__(self, dct: dict): + if isinstance(dct, str): + dct = json.loads(dct) + self.id: int = dct.get('id', 0) + self.url: str = dct.get('url', '') + self.width: int = dct.get('width', 0) + self.height: int = dct.get('height', 0) + self.type: str = dct.get('type', 'Unknown') + self.dct: dict = dct if full_dct else {} + + def __str__(self): + return f'ModelImage(id={self.id} url="{self.url}" width={self.width} height={self.height} type="{self.type}")' + +@dataclass +class ModelFile(): + def __init__(self, dct: dict): + if isinstance(dct, str): + dct = json.loads(dct) + self.id: int = dct.get('id', 0) + self.size: int = int(1024 * dct.get('sizeKB', 0)) + self.name: str = dct.get('name', 'Unknown') + self.type: str = dct.get('type', 'Unknown') + self.hashes: list[str] = dct.get('hashes', {}).values() + self.url: str = dct.get('downloadUrl', '') + self.dct: dict = dct if full_dct else {} + + def __str__(self): + return f'ModelFile(id={self.id} name="{self.name}" size={self.size} type="{self.type}" url="{self.url}")' + + +@dataclass +class ModelVersion(): + def __init__(self, dct: dict): + import bs4 + if isinstance(dct, str): + dct = json.loads(dct) + self.id: int = dct.get('id', 0) + self.name: str = dct.get('name', 'Unknown') + self.base: str = dct.get('baseModel', 'Unknown') + self.mtime: str = dct.get('publishedAt', '') + self.downloads: int = dct.get('stats', {}).get('downloadCount', 0) + self.availability: str = dct.get('availability', 'Unknown') + self.html: str = dct.get('description', '') or '' if full_html else '' + self.desc: str = bs4.BeautifulSoup(dct.get('description', '') or '', features="html.parser").get_text() + self.files = [ModelFile(f) for f in dct.get('files', [])] + self.images = [ModelImage(i) for i in dct.get('images', [])] + self.dct: dict = dct if full_dct else {} + + def __str__(self): + return f'ModelVersion(id={self.id} name="{self.name}" base="{self.base}" mtime="{self.mtime}" downloads={self.downloads} availability={self.availability} desc="{self.desc[:30]}...")' + + +@dataclass +class Model(): + def __init__(self, dct: dict): + import bs4 + if isinstance(dct, str): + dct = json.loads(dct) + self.id: int = dct.get('id', 0) + self.url: str = f'https://civitai.com/models/{self.id}' + self.type: str = dct.get('type', 'Unknown') + self.name: str = dct.get('name', 'Unknown') + self.html: str = dct.get('description', '') or '' if full_html else '' + self.desc: str = bs4.BeautifulSoup(dct.get('description', '') or '', features="html.parser").get_text() + self.tags: list[str] = dct.get('tags', []) + self.nsfw: bool = dct.get('nsfw', False) + self.level: str = dct.get('nsfwLevel', 0) + self.availability: str = dct.get('availability', 'Unknown') + self.downloads: int = dct.get('stats', {}).get('downloadCount', 0) + self.creator: str = dct.get('creator', {}).get('username', 'Unknown') + self.versions: list[ModelVersion] = [ModelVersion(v) for v in dct.get('modelVersions', [])] + self.dct: dict = dct if full_dct else {} + + def __str__(self): + return f'Model(id={self.id} type={self.type} name="{self.name}" versions={len(self.versions)} nsfw={self.nsfw}/{self.level} downloads={self.downloads} author="{self.creator}" tags={self.tags} desc="{self.desc[:30]}...")' + + +def search_civitai( + query:str, + tag:str = '', # optional:tag name + types:str = '', # (Checkpoint, TextualInversion, Hypernetwork, AestheticGradient, LORA, Controlnet, Poses) + sort:str = '', # (Highest Rated, Most Downloaded, Newest) + period:str = '', # (AllTime, Year, Month, Week, Day) + nsfw:bool = None, # optional:bool + limit:int = 0, + base:list[str] = [], # list + token:str = None, + exact:bool = True, +): + import requests + from urllib.parse import urlencode + + if len(query) == 0: + log.error('CivitAI: empty query') + return [] + + t0 = time.time() + dct = { 'query': query } + if len(tag) > 0: + dct['tag'] = tag + if nsfw is not None: + dct['nsfw'] = 'true' if nsfw else 'false' + if limit > 0: + dct['limit'] = limit + if len(types) > 0: + dct['types'] = types + if len(sort) > 0: + dct['sort'] = sort + if len(period) > 0: + dct['period'] = period + if len(base) > 0: + dct['baseModels'] = ','.join(base) + encoded = urlencode(dct) + + headers = {} + if token is None: + token = os.environ.get('CIVITAI_TOKEN', None) + if token is not None and len(token) > 0: + headers['Authorization'] = f'Bearer {token}' + + url = 'https://civitai.com/api/v1/models' + uri = f'{url}?{encoded}' + log.info(f'CivitAI request: uri="{uri}" dct={dct} token={token is not None}') + result = requests.get(uri, headers=headers, timeout=60) + + if result.status_code != 200: + log.error(f'CivitAI: code={result.status_code} reason={result.reason} uri={result.url}') + return [] + + models: list[Model] = [] + exact_models: list[Model] = [] + items = result.json().get('items', []) + for item in items: + models.append(Model(item)) + + if exact: + for model in models: + model_names = [model.name.lower()] + version_names = [v.name.lower() for v in model.versions] + file_names = [f.name.lower() for v in model.versions for f in v.files] + if any([query.lower() in name for name in model_names + version_names + file_names]): # noqa: C419 + exact_models.append(model) + + t1 = time.time() + log.info(f'CivitAI result: code={result.status_code} exact={len(exact_models)} total={len(models)} time={t1-t0:.2f}') + return exact_models if len(exact_models) > 0 else models + + +def models_to_dct(all_models:list, model_id:int=None): + dct = [] + for model in all_models: + if model_id is not None and model.id != model_id: + continue + model_dct = model.__dict__.copy() + versions_dct = [] + for version in model.versions: + version_dct = version.__dict__.copy() + version_dct['files'] = [f.__dict__.copy() for f in version.files] + version_dct['images'] = [i.__dict__.copy() for i in version.images] + versions_dct.append(version_dct) + model_dct['versions'] = versions_dct + dct.append(model_dct) + return dct + + +def print_models(models: list[Model]): + if debug: + from rich import print as dbg + else: + dbg = lambda *args, **kwargs: None # pylint: disable=unnecessary-lambda-assignment + for model in models: + log.info(f' {model}') + dbg('Model', model.dct) + for version in model.versions: + log.info(f' {version}') + dbg('ModelVersion', version.dct) + for file in version.files: + log.info(f' {file}') + dbg('ModelFile', file.dct) + for image in version.images: + log.info(f' {image}') + dbg('ModelImage', image.dct) + + +if __name__ == "__main__": + sys.argv.pop(0) + txt = ' '.join(sys.argv) + res = search_civitai( + query=txt, + # tag = '', + # types = '', + # sort = 'Most Downloaded', + # period = 'Year', + # nsfw = True, + # base = [], + # exact= True, + # limit=100, + ) + print_models(res) diff --git a/cli/image-exif.py b/cli/image-exif.py index 9a48d2dd7..fc2573220 100755 --- a/cli/image-exif.py +++ b/cli/image-exif.py @@ -64,7 +64,10 @@ class Exif: # pylint: disable=single-string-used-for-slots def decode(self, s: bytes): remove_prefix = lambda text, prefix: text[len(prefix):] if text.startswith(prefix) else text # pylint: disable=unnecessary-lambda-assignment - for encoding in ['utf-8', 'utf-16', 'ascii', 'latin_1', 'cp1252', 'cp437']: # try different encodings + # from encodings.aliases import aliases + # cp = list(set(aliases.values())) + for encoding in ['utf_16_be', 'utf-8', 'utf-16', 'ascii', 'latin_1', 'cp1252', 'cp437']: # try different encodings + # for encoding in cp: try: s = remove_prefix(s, b'UNICODE') s = remove_prefix(s, b'ASCII') diff --git a/configs/Dockerfile.cuda b/configs/Dockerfile.cuda index ca117bd4c..4e5a496e1 100644 --- a/configs/Dockerfile.cuda +++ b/configs/Dockerfile.cuda @@ -2,7 +2,7 @@ # docs: # base image -FROM pytorch/pytorch:2.7.0-cuda12.8-cudnn9-runtime +FROM pytorch/pytorch:2.8.0-cuda12.8-cudnn9-runtime # metadata LABEL org.opencontainers.image.vendor="SD.Next" @@ -13,7 +13,7 @@ LABEL org.opencontainers.image.source="https://github.com/vladmandic/sdnext/" LABEL org.opencontainers.image.licenses="AGPL-3.0" LABEL org.opencontainers.image.title="SD.Next" LABEL org.opencontainers.image.description="SD.Next: Advanced Implementation of Stable Diffusion and other Diffusion-based generative image models" -LABEL org.opencontainers.image.base.name="https://hub.docker.com/pytorch/pytorch:2.7.0-cuda12.8-cudnn9-runtime" +LABEL org.opencontainers.image.base.name="https://hub.docker.com/pytorch/pytorch:2.8.0-cuda12.8-cudnn9-runtime" LABEL org.opencontainers.image.version="latest" # minimum install diff --git a/extensions-builtin/sd-extension-system-info b/extensions-builtin/sd-extension-system-info index 0760f3bce..615d2f810 160000 --- a/extensions-builtin/sd-extension-system-info +++ b/extensions-builtin/sd-extension-system-info @@ -1 +1 @@ -Subproject commit 0760f3bcee4cd8448089e749dd7c22cdfebf15c3 +Subproject commit 615d2f8103fb61779037bd81523906457fbf7277 diff --git a/extensions-builtin/sdnext-modernui b/extensions-builtin/sdnext-modernui index 43ed2ea51..574fcf4e8 160000 --- a/extensions-builtin/sdnext-modernui +++ b/extensions-builtin/sdnext-modernui @@ -1 +1 @@ -Subproject commit 43ed2ea51049a5926fe553aed3e9a717657352a6 +Subproject commit 574fcf4e8790e6faf3a3a500e4aedf399d0b0e4a diff --git a/html/locale_en.json b/html/locale_en.json index b222c2ed5..d35d8a96d 100644 --- a/html/locale_en.json +++ b/html/locale_en.json @@ -19,6 +19,7 @@ {"id":"","label":"🖌️","localized":"","hint":"LaMa remove selected object from image"}, {"id":"","label":"🖼️","localized":"","hint":"Show preview"}, {"id":"","label":"♻","localized":"","hint":"Interrogate image"}, + {"id":"","label":"⁜","localized":"","hint":"Cycle image fit method"}, {"id":"","label":"↶","localized":"","hint":"Apply selected style to prompt"}, {"id":"","label":"↷","localized":"","hint":"Save current prompt to style"}, {"id":"","label":"","localized":"","hint":"Sort by name, ascending"}, @@ -92,7 +93,7 @@ {"id":"","label":"Denoise","localized":"","hint":"Denoising settings. Higher denoise means that more of existing image content is allowed to change during generate"}, {"id":"","label":"Mask","localized":"","hint":"Image masking and mask options"}, {"id":"","label":"Input","localized":"","hint":"Selection of input media"}, - {"id":"","label":"Video","localized":"","hint":"Settings related to video generation"}, + {"id":"","label":"Video","localized":"","hint":"Create video using guidance"}, {"id":"","label":"Control elements","localized":"","hint":"Control elements are advanced models that can guide generation towards desired outcome"}, {"id":"","label":"IP adapter","localized":"","hint":"Guide generation towards desired outcome using IP adapters plugin models"}, {"id":"","label":"IP adapters","localized":"","hint":"IP adapters are plugin models that can guide generation towards desired outcome"}, diff --git a/html/previews.json b/html/previews.json index 3b59d376c..aabb38822 100644 --- a/html/previews.json +++ b/html/previews.json @@ -14,5 +14,8 @@ "Efficient-Large-Model--Sana_1600M_4Kpx_BF16_diffusers": "models/Reference/Efficient-Large-Model--Sana_1600M_1024px_diffusers.jpg", "Efficient-Large-Model--Sana_600M_1024px_diffusers": "models/Reference/Efficient-Large-Model--Sana_1600M_1024px_diffusers.jpg", "stabilityai--stable-video-diffusion-img2vid-xt-1-1": "models/Reference/stabilityai--stable-video-diffusion-img2vid-xt.jpg", - "shuttleai--shuttle-3-diffusion": "models/Reference/shuttleai--shuttle-3-diffusion.jpg" + "shuttleai--shuttle-3-diffusion": "models/Reference/shuttleai--shuttle-3-diffusion.jpg", + "HiDream-ai/HiDream-I1-Full": "models/Reference/HiDream-I1 Full", + "vladmandic/chroma-unlocked-v50": "models/Reference/lodestones Chroma Unlocked HD", + "vladmandic/chroma-unlocked-v50-annealed": "models/Reference/lodestones Chroma Unlocked HD" } diff --git a/html/reference.json b/html/reference.json index 4f21017ac..d2d1e6e6d 100644 --- a/html/reference.json +++ b/html/reference.json @@ -1,25 +1,26 @@ + { - "Tempest SD-XL v0.1": { - "path": "TempestV0.1-Artistic.safetensors@https://huggingface.co/dataautogpt3/TempestV0.1/resolve/main/TempestV0.1-Artistic.safetensors?download=true", - "preview": "TempestV0.1-Artistic.jpg", - "desc": "The TempestV0.1 Initiative is a powerhouse in image generation, leveraging an unparalleled dataset of over 6 million images. The collection's vast scale, with resolutions from 1400x2100 to 4800x7200, encompasses 200GB of high-quality content.", - "extras": "width: 2048, height: 1024, sampler: DEIS, steps: 40, cfg_scale: 6.0" + "Tempest-by-Vlad XL": { + "path": "tempestByVlad_baseV01.safetensors@https://civitai.com/api/download/models/1301775", + "preview": "tempest-by-vlad-base.jpg", + "desc": "Flexible SDXL model with custom encoder and finetuned for larger landscape resolutions with high details and high contrast.", + "extras": "" + }, + "Tempest-by-Vlad XL Hyper": { + "path": "tempestByVlad_hyperV01.safetensors@https://civitai.com/api/download/models/1343512", + "preview": "tempest-by-vlad-hyper.jpg", + "desc": "Custom distilled variant with goal to get as-normal-as-possible model that works with low steps and guidance-free", + "extras": "" }, - "Juggernaut SD-XL XI": { + "Juggernaut XL XI": { "path": "juggernautXL_juggXIByRundiffusion.safetensors@https://civitai.com/api/download/models/782002", "preview": "juggernautXL_v9Rundiffusionphoto2.jpg", "desc": "Showcase finetuned model based on Stable diffusion XL", "extras": "sampler: DEIS, steps: 20, cfg_scale: 6.0" }, - "Juggernaut SD-XL X Hyper": { - "path": "Juggernaut_X_RunDiffusion_Hyper.safetensors@https://civitai.com/api/download/models/471120", - "preview": "juggernautXL_v9Rundiffusionphoto2.jpg", - "desc": "Showcase finetuned model based on Stable diffusion XL", - "extras": "sampler: DEIS, steps: 20, cfg_scale: 6.0" - }, - "Juggernaut SD-XL IX Lightning": { - "path": "juggernautXL_v9Rdphoto2Lightning.safetensors@https://civitai.com/api/download/models/357609", + "Juggernaut XL XI Lightning": { + "path": "juggernautXL_juggXILightningByRD.safetensors@https://civitai.com/api/download/models/920957", "preview": "juggernautXL_v9Rdphoto2Lightning.jpg", "desc": "Showcase finetuned model based on Stable diffusion XL", "extras": "sampler: DPM SDE, steps: 6, cfg_scale: 2.0" @@ -32,40 +33,6 @@ "extras": "width: 512, height: 512, sampler: DEIS, steps: 20, cfg_scale: 6.0" }, - "DreamShaper SD v8": { - "original": true, - "path": "dreamshaper_8.safetensors@https://civitai.com/api/download/models/128713", - "preview": "dreamshaper_8.jpg", - "desc": "Showcase finetuned model based on Stable diffusion 1.5", - "extras": "width: 512, height: 512, sampler: DEIS, steps: 20, cfg_scale: 6.0" - }, - "Dreamshaper SD v7 LCM": { - "path": "SimianLuo/LCM_Dreamshaper_v7", - "preview": "SimianLuo--LCM_Dreamshaper_v7.jpg", - "desc": "Latent Consistencey Models enable swift inference with minimal steps on any pre-trained LDMs, including Stable Diffusion. By distilling classifier-free guidance into the model's input, LCM can generate high-quality images in very short inference time. LCM can generate quality images in as few as 3-4 steps, making it blazingly fast.", - "extras": "width: 512, height: 512, sampler: LCM, steps: 4, cfg_scale: 0.0" - }, - "DreamShaper SD-XL Turbo": { - "path": "dreamshaperXL_v21TurboDPMSDE.safetensors@https://civitai.com/api/download/models/351306", - "preview": "dreamshaperXL_v21TurboDPMSDE.jpg", - "desc": "Showcase finetuned model based on Stable diffusion XL", - "extras": "sampler: DPM SDE, steps: 8, cfg_scale: 2.0" - }, - - "SDXS DreamShaper 512": { - "path": "IDKiro/sdxs-512-dreamshaper", - "preview": "IDKiro--sdxs-512-dreamshaper.jpg", - "desc": "SDXS: Real-Time One-Step Latent Diffusion Models with Image Conditions", - "extras": "width: 512, height: 512, sampler: CMSI, steps: 1, cfg_scale: 0.0" - }, - "SDXL Flash Mini": { - "path": "SDXL-Flash_Mini.safetensors@https://huggingface.co/sd-community/sdxl-flash-mini/resolve/main/SDXL-Flash_Mini.safetensors?download=true", - "preview": "SDXL-Flash_Mini.jpg", - "desc": "Introducing the new fast model SDXL Flash (Mini), we learned that all fast XL models work fast, but the quality decreases, and we also made a fast model, but it is not as fast as LCM, Turbo, Lightning and Hyper, but the quality is higher.", - "extras": "width: 2048, height: 1024, sampler: DEIS, steps: 40, cfg_scale: 6.0", - "experimental": true - }, - "RunwayML StableDiffusion 1.5": { "original": true, "path": "v1-5-pruned-fp16-emaonly.safetensors@https://huggingface.co/Aptronym/SDNext/resolve/main/Reference/v1-5-pruned-fp16-emaonly.safetensors?download=true", @@ -173,28 +140,57 @@ "extras": "sampler: Default, cfg_scale: 4.5" }, + "lodestones Chroma Unlocked HD": { + "path": "lodestones/Chroma1-HD", + "preview": "lodestones--Chroma.jpg", + "desc": "Chroma is a 8.9B parameter model based on FLUX.1-schnell. It’s fully Apache 2.0 licensed, ensuring that anyone can use, modify, and build on top of it—no corporate gatekeeping. The model is still training right now, and I’d love to hear your thoughts! Your input and feedback are really appreciated.", + "skip": true, + "extras": "sampler: Default, cfg_scale: 3.5" + }, + "lodestones Chroma Unlocked HD Annealed": { + "path": "vladmandic/chroma-unlocked-v50-annealed", + "preview": "lodestones--Chroma.jpg", + "desc": "Chroma is a 8.9B parameter model based on FLUX.1-schnell. It’s fully Apache 2.0 licensed, ensuring that anyone can use, modify, and build on top of it—no corporate gatekeeping. The model is still training right now, and I’d love to hear your thoughts! Your input and feedback are really appreciated.", + "skip": true, + "extras": "sampler: Default, cfg_scale: 3.5" + }, + "lodestones Chroma Unlocked HD Flash": { + "path": "lodestones/Chroma1-Flash", + "preview": "lodestones--Chroma.jpg", + "desc": "Chroma is a 8.9B parameter model based on FLUX.1-schnell. It’s fully Apache 2.0 licensed, ensuring that anyone can use, modify, and build on top of it—no corporate gatekeeping. The model is still training right now, and I’d love to hear your thoughts! Your input and feedback are really appreciated.", + "skip": true, + "extras": "sampler: Default, cfg_scale: 1.0" + }, "lodestones Chroma Unlocked v48": { "path": "vladmandic/chroma-unlocked-v48", "preview": "lodestones--Chroma.jpg", "desc": "Chroma is a 8.9B parameter model based on FLUX.1-schnell. It’s fully Apache 2.0 licensed, ensuring that anyone can use, modify, and build on top of it—no corporate gatekeeping. The model is still training right now, and I’d love to hear your thoughts! Your input and feedback are really appreciated.", "skip": true, - "extras": "sampler: Default, cfg_scale: 3.5" + "extras": "sampler: Default, cfg_scale: 1.0" }, "lodestones Chroma Unlocked v48 Detail Calibrated": { "path": "vladmandic/chroma-unlocked-v48-detail-calibrated", "preview": "lodestones--Chroma.jpg", "desc": "Chroma is a 8.9B parameter model based on FLUX.1-schnell. It’s fully Apache 2.0 licensed, ensuring that anyone can use, modify, and build on top of it—no corporate gatekeeping. The model is still training right now, and I’d love to hear your thoughts! Your input and feedback are really appreciated.", "skip": true, - "extras": "sampler: Default, cfg_scale: 3.5" - }, - "lodestones Chroma Unlocked v48 Flash": { - "path": "vladmandic/chroma-unlocked-v46-flash", - "preview": "lodestones--Chroma.jpg", - "desc": "Chroma is a 8.9B parameter model based on FLUX.1-schnell. It’s fully Apache 2.0 licensed, ensuring that anyone can use, modify, and build on top of it—no corporate gatekeeping. The model is still training right now, and I’d love to hear your thoughts! Your input and feedback are really appreciated.", - "skip": true, "extras": "sampler: Default, cfg_scale: 1.0" }, + "Qwen-Image": { + "path": "Qwen/Qwen-Image", + "preview": "Qwen--Qwen-Image.jpg", + "desc": " Qwen-Image, an image generation foundation model in the Qwen series that achieves significant advances in complex text rendering and precise image editing.", + "skip": true, + "extras": "" + }, + "Qwen-Lightning": { + "path": "vladmandic/Qwen-Lightning", + "preview": "Qwen--Qwen-Image.jpg", + "desc": " Qwen-Lightning is step-distilled from Qwen-Image to allow for generation in 8 steps.", + "skip": true, + "extras": "steps: 8" + }, + "Ostris Flex.2 Preview": { "path": "ostris/Flex.2-preview", "preview": "ostris--Flex.2-preview.jpg", @@ -261,6 +257,20 @@ "extras": "sampler: Default, cfg_scale: 3.5" }, + "SDXS DreamShaper 512": { + "path": "IDKiro/sdxs-512-dreamshaper", + "preview": "IDKiro--sdxs-512-dreamshaper.jpg", + "desc": "SDXS: Real-Time One-Step Latent Diffusion Models with Image Conditions", + "extras": "width: 512, height: 512, sampler: CMSI, steps: 1, cfg_scale: 0.0" + }, + "SDXL Flash Mini": { + "path": "SDXL-Flash_Mini.safetensors@https://huggingface.co/sd-community/sdxl-flash-mini/resolve/main/SDXL-Flash_Mini.safetensors?download=true", + "preview": "SDXL-Flash_Mini.jpg", + "desc": "Introducing the new fast model SDXL Flash (Mini), we learned that all fast XL models work fast, but the quality decreases, and we also made a fast model, but it is not as fast as LCM, Turbo, Lightning and Hyper, but the quality is higher.", + "extras": "width: 2048, height: 1024, sampler: DEIS, steps: 40, cfg_scale: 6.0", + "experimental": true + }, + "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.", @@ -419,6 +429,24 @@ "preview": "Tencent-Hunyuan--HunyuanDiT-v1.2-Diffusers.jpg", "extras": "sampler: Default, cfg_scale: 2.0" }, + "Tencent HunyuanDiT 1.2 Distilled": { + "path": "Tencent-Hunyuan/HunyuanDiT-v1.2-Diffusers-Distilled", + "desc": "Hunyuan-DiT : A Powerful Multi-Resolution Diffusion Transformer with Fine-Grained Chinese Understanding.", + "preview": "Tencent-Hunyuan--HunyuanDiT-v1.2-Diffusers.jpg", + "extras": "sampler: Default, cfg_scale: 2.0" + }, + "Tencent HunyuanDiT 1.1": { + "path": "Tencent-Hunyuan/HunyuanDiT-v1.1-Diffusers", + "desc": "Hunyuan-DiT : A Powerful Multi-Resolution Diffusion Transformer with Fine-Grained Chinese Understanding.", + "preview": "Tencent-Hunyuan--HunyuanDiT-v1.2-Diffusers.jpg", + "extras": "sampler: Default, cfg_scale: 2.0" + }, + "Tencent HunyuanDiT 1.1 Distilled": { + "path": "Tencent-Hunyuan/HunyuanDiT-v1.1-Diffusers-Distilled", + "desc": "Hunyuan-DiT : A Powerful Multi-Resolution Diffusion Transformer with Fine-Grained Chinese Understanding.", + "preview": "Tencent-Hunyuan--HunyuanDiT-v1.2-Diffusers.jpg", + "extras": "sampler: Default, cfg_scale: 2.0" + }, "AlphaVLLM Lumina Next SFT": { "path": "Alpha-VLLM/Lumina-Next-SFT-diffusers", @@ -480,7 +508,7 @@ }, "Kandinsky 2.2": { "path": "kandinsky-community/kandinsky-2-2-decoder", - "desc": "Kandinsky 2.2 is a text-conditional diffusion model (+0.1!) based on unCLIP and latent diffusion, composed of a transformer-based image prior model, a unet diffusion model, and a decoder. Kandinsky 2.1 inherits best practices from Dall-E 2 and Latent diffusion while introducing some new ideas. It uses the CLIP model as a text and image encoder, and diffusion image prior (mapping) between latent spaces of CLIP modalities. This approach increases the visual performance of the model and unveils new horizons in blending images and text-guided image manipulation.", + "desc": "Kandinsky 2.2 is a text-conditional diffusion model (+0.1!) based on unCLIP and latent diffusion, composed of a transformer-based image prior model, a unet diffusion model, and a decoder. Kandinsky 2.2 inherits best practices from Dall-E 2 and Latent diffusion while introducing some new ideas. It uses the CLIP model as a text and image encoder, and diffusion image prior (mapping) between latent spaces of CLIP modalities. This approach increases the visual performance of the model and unveils new horizons in blending images and text-guided image manipulation.", "preview": "kandinsky-community--kandinsky-2-2-decoder.jpg", "extras": "width: 768, height: 768, sampler: Default" }, @@ -524,13 +552,13 @@ }, "CogView 4": { - "path": "THUDM/CogView4-6B", + "path": "zai-org/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", + "path": "zai-org/CogView3-Plus-3B", "desc": "An innovative cascaded framework that enhances the performance of text-to-image diffusion. CogView is the first model implementing relay diffusion in the realm of text-to-image generation, executing the task by first creating low-resolution images and subsequently applying relay-based super-resolution.", "preview": "THUDM--CogView3-Plus-3B.jpg", "skip": true diff --git a/installer.py b/installer.py index a0bbc2828..5e796f119 100644 --- a/installer.py +++ b/installer.py @@ -593,7 +593,7 @@ def check_diffusers(): t_start = time.time() if args.skip_all or args.skip_git: return - sha = '0c71189abeaa8ab4b28dd7e5a309ac75c64968a2' # diffusers commit hash + sha = '58bf2682612bc29b7cdb8a10ba6eee28a024d6d3' # 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 -1) cur = opts.get('diffusers_version', '') if minor > -1 else '' @@ -618,7 +618,7 @@ def check_transformers(): if args.use_directml: target = '4.52.4' else: - target = '4.54.1' + target = '4.55.2' if (pkg is None) or ((pkg.version != target) and (not args.experimental)): if pkg is None: log.info(f'Transformers install: version={target}') @@ -648,8 +648,7 @@ def install_cuda(): if args.use_nightly: cmd = os.environ.get('TORCH_COMMAND', '--upgrade --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cu128 --extra-index-url https://download.pytorch.org/whl/nightly/cu126') else: - # cmd = os.environ.get('TORCH_COMMAND', 'torch==2.6.0+cu126 torchvision==0.21.0+cu126 --index-url https://download.pytorch.org/whl/cu126') - cmd = os.environ.get('TORCH_COMMAND', 'torch==2.7.1+cu128 torchvision==0.22.1+cu128 --index-url https://download.pytorch.org/whl/cu128') + cmd = os.environ.get('TORCH_COMMAND', 'torch==2.8.0+cu128 torchvision==0.23.0+cu128 --index-url https://download.pytorch.org/whl/cu128') return cmd @@ -738,13 +737,13 @@ def install_rocm_zluda(): if args.use_nightly: if rocm.version is None or float(rocm.version) >= 6.4: # assume the latest if version check fails torch_command = os.environ.get('TORCH_COMMAND', '--upgrade --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/rocm6.4') - elif rocm.version == "6.3": + else: # oldest rocm version on nightly is 6.3 torch_command = os.environ.get('TORCH_COMMAND', '--upgrade --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/rocm6.3') - else: # oldest rocm version on nightly is 6.2.4 - torch_command = os.environ.get('TORCH_COMMAND', '--upgrade --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/rocm6.2.4') else: - if rocm.version is None or float(rocm.version) >= 6.3: # assume the latest if version check fails - torch_command = os.environ.get('TORCH_COMMAND', 'torch==2.7.1+rocm6.3 torchvision==0.22.1+rocm6.3 --index-url https://download.pytorch.org/whl/rocm6.3') + if rocm.version is None or float(rocm.version) >= 6.4: # assume the latest if version check fails + torch_command = os.environ.get('TORCH_COMMAND', 'torch==2.8.0+rocm6.4 torchvision==0.23.0+rocm6.4 --index-url https://download.pytorch.org/whl/rocm6.4') + elif rocm.version == "6.3": + torch_command = os.environ.get('TORCH_COMMAND', 'torch==2.8.0+rocm6.3 torchvision==0.23.0+rocm6.3 --index-url https://download.pytorch.org/whl/rocm6.3') elif rocm.version == "6.2": # use rocm 6.2.4 instead of 6.2 as torch==2.7.1+rocm6.2 doesn't exists torch_command = os.environ.get('TORCH_COMMAND', 'torch==2.7.1+rocm6.2.4 torchvision==0.22.1+rocm6.2.4 --index-url https://download.pytorch.org/whl/rocm6.2.4') @@ -804,6 +803,7 @@ def install_ipex(): if args.use_nightly: torch_command = os.environ.get('TORCH_COMMAND', '--upgrade --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/xpu') else: + # torch 2.8 segfaults with torch.compile: https://github.com/pytorch/pytorch/issues/159974 torch_command = os.environ.get('TORCH_COMMAND', 'torch==2.7.1+xpu torchvision==0.22.1+xpu --index-url https://download.pytorch.org/whl/xpu') ts('ipex', t_start) @@ -815,9 +815,9 @@ def install_openvino(): #check_python(supported_minors=[9, 10, 11, 12, 13], reason='OpenVINO backend requires a Python version between 3.9 and 3.13') log.info('OpenVINO: selected') if sys.platform == 'darwin': - torch_command = os.environ.get('TORCH_COMMAND', 'torch==2.7.1 torchvision==0.22.1') + torch_command = os.environ.get('TORCH_COMMAND', 'torch==2.8.0 torchvision==0.23.0') else: - torch_command = os.environ.get('TORCH_COMMAND', 'torch==2.7.1+cpu torchvision==0.22.1+cpu --index-url https://download.pytorch.org/whl/cpu') + torch_command = os.environ.get('TORCH_COMMAND', 'torch==2.8.0+cpu torchvision==0.23.0 --index-url https://download.pytorch.org/whl/cpu') install(os.environ.get('OPENVINO_COMMAND', 'openvino==2025.2.0'), 'openvino') install(os.environ.get('NNCF_COMMAND', 'nncf==2.17.0'), 'nncf') @@ -1310,6 +1310,7 @@ def install_requirements(): # set environment variables controling the behavior of various libraries def set_environment(): + from modules.paths import models_path log.debug('Setting environment tuning') os.environ.setdefault('ACCELERATE', 'True') os.environ.setdefault('ATTN_PRECISION', 'fp16') @@ -1336,7 +1337,7 @@ def set_environment(): os.environ.setdefault('DO_NOT_TRACK', '1') os.environ.setdefault('UV_INDEX_STRATEGY', 'unsafe-any-match') os.environ.setdefault('UV_NO_BUILD_ISOLATION', '1') - os.environ.setdefault('HF_HUB_CACHE', opts.get('hfcache_dir', os.path.join(os.path.expanduser('~'), '.cache', 'huggingface', 'hub'))) + os.environ.setdefault('HF_HUB_CACHE', opts.get('hfcache_dir', os.path.join(models_path, 'huggingface'))) allocator = f'garbage_collection_threshold:{opts.get("torch_gc_threshold", 80)/100:0.2f},max_split_size_mb:512' if opts.get("torch_malloc", "native") == 'cudaMallocAsync': allocator += ',backend:cudaMallocAsync' diff --git a/javascript/base.css b/javascript/base.css index 75ca46f55..f6a7f7d09 100644 --- a/javascript/base.css +++ b/javascript/base.css @@ -114,7 +114,7 @@ table.settings-value-table td { padding: 0.4em; border: 1px solid #ccc; max-widt /* custom component */ .folder-selector textarea { height: 2em !important; padding: 6px !important; } -.nvml { position: fixed; bottom: 10px; right: 10px; background: var(--background-fill-primary); border: 1px solid var(--button-primary-border-color); padding: 6px; color: var(--button-primary-text-color); +.gpu { position: fixed; bottom: 10px; right: 10px; background: var(--background-fill-primary); border: 1px solid var(--button-primary-border-color); padding: 6px; color: var(--button-primary-text-color); font-size: 0.7em; z-index: 50; font-family: monospace; display: none; } /* image browser */ diff --git a/javascript/black-teal-reimagined.css b/javascript/black-teal-reimagined.css index 3f1cdf9a9..1e7d4dc0b 100644 --- a/javascript/black-teal-reimagined.css +++ b/javascript/black-teal-reimagined.css @@ -862,19 +862,6 @@ svg.feather.feather-image, border-radius: var(--radius-md); } -/* Overlay Name Styles */ -.extra-network-cards .card .overlay .name { - font-size: var(--text-lg); - font-weight: bold; - text-shadow: 1px 1px black; - color: white; - overflow-wrap: anywhere; - position: absolute; - bottom: 0; - padding: 0.2em; - z-index: 10; -} - /* Preview Styles */ .extra-network-cards .card .preview { box-shadow: var(--button-shadow); @@ -1124,9 +1111,7 @@ svg.feather.feather-image, } /* Based on Gradio Built-in Dark Theme */ -:root, -.light, -.dark { +:root, .light, .dark { --body-background-fill: var(--background-color); --color-accent-soft: var(--neutral-700); --background-fill-secondary: none; diff --git a/javascript/civitai.js b/javascript/civitai.js new file mode 100644 index 000000000..31f1fc890 --- /dev/null +++ b/javascript/civitai.js @@ -0,0 +1,147 @@ +String.prototype.format = function (args) { // eslint-disable-line no-extend-native, func-names + let thisString = ''; + for (let charPos = 0; charPos < this.length; charPos++) thisString += this[charPos]; + for (const key in args) { // eslint-disable-line guard-for-in + const stringKey = `{${key}}`; + thisString = thisString.replace(new RegExp(stringKey, 'g'), args[key]); + } + return thisString; +}; + +let selectedURL = ''; +let selectedName = ''; +let selectedType = ''; + +function clearModelDetails() { + const el = gradioApp().getElementById('model-details') || gradioApp().getElementById('civitai_models_output') || gradioApp().getElementById('models_outcome'); + if (!el) return; + el.innerHTML = ''; +} + +const modelDetailsHTML = ` +
+ + + + + + + + + + + + +
Name{name}
Type{type}
Tags
{tags}
NSFW{nsfw} | {level}
Availability{availability}
Downloads{downloads}
Author{creator}
Description
{desc}
Download
All variants
+
+ + + + + + + + + + + + + + + + {versions} + +
VersionTypeBaseFileUpdatedSizeAvailabilityDescription
+
+`; + +const modelVersionsHTML = ` + + {url} + {name} + {type} + {base} + {file} + {mtime} + {size} + {availability} +
{desc}
+ +`; + +async function modelCardClick(id) { + log('modelCardClick id', id); + const el = gradioApp().getElementById('model-details') || gradioApp().getElementById('civitai_models_output') || gradioApp().getElementById('models_outcome'); + if (!el) return; + const res = await fetch(`${window.api}/civitai?model_id=${encodeURI(id)}`); + if (!res || res.status !== 200) { + error(`modelCardClick: id=${id} status=${res ? res.status : 'unknown'}`); + return; + } + let data = await res.json(); + log('modelCardClick data', data); + if (!data || data.length === 0) return; + data = data[0]; // assuming the first item is the one we want + + const versionsHTML = data.versions.map((v) => modelVersionsHTML.format({ + url: `
\udb80\uddda
`, + name: v.name || 'unknown', + type: v.files[0]?.type || 'unknown', + base: v.base || 'unknown', + mtime: (new Date(v.mtime)).toLocaleDateString(), + availability: v.availability || 'unknown', + size: v.files[0]?.size ? `${(v.files[0].size / 1024 / 1024).toFixed(2)} MB` : 'unknown', + file: `${v.files[0]?.name || 'unknown'}`, + desc: v.desc || 'no description available', + })).join(''); + const url = `${data.name || 'unknown'}`; + const creator = `${data.creator || 'unknown'}`; + const images = data.versions.map((v) => v.images).flat().map((i) => i.url); // TODO image gallery + const modelHTML = modelDetailsHTML.format({ + name: url, + type: data.type || 'unknown', + tags: data.tags?.join(', ') || '', + nsfw: data.nsfw ? 'yes' : 'no', + level: data.level?.toString() || '', + availability: data.availability || 'unknown', + downloads: data.downloads?.toString() || '', + creator, + desc: data.desc || 'no description available', + image: images.length > 0 ? images[0] : '/sdapi/v1/network/thumb?filename=html/card-no-preview.png', + versions: versionsHTML || '', + }); + el.innerHTML = modelHTML; +} + +function startCivitDownload(url, name, type) { + log('startCivitDownload', { url, name, type }); + selectedURL = [url]; + selectedName = [name]; + selectedType = [type]; + const civitDownloadBtn = gradioApp().getElementById('civitai_download_btn'); + if (civitDownloadBtn) civitDownloadBtn.click(); +} + +function startCivitAllDownload(evt) { + log('startCivitAllDownload', evt); + const versions = gradioApp().getElementById('model-versions-table').querySelectorAll('tr'); + selectedURL = []; + selectedName = []; + selectedType = []; + for (const version of versions) { + const parsed = version.querySelector('td:nth-child(1) div')?.getAttribute('onclick')?.match(/startCivitDownload\('([^']+)', '([^']+)', '([^']+)'\)/); + if (!parsed || parsed.length < 4) continue; + selectedURL.push(parsed[1]); + selectedName.push(parsed[2]); + selectedType.push(parsed[3]); + } + const civitDownloadBtn = gradioApp().getElementById('civitai_download_btn'); + if (civitDownloadBtn) civitDownloadBtn.click(); +} + +function downloadCivitModel(modelUrl, modelName, modelType, modelPath, civitToken, innerHTML) { + log('downloadCivitModel', { modelUrl, modelName, modelType, modelPath, civitToken }); + const el = gradioApp().getElementById('civitai_models_output') || gradioApp().getElementById('models_outcome'); + const currentHTML = el?.innerHTML || ''; + return [selectedURL, selectedName, selectedType, modelPath, civitToken, currentHTML]; +} diff --git a/javascript/contextMenus.js b/javascript/contextMenus.js index 271f5a9e4..96c57a231 100644 --- a/javascript/contextMenus.js +++ b/javascript/contextMenus.js @@ -149,7 +149,6 @@ async function initContextMenu() { appendContextMenuOption(id, 'Generate forever', () => generateForever(`#${tab}_generate`)); appendContextMenuOption(id, 'Apply selected style', quickApplyStyle); appendContextMenuOption(id, 'Quick save style', quickSaveStyle); - appendContextMenuOption(id, 'nVidia overlay', initNVML); id = `#${tab}_reprocess`; appendContextMenuOption(id, 'Decode full quality', () => reprocessClick(`${tab}`, 'reprocess_decode'), true); appendContextMenuOption(id, 'Refine & HiRes pass', () => reprocessClick(`${tab}`, 'reprocess_refine'), true); diff --git a/javascript/extraNetworks.js b/javascript/extraNetworks.js index 1a9ae4fe6..e13323307 100644 --- a/javascript/extraNetworks.js +++ b/javascript/extraNetworks.js @@ -1,32 +1,38 @@ const activePromptTextarea = {}; let sortVal = -1; let totalCards = -1; +let lastTab = 'control'; // helpers const getENActiveTab = () => { let tabName = ''; - if (gradioApp().getElementById('txt2img_prompt')?.checkVisibility()) return 'txt2img'; - if (gradioApp().getElementById('img2img_prompt')?.checkVisibility()) return 'img2img'; - if (gradioApp().getElementById('control_prompt')?.checkVisibility()) return 'control'; - if (gradioApp().getElementById('video_prompt')?.checkVisibility()) return 'video'; - if (gradioApp().getElementById('framepack_prompt_row')?.checkVisibility()) return 'framepack'; + if (gradioApp().getElementById('txt2img_prompt')?.checkVisibility()) tabName = 'txt2img'; + else if (gradioApp().getElementById('img2img_prompt')?.checkVisibility()) tabName = 'img2img'; + else if (gradioApp().getElementById('control_prompt')?.checkVisibility()) tabName = 'control'; + else if (gradioApp().getElementById('video_prompt')?.checkVisibility()) tabName = 'video'; + else if (gradioApp().getElementById('extras_image')?.checkVisibility()) tabName = 'process'; + else if (gradioApp().getElementById('interrogate_image')?.checkVisibility()) tabName = 'caption'; + else if (gradioApp().getElementById('tab-gallery-search')?.checkVisibility()) tabName = 'gallery'; + if (tabName in ['process', 'caption', 'gallery']) tabName = lastTab; + else lastTab = tabName; + if (tabName !== '') return tabName; // legacy method 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'; - else if (gradioApp().getElementById('tab_framepack_tab')?.style.display === 'block') tabName = 'framepack'; + else tabName = 'control'; // log('getENActiveTab', tabName); return tabName; }; const getENActivePage = () => { - const tabname = getENActiveTab(); - let page = gradioApp().querySelector(`#${tabname}_extra_networks > .tabs > .tab-nav > .selected`); - if (!page) page = gradioApp().querySelector(`#${tabname}_extra_tabs > .tab-nav > .selected`); + const tabName = getENActiveTab(); + let page = gradioApp().querySelector(`#${tabName}_extra_networks > .tabs > .tab-nav > .selected`); + if (!page) page = gradioApp().querySelector(`#${tabName}_extra_tabs > .tab-nav > .selected`); const pageName = page ? page.innerText : ''; - const btnApply = gradioApp().getElementById(`${tabname}_extra_apply`); + const btnApply = gradioApp().getElementById(`${tabName}_extra_apply`); if (btnApply) btnApply.style.display = pageName === 'Style' ? 'inline-flex' : 'none'; // log('getENActivePage', pageName); return pageName; @@ -48,8 +54,8 @@ const setENState = (state) => { function showCardDetails(event) { // log('showCardDetails', event); - const tabname = getENActiveTab(); - const btn = gradioApp().getElementById(`${tabname}_extra_details_btn`); + const tabName = getENActiveTab(); + const btn = gradioApp().getElementById(`${tabName}_extra_details_btn`); btn.click(); event.stopPropagation(); event.preventDefault(); @@ -96,9 +102,9 @@ function readCardTags(el, tags) { } function readCardDescription(page, item) { - xhrGet('/sd_extra_networks/description', { page, item }, (data) => { - const tabname = getENActiveTab(); - const description = gradioApp().querySelector(`#${tabname}_description > label > textarea`); + xhrGet('/sdapi/v1/network/desc', { page, item }, (data) => { + const tabName = getENActiveTab(); + const description = gradioApp().querySelector(`#${tabName}_description > label > textarea`); if (description) { description.value = data?.description?.trim() || ''; updateInput(description); @@ -108,10 +114,10 @@ function readCardDescription(page, item) { } function getCardsForActivePage() { - const pagename = getENActivePage(); - if (!pagename) return []; - const allCards = Array.from(gradioApp().querySelectorAll('.extra-network-cards > .card')); - const cards = allCards.filter((el) => el.dataset.page.toLowerCase().includes(pagename.toLowerCase())); + const pageName = getENActivePage(); + if (!pageName) return []; + let allCards = Array.from(gradioApp().querySelectorAll('.extra-network-cards > .card')); + allCards = allCards.filter((el) => el.dataset.page?.toLowerCase().includes(pageName.toLowerCase())); // log('getCardsForActivePage', pagename, cards.length); return allCards; } @@ -220,8 +226,8 @@ function sortExtraNetworks(fixed = 'no') { case 0: return 0; case 1: return a.dataset.name ? a.dataset.name.localeCompare(b.dataset.name) : 0; case 2: return b.dataset.name ? b.dataset.name.localeCompare(a.dataset.name) : 0; - case 3: return a.dataset.mtime && !isNaN(a.dataset.mtime) ? parseFloat(b.dataset.mtime) - parseFloat(a.dataset.mtime) : 0; - case 4: return b.dataset.mtime && !isNaN(b.dataset.mtime) ? parseFloat(a.dataset.mtime) - parseFloat(b.dataset.mtime) : 0; + case 3: return a.dataset.mtime ? (new Date(b.dataset.mtime)).getTime() - (new Date(a.dataset.mtime)).getTime() : 0; + case 4: return b.dataset.mtime ? (new Date(a.dataset.mtime)).getTime() - (new Date(b.dataset.mtime)).getTime() : 0; case 5: return a.dataset.size && !isNaN(a.dataset.size) ? parseFloat(b.dataset.size) - parseFloat(a.dataset.size) : 0; case 6: return b.dataset.size && !isNaN(b.dataset.size) ? parseFloat(a.dataset.size) - parseFloat(b.dataset.size) : 0; } @@ -234,25 +240,42 @@ function sortExtraNetworks(fixed = 'no') { return desc; } -function refreshENInput(tabname) { - log('refreshNetworks', tabname, gradioApp().querySelector(`#${tabname}_extra_networks textarea`)?.value); - gradioApp().querySelector(`#${tabname}_extra_networks textarea`)?.dispatchEvent(new Event('input')); +function refreshENInput(tabName) { + log('refreshNetworks', tabName, gradioApp().querySelector(`#${tabName}_extra_networks textarea`)?.value); + gradioApp().querySelector(`#${tabName}_extra_networks textarea`)?.dispatchEvent(new Event('input')); } -function cardClicked(textToAdd, allowNegativePrompt) { - // log('cardClicked', textToAdd, allowNegativePrompt); - const tabname = getENActiveTab(); - log('cardClicked', tabname, textToAdd); - const textarea = allowNegativePrompt ? activePromptTextarea[tabname] : gradioApp().querySelector(`#${tabname}_prompt > label > textarea`); +async function markSelectedCards(selected, page = '') { + log('markSelectedCards', selected, page); + gradioApp().querySelectorAll('.extra-network-cards .card').forEach((el) => { + if (page.length > 0 && el.dataset.page !== page) return; // filter by page + if (selected.includes(el.dataset.name) || selected.includes(el.dataset.short)) el.classList.add('card-selected'); + else el.classList.remove('card-selected'); + }); +} + +function extractLoraNames(prompt) { + const regex = /]+)(?::[\d.]+)?>/g; + const names = []; + let match; + while ((match = regex.exec(prompt)) !== null) names.push(match[1]); // eslint-disable-line no-cond-assign + return names; +} + +function cardClicked(textToAdd) { + const tabName = getENActiveTab(); + log('cardClicked', tabName, textToAdd); + const textarea = activePromptTextarea[tabName]; if (textarea.value.indexOf(textToAdd) !== -1) textarea.value = textarea.value.replace(textToAdd, ''); else textarea.value += textToAdd; updateInput(textarea); + markSelectedCards(extractLoraNames(textarea.value), 'lora'); } function extraNetworksSearchButton(event) { // log('extraNetworksSearchButton', event); - const tabname = getENActiveTab(); - const searchTextarea = gradioApp().querySelector(`#${tabname}_extra_search textarea`); + const tabName = getENActiveTab(); + const searchTextarea = gradioApp().querySelector(`#${tabName}_extra_search textarea`); const button = event.target; searchTextarea.value = `${button.textContent.trim()}/`; updateInput(searchTextarea); @@ -261,8 +284,8 @@ function extraNetworksSearchButton(event) { let desiredStyle = ''; function selectStyle(name) { desiredStyle = name; - const tabname = getENActiveTab(); - const button = gradioApp().querySelector(`#${tabname}_styles_select`); + const tabName = getENActiveTab(); + const button = gradioApp().querySelector(`#${tabName}_styles_select`); button.click(); } @@ -271,31 +294,28 @@ function applyStyles(styles) { if (styles) { newStyles = Array.isArray(styles) ? styles : [styles]; } else { - const tabname = getENActiveTab(); - styles = gradioApp().querySelectorAll(`#${tabname}_styles .token span`); + const tabName = getENActiveTab(); + styles = gradioApp().querySelectorAll(`#${tabName}_styles .token span`); newStyles = Array.from(styles).map((el) => el.textContent).filter((el) => el.length > 0); } const index = newStyles.indexOf(desiredStyle); if (index > -1) newStyles.splice(index, 1); else newStyles.push(desiredStyle); - gradioApp().querySelectorAll('.extra-network-cards .card').forEach((el) => { - if (newStyles.includes(el.getAttribute('data-name'))) el.style.boxShadow = '0 0 2px 4px var(--button-primary-border-color)'; - else el.style.boxShadow = 'none'; - }); + markSelectedCards(newStyles, 'style'); return newStyles.join('|'); } function quickApplyStyle() { - const tabname = getENActiveTab(); - const btnApply = gradioApp().getElementById(`${tabname}_extra_apply`); + const tabName = getENActiveTab(); + const btnApply = gradioApp().getElementById(`${tabName}_extra_apply`); if (btnApply) btnApply.click(); } function quickSaveStyle() { - const tabname = getENActiveTab(); - const btnSave = gradioApp().getElementById(`${tabname}_extra_quicksave`); + const tabName = getENActiveTab(); + const btnSave = gradioApp().getElementById(`${tabName}_extra_quicksave`); if (btnSave) btnSave.click(); - const btnRefresh = gradioApp().getElementById(`${tabname}_extra_refresh`); + const btnRefresh = gradioApp().getElementById(`${tabName}_extra_refresh`); if (btnRefresh) { setTimeout(() => btnRefresh.click(), 100); // setTimeout(() => sortExtraNetworks('fixed'), 500); @@ -313,18 +333,18 @@ let enDirty = false; function closeDetailsEN(...args) { // log('closeDetailsEN'); enDirty = true; - const tabname = getENActiveTab(); - const btnClose = gradioApp().getElementById(`${tabname}_extra_details_close`); + const tabName = getENActiveTab(); + const btnClose = gradioApp().getElementById(`${tabName}_extra_details_close`); if (btnClose) setTimeout(() => btnClose.click(), 100); - const btnRefresh = gradioApp().getElementById(`${tabname}_extra_refresh`); + const btnRefresh = gradioApp().getElementById(`${tabName}_extra_refresh`); if (btnRefresh && enDirty) setTimeout(() => btnRefresh.click(), 100); return [...args]; } function refeshDetailsEN(args) { // log(`refeshDetailsEN: ${enDirty}`); - const tabname = getENActiveTab(); - const btnRefresh = gradioApp().getElementById(`${tabname}_extra_refresh`); + const tabName = getENActiveTab(); + const btnRefresh = gradioApp().getElementById(`${tabName}_extra_refresh`); if (btnRefresh && enDirty) setTimeout(() => btnRefresh.click(), 100); enDirty = false; return args; @@ -334,30 +354,30 @@ function refeshDetailsEN(args) { function refreshENpage() { if (getCardsForActivePage().length === 0) { // log('refreshENpage'); - const tabname = getENActiveTab(); - const btnRefresh = gradioApp().getElementById(`${tabname}_extra_refresh`); + const tabName = getENActiveTab(); + const btnRefresh = gradioApp().getElementById(`${tabName}_extra_refresh`); if (btnRefresh) btnRefresh.click(); } } // init -function setupExtraNetworksForTab(tabname) { - let tabs = gradioApp().querySelector(`#${tabname}_extra_tabs`); +function setupExtraNetworksForTab(tabName) { + let tabs = gradioApp().querySelector(`#${tabName}_extra_tabs`); if (tabs) tabs.classList.add('extra-networks'); - const en = gradioApp().getElementById(`${tabname}_extra_networks`); - tabs = gradioApp().querySelector(`#${tabname}_extra_tabs > div`); + const en = gradioApp().getElementById(`${tabName}_extra_networks`); + tabs = gradioApp().querySelector(`#${tabName}_extra_tabs > div`); if (!tabs) return; // buttons - const btnShow = gradioApp().getElementById(`${tabname}_extra_networks_btn`); - const btnRefresh = gradioApp().getElementById(`${tabname}_extra_refresh`); - const btnScan = gradioApp().getElementById(`${tabname}_extra_scan`); - const btnSave = gradioApp().getElementById(`${tabname}_extra_save`); - const btnClose = gradioApp().getElementById(`${tabname}_extra_close`); - const btnSort = gradioApp().getElementById(`${tabname}_extra_sort`); - const btnView = gradioApp().getElementById(`${tabname}_extra_view`); - const btnModel = gradioApp().getElementById(`${tabname}_extra_model`); - const btnApply = gradioApp().getElementById(`${tabname}_extra_apply`); + const btnShow = gradioApp().getElementById(`${tabName}_extra_networks_btn`); + const btnRefresh = gradioApp().getElementById(`${tabName}_extra_refresh`); + const btnScan = gradioApp().getElementById(`${tabName}_extra_scan`); + const btnSave = gradioApp().getElementById(`${tabName}_extra_save`); + const btnClose = gradioApp().getElementById(`${tabName}_extra_close`); + const btnSort = gradioApp().getElementById(`${tabName}_extra_sort`); + const btnView = gradioApp().getElementById(`${tabName}_extra_view`); + const btnModel = gradioApp().getElementById(`${tabName}_extra_model`); + const btnApply = gradioApp().getElementById(`${tabName}_extra_apply`); const buttons = document.createElement('span'); buttons.classList.add('buttons'); if (btnRefresh) buttons.appendChild(btnRefresh); @@ -373,8 +393,8 @@ function setupExtraNetworksForTab(tabname) { tabs.appendChild(buttons); // details - const detailsImg = gradioApp().getElementById(`${tabname}_extra_details_img`); - const detailsClose = gradioApp().getElementById(`${tabname}_extra_details_close`); + const detailsImg = gradioApp().getElementById(`${tabName}_extra_details_img`); + const detailsClose = gradioApp().getElementById(`${tabName}_extra_details_close`); if (detailsImg && detailsClose) { detailsImg.title = 'Close details'; detailsImg.onclick = () => detailsClose.click(); @@ -384,9 +404,9 @@ function setupExtraNetworksForTab(tabname) { const div = document.createElement('div'); div.classList.add('second-line'); tabs.appendChild(div); - const txtSearch = gradioApp().querySelector(`#${tabname}_extra_search`); - const txtSearchValue = gradioApp().querySelector(`#${tabname}_extra_search textarea`); - const txtDescription = gradioApp().getElementById(`${tabname}_description`); + const txtSearch = gradioApp().querySelector(`#${tabName}_extra_search`); + const txtSearchValue = gradioApp().querySelector(`#${tabName}_extra_search textarea`); + const txtDescription = gradioApp().getElementById(`${tabName}_description`); txtSearch.classList.add('search'); txtDescription.classList.add('description'); div.appendChild(txtSearch); @@ -404,7 +424,7 @@ function setupExtraNetworksForTab(tabname) { let hoverTimer = null; let previousCard = null; if (window.opts.extra_networks_fetch) { - gradioApp().getElementById(`${tabname}_extra_tabs`).onmouseover = async (e) => { + gradioApp().getElementById(`${tabName}_extra_tabs`).onmouseover = async (e) => { const el = e.target.closest('.card'); // bubble-up to card if (!el || (el.title === previousCard)) return; if (!hoverTimer) { @@ -424,7 +444,7 @@ function setupExtraNetworksForTab(tabname) { // auto-resize networks sidebar const resizeObserver = new ResizeObserver((entries) => { for (const entry of entries) { - for (const el of Array.from(gradioApp().getElementById(`${tabname}_extra_tabs`).querySelectorAll('.extra-networks-page'))) { + for (const el of Array.from(gradioApp().getElementById(`${tabName}_extra_tabs`).querySelectorAll('.extra-networks-page'))) { const h = Math.trunc(entry.contentRect.height); if (h <= 0) return; const vh = opts.logmonitor_show ? '55vh' : '68vh'; @@ -432,12 +452,12 @@ function setupExtraNetworksForTab(tabname) { 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}`); + // log(`${tabName} height: ${entry.target.id}=${h} ${el.id}=${el.clientHeight}`); } } }); - const settingsEl = gradioApp().getElementById(`${tabname}_settings`); - const interfaceEl = gradioApp().getElementById(`${tabname}_interface`); + const settingsEl = gradioApp().getElementById(`${tabName}_settings`); + const interfaceEl = gradioApp().getElementById(`${tabName}_interface`); if (settingsEl) resizeObserver.observe(settingsEl); if (interfaceEl) resizeObserver.observe(interfaceEl); @@ -452,7 +472,7 @@ function setupExtraNetworksForTab(tabname) { const target = window.opts.extra_networks_card_cover === 'sidebar' ? 0 : window.opts.extra_networks_height; if (window.opts.theme_type === 'Standard') h = target > 0 ? target : 55; else h = target > 0 ? target : 87; - for (const el of Array.from(gradioApp().getElementById(`${tabname}_extra_tabs`).querySelectorAll('.extra-networks-page'))) { + for (const el of Array.from(gradioApp().getElementById(`${tabName}_extra_tabs`).querySelectorAll('.extra-networks-page'))) { if (h > 0) el.style.height = `${h}vh`; el.parentElement.style.width = '-webkit-fill-available'; } @@ -476,7 +496,7 @@ function setupExtraNetworksForTab(tabname) { en.style.top = '13em'; en.style.transition = ''; en.style.zIndex = 100; - gradioApp().getElementById(`${tabname}_settings`).parentNode.style.width = 'unset'; + gradioApp().getElementById(`${tabName}_settings`).parentNode.style.width = 'unset'; } else if (window.opts.extra_networks_card_cover === 'sidebar') { en.style.position = 'absolute'; en.style.height = 'auto'; @@ -487,7 +507,7 @@ function setupExtraNetworksForTab(tabname) { en.style.top = '13em'; en.style.transition = 'width 0.3s ease'; en.style.zIndex = 100; - gradioApp().getElementById(`${tabname}_settings`).parentNode.style.width = `calc(100vw - 2em - min(${window.opts.extra_networks_sidebar_width}vw, 50vw))`; + gradioApp().getElementById(`${tabName}_settings`).parentNode.style.width = `calc(100vw - 2em - min(${window.opts.extra_networks_sidebar_width}vw, 50vw))`; } else { en.style.position = 'relative'; en.style.height = 'unset'; @@ -498,15 +518,15 @@ function setupExtraNetworksForTab(tabname) { en.style.top = 0; en.style.transition = ''; en.style.zIndex = 0; - gradioApp().getElementById(`${tabname}_settings`).parentNode.style.width = 'unset'; + gradioApp().getElementById(`${tabName}_settings`).parentNode.style.width = 'unset'; } } else { if (window.opts.extra_networks_card_cover === 'sidebar') en.style.width = 0; - gradioApp().getElementById(`${tabname}_settings`).parentNode.style.width = 'unset'; + gradioApp().getElementById(`${tabName}_settings`).parentNode.style.width = 'unset'; } - if (tabname === 'video') { - gradioApp().getElementById('framepack_settings').parentNode.style.width = gradioApp().getElementById(`${tabname}_settings`).parentNode.style.width; - gradioApp().getElementById('ltx_settings').parentNode.style.width = gradioApp().getElementById(`${tabname}_settings`).parentNode.style.width; + if (tabName === 'video') { + gradioApp().getElementById('framepack_settings').parentNode.style.width = gradioApp().getElementById(`${tabName}_settings`).parentNode.style.width; + gradioApp().getElementById('ltx_settings').parentNode.style.width = gradioApp().getElementById(`${tabName}_settings`).parentNode.style.width; } } }); @@ -514,8 +534,8 @@ function setupExtraNetworksForTab(tabname) { } async function showNetworks() { - for (const tabname of ['txt2img', 'img2img', 'control', 'video']) { - if (window.opts.extra_networks_show) gradioApp().getElementById(`${tabname}_extra_networks_btn`).click(); + for (const tabName of ['txt2img', 'img2img', 'control', 'video']) { + if (window.opts.extra_networks_show) gradioApp().getElementById(`${tabName}_extra_networks_btn`).click(); } log('showNetworks'); } @@ -526,11 +546,11 @@ async function setupExtraNetworks() { setupExtraNetworksForTab('control'); setupExtraNetworksForTab('video'); - function registerPrompt(tabname, id) { + function registerPrompt(tabName, id) { const textarea = gradioApp().querySelector(`#${id} > label > textarea`); if (!textarea) return; - if (!activePromptTextarea[tabname]) activePromptTextarea[tabname] = textarea; - textarea.addEventListener('focus', () => { activePromptTextarea[tabname] = textarea; }); + if (!activePromptTextarea[tabName]) activePromptTextarea[tabName] = textarea; + textarea.addEventListener('focus', () => { activePromptTextarea[tabName] = textarea; }); } registerPrompt('txt2img', 'txt2img_prompt'); diff --git a/javascript/gallery.js b/javascript/gallery.js index 2de1d70a3..38797551b 100644 --- a/javascript/gallery.js +++ b/javascript/gallery.js @@ -1,12 +1,4 @@ /* eslint-disable max-classes-per-file */ -// Known issues -// Images flash on the screen before they get processed and separator is properly closed, especially when root/subfolder has large amount of files -// Search is a bit wonky, I tried to get the separators to hide if 0 hits in seperator are found, but no luck so -// Sorting huge amount of images is slow, might look at optimising, I don't think it's a regression. - -// TODO -// Setting to enable or disable separator state persistence - let ws; let url; let currentImage; @@ -192,7 +184,8 @@ async function addSeparators() { async function delayFetchThumb(fn) { while (outstanding > 16) await new Promise((resolve) => setTimeout(resolve, 50)); // eslint-disable-line no-promise-executor-return outstanding++; - const res = await fetch(`${window.api}/browser/thumb?file=${encodeURI(fn)}`, { priority: 'low' }); + const ts = Date.now().toString(); + const res = await fetch(`${window.api}/browser/thumb?file=${encodeURI(fn)}&ts=${ts}`, { priority: 'low' }); if (!res.ok) { error(`fetchThumb: ${res.statusText}`); outstanding--; @@ -662,6 +655,27 @@ async function galleryHidden() { if (pruneImagesTimer) clearInterval(pruneImagesTimer); } +async function monitorGalleries() { + async function galleryMutation(mutations) { + const galleries = mutations.filter((m) => m.target?.classList?.contains('preview')); + for (const gallery of galleries) { + const links = gallery.target.querySelectorAll('a'); + for (const link of links) { + const href = link.getAttribute('href'); + if (!href) continue; + const fn = href.split('/').pop().split('\\').pop(); + link.setAttribute('download', fn); + } + } + } + + const galleryElements = gradioApp().querySelectorAll('.gradio-gallery'); + for (const gallery of galleryElements) { + const galleryObserver = new MutationObserver(galleryMutation); + galleryObserver.observe(gallery, { childList: true, subtree: true, attributes: true }); + } +} + async function initGallery() { // triggered on gradio change to monitor when ui gets sufficiently constructed log('initGallery'); el.folders = gradioApp().getElementById('tab-gallery-folders'); @@ -681,6 +695,7 @@ async function initGallery() { // triggered on gradio change to monitor when ui if (entries[0].intersectionRatio > 0) galleryVisible(); }); intersectionObserver.observe(el.folders); + monitorGalleries(); } // register on startup diff --git a/javascript/gpu.js b/javascript/gpu.js new file mode 100644 index 000000000..d95542f3d --- /dev/null +++ b/javascript/gpu.js @@ -0,0 +1,75 @@ +let gpuInterval = null; // eslint-disable-line prefer-const +const chartData = { mem: [], load: [] }; + +async function updateGPUChart(mem, load) { + const maxLen = 120; + const colorRangeMap = $.range_map({ // eslint-disable-line no-undef + '0:5': '#fffafa', + '6:10': '#fff7ed', + '11:20': '#fed7aa', + '21:30': '#fdba74', + '31:40': '#fb923c', + '41:50': '#f97316', + '51:60': '#ea580c', + '61:70': '#c2410c', + '71:80': '#9a3412', + '81:90': '#7c2d12', + '91:100': '#6c2e12', + }); + const sparklineConfigLOAD = { type: 'bar', height: '128px', barWidth: '3px', barSpacing: '1px', chartRangeMin: 0, chartRangeMax: 100, barColor: '#89007D' }; + const sparklineConfigMEM = { type: 'bar', height: '128px', barWidth: '3px', barSpacing: '1px', chartRangeMin: 0, chartRangeMax: 100, colorMap: colorRangeMap, composite: true }; + if (chartData.load.length > maxLen) chartData.load.shift(); + chartData.load.push(load); + if (chartData.mem.length > maxLen) chartData.mem.shift(); + chartData.mem.push(mem); + $('#gpuChart').sparkline(chartData.load, sparklineConfigLOAD); // eslint-disable-line no-undef + $('#gpuChart').sparkline(chartData.mem, sparklineConfigMEM); // eslint-disable-line no-undef +} + +async function updateGPU() { + const gpuEl = document.getElementById('gpu'); + const gpuTable = document.getElementById('gpu-table'); + try { + const res = await fetch(`${window.api}/gpu`); + if (!res.ok) { + clearInterval(gpuInterval); + gpuEl.style.display = 'none'; + return; + } + const data = await res.json(); + if (!data) { + clearInterval(gpuInterval); + gpuEl.style.display = 'none'; + return; + } + const gpuTbody = gpuTable.querySelector('tbody'); + for (const gpu of data) { + console.log(gpu); + let rows = `GPU${gpu.name}`; + for (const item of Object.entries(gpu.data)) rows += `${item[0]}${item[1]}`; + gpuTbody.innerHTML = rows; + if (gpu.chart && gpu.chart.length === 2) updateGPUChart(gpu.chart); + } + gpuEl.style.display = 'block'; + } catch (e) { + error('updateGPU', e); + clearInterval(gpuInterval); + gpuEl.style.display = 'none'; + } +} + +async function startGPU() { + const gpuEl = document.getElementById('gpu'); + gpuEl.style.display = 'block'; + if (gpuInterval) clearInterval(gpuInterval); + const interval = window.opts?.gpu_monitor || 3000; + log('startGPU', interval); + gpuInterval = setInterval(updateGPU, interval); + updateGPU(); +} + +async function disableGPU() { + clearInterval(gpuInterval); + const gpuEl = document.getElementById('gpu'); + gpuEl.style.display = 'none'; +} diff --git a/javascript/imageViewer.js b/javascript/imageViewer.js index c8a2dfd65..5631c173b 100644 --- a/javascript/imageViewer.js +++ b/javascript/imageViewer.js @@ -3,6 +3,18 @@ let previewDrag = false; let modalPreviewZone; let previewInstance; +function cycleImageFit() { + const root = document.documentElement; + const current = getComputedStyle(root).getPropertyValue('--sd-image-fit').trim(); + let next = 'contain'; + if (current === 'contain') next = 'cover'; + else if (current === 'cover') next = 'fill'; + else if (current === 'fill') next = 'scale-down'; + else if (current === 'scale-down') next = 'none'; + root.style.setProperty('--sd-image-fit', next); + log('cycleImageFit', current, next); +} + function closeModal(evt, force = false) { if (force) gradioApp().getElementById('lightboxModal').style.display = 'none'; if (previewDrag) return; diff --git a/javascript/light-teal.css b/javascript/light-teal.css index df8a3ab51..5d0ccfc22 100644 --- a/javascript/light-teal.css +++ b/javascript/light-teal.css @@ -197,6 +197,7 @@ svg.feather.feather-image, .feather .feather-image { display: none } --checkbox-border-radius: var(--radius-sm); --checkbox-border-width: var(--input-border-width); --checkbox-check: url("data:image/svg+xml,%3csvg viewBox='0 0 16 16' fill='white' xmlns='http://www.w3.org/2000/svg'%3e%3cpath d='M12.207 4.793a1 1 0 010 1.414l-5 5a1 1 0 01-1.414 0l-2-2a1 1 0 011.414-1.414L6.5 9.086l4.293-4.293a1 1 0 011.414 0z'/%3e%3c/svg%3e"); + --radio-circle: url("data:image/svg+xml,%3csvg viewBox='0 0 16 16' fill='white' xmlns='http://www.w3.org/2000/svg'%3e%3ccircle cx='8' cy='8' r='3'/%3e%3c/svg%3e"); --checkbox-label-background-fill-hover: None; --checkbox-label-background-fill-selected: var(--checkbox-label-background-fill); --checkbox-label-background-fill: None; diff --git a/javascript/nvml.js b/javascript/nvml.js deleted file mode 100644 index cf55bb196..000000000 --- a/javascript/nvml.js +++ /dev/null @@ -1,105 +0,0 @@ -let nvmlInterval = null; // eslint-disable-line prefer-const -let nvmlEl = null; -let nvmlTable = null; -const chartData = { mem: [], load: [] }; - -async function updateNVMLChart(mem, load) { - const maxLen = 120; - const colorRangeMap = $.range_map({ // eslint-disable-line no-undef - '0:5': '#fffafa', - '6:10': '#fff7ed', - '11:20': '#fed7aa', - '21:30': '#fdba74', - '31:40': '#fb923c', - '41:50': '#f97316', - '51:60': '#ea580c', - '61:70': '#c2410c', - '71:80': '#9a3412', - '81:90': '#7c2d12', - '91:100': '#6c2e12', - }); - const sparklineConfigLOAD = { type: 'bar', height: '100px', barWidth: '2px', barSpacing: '1px', chartRangeMin: 0, chartRangeMax: 100, barColor: '#89007D' }; - const sparklineConfigMEM = { type: 'bar', height: '100px', barWidth: '2px', barSpacing: '1px', chartRangeMin: 0, chartRangeMax: 100, colorMap: colorRangeMap, composite: true }; - if (chartData.load.length > maxLen) chartData.load.shift(); - chartData.load.push(load); - if (chartData.mem.length > maxLen) chartData.mem.shift(); - chartData.mem.push(mem); - $('#nvmlChart').sparkline(chartData.load, sparklineConfigLOAD); // eslint-disable-line no-undef - $('#nvmlChart').sparkline(chartData.mem, sparklineConfigMEM); // eslint-disable-line no-undef -} - -async function updateNVML() { - try { - const res = await fetch(`${window.api}/nvml`); - if (!res.ok) { - clearInterval(nvmlInterval); - nvmlEl.style.display = 'none'; - return; - } - const data = await res.json(); - if (!data) { - clearInterval(nvmlInterval); - nvmlEl.style.display = 'none'; - return; - } - const nvmlTbody = nvmlTable.querySelector('tbody'); - for (const gpu of data) { - const rows = ` - GPU${gpu.name} - Driver${gpu.version.driver} - VBIOS${gpu.version.vbios} - ROM${gpu.version.rom} - Driver${gpu.version.driver} - PCIGen.${gpu.pci.link} x${gpu.pci.width} - Memory${gpu.memory.used}Mb / ${gpu.memory.total}Mb - Clock${gpu.clock.gpu[0]}Mhz / ${gpu.clock.gpu[1]}Mhz - Power${gpu.power[0]}W / ${gpu.power[1]}W - Load GPU${gpu.load.gpu}% - Load Memory${gpu.load.memory}% - Temperature${gpu.load.temp}°C - Fans${gpu.load.fan}% - State${gpu.state} - `; - nvmlTbody.innerHTML = rows; - updateNVMLChart(gpu.load.memory, gpu.load.gpu); - } - nvmlEl.style.display = 'block'; - } catch (e) { - clearInterval(nvmlInterval); - nvmlEl.style.display = 'none'; - } -} - -async function initNVML() { - nvmlEl = document.getElementById('nvml'); - if (!nvmlEl) { - nvmlEl = document.createElement('div'); - nvmlEl.className = 'nvml'; - nvmlEl.id = 'nvml'; - nvmlTable = document.createElement('table'); - nvmlTable.className = 'nvml-table'; - nvmlTable.id = 'nvml-table'; - nvmlTable.innerHTML = ` - - - `; - const nvmlChart = document.createElement('div'); - nvmlChart.id = 'nvmlChart'; - nvmlEl.appendChild(nvmlTable); - nvmlEl.appendChild(nvmlChart); - gradioApp().appendChild(nvmlEl); - log('initNVML'); - } - if (nvmlInterval) { - clearInterval(nvmlInterval); - nvmlInterval = null; - nvmlEl.style.display = 'none'; - } else { - nvmlInterval = setInterval(updateNVML, 1000); - } -} - -async function disableNVML() { - clearInterval(nvmlInterval); - nvmlEl.style.display = 'none'; -} diff --git a/javascript/sdnext.css b/javascript/sdnext.css index 4ac257c89..68255a8af 100644 --- a/javascript/sdnext.css +++ b/javascript/sdnext.css @@ -14,6 +14,7 @@ --color-trace: #666666; --color-warning: #FF9900; --left-column: 530px; + --sd-image-fit: contain; } a { @@ -66,6 +67,10 @@ button { min-width: unset !important; } +h4 { + margin: 0.2em 0em 0.2em 0em; +} + input[type='color'] { height: 32px; width: 64px; @@ -122,6 +127,17 @@ input::-webkit-outer-spin-button, input::-webkit-inner-spin-button { overflow: auto; } +.link { + background-color: var(--background-fill-primary); + cursor: pointer; + border-radius: var(--input-radius); + width: 2em; +} + +.link:hover { + background-color: var(--button-primary-background-fill); +} + .gradio-dropdown, .block.gradio-slider, .block.gradio-checkbox, .block.gradio-textbox, .block.gradio-radio, .block.gradio-checkboxgroup, .block.gradio-number, .block.gradio-colorpicker { border-width: 0 !important; box-shadow: none !important; @@ -500,17 +516,36 @@ color: var(--primary-500) !important color: var(--body-text-color-subdued) !important } +.gradio-gallery img, .image-container img { + max-width: 100%; + object-position: top; + width: 100%; + height: 100%; + object-fit: var(--sd-image-fit) !important; +} + .interrogate { background: none !important; font-size: 1.5em !important; max-width: fit-content; position: absolute; right: 2.8em; - top: 0.2em; + top: 0.1em; z-index: 50; } -.interrogate:hover { +.image-fit { + background: none !important; + font-size: 1.5em !important; + max-width: fit-content; + position: absolute; + right: 4.0em; + top: 0.1em; + z-index: 50; +} + +.interrogate:hover, +.image-fit:hover { background: var(--button-primary-background-fill-hover) !important; } @@ -1226,11 +1261,18 @@ table.settings-value-table td { } .extra-network-cards .card { - height: fit-content; - margin: 0 0 0.5em 0.5em; - position: relative; - scroll-margin-top: 0; - scroll-snap-align: start; + margin: 0 0 0.5em 0.5em; + position: relative; + scroll-margin-top: 0; + scroll-snap-align: start; + height: var(--card-size); + width: var(--card-size); + contain: strict; +} + +*.extra-network-cards .card-selected { + transform: scale(0.9); + box-shadow: 0 0 2em var(--button-primary-background-fill); } .extra-network-cards .card .overlay { @@ -1243,17 +1285,16 @@ table.settings-value-table td { bottom: 0; color: white; font-size: var(--text-lg); - font-weight: bold; - overflow-wrap: anywhere; + overflow-wrap: break-word; padding: 0.2em; position: absolute; - text-shadow: 1px 1px black; z-index: 10; + text-shadow: 2px 2px 2px black; + filter: drop-shadow(0px 0px 4px black); } .extra-network-cards .card .overlay .reference { - - color: var(--body-text-color-subdued) + background-color: rgba(0, 0, 0, 0.2); } .extra-network-cards .card .preview { @@ -1515,12 +1556,6 @@ background: var(--background-color) min-height: 0; } -#models_error { - font-family: monospace; - -color: var(--body-text-color-subdued) -} - #model_loader_df button { display: none !important; } @@ -1580,14 +1615,13 @@ color: var(--body-text-color-subdued) padding: 6px !important; } -.nvml { +.gpu { background: var(--background-fill-primary); border: 1px solid var(--button-primary-border-color); bottom: 10px; color: var(--button-primary-text-color); display: none; font-family: monospace; - font-size: var(--text-xxs); padding: 6px; position: fixed; right: 10px; @@ -1929,10 +1963,6 @@ div:has(>#tab-gallery-folders) { padding: 0.2em; } -.docs-results { - background-color: var(--sd-group-background-color); -} - .docs-card { margin: 1em 0; background-color: var(--background-fill-primary); @@ -1965,6 +1995,86 @@ div:has(>#tab-gallery-folders) { margin-top: 0.2em; } +#model_desc { + overflow: auto; +} + +#model_list_table { + overflow: auto; + max-height: 50vh; +} + +#civit_metadata { + overflow: auto; +} + +.model-config { + font-size: 0.8em !important; + opacity: 0.8; + max-height: 6em; + overflow-y: auto; +} + +.simple-table tr { + vertical-align: baseline; +} + +.simple-table td { + padding: 0.2em !important; +} + +.simple-table tr { + vertical-align: baseline; +} + +.simple-table thead tr { + background-color: var(--button-primary-border-color) !important; +} + +.simple-table tr:nth-child(odd) { + background-color: var(--neutral-900); +} + +.simple-table td { + padding: 0.2em !important; + white-space: pre-wrap; +} + +.simple-table td div { + padding: 0.2em !important; + white-space: pre-wrap; + max-height: 7em; + overflow-x: hidden; + overflow-y: auto; +} + +.simple-table td:nth-child(1) { + color: var(--button-primary-border-color); + font-weight: bold; +} + +.div-link { + cursor: pointer; +} + +.div-link:hover { + background-color: var(--button-primary-background-fill); +} + +.video-model-link { + color: var(--button-primary-background-fill); + font-weight: normal; +} + +#civitai_token textarea, #hf_token textarea, #setting_huggingface_token textarea { + filter: blur(4px); +} + +#civitai_token textarea:hover, #hf_token textarea:hover, #setting_huggingface_token textarea:hover, +#civitai_token textarea:focus, #hf_token textarea:focus, #setting_huggingface_token textarea:focus { + filter: blur(0); +} + @keyframes move { from { background-position-x: 0, -40px; diff --git a/javascript/ui.js b/javascript/ui.js index 4df527ceb..6fe2f954b 100644 --- a/javascript/ui.js +++ b/javascript/ui.js @@ -32,6 +32,12 @@ function clip_gallery_urls(gallery) { ); } +function isVisible(el) { + const rect = el.getBoundingClientRect(); + if (rect.width === 0 && rect.height === 0) return false; + return rect.top >= 0 && rect.left >= 0 && rect.bottom <= (window.innerHeight || document.documentElement.clientHeight) && rect.right <= (window.innerWidth || document.documentElement.clientWidth); +} + function all_gallery_buttons() { let allGalleryButtons = gradioApp().querySelectorAll('[style="display: block;"].tabitem div[id$=_gallery].gradio-gallery .thumbnails > .thumbnail-item.thumbnail-small'); if (allGalleryButtons.length === 0) allGalleryButtons = gradioApp().querySelectorAll('.gradio-gallery .thumbnails > .thumbnail-item.thumbnail-small'); @@ -66,6 +72,7 @@ function selected_gallery_files() { let allCurrentButtons = gradioApp().querySelectorAll('[style="display: block;"].tabitem div[id$=_gallery].gradio-gallery .thumbnail-item.thumbnail-small'); if (allCurrentButtons.length === 0) allCurrentButtons = gradioApp().querySelectorAll('.gradio-gallery .thumbnails > .thumbnail-item.thumbnail-small'); allImages = Array.from(allCurrentButtons).map((v) => v.querySelector('img')?.src); + allImages = allImages.filter((el) => isVisible(el)); } catch { /**/ } const selectedIndex = selected_gallery_index(); return [allImages, selectedIndex]; @@ -178,7 +185,7 @@ function switch_to_caption(...args) { function get_tab_index(tabId) { let res = 0; - gradioApp().getElementById(tabId).querySelector('div').querySelectorAll('button') + gradioApp().getElementById(tabId)?.querySelector('div').querySelectorAll('button') .forEach((button, i) => { if (button.className.indexOf('selected') !== -1) res = i; }); @@ -472,25 +479,28 @@ function updateInput(target) { let desiredCheckpointName = null; function selectCheckpoint(name) { desiredCheckpointName = name; - const tabname = getENActiveTab(); - const btnModel = gradioApp().getElementById(`${tabname}_extra_model`); + const tabName = getENActiveTab(); + const btnModel = gradioApp().getElementById(`${tabName}_extra_model`); const isRefiner = btnModel && btnModel.classList.contains('toolbutton-selected'); if (isRefiner) gradioApp().getElementById('change_refiner').click(); else gradioApp().getElementById('change_checkpoint').click(); - log(`Change ${isRefiner ? 'refiner' : 'model'}: ${desiredCheckpointName}`); + log(`selectCheckpoint ${isRefiner ? 'refiner' : 'model'}: ${desiredCheckpointName}`); + markSelectedCards([desiredCheckpointName], 'model'); } let desiredVAEName = null; function selectVAE(name) { desiredVAEName = name; gradioApp().getElementById('change_vae').click(); - log(`Change VAE: ${desiredVAEName}`); + log(`selectVAE: ${desiredVAEName}`); + markSelectedCards([desiredVAEName], 'vae'); } function selectReference(name) { - log(`Select reference: ${name}`); + log(`selectReference: ${name}`); desiredCheckpointName = name; gradioApp().getElementById('change_reference').click(); + markSelectedCards([desiredCheckpointName], 'model'); } function currentImageResolutionimg2img(_a, _b, scaleBy) { diff --git a/launch.py b/launch.py index 7f7d3c335..423842b1b 100755 --- a/launch.py +++ b/launch.py @@ -45,9 +45,9 @@ def init_args(): def init_paths(): global script_path, extensions_dir # pylint: disable=global-statement import modules.paths - modules.paths.register_paths() script_path = modules.paths.script_path extensions_dir = modules.paths.extensions_dir + sys.path.insert(0, script_path) rec('paths') @@ -221,10 +221,7 @@ def start_server(immediate=True, server=None): installer.log.trace('Logging: level=trace') server.wants_restart = False else: - if args.api_only: - uvicorn = server.api_only() - else: - uvicorn = server.webui(restart=not immediate) + uvicorn = server.webui(restart=not immediate) if args.profile: pr.disable() installer.print_profile(pr, 'WebUI') diff --git a/models/Reference/Qwen--Qwen-Image.jpg b/models/Reference/Qwen--Qwen-Image.jpg new file mode 100644 index 000000000..b55838835 Binary files /dev/null and b/models/Reference/Qwen--Qwen-Image.jpg differ diff --git a/models/Reference/tempest-by-vlad-base.jpg b/models/Reference/tempest-by-vlad-base.jpg new file mode 100644 index 000000000..0d48f0a62 Binary files /dev/null and b/models/Reference/tempest-by-vlad-base.jpg differ diff --git a/models/Reference/tempest-by-vlad-hyper.jpg b/models/Reference/tempest-by-vlad-hyper.jpg new file mode 100644 index 000000000..c2ffc23e8 Binary files /dev/null and b/models/Reference/tempest-by-vlad-hyper.jpg differ diff --git a/modules/api/api.py b/modules/api/api.py index dfc4cacb3..a260c22ae 100644 --- a/modules/api/api.py +++ b/modules/api/api.py @@ -5,7 +5,7 @@ from fastapi import FastAPI, APIRouter, Depends, Request from fastapi.security import HTTPBasic, HTTPBasicCredentials from fastapi.exceptions import HTTPException from modules import errors, shared, postprocessing -from modules.api import models, endpoints, script, helpers, server, nvml, generate, process, control, docs +from modules.api import models, endpoints, script, helpers, server, generate, process, control, docs, gpu errors.install() @@ -54,7 +54,7 @@ class Api: self.add_api_route("/sdapi/v1/options", server.get_config, methods=["GET"], response_model=models.OptionsModel) self.add_api_route("/sdapi/v1/options", server.set_config, methods=["POST"]) self.add_api_route("/sdapi/v1/cmd-flags", server.get_cmd_flags, methods=["GET"], response_model=models.FlagsModel) - self.add_api_route("/sdapi/v1/nvml", nvml.get_nvml, methods=["GET"], response_model=List[models.ResNVML]) + self.add_api_route("/sdapi/v1/gpu", gpu.get_gpu_status, methods=["GET"], response_model=List[models.ResGPU]) # core api using locking self.add_api_route("/sdapi/v1/txt2img", self.generate.post_text2img, methods=["POST"], response_model=models.ResTxt2Img) @@ -91,6 +91,7 @@ class Api: self.add_api_route("/sdapi/v1/interrogate", endpoints.post_interrogate, methods=["POST"]) self.add_api_route("/sdapi/v1/vqa", endpoints.post_vqa, methods=["POST"]) self.add_api_route("/sdapi/v1/checkpoint", endpoints.get_checkpoint, methods=["GET"]) + self.add_api_route("/sdapi/v1/checkpoint", endpoints.set_checkpoint, methods=["POST"]) self.add_api_route("/sdapi/v1/refresh-checkpoints", endpoints.post_refresh_checkpoints, methods=["POST"]) self.add_api_route("/sdapi/v1/unload-checkpoint", endpoints.post_unload_checkpoint, methods=["POST"]) self.add_api_route("/sdapi/v1/reload-checkpoint", endpoints.post_reload_checkpoint, methods=["POST"]) @@ -98,6 +99,7 @@ class Api: self.add_api_route("/sdapi/v1/refresh-vae", endpoints.post_refresh_vae, methods=["POST"]) self.add_api_route("/sdapi/v1/latents", endpoints.get_latent_history, methods=["GET"], response_model=List[str]) self.add_api_route("/sdapi/v1/latents", endpoints.post_latent_history, methods=["POST"], response_model=int) + self.add_api_route("/sdapi/v1/modules", endpoints.get_modules, methods=["GET"]) # lora api from modules.api import loras @@ -111,10 +113,12 @@ class Api: from modules.api import nudenet nudenet.register_api() + # civitai api + from modules.civitai import api_civitai + api_civitai.register_api() + def add_api_route(self, path: str, endpoint, **kwargs): - if (shared.cmd_opts.auth or shared.cmd_opts.auth_file) and shared.cmd_opts.api_only: - kwargs['dependencies'] = [Depends(self.auth)] if shared.opts.subpath is not None and len(shared.opts.subpath) > 0: self.app.add_api_route(f'{shared.opts.subpath}{path}', endpoint, **kwargs) self.app.add_api_route(path, endpoint, **kwargs) diff --git a/modules/api/endpoints.py b/modules/api/endpoints.py index 28c573559..44120e7b5 100644 --- a/modules/api/endpoints.py +++ b/modules/api/endpoints.py @@ -140,6 +140,17 @@ def get_checkpoint(): checkpoint['hash'] = shared.sd_model.sd_checkpoint_info.shorthash return checkpoint +def set_checkpoint(sd_model_checkpoint: str, dtype:str=None, force:bool=False): + from modules import sd_models, devices + if force: + sd_models.unload_model_weights(op='model') + if dtype is not None: + shared.opts.cuda_dtype = dtype + devices.set_dtype() + shared.opts.sd_model_checkpoint = sd_model_checkpoint + model = sd_models.reload_model_weights() + return { 'ok': model is not None } + def post_refresh_checkpoints(): shared.refresh_checkpoints() return {} @@ -148,6 +159,30 @@ def post_refresh_vae(): shared.refresh_vaes() return {} +def get_modules(): + from modules import modelstats + model = modelstats.analyze() + if model is None: + return {} + model_obj = { + 'model': model.name, + 'type': model.type, + 'class': model.cls, + 'size': model.size, + 'mtime': str(model.mtime), + 'modules': [] + } + for m in model.modules: + model_obj['modules'].append({ + 'class': m.cls, + 'params': m.params, + 'modules': m.modules, + 'quant': m.quant, + 'device': str(m.device), + 'dtype': str(m.dtype) + }) + return model_obj + def get_extensions_list(): from modules import extensions extensions.list_extensions() diff --git a/modules/api/gallery.py b/modules/api/gallery.py index 6510cd673..8b74e2c04 100644 --- a/modules/api/gallery.py +++ b/modules/api/gallery.py @@ -9,7 +9,7 @@ from fastapi.responses import JSONResponse from starlette.websockets import WebSocket, WebSocketState from pydantic import BaseModel, Field # pylint: disable=no-name-in-module from PIL import Image -from modules import shared, images, files_cache +from modules import shared, images, files_cache, modelstats debug = shared.log.debug if os.environ.get('SD_BROWSER_DEBUG', None) is not None else lambda *args, **kwargs: None @@ -76,7 +76,7 @@ def register_api(app: FastAPI): # register api def get_video_thumbnail(filepath): from modules.video import get_video_params try: - stat = os.stat(filepath) + stat_size, stat_mtime = modelstats.stat(filepath) frames, fps, duration, width, height, codec, frame = get_video_params(filepath, capture=True) h = shared.opts.extra_networks_card_size w = shared.opts.extra_networks_card_size if shared.opts.browser_fixed_width else width * h // height @@ -91,8 +91,8 @@ def register_api(app: FastAPI): # register api 'data': data_url, 'width': width, 'height': height, - 'size': stat.st_size, - 'mtime': stat.st_mtime, + 'size': stat_size, + 'mtime': stat_mtime.timestamp(), } return content except Exception as e: @@ -101,7 +101,7 @@ def register_api(app: FastAPI): # register api def get_image_thumbnail(filepath): try: - stat = os.stat(filepath) + stat_size, stat_mtime = modelstats.stat(filepath) image = Image.open(filepath) geninfo, _items = images.read_info_from_image(image) h = shared.opts.extra_networks_card_size @@ -118,8 +118,8 @@ def register_api(app: FastAPI): # register api 'data': data_url, 'width': width, 'height': height, - 'size': stat.st_size, - 'mtime': stat.st_mtime, + 'size': stat_size, + 'mtime': stat_mtime.timestamp(), } return content except Exception as e: diff --git a/modules/api/gpu.py b/modules/api/gpu.py new file mode 100644 index 000000000..7f33156c4 --- /dev/null +++ b/modules/api/gpu.py @@ -0,0 +1,37 @@ +import torch +from installer import log + + +device = None + + +def get_gpu_status(): + global device # pylint: disable=global-statement + if device is None: + try: + device = torch.cuda.get_device_name(torch.cuda.current_device()) + log.info(f'GPU monitoring: device={device}') + except Exception: + device = '' + # per vendor modules + if 'nvidia' in device.lower(): + from modules.api import nvml + return nvml.get_nvml() + elif 'amd' in device.lower(): + from modules.api import rocm_smi + return rocm_smi.get_rocm_smi() + return [] + + +""" +Resut should always be: list[ResGPU] +class ResGPU(BaseModel): + name: str = Field(title="GPU Name") + data: dict = Field(title="Name/Value data") + chart: list[float, float] = Field(title="Exactly two items to place on chart") +""" + +if __name__ == '__main__': + from rich import print as rprint + for gpu in get_gpu_status(): + rprint(gpu) diff --git a/modules/api/loras.py b/modules/api/loras.py index 4fbae29c5..c192ec62d 100644 --- a/modules/api/loras.py +++ b/modules/api/loras.py @@ -7,16 +7,16 @@ def get_lora(lora: str) -> dict: if lora not in lora_load.available_networks: raise HTTPException(status_code=404, detail=f"Lora '{lora}' not found") obj = lora_load.available_networks[lora] - obj.info = obj.get_info() - obj.desc = obj.get_desc() return obj.__dict__ + def get_loras(): 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 lora_load.available_networks.values()] + def post_refresh_loras(): from modules.lora import lora_load return lora_load.list_available_networks() diff --git a/modules/api/middleware.py b/modules/api/middleware.py index 5c72f2204..d276bd46b 100644 --- a/modules/api/middleware.py +++ b/modules/api/middleware.py @@ -15,6 +15,7 @@ import modules.errors as errors errors.install() +ignore_endpoints = ['/sdapi/v1/log', '/sdapi/v1/browser', '/sdapi/v1/gpu', '/sdapi/v1/network/thumb'] def setup_middleware(app: FastAPI, cmd_opts): @@ -42,8 +43,8 @@ def setup_middleware(app: FastAPI, cmd_opts): res.headers["X-Process-Time"] = duration endpoint = req.scope.get('path', 'err') token = req.cookies.get("access-token") or req.cookies.get("access-token-unsecure") - if (cmd_opts.api_log or cmd_opts.api_only) and endpoint.startswith('/sdapi'): - if '/sdapi/v1/log' in endpoint or '/sdapi/v1/browser' in endpoint: + if (cmd_opts.api_log) and endpoint.startswith('/sdapi'): + if any([endpoint.startswith(x) for x in ignore_endpoints]): # noqa C419 # pylint: disable=use-a-generator return res log.info('API user={user} code={code} {prot}/{ver} {method} {endpoint} {cli} {duration}'.format( # pylint: disable=consider-using-f-string, logging-format-interpolation user = app.tokens.get(token) if hasattr(app, 'tokens') else None, diff --git a/modules/api/models.py b/modules/api/models.py index 010467f66..a97ddff0c 100644 --- a/modules/api/models.py +++ b/modules/api/models.py @@ -424,16 +424,10 @@ class ResScripts(BaseModel): img2img: list = Field(default=None, title="Img2img", description="Titles of scripts (img2img)") control: list = Field(default=None, title="Control", description="Titles of scripts (control)") -class ResNVML(BaseModel): # definition of http response - name: str = Field(title="Name") - version: dict = Field(title="Version") - pci: dict = Field(title="Version") - memory: dict = Field(title="Version") - clock: dict = Field(title="Version") - load: dict = Field(title="Version") - power: list = [] - state: str = Field(title="State") - +class ResGPU(BaseModel): # definition of http response + name: str = Field(title="GPU Name") + data: dict = Field(title="Name/Value data") + chart: list[float, float] = Field(title="Exactly two items to place on chart") # helper function diff --git a/modules/api/nudenet.py b/modules/api/nudenet.py index 4756d42f3..a1243592a 100644 --- a/modules/api/nudenet.py +++ b/modules/api/nudenet.py @@ -58,7 +58,7 @@ def banned_words( def register_api(): from modules.shared import api as api_instance - api_instance.add_api_route("/sdapi/v1//nudenet", nudenet_censor, methods=["POST"], response_model=dict) - api_instance.add_api_route("/sdapi/v1//prompt-lang", prompt_check, methods=["POST"], response_model=dict) - api_instance.add_api_route("/sdapi/v1//image-guard", image_guard, methods=["POST"], response_model=dict) - api_instance.add_api_route("/sdapi/v1//prompt-banned", banned_words, methods=["POST"], response_model=list) + api_instance.add_api_route("/sdapi/v1/nudenet", nudenet_censor, methods=["POST"], response_model=dict) + api_instance.add_api_route("/sdapi/v1/prompt-lang", prompt_check, methods=["POST"], response_model=dict) + api_instance.add_api_route("/sdapi/v1/image-guard", image_guard, methods=["POST"], response_model=dict) + api_instance.add_api_route("/sdapi/v1/prompt-banned", banned_words, methods=["POST"], response_model=list) diff --git a/modules/api/nvml.py b/modules/api/nvml.py index ba9d9b12f..177c4fc54 100644 --- a/modules/api/nvml.py +++ b/modules/api/nvml.py @@ -25,9 +25,11 @@ def get_reason(val): reason = ', '.join([throttle[i] for i in throttle if i & val]) return reason if len(reason) > 0 else 'ok' + def get_nvml(): global nvml_initialized # pylint: disable=global-statement try: + from modules.memstats import ram_stats if not nvml_initialized: install('pynvml', quiet=True) import pynvml # pylint: disable=redefined-outer-name @@ -43,41 +45,29 @@ def get_nvml(): name = pynvml.nvmlDeviceGetName(dev) except Exception: name = '' - device = { - 'name': name, - 'version': { - 'cuda': pynvml.nvmlSystemGetCudaDriverVersion(), - 'driver': pynvml.nvmlSystemGetDriverVersion(), - 'vbios': pynvml.nvmlDeviceGetVbiosVersion(dev), - 'rom': pynvml.nvmlDeviceGetInforomImageVersion(dev), - 'capabilities': pynvml.nvmlDeviceGetCudaComputeCapability(dev), - }, - 'pci': { - 'link': pynvml.nvmlDeviceGetCurrPcieLinkGeneration(dev), - 'width': pynvml.nvmlDeviceGetCurrPcieLinkWidth(dev), - 'busid': pynvml.nvmlDeviceGetPciInfo(dev).busId, - 'deviceid': pynvml.nvmlDeviceGetPciInfo(dev).pciDeviceId, - }, - 'memory': { - 'total': round(pynvml.nvmlDeviceGetMemoryInfo(dev).total/1024/1024, 2), - 'free': round(pynvml.nvmlDeviceGetMemoryInfo(dev).free/1024/1024,2), - 'used': round(pynvml.nvmlDeviceGetMemoryInfo(dev).used/1024/1024,2), - }, - 'clock': { # gpu, sm, memory - 'gpu': [pynvml.nvmlDeviceGetClockInfo(dev, 0), pynvml.nvmlDeviceGetMaxClockInfo(dev, 0)], - 'sm': [pynvml.nvmlDeviceGetClockInfo(dev, 1), pynvml.nvmlDeviceGetMaxClockInfo(dev, 1)], - 'memory': [pynvml.nvmlDeviceGetClockInfo(dev, 2), pynvml.nvmlDeviceGetMaxClockInfo(dev, 2)], - }, - 'load': { - 'gpu': round(pynvml.nvmlDeviceGetUtilizationRates(dev).gpu), - 'memory': round(pynvml.nvmlDeviceGetUtilizationRates(dev).memory), - 'temp': pynvml.nvmlDeviceGetTemperature(dev, 0), - 'fan': pynvml.nvmlDeviceGetFanSpeed(dev), - }, - 'power': [round(pynvml.nvmlDeviceGetPowerUsage(dev)/1000, 2), round(pynvml.nvmlDeviceGetEnforcedPowerLimit(dev)/1000, 2)], - 'state': get_reason(pynvml.nvmlDeviceGetCurrentClocksThrottleReasons(dev)), + load = pynvml.nvmlDeviceGetUtilizationRates(dev) + mem = pynvml.nvmlDeviceGetMemoryInfo(dev) + ram = ram_stats() + data = { + "CUDA": f'Version {pynvml.nvmlSystemGetCudaDriverVersion()} Compute {pynvml.nvmlDeviceGetCudaComputeCapability(dev)}', + "Driver": pynvml.nvmlSystemGetDriverVersion(), + "Hardware": f'VBIOS {pynvml.nvmlDeviceGetVbiosVersion(dev)} ROM {pynvml.nvmlDeviceGetInforomImageVersion(dev)}', + "PCI link": f'Gen.{pynvml.nvmlDeviceGetCurrPcieLinkGeneration(dev)} x{pynvml.nvmlDeviceGetCurrPcieLinkWidth(dev)}', + "Power": f'{round(pynvml.nvmlDeviceGetPowerUsage(dev)/1000, 2)} W / {round(pynvml.nvmlDeviceGetEnforcedPowerLimit(dev)/1000, 2)} W', + "GPU clock": f'{pynvml.nvmlDeviceGetClockInfo(dev, 0)} Mhz / {pynvml.nvmlDeviceGetMaxClockInfo(dev, 0)} Mhz', + "SM clock": f'{pynvml.nvmlDeviceGetClockInfo(dev, 1)} Mhz / {pynvml.nvmlDeviceGetMaxClockInfo(dev, 1)} Mhz', + "VRAM clock": f'{pynvml.nvmlDeviceGetClockInfo(dev, 2)} Mhz / {pynvml.nvmlDeviceGetMaxClockInfo(dev, 2)} Mhz', + "VRAM usage": f'{round(100 * mem.used / mem.total)}% | {round(mem.used / 1024 / 1024)} MB used | {round(mem.free / 1024 / 1024)} MB free | {round(mem.total / 1024 / 1024)} MB total', + "RAM usage": f'{round(100 * ram["used"] / ram["total"])}% | {round(1024 * ram["used"])} MB used | {round(1024 * ram["free"])} MB free | {round(1024 * ram["total"])} MB total', + "System load": f'GPU {load.gpu}% | VRAM {load.memory}% | Temp {pynvml.nvmlDeviceGetTemperature(dev, 0)}C | Fan {pynvml.nvmlDeviceGetFanSpeed(dev)}%', + 'State': get_reason(pynvml.nvmlDeviceGetCurrentClocksThrottleReasons(dev)), } - devices.append(device) + chart = [load.memory, load.gpu] + devices.append({ + 'name': name, + 'data': data, + 'chart': chart, + }) # log.debug(f'nmvl: {devices}') return devices except Exception as e: diff --git a/modules/api/rocm_smi.py b/modules/api/rocm_smi.py new file mode 100644 index 000000000..bbf7f1ba7 --- /dev/null +++ b/modules/api/rocm_smi.py @@ -0,0 +1,114 @@ +import math +import json +import subprocess as sp +from enum import IntFlag + + +try: + from installer import log +except Exception: + import logging + log = logging.getLogger(__name__) + + +try: + from modules.rocm import version as rocm_version +except Exception: + rocm_version = "unknown" + + +# ThrottleStatus is from leuc/amdgpu_metrics.py +class ThrottleStatus(IntFlag): + # linux/drivers/gpu/drm/amd/pm/inc/amdgpu_smu.h + PPT0 = 1 << 0 + PPT1 = 1 << 1 + PPT2 = 1 << 2 + PPT3 = 1 << 3 + SPL = 1 << 4 + FPPT = 1 << 5 + SPPT = 1 << 6 + SPPT_APU = 1 << 7 + TDC_GFX = 1 << 16 + TDC_SOC = 1 << 17 + TDC_MEM = 1 << 18 + TDC_VDD = 1 << 19 + TDC_CVIP = 1 << 20 + EDC_CPU = 1 << 21 + EDC_GFX = 1 << 22 + APCC = 1 << 23 + TEMP_GPU = 1 << 32 + TEMP_CORE = 1 << 33 + TEMP_MEM = 1 << 34 + TEMP_EDGE = 1 << 35 + TEMP_HOTSPOT = 1 << 36 + TEMP_SOC = 1 << 37 + TEMP_VR_GFX = 1 << 38 + TEMP_VR_SOC = 1 << 39 + TEMP_VR_MEM0 = 1 << 40 + TEMP_VR_MEM1 = 1 << 41 + TEMP_LIQUID0 = 1 << 42 + TEMP_LIQUID1 = 1 << 43 + VRHOT0 = 1 << 44 + VRHOT1 = 1 << 45 + PROCHOT_CPU = 1 << 46 + PROCHOT_GFX = 1 << 47 + PPM = 1 << 56 + FIT = 1 << 57 + + def active(self): + members = self.__class__.__members__ + return (m for m in members if getattr(self, m)._value_ & self.value != 0) # pylint: disable=protected-access + + def __iter__(self): + return self.active() + + def __str__(self): + return ', '.join(self.active()) + + +def get_rocm_smi(): + try: + rocm_smi_data = json.loads(sp.check_output(("rocm-smi", "-a", "--json"))) + driver_version = rocm_smi_data.pop("system", {"Driver version": "unknown"}).get("Driver version") + + devices = [] + for key in rocm_smi_data.keys(): + load = { + 'gpu': rocm_smi_data[key].get('GPU use (%)', 'unknown'), + 'memory': rocm_smi_data[key].get("GPU Memory Allocated (VRAM%)", "unknown"), + 'temp': rocm_smi_data[key].get('Temperature (Sensor edge) (C)', 'unknown'), + 'temp_junction': rocm_smi_data[key].get('Temperature (Sensor junction) (C)', 'unknown'), + 'temp_memory': rocm_smi_data[key].get('Temperature (Sensor memory) (C)', 'unknown'), + 'fan': rocm_smi_data[key].get('Fan speed (%)', 'unknown'), + } + + data = { + "ROCm": f'version {rocm_version} agent {rocm_smi_data[key].get("GFX Version", "unknown")}', + "Driver": driver_version, + "Hardware": f'VBIOS {rocm_smi_data[key].get("VBIOS version", "unknown")}', + "PCI link": f'Gen.{int(math.log2(float(rocm_smi_data[key].get("pcie_link_speed (0.1 GT/s)", 10)) / 10))} x{rocm_smi_data[key].get("pcie_link_width (Lanes)", "unknown")}', + "Power": f'{round(float(rocm_smi_data[key].get("Average Graphics Package Power (W)", 0)), 2)} W / {round(float(rocm_smi_data[key].get("Max Graphics Package Power (W)", 0)), 2)} W', + "GPU clock": f'{rocm_smi_data[key].get("average_gfxclk_frequency (MHz)", 0)} Mhz / {rocm_smi_data[key].get("Valid sclk range", "0").split(" - ")[-1].removesuffix("Mhz")} Mhz', + "VRAM clock": f'{rocm_smi_data[key].get("current_uclk (MHz)", 0)} Mhz / {rocm_smi_data[key].get("Valid mclk range", "0").split(" - ")[-1].removesuffix("Mhz")} Mhz', + "VRAM usage": f'{load["memory"]}% Used | {rocm_smi_data[key].get("GPU Memory Read/Write Activity (%)", "unknown")}% Activity', + "GPU usage": f'GPU {load["gpu"]}% | Fan {load["fan"]}%', + "GPU temp": f'Edge {load["temp"]}C | Junction {load["temp_junction"]}C | Memory {load["temp_memory"]}C', + 'Throttle reason': str(ThrottleStatus(int(rocm_smi_data[key].get("throttle_status", 0)))), + } + name = rocm_smi_data[key].get('Device Name', 'unknown') + chart = [load["memory"], load["gpu"]] + devices.append({ + 'name': name, + 'data': data, + 'chart': chart, + }) + return devices + except Exception as e: + log.error(f'ROCm SMI: {e}') + return [] + + +if __name__ == '__main__': + from rich import print as rprint + for gpu in get_rocm_smi(): + rprint(gpu) diff --git a/modules/civitai/api_civitai.py b/modules/civitai/api_civitai.py new file mode 100644 index 000000000..92add902f --- /dev/null +++ b/modules/civitai/api_civitai.py @@ -0,0 +1,59 @@ +from starlette.responses import JSONResponse + + +def models_to_json(all_models:list, model_id:int=None): + dct = [] + for model in all_models: + if model_id is not None and model.id != model_id: + continue + model_dct = model.__dict__.copy() + versions_dct = [] + for version in model.versions: + version_dct = version.__dict__.copy() + version_dct['files'] = [f.__dict__.copy() for f in version.files] + version_dct['images'] = [i.__dict__.copy() for i in version.images] + versions_dct.append(version_dct) + model_dct['versions'] = versions_dct + dct.append(model_dct) + # obj = json.dumps(dct, indent=2, ensure_ascii=False) + return dct + + +def get_civitai( + model_id:int=None, # if model_id is provided assume fetch-from-cache + query:str = '', # search query or tag is required + tag:str = '', # search query or tag is required + types:str = '', # Checkpoint, TextualInversion, Hypernetwork, AestheticGradient, LORA, Controlnet, Poses + sort:str = '', # Highest Rated, Most Downloaded, Newest + period:str = '', # AllTime, Year, Month, Week, Day + nsfw:bool = None, # optional:bool + limit:int = 0, + base:str = '', + token:str = None, + exact:bool = True, +): + from modules.civitai import search_civitai + if model_id is not None: + dct = models_to_json(search_civitai.models, model_id=model_id) + return JSONResponse(content=dct, status_code=200) + if len(query) > 0 or len(tag) > 0: + models = search_civitai.search_civitai( + query=query, + tag=tag, + types=types, + sort=sort, + period=period, + nsfw=nsfw, + limit=limit, + base=base, + token=token, + exact=exact + ) + dct = models_to_json(models) + return JSONResponse(content=dct, status_code=200) + return JSONResponse(content=[], status_code=200) + + +def register_api(): + from modules.shared import api + api.add_api_route("/sdapi/v1/civitai", get_civitai, methods=["GET"], response_model=list) diff --git a/modules/civitai/download_civitai.py b/modules/civitai/download_civitai.py new file mode 100644 index 000000000..932f8ce7e --- /dev/null +++ b/modules/civitai/download_civitai.py @@ -0,0 +1,190 @@ +import os +import json +import rich.progress as p +from PIL import Image +from modules import shared, errors, paths + + +pbar = None + + +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_meta(model_path: str, model_id): + fn = os.path.splitext(model_path)[0] + '.json' + url = f'https://civitai.com/api/v1/models/{model_id}' + r = shared.req(url) + if r.status_code == 200: + try: + data = r.json() + shared.writefile(data, filename=fn, mode='w', silent=True) + shared.log.info(f'CivitAI download: id={model_id} url={url} file="{fn}"') + return r.status_code, len(data), '' # code/size/note + except Exception as e: + errors.display(e, 'civitai meta') + shared.log.error(f'CivitAI meta: id={model_id} url={url} file="{fn}" {e}') + return r.status_code, '', str(e) + return r.status_code, '', '' + + +def download_civit_preview(model_path: str, preview_url: str): + global pbar # pylint: disable=global-statement + if model_path is None: + pbar = None + return 500, '', '' + 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 500, '', 'exepected preview image got json' + if os.path.exists(preview_file): + return 304, '', 'already exists' + # res = f'CivitAI download: url={preview_url} file="{preview_file}"' + r = shared.req(preview_url, stream=True) + total_size = int(r.headers.get('content-length', 0)) + block_size = 16384 # 16KB blocks + 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 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) + pbar.update(task, advance=block_size) + if written < 1024: # min threshold + os.remove(preview_file) + return 400, '', 'removed invalid download' + if is_video: + img = save_video_frame(preview_file) + else: + img = Image.open(preview_file) + except Exception as e: + shared.log.error(f'CivitAI download error: url={preview_url} file="{preview_file}" written={written} {e}') + return 500, '', str(e) + shared.state.end() + if img is None: + return 500, '', 'image is none' + shared.log.info(f'CivitAI download: url={preview_url} file="{preview_file}" size={total_size} image={img.size}') + img.close() + return 200, str(total_size), '' # code/size/note + + +def download_civit_model_thread(model_name: str, model_url: str, model_path: str = "", model_type: str = "Model", token: str = None): + import hashlib + sha256 = hashlib.sha256() + sha256.update(model_url.encode('utf-8')) + temp_file = sha256.hexdigest()[:8] + '.tmp' + + headers = {} + starting_pos = 0 + if os.path.isfile(temp_file): + starting_pos = os.path.getsize(temp_file) + headers['Range'] = f'bytes={starting_pos}-' + if token is None or len(token) == 0: + token = shared.opts.civitai_token + if token is not None and len(token) > 0: + headers['Authorization'] = f'Bearer {token}' + + r = shared.req(model_url, headers=headers, stream=True) + total_size = int(r.headers.get('content-length', 0)) + if model_name is None or len(model_name) == 0: + cn = r.headers.get('content-disposition', '') + model_name = cn.split('filename=')[-1].strip('"') + + model_path = model_path.strip() + if len(model_path) > 0: + if os.path.isabs(model_path): + pass + else: + model_path = os.path.join(paths.models_path, model_path) + elif model_type.lower() == 'lora': + model_path = shared.opts.lora_dir + elif model_type.lower() == 'embedding': + model_path = shared.opts.embeddings_dir + elif model_type.lower() == 'vae': + model_path = shared.opts.vae_dir + else: + model_path = shared.opts.ckpt_dir + model_file = os.path.join(model_path, model_name) + temp_file = os.path.join(model_path, temp_file) + + res = f'Model download: name="{model_name}" url="{model_url}" path="{model_path}" temp="{temp_file}"' + if os.path.isfile(model_file): + res += ' already exists' + shared.log.warning(res) + return res + + res += f' size={round((starting_pos + total_size)/1024/1024, 2)}Mb' + shared.log.info(res) + shared.state.begin('CivitAI') + block_size = 16384 # 16KB blocks + written = starting_pos + global pbar # pylint: disable=global-statement + if pbar is None: + pbar = p.Progress(p.TextColumn('[cyan]{task.description}'), p.DownloadColumn(), p.BarColumn(), p.TaskProgressColumn(), p.TimeRemainingColumn(), p.TimeElapsedColumn(), p.TransferSpeedColumn(), p.TextColumn('[cyan]{task.fields[name]}'), console=shared.console) + with pbar: + task = pbar.add_task(description="Download starting", total=starting_pos+total_size, name=model_name) + try: + with open(temp_file, 'ab') as f: + for data in r.iter_content(block_size): + if written == 0: + try: # check if response is JSON message instead of bytes + shared.log.error(f'Model download: response={json.loads(data.decode("utf-8"))}') + raise ValueError('response: type=json expected=bytes') + except Exception: # this is good + pass + written = written + len(data) + f.write(data) + pbar.update(task, description="Download", completed=written) + if written < 1024: # min threshold + os.remove(temp_file) + raise ValueError(f'removed invalid download: bytes={written}') + except Exception as e: + shared.log.error(f'{res} {e}') + finally: + pbar.stop_task(task) + pbar.remove_task(task) + if starting_pos+total_size != written: + shared.log.warning(f'{res} written={round(written/1024/1024)}Mb incomplete download') + elif os.path.exists(temp_file): + shared.log.debug(f'Model download complete: temp="{temp_file}" path="{model_file}"') + os.rename(temp_file, model_file) + shared.state.end() + if os.path.exists(model_file): + return model_file + else: + return None + + +def download_civit_model(model_url: str, model_name: str = '', model_path: str = '', model_type: str = '', token: str = None): + import threading + if model_url is None or len(model_url) == 0: + shared.log.error('Model download: no url provided') + return + thread = threading.Thread(target=download_civit_model_thread, args=(model_name, model_url, model_path, model_type, token)) + thread.start() + thread.join() + from modules.sd_models import list_models # pylint: disable=W0621 + list_models() diff --git a/modules/civitai/metadata_civitai.py b/modules/civitai/metadata_civitai.py new file mode 100644 index 000000000..93bc143be --- /dev/null +++ b/modules/civitai/metadata_civitai.py @@ -0,0 +1,285 @@ +import os +import re +import time +import gradio as gr +from modules.shared import log, opts, req, readfile, max_workers + + +data = [] +selected_model = None + + +class CivitModel: + def __init__(self, name, fn, sha = None, meta = {}): + self.name = name + self.file = name + self.id = meta.get('id', 0) + self.fn = fn + self.sha = sha + self.meta = meta + self.versions = 0 + self.vername = '' + self.latest = '' + self.latest_hashes = [] + self.latest_name = '' + self.url = None + self.status = 'Not found' + + +def civit_update_metadata(): + def create_update_metadata_table(rows: list[CivitModel]): + html = """ + + + + + + {tbody} + +
FileIDNameHashVersionsLatestStatus
+ """ + tbody = '' + for row in rows: + try: + tbody += f""" + + {row.file} + {row.id} + {row.name} + {row.sha} + {row.versions} + {row.latest} + {row.status} + + """ + except Exception as e: + log.error(f'Model list: row={row} {e}') + return html.format(tbody=tbody) + + log.debug('CivitAI update metadata: models') + from modules import ui_extra_networks + from modules.civitai.download_civitai import download_civit_meta + pages = ui_extra_networks.get_pages('Model') + if len(pages) == 0: + return 'CivitAI update metadata: no models found' + page: ui_extra_networks.ExtraNetworksPage = pages[0] + results = [] + all_hashes = [(item.get('hash', None) or 'XXXXXXXX').upper()[:8] for item in page.list_items()] + for item in page.list_items(): + model = CivitModel(name=item['name'], fn=item['filename'], sha=item.get('hash', None), meta=item.get('metadata', {})) + if model.sha is None or len(model.sha) == 0: + log.debug(f'CivitAI skip search: name="{model.name}" hash=None') + else: + r = req(f'https://civitai.com/api/v1/model-versions/by-hash/{model.sha}') + log.debug(f'CivitAI search: name="{model.name}" hash={model.sha} status={r.status_code}') + if r.status_code == 200: + d = r.json() + model.id = d['modelId'] + download_civit_meta(model.fn, model.id) + fn = os.path.splitext(item['filename'])[0] + '.json' + model.meta = readfile(fn, silent=True) + model.name = model.meta.get('name', model.name) + model.versions = len(model.meta.get('modelVersions', [])) + versions = model.meta.get('modelVersions', []) + if len(versions) > 0: + model.latest = versions[0].get('name', '') + model.latest_hashes.clear() + for v in versions[0].get('files', []): + for h in v.get('hashes', {}).values(): + model.latest_hashes.append(h[:8].upper()) + for ver in versions: + for f in ver.get('files', []): + for h in f.get('hashes', {}).values(): + if h[:8].upper() == model.sha[:8].upper(): + model.vername = ver.get('name', '') + model.url = f.get('downloadUrl', None) + model.latest_name = f.get('name', '') + if model.vername == model.latest: + model.status = 'Latest version' + elif any(map(lambda v: v in model.latest_hashes, all_hashes)): # pylint: disable=cell-var-from-loop # noqa: C417 + model.status = 'Update downloaded' + else: + model.status = 'Update available' + break + results.append(model) + yield create_update_metadata_table(results) + return create_update_metadata_table(results) + + +def civit_search_model(name, tag, model_type): + # types = 'LORA' if model_type == 'LoRA' else 'Checkpoint' + url = 'https://civitai.com/api/v1/models?limit=25&Sort=Newest' + if model_type == 'Model': + url += '&types=Checkpoint' + elif model_type == 'LoRA': + url += '&types=LORA&types=DoRA&types=LoCon' + elif model_type == 'Embedding': + url += '&types=TextualInversion' + elif model_type == 'VAE': + url += '&types=VAE' + if name is not None and len(name) > 0: + url += f'&query={name}' + if tag is not None and len(tag) > 0: + url += f'&tag={tag}' + r = req(url) + log.debug(f'CivitAI search: type={model_type} name="{name}" tag={tag or "none"} url="{url}" status={r.status_code}') + if r.status_code != 200: + log.warning(f'CivitAI search: name="{name}" tag={tag} status={r.status_code}') + return [], gr.update(visible=False, value=[]), gr.update(visible=False, value=None), gr.update(visible=False, value=None) + try: + body = r.json() + except Exception as e: + log.error(f'CivitAI search: name="{name}" tag={tag} {e}') + return [], gr.update(visible=False, value=[]), gr.update(visible=False, value=None), gr.update(visible=False, value=None) + global data # pylint: disable=global-statement + data = body.get('items', []) + data1 = [] + for model in data: + found = 0 + if model_type == 'LoRA' and model['type'].lower() in ['lora', 'locon', 'dora', 'lycoris']: + found += 1 + elif model_type == 'Embedding' and model['type'].lower() in ['textualinversion', 'embedding']: + found += 1 + elif model_type == 'Model' and model['type'].lower() in ['checkpoint']: + found += 1 + elif model_type == 'VAE' and model['type'].lower() in ['vae']: + found += 1 + elif model_type == 'Other': + found += 1 + if found > 0: + data1.append([ + model['id'], + model['name'], + ', '.join(model['tags']), + model['stats']['downloadCount'], + model['stats']['rating'] + ]) + res = f'Search result: name={name} tag={tag or "none"} type={model_type} models={len(data1)}' + return res, gr.update(visible=len(data1) > 0, value=data1 if len(data1) > 0 else []), gr.update(visible=False, value=None), gr.update(visible=False, value=None) + + +def atomic_civit_search_metadata(item, results): + from modules.civitai.download_civitai import download_civit_preview, download_civit_meta + if item is None: + return + try: + meta = os.path.splitext(item['filename'])[0] + '.json' + except Exception: + # log.error(f'CivitAI search metadata: item={item} {e}') + return + has_meta = os.path.isfile(meta) and os.stat(meta).st_size > 0 + if ('card-no-preview.png' in item['preview'] or not has_meta) and os.path.isfile(item['filename']): + sha = item.get('hash', None) + found = False + result = { + 'id': '', + 'name': item['name'], + 'type': '', + 'hash': '', + 'code': '', + 'size': '', + 'note': '', + } + if sha is not None and len(sha) > 0: + r = req(f'https://civitai.com/api/v1/model-versions/by-hash/{sha}') + log.debug(f'CivitAI search: name="{item["name"]}" hash={sha} status={r.status_code}') + result['hash'] = sha + result['code'] = r.status_code + if r.status_code == 200: + d = r.json() + result['code'], result['size'], result['note'] = download_civit_meta(item['filename'], d['modelId']) + result['id'] = d['modelId'] + result['type'] = 'metadata' + results.append(result) + if d.get('images') is not None: + for i in d['images']: + result['code'], result['size'], result['note'] = download_civit_preview(item['filename'], i['url']) + if result['code'] == 200: + result['type'] = 'preview' + results.append(result) + found = True + break + if not found and os.stat(item['filename']).st_size < (1024 * 1024 * 1024): + from modules import hashes + sha = hashes.calculate_sha256(item['filename'], quiet=True)[:10] + r = req(f'https://civitai.com/api/v1/model-versions/by-hash/{sha}') + log.debug(f'CivitAI search: name="{item["name"]}" hash={sha} status={r.status_code}') + result['hash'] = sha + result['code'] = r.status_code + if r.status_code == 200: + d = r.json() + result['code'], result['size'], result['note'] = download_civit_meta(item['filename'], d['modelId']) + result['id'] = d['modelId'] + result['type'] = 'metadata' + results.append(result) + if d.get('images') is not None: + for i in d['images']: + result['code'], result['size'], result['note'] = download_civit_preview(item['filename'], i['url']) + if result['code'] == 200: + result['type'] = 'preview' + results.append(result) + found = True + break + if not found: + results.append(result) + + +def civit_search_metadata(title: str = None): + def create_search_metadata_table(rows): + html = """ + + + {tbody} +
NameIDTypeCodeHashSizeNote
+ """ + tbody = '' + for row in rows: + try: + tbody += f""" + + {row['name']} + {row['id']} + {row['type']} + {row['code']} + {row['hash']} + {row['size']} + {row['note']} + + """ + except Exception as e: + log.error(f'Model list: row={row} {e}') + return html.format(tbody=tbody) + + from modules.ui_extra_networks import get_pages + results = [] + scanned, skipped = 0, 0 + t0 = time.time() + candidates = [] + re_skip = [r.strip() for r in opts.extra_networks_scan_skip.split(',') if len(r.strip()) > 0] + for page in get_pages(): + if type(title) == str: + if page.title != title: + continue + if page.name == 'style': + continue + for item in page.list_items(): + if item is None: + continue + if any(re.search(re_str, item.get('name', '') + item.get('filename', '')) for re_str in re_skip): + skipped += 1 + continue + scanned += 1 + candidates.append(item) + log.debug(f'CivitAI search metadata: type={title if type(title) == str else "all"} workers={max_workers} skip={len(re_skip)} items={len(candidates)}') + import concurrent + with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor: + future_items = {} + for fn in candidates: + future_items[executor.submit(atomic_civit_search_metadata, fn, results)] = fn + for future in concurrent.futures.as_completed(future_items): + future.result() + yield create_search_metadata_table(results) + + t1 = time.time() + log.debug(f'CivitAI search metadata: scanned={scanned} skipped={skipped} time={t1-t0:.2f}') + yield create_search_metadata_table(results) diff --git a/modules/civitai/search_civitai.py b/modules/civitai/search_civitai.py new file mode 100644 index 000000000..226925bf9 --- /dev/null +++ b/modules/civitai/search_civitai.py @@ -0,0 +1,221 @@ +from dataclasses import dataclass +import os +import json +import time +from installer import install, log + + +full_dct = False +full_html = False +base_models = ['', 'ODOR', 'SD 1.4', 'SD 1.5', 'SD 1.5 LCM', 'SD 1.5 Hyper', 'SD 2.0', 'SD 2.0 768', 'SD 2.1', 'SD 2.1 768', 'SD 2.1 Unclip', 'SDXL 0.9', 'SDXL 1.0', 'SD 3', 'SD 3.5', 'SD 3.5 Medium', 'SD 3.5 Large', 'SD 3.5 Large Turbo', 'Pony', 'Flux.1 S', 'Flux.1 D', 'Flux.1 Kontext', 'AuraFlow', 'SDXL 1.0 LCM', 'SDXL Distilled', 'SDXL Turbo', 'SDXL Lightning', 'SDXL Hyper', 'Stable Cascade', 'SVD', 'SVD XT', 'Playground v2', 'PixArt a', 'PixArt E', 'Hunyuan 1', 'Hunyuan Video', 'Lumina', 'Kolors', 'Illustrious', 'Mochi', 'LTXV', 'CogVideoX', 'NoobAI', 'Wan Video', 'Wan Video 1.3B t2v', 'Wan Video 14B t2v', 'Wan Video 14B i2v 480p', 'Wan Video 14B i2v 720p', 'HiDream', 'OpenAI', 'Imagen4', 'Other'] + + +@dataclass +class ModelImage(): + def __init__(self, dct: dict): + if isinstance(dct, str): + dct = json.loads(dct) + self.id: int = dct.get('id', 0) + self.url: str = dct.get('url', '') + self.width: int = dct.get('width', 0) + self.height: int = dct.get('height', 0) + self.type: str = dct.get('type', 'Unknown') + self.dct: dict = dct if full_dct else {} + + def __str__(self): + return f'ModelImage(id={self.id} url="{self.url}" width={self.width} height={self.height} type="{self.type}")' + + +@dataclass +class ModelFile(): + def __init__(self, dct: dict): + if isinstance(dct, str): + dct = json.loads(dct) + self.id: int = dct.get('id', 0) + self.size: int = int(1024 * dct.get('sizeKB', 0)) + self.name: str = dct.get('name', 'Unknown') + self.type: str = dct.get('type', 'Unknown') + self.hashes: list[str] = [str(h) for h in dct.get('hashes', {}).values()] + self.url: str = dct.get('downloadUrl', '') + self.dct: dict = dct if full_dct else {} + + def __str__(self): + return f'ModelFile(id={self.id} name="{self.name}" size={self.size} type="{self.type}" url="{self.url}")' + + +@dataclass +class ModelVersion(): + def __init__(self, dct: dict): + import bs4 + if isinstance(dct, str): + dct = json.loads(dct) + self.id: int = dct.get('id', 0) + self.name: str = dct.get('name', 'Unknown') + self.base: str = dct.get('baseModel', 'Unknown') + self.mtime: str = dct.get('publishedAt', '') + self.downloads: int = dct.get('stats', {}).get('downloadCount', 0) + self.availability: str = dct.get('availability', 'Unknown') + self.html: str = dct.get('description', '') or '' if full_html else '' + self.desc: str = bs4.BeautifulSoup(dct.get('description', '') or '', features="html.parser").get_text() + self.files = [ModelFile(f) for f in dct.get('files', [])] + self.images = [ModelImage(i) for i in dct.get('images', [])] + self.dct: dict = dct if full_dct else {} + + def __str__(self): + return f'ModelVersion(id={self.id} name="{self.name}" base="{self.base}" mtime="{self.mtime}" downloads={self.downloads} availability={self.availability} desc="{self.desc[:30]}...")' + + +@dataclass +class Model(): + def __init__(self, dct: dict): + import bs4 + if isinstance(dct, str): + dct = json.loads(dct) + self.id: int = dct.get('id', 0) + self.url: str = f'https://civitai.com/models/{self.id}' + self.type: str = dct.get('type', 'Unknown') + self.name: str = dct.get('name', 'Unknown') + self.html: str = dct.get('description', '') or '' if full_html else '' + self.desc: str = bs4.BeautifulSoup(dct.get('description', '') or '', features="html.parser").get_text() + self.tags: list[str] = dct.get('tags', []) + self.nsfw: bool = dct.get('nsfw', False) + self.level: str = dct.get('nsfwLevel', 0) + self.availability: str = dct.get('availability', 'Unknown') + self.downloads: int = dct.get('stats', {}).get('downloadCount', 0) + self.creator: str = dct.get('creator', {}).get('username', 'Unknown') + self.versions: list[ModelVersion] = [ModelVersion(v) for v in dct.get('modelVersions', [])] + self.dct: dict = dct if full_dct else {} + + def __str__(self): + return f'Model(id={self.id} type={self.type} name="{self.name}" versions={len(self.versions)} nsfw={self.nsfw}/{self.level} downloads={self.downloads} author="{self.creator}" tags={self.tags} desc="{self.desc[:30]}...")' + + +models: list[Model] = [] # global cache for civitai search results + + +def search_civitai( + query:str, + tag:str = '', # optional:tag name + types:str = '', # (Checkpoint, TextualInversion, Hypernetwork, AestheticGradient, LORA, Controlnet, Poses) + sort:str = '', # (Highest Rated, Most Downloaded, Newest) + period:str = '', # (AllTime, Year, Month, Week, Day) + nsfw:bool = None, # optional:bool + limit:int = 0, + base:str = '', # list + token:str = None, + exact:bool = True, +): + global models # pylint: disable=global-statement + import requests + from urllib.parse import urlencode + install('bs4') # Ensure BeautifulSoup is installed + + if len(query) == 0: + log.error('CivitAI: empty query') + return [] + + t0 = time.time() + dct = { 'query': query } + if len(tag) > 0: + dct['tag'] = tag + if nsfw is not None: + dct['nsfw'] = 'true' if nsfw else 'false' + if limit > 0: + dct['limit'] = limit + if len(types) > 0: + dct['types'] = types + if len(sort) > 0: + dct['sort'] = sort + if len(period) > 0: + dct['period'] = period + if len(base) > 0: + dct['baseModels'] = base + encoded = urlencode(dct) + + headers = {} + if token is None: + token = os.environ.get('CIVITAI_TOKEN', None) + if token is not None and len(token) > 0: + headers['Authorization'] = f'Bearer {token}' + + url = 'https://civitai.com/api/v1/models' + if query.isnumeric(): + uri = f'{url}/{query}' + else: + uri = f'{url}?{encoded}' + + log.info(f'CivitAI request: uri="{uri}" dct={dct} token={token is not None}') + result = requests.get(uri, headers=headers, timeout=60) + + if result.status_code != 200: + log.error(f'CivitAI: code={result.status_code} reason={result.reason} uri={result.url}') + return [] + + all_models: list[Model] = [] + exact_models: list[Model] = [] + dct = result.json() + if 'items' not in dct: + items = [dct] # single model + else: + items = dct.get('items', []) + for item in items: + all_models.append(Model(item)) + + if exact: + for model in all_models: + model_names = [model.name.lower()] + version_names = [v.name.lower() for v in model.versions] + file_names = [f.name.lower() for v in model.versions for f in v.files] + if any([query.lower() in name for name in model_names + version_names + file_names]): # noqa: C419 # pylint: disable=use-a-generator + exact_models.append(model) + + t1 = time.time() + log.info(f'CivitAI result: code={result.status_code} exact={len(exact_models)} total={len(models)} time={t1-t0:.2f}') + models = exact_models if len(exact_models) > 0 else all_models + return models + + +def create_model_cards(all_models: list[Model]) -> str: + details = """ +
+
+ """ + cards = """ +
+ {cards} +
+ """ + card = """ +
+
{name}
+
{type}
+ {name} +
+ """ + all_cards = '' + for model in all_models: + previews = [] + for version in model.versions: + for image in version.images: + if image.url and len(image.url) > 0 and not image.url.lower().endswith('.mp4'): + previews.append(image.url) + if len(previews) == 0: + previews = ['/sdapi/v1/network/thumb?filename=html/card-no-preview.png'] + all_cards += card.format(id=model.id, name=model.name, type=model.type, preview=previews[0]) + html = details + cards.format(cards=all_cards) + return html + + +def print_models(all_models: list[Model]): + for model in all_models: + log.info(f' {model}') + log.trace('Model', model.dct) + for version in model.versions: + log.info(f' {version}') + log.trace('ModelVersion', version.dct) + for file in version.files: + log.info(f' {file}') + log.trace('ModelFile', file.dct) + for image in version.images: + log.info(f' {image}') + log.trace('ModelImage', image.dct) diff --git a/modules/cmd_args.py b/modules/cmd_args.py index 84da8f5f6..d368babb5 100644 --- a/modules/cmd_args.py +++ b/modules/cmd_args.py @@ -62,7 +62,6 @@ def main_args(): group_http.add_argument('--docs', default=os.environ.get("SD_DOCS", False), action='store_true', help = "Mount API docs, default: %(default)s") group_http.add_argument("--auth", type=str, default=os.environ.get("SD_AUTH", None), help='Set access authentication like "user:pwd,user:pwd""') group_http.add_argument("--auth-file", type=str, default=os.environ.get("SD_AUTHFILE", None), help='Set access authentication using file, default: %(default)s') - group_http.add_argument('--api-only', default=os.environ.get("SD_APIONLY", False), action='store_true', help = "Run in API only mode without starting UI") group_http.add_argument("--allowed-paths", nargs='+', default=[], type=str, required=False, help="add additional paths to paths allowed for web access") group_http.add_argument("--share", default=os.environ.get("SD_SHARE", False), action='store_true', help="Enable UI accessible through Gradio site, default: %(default)s") group_http.add_argument("--insecure", default=os.environ.get("SD_INSECURE", False), action='store_true', help="Enable extensions tab regardless of other options, default: %(default)s") @@ -89,6 +88,7 @@ def compatibility_args(): group_compat.add_argument("--disable-extension-access", default=False, action='store_true', help=argparse.SUPPRESS) group_compat.add_argument("--api", action='store_true', help=argparse.SUPPRESS, default=True) group_compat.add_argument("--api-auth", type=str, help=argparse.SUPPRESS, default=None) + group_compat.add_argument('--api-only', default=False, help=argparse.SUPPRESS) group_compat.add_argument("--disable-queue", default=os.environ.get("SD_DISABLEQUEUE", False), action='store_true', help=argparse.SUPPRESS) diff --git a/modules/control/run.py b/modules/control/run.py index 57cd402f0..bc25d27bb 100644 --- a/modules/control/run.py +++ b/modules/control/run.py @@ -13,7 +13,7 @@ from modules.control.units import lite # Kohya ControlLLLite from modules.control.units import t2iadapter # TencentARC T2I-Adapter from modules.control.units import reference # ControlNet-Reference from modules.control.processor import preprocess_image -from modules import devices, shared, errors, processing, images, sd_models, scripts_manager, masking +from modules import devices, shared, errors, processing, images, sd_models, sd_vae, scripts_manager, masking from modules.processing_class import StableDiffusionProcessingControl from modules.ui_common import infotext_to_html from modules.api import script @@ -384,10 +384,11 @@ def control_run(state: str = '', # pylint: disable=keyword-arg-before-vararg p.selected_scale_tab_mask = 1 # hires/refine defined outside of main init + vae_scale_factor = sd_vae.get_vae_scale_factor() if p.enable_hr and (p.hr_resize_x == 0 or p.hr_resize_y == 0): - p.hr_upscale_to_x, p.hr_upscale_to_y = 8 * int(p.width_before * p.hr_scale / 8), 8 * int(p.height_before * p.hr_scale / 8) + p.hr_upscale_to_x, p.hr_upscale_to_y = vae_scale_factor * int(p.width_before * p.hr_scale / vae_scale_factor), vae_scale_factor * int(p.height_before * p.hr_scale / vae_scale_factor) elif p.enable_hr and (p.hr_upscale_to_x == 0 or p.hr_upscale_to_y == 0): - p.hr_upscale_to_x, p.hr_upscale_to_y = 8 * int(p.hr_resize_x / 8), 8 * int(p.hr_resize_y / 8) + p.hr_upscale_to_x, p.hr_upscale_to_y = 8 * int(p.hr_resize_x / vae_scale_factor), vae_scale_factor * int(p.hr_resize_y / vae_scale_factor) global p_extra_args # pylint: disable=global-statement for k, v in p_extra_args.items(): diff --git a/modules/control/tile.py b/modules/control/tile.py index de9df1131..1d6478edc 100644 --- a/modules/control/tile.py +++ b/modules/control/tile.py @@ -1,6 +1,6 @@ import time from PIL import Image -from modules import shared, processing, images, sd_models +from modules import shared, processing, images, sd_models, sd_vae def get_tile(image: Image.Image, x: int, y: int, sx: int, sy: int) -> Image.Image: @@ -23,17 +23,18 @@ def run_tiling(p: processing.StableDiffusionProcessing, input_image: Image.Image sx, sy = p.control_tile.split('x') sx = int(sx) sy = int(sy) + vae_scale_factor = sd_vae.get_vae_scale_factor() if sx <= 0 or sy <= 0: raise ValueError('Control Tile: invalid tile size') control_image = p.task_args.get('control_image', None) or p.task_args.get('image', None) control_upscaled = None if isinstance(control_image, list) and len(control_image) > 0: - w, h = 8 * int(sx * control_image[0].width) // 8, 8 * int(sy * control_image[0].height) // 8 + w, h = vae_scale_factor * int(sx * control_image[0].width) // vae_scale_factor, vae_scale_factor * int(sy * control_image[0].height) // vae_scale_factor control_upscaled = images.resize_image(resize_mode=1 if sx==sy else 5, im=control_image[0], width=w, height=h, context='add with forward') init_image = p.override or input_image init_upscaled = None if init_image is not None: - w, h = 8 * int(sx * init_image.width) // 8, 8 * int(sy * init_image.height) // 8 + w, h = vae_scale_factor * int(sx * init_image.width) // vae_scale_factor, vae_scale_factor * int(sy * init_image.height) // vae_scale_factor init_upscaled = images.resize_image(resize_mode=1 if sx==sy else 5, im=init_image, width=w, height=h, context='add with forward') t1 = time.time() shared.log.debug(f'Control Tile: scale={sx}x{sy} resize={"fixed" if sx==sy else "context"} control={control_upscaled} init={init_upscaled} time={t1-t0:.3f}') diff --git a/modules/images.py b/modules/images.py index c8d13f92c..1b8eb6cce 100644 --- a/modules/images.py +++ b/modules/images.py @@ -197,7 +197,7 @@ def save_image(image, dirname = os.path.dirname(params.filename) if dirname is not None and len(dirname) > 0: os.makedirs(dirname, exist_ok=True) - params.filename = namegen.sequence(params.filename, dirname, basename) + params.filename = namegen.sequence(params.filename) params.filename = namegen.sanitize(params.filename) # callbacks script_callbacks.before_image_saved_callback(params) @@ -218,7 +218,7 @@ def save_image(image, def safe_decode_string(s: bytes): remove_prefix = lambda text, prefix: text[len(prefix):] if text.startswith(prefix) else text # pylint: disable=unnecessary-lambda-assignment - for encoding in ['utf-8', 'utf-16', 'ascii', 'latin_1', 'cp1252', 'cp437']: # try different encodings + for encoding in ['utf_16_be', 'utf-8', 'utf-16', 'ascii', 'latin_1', 'cp1252', 'cp437']: # try different encodings try: s = remove_prefix(s, b'UNICODE') s = remove_prefix(s, b'ASCII') diff --git a/modules/images_namegen.py b/modules/images_namegen.py index 8cf565b67..c8971eff2 100644 --- a/modules/images_namegen.py +++ b/modules/images_namegen.py @@ -15,6 +15,7 @@ re_pattern_arg = re.compile(r"(.*)<([^>]*)>$") re_attention = re.compile(r'[\(*\[*](\w+)(:\d+(\.\d+))?[\)*\]*]|') re_network = re.compile(r'\<\w+:(\w+)(:\d+(\.\d+))?\>|') re_brackets = re.compile(r'[\([{})\]]') +seq = 0 NOTHING = object() @@ -49,6 +50,7 @@ class FilenameGenerator: 'seed': lambda self: (self.seed and str(self.seed)) or '', 'steps': lambda self: self.p and getattr(self.p, 'steps', 0), 'cfg': lambda self: self.p and getattr(self.p, 'cfg_scale', 0), + 'pag': lambda self: self.p and getattr(self.p, 'pag_scale', 0), 'clip_skip': lambda self: self.p and getattr(self.p, 'clip_skip', 0), 'denoising': lambda self: self.p and getattr(self.p, 'denoising_strength', 0), 'styles': lambda self: (self.p and ", ".join([style for style in self.p.styles if not style == "None"])) or "None", @@ -181,18 +183,20 @@ class FilenameGenerator: debug(f'Filename sanitize: input="{filename}" parts={parts} output="{fn}" ext={ext} max={max_length} len={len(fn)}') return fn - def sequence(self, fn, dirname, basename): + def sequence(self, fn): + global seq # pylint: disable=global-statement x = fn + dirname = os.path.dirname(fn) + if seq == 0: + seq = len(os.listdir(dirname)) if os.path.exists(dirname) and os.path.isdir(dirname) else 0 if shared.opts.save_images_add_number or '[seq]' in fn: if '[seq]' not in fn: fn = os.path.join(os.path.dirname(fn), f"[seq]-{os.path.basename(fn)}") - basecount = get_next_sequence_number(dirname, basename) - for i in range(9999): - seq = f"{basecount + i:05}" - filename = fn.replace('[seq]', seq) - if not os.path.exists(filename): - debug(f'Prompt sequence: input="{fn}" seq={seq} output="{filename}"') - x = filename + for _i in range(99999): # 99999/000001 + seq += 1 + dst = fn.replace('[seq]', f'{seq:05}') + if not os.path.exists(dst): + x = dst break return x @@ -219,7 +223,7 @@ class FilenameGenerator: replacement = fun(self, *pattern_args) except Exception as e: replacement = None - errors.display(e, 'Filename apply pattern') + errors.display(e, 'namegen') shared.log.error(f'Filename apply pattern: {x} {e}') if replacement == NOTHING: continue @@ -231,21 +235,7 @@ class FilenameGenerator: return res -def get_next_sequence_number(path, basename): - """ - Determines and returns the next sequence number to use when saving an image in the specified directory. - """ - result = -1 - if basename != '': - basename = f"{basename}-" - prefix_length = len(basename) - if not os.path.isdir(path): - return 0 - for p in os.listdir(path): - if p.startswith(basename): - parts = os.path.splitext(p[prefix_length:])[0].split('-') # splits the filename (removing the basename first if one is defined, so the sequence number is always the first element) - try: - result = max(int(parts[0]), result) - except ValueError: - pass - return result + 1 +def get_next_sequence_number(path, basename): # pylint: disable=unused-argument + global seq # pylint: disable=global-statement + seq += 1 + return seq # unused diff --git a/modules/img2img.py b/modules/img2img.py index d91daaca1..b3c0880aa 100644 --- a/modules/img2img.py +++ b/modules/img2img.py @@ -138,10 +138,10 @@ def process_batch(p, input_files, input_dir, output_dir, inpaint_mask_dir, args) if output_dir == '': output_dir = shared.opts.outdir_img2img_samples os.makedirs(output_dir, exist_ok=True) - geninfo, items = images.read_info_from_image(image) + info, items = images.read_info_from_image(image) for k, v in items.items(): image.info[k] = v - images.save_image(image, path=output_dir, basename=basename, seed=None, prompt=None, extension=ext, info=geninfo, grid=False, pnginfo_section_name="extras", existing_info=image.info, forced_filename=forced_filename) + images.save_image(image, path=output_dir, basename=basename, seed=None, prompt=None, extension=ext, info=info, grid=False, pnginfo_section_name="extras", existing_info=image.info, forced_filename=forced_filename) processed = scripts_manager.scripts_img2img.after(p, processed, *args) shared.log.debug(f'Processed: images={len(batch_image_files)} memory={memory_stats()} batch') diff --git a/modules/intel/ipex/__init__.py b/modules/intel/ipex/__init__.py index d165f050e..f590da35e 100644 --- a/modules/intel/ipex/__init__.py +++ b/modules/intel/ipex/__init__.py @@ -125,14 +125,21 @@ def ipex_init(): # pylint: disable=too-many-statements torch.cuda.Tuple = torch.xpu.Tuple torch.cuda.List = torch.xpu.List + if torch_version < 2.8: + if has_ipex: + torch.cuda.memory_summary = torch.xpu.memory_summary + torch.cuda.memory_snapshot = torch.xpu.memory_snapshot + if torch_version < 2.9: - # torch._int_mm via onednn is supposed to land on pytorch with torch 2.8 or 2.9 - # ipex 2.7+ has experimental torch._int_mm support but uses the cpu with torch.compile and also runs as slow as onednn.qlinear - if (not has_ipex or torch_version <= 2.7) and hasattr(torch.ops, "onednn") and hasattr(torch.ops.onednn, "qlinear_pointwise"): - def onednn_mm(x: torch.Tensor, y: torch.Tensor, output_dtype=torch.float32): - # supports int8, fp32, fp16, and bf16 matmul with accumulation using a different float dtype + # torch._int_mm via onednn quantized matmul is supported with torch 2.9 + # ipex 2.7+ has the same torch._int_mm support as torch 2.9 but doesn't support torch.compile + # torch._int_mm directly uses onednn quantized matmul + # onednn qlinear is a wrapper around onednn quantized matmul + if hasattr(torch.ops, "onednn") and hasattr(torch.ops.onednn, "qlinear_pointwise"): + def onednn_mm(x: torch.Tensor, y: torch.Tensor): + # supports int8, fp32, fp16, and bf16 matmul with accumulation using a different dtype # int8 matmul with onednn is slower than 16 bit with dim_size < 4096 - return torch.ops.onednn.qlinear_pointwise(x, 1.0, 0, y, torch.ones(1, device=y.device), torch.zeros(1, device=y.device), None, 1.0, 0, output_dtype, "none", [], "none") + return torch.ops.onednn.qlinear_pointwise.default(x, 1.0, 0, y, torch.ones(1, device=y.device), torch.zeros(1, device=y.device), None, 1.0, 0, torch.float32, "none", [], "none") torch._int_mm = onednn_mm try: # torch.compile fix @@ -146,9 +153,6 @@ def ipex_init(): # pylint: disable=too-many-statements torch.xpu.empty_cache = lambda: None torch.cuda.empty_cache = torch.xpu.empty_cache - if has_ipex: - torch.cuda.memory_summary = torch.xpu.memory_summary - torch.cuda.memory_snapshot = torch.xpu.memory_snapshot torch.cuda.memory = torch.xpu.memory torch.cuda.memory_stats = torch.xpu.memory_stats torch.cuda.memory_allocated = torch.xpu.memory_allocated diff --git a/modules/intel/ipex/int_mm.py b/modules/intel/ipex/int_mm.py index 9c5fab093..4256937ac 100644 --- a/modules/intel/ipex/int_mm.py +++ b/modules/intel/ipex/int_mm.py @@ -49,7 +49,7 @@ def qlinear_unary( # GEMM template needs 2D input, normalize input shape here x = view(x, [-1, x_size[-1]]) if not isinstance(x_scale, ir.TensorBox): - assert type(x_scale) == float + assert isinstance(x_scale, float) x_scale = V.graph.add_tensor_constant( torch.tensor(x_scale, dtype=torch.float32), name="x_scale" ) @@ -71,7 +71,7 @@ def qlinear_unary( torch.tensor(0, dtype=torch.int32), name="x_zp" ) if not isinstance(x_zp, ir.TensorBox): - assert type(x_zp) == int + assert isinstance(x_zp, int) x_zp = V.graph.add_tensor_constant( torch.tensor(x_zp, dtype=torch.int32), name="x_zp" ) diff --git a/modules/interrogate/vqa.py b/modules/interrogate/vqa.py index 07f7e7a15..6f419a4e5 100644 --- a/modules/interrogate/vqa.py +++ b/modules/interrogate/vqa.py @@ -29,8 +29,9 @@ vlm_models = { "Google Gemma 3n E4B": "google/gemma-3n-E4B-it", # 1.5GB "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", + "Alibaba Qwen 2.0 VL 2B": "Qwen/Qwen2-VL-2B-Instruct", "Alibaba Qwen 2.5 Omni 3B": "Qwen/Qwen2.5-Omni-3B", + "Alibaba Qwen 2.5 VL 4B": "Qwen/Qwen2.5-VL-3B-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 @@ -122,7 +123,11 @@ def qwen(question: str, image: Image.Image, repo: str = None, system_prompt: str if model is None or loaded != repo: shared.log.debug(f'Interrogate load: vlm="{repo}"') model = None - model = transformers.Qwen2VLForConditionalGeneration.from_pretrained( + if '2.5' in repo: + cls_name = transformers.Qwen2_5_VLForConditionalGeneration + else: + cls_name = transformers.Qwen2VLForConditionalGeneration + model = cls_name.from_pretrained( repo, torch_dtype=devices.dtype, cache_dir=shared.opts.hfcache_dir, diff --git a/modules/ldsr/ldsr_model_arch.py b/modules/ldsr/ldsr_model_arch.py deleted file mode 100644 index 7494fcfcc..000000000 --- a/modules/ldsr/ldsr_model_arch.py +++ /dev/null @@ -1,236 +0,0 @@ -import os -import time -import numpy as np -import torch -import torchvision -from PIL import Image -from einops import rearrange, repeat -from omegaconf import OmegaConf -import safetensors.torch -from ldm.models.diffusion.ddim import DDIMSampler -from ldm.util import instantiate_from_config, ismap -from modules import devices, shared, sd_hijack -from modules.upscaler import compile_upscaler - -cached_ldsr_model: torch.nn.Module = None - - -# Create LDSR Class -class LDSR: - def load_model_from_config(self, half_attention): - global cached_ldsr_model # pylint: disable=global-statement - - if cached_ldsr_model is not None: - shared.log.info(f"Upscaler cached: type=LDSR model={self.modelPath}") - model: torch.nn.Module = cached_ldsr_model - else: - _, extension = os.path.splitext(self.modelPath) - if extension.lower() == ".safetensors": - pl_sd = safetensors.torch.load_file(self.modelPath, device="cpu") - else: - pl_sd = torch.load(self.modelPath, map_location="cpu") - shared.log.info(f"Upscaler loaded: type=LDSR model={self.modelPath}") - sd = pl_sd["state_dict"] if "state_dict" in pl_sd else pl_sd - config = OmegaConf.load(self.yamlPath) - config.model.target = "ldm.models.diffusion.ddpm.LatentDiffusionV1" - model: torch.nn.Module = instantiate_from_config(config.model) - model.load_state_dict(sd, strict=False) - model = model.to(devices.device) - if half_attention: - model = model.half() - if shared.cmd_opts.opt_channelslast: - model = model.to(memory_format=torch.channels_last) - sd_hijack.model_hijack.hijack(model) # apply optimization - model.eval() - model = compile_upscaler(model) - cached_ldsr_model = model - return {"model": model} - - def __init__(self, model_path, yaml_path): - self.modelPath = model_path - self.yamlPath = yaml_path - - @staticmethod - def run(model, selected_path, custom_steps, eta): - example = get_cond(selected_path) - n_runs = 1 - guider = None - ckwargs = None - ddim_use_x0_pred = False - temperature = 1. - eta = eta # pylint: disable=self-assigning-variable - custom_shape = None - height, width = example["image"].shape[1:3] - split_input = height >= 128 and width >= 128 - if split_input: - ks = 128 - stride = 64 - vqf = 4 # - model.split_input_params = {"ks": (ks, ks), "stride": (stride, stride), - "vqf": vqf, - "patch_distributed_vq": True, - "tie_braker": False, - "clip_max_weight": 0.5, - "clip_min_weight": 0.01, - "clip_max_tie_weight": 0.5, - "clip_min_tie_weight": 0.01} - else: - if hasattr(model, "split_input_params"): - delattr(model, "split_input_params") - - x_t = None - logs = None - for _ in range(n_runs): - if custom_shape is not None: - x_t = torch.randn(1, custom_shape[1], custom_shape[2], custom_shape[3]).to(model.device) - x_t = repeat(x_t, '1 c h w -> b c h w', b=custom_shape[0]) - - logs = make_convolutional_sample(example, model, - custom_steps=custom_steps, - eta=eta, quantize_x0=False, - custom_shape=custom_shape, - temperature=temperature, noise_dropout=0., - corrector=guider, corrector_kwargs=ckwargs, x_T=x_t, - ddim_use_x0_pred=ddim_use_x0_pred - ) - return logs - - def super_resolution(self, image, steps=100, target_scale=2, half_attention=False): - model = self.load_model_from_config(half_attention) - # Run settings - diffusion_steps = int(steps) - eta = 1.0 - im_og = image - width_og, height_og = im_og.size - # If we can adjust the max upscale size, then the 4 below should be our variable - down_sample_rate = target_scale / 4 - wd = width_og * down_sample_rate - hd = height_og * down_sample_rate - width_downsampled_pre = int(np.ceil(wd)) - height_downsampled_pre = int(np.ceil(hd)) - if down_sample_rate != 1: - shared.log.info(f'LDSR Downsampling from [{width_og}, {height_og}] to [{width_downsampled_pre}, {height_downsampled_pre}]') - im_og = im_og.resize((width_downsampled_pre, height_downsampled_pre), Image.LANCZOS) - else: - shared.log.info(f"LDSR Downsample rate is 1 from {target_scale} / 4 (Not downsampling)") - - # pad width and height to multiples of 64, pads with the edge values of image to avoid artifacts - pad_w, pad_h = np.max(((2, 2), np.ceil(np.array(im_og.size) / 64).astype(int)), axis=0) * 64 - im_og.size - im_padded = Image.fromarray(np.pad(np.array(im_og), ((0, pad_h), (0, pad_w), (0, 0)), mode='edge')) - - logs = self.run(model["model"], im_padded, diffusion_steps, eta) - - sample = logs["sample"] - sample = sample.detach().cpu() - sample = torch.clamp(sample, -1., 1.) - sample = (sample + 1.) / 2. * 255 - sample = sample.numpy().astype(np.uint8) - sample = np.transpose(sample, (0, 2, 3, 1)) - a = Image.fromarray(sample[0]) - # remove padding - a = a.crop((0, 0) + tuple(np.array(im_og.size) * 4)) - - if shared.opts.upscaler_unload: - del model - global cached_ldsr_model # pylint: disable=global-statement - cached_ldsr_model = None - shared.log.debug(f"Upscaler unloaded: type=LDSR model={self.modelPath}") - devices.torch_gc(force=True) - - return a - - -def get_cond(selected_path): - example = {} - up_f = 4 - c = selected_path.convert('RGB') - c = torch.unsqueeze(torchvision.transforms.ToTensor()(c), 0) - c_up = torchvision.transforms.functional.resize(c, size=[up_f * c.shape[2], up_f * c.shape[3]], - antialias=True) - c_up = rearrange(c_up, '1 c h w -> 1 h w c') - c = rearrange(c, '1 c h w -> 1 h w c') - c = 2. * c - 1. - - c = c.to(devices.device) - example["LR_image"] = c - example["image"] = c_up - - return example - - -@torch.no_grad() -def convsample_ddim(model, cond, steps, shape, eta=1.0, callback=None, normals_sequence=None, - mask=None, x0=None, quantize_x0=False, temperature=1., score_corrector=None, - corrector_kwargs=None, x_t=None - ): - ddim = DDIMSampler(model) - bs = shape[0] - shape = shape[1:] - shared.log.info(f"LDSR Sampling with eta = {eta}; steps: {steps}") - samples, intermediates = ddim.sample(steps, batch_size=bs, shape=shape, conditioning=cond, callback=callback, - normals_sequence=normals_sequence, quantize_x0=quantize_x0, eta=eta, - mask=mask, x0=x0, temperature=temperature, verbose=False, - score_corrector=score_corrector, - corrector_kwargs=corrector_kwargs, x_t=x_t) - - return samples, intermediates - - -@torch.no_grad() -def make_convolutional_sample(batch, model, custom_steps=None, eta=1.0, quantize_x0=False, custom_shape=None, temperature=1., noise_dropout=0., corrector=None, - corrector_kwargs=None, x_T=None, ddim_use_x0_pred=False): - log = {} - - z, c, x, xrec, xc = model.get_input(batch, model.first_stage_key, - return_first_stage_outputs=True, - force_c_encode=not (hasattr(model, 'split_input_params') - and model.cond_stage_key == 'coordinates_bbox'), - return_original_cond=True) - - if custom_shape is not None: - z = torch.randn(custom_shape) - shared.log.info(f"LDSR Generating {custom_shape[0]} samples of shape {custom_shape[1:]}") - - z0 = None - - log["input"] = x - log["reconstruction"] = xrec - - if ismap(xc): - log["original_conditioning"] = model.to_rgb(xc) - if hasattr(model, 'cond_stage_key'): - log[model.cond_stage_key] = model.to_rgb(xc) - - else: - log["original_conditioning"] = xc if xc is not None else torch.zeros_like(x) - if model.cond_stage_model: - log[model.cond_stage_key] = xc if xc is not None else torch.zeros_like(x) - if model.cond_stage_key == 'class_label': - log[model.cond_stage_key] = xc[model.cond_stage_key] - - with model.ema_scope("Plotting"): - t0 = time.time() - - sample, intermediates = convsample_ddim(model, c, steps=custom_steps, shape=z.shape, - eta=eta, - quantize_x0=quantize_x0, mask=None, x0=z0, - temperature=temperature, score_corrector=corrector, corrector_kwargs=corrector_kwargs, - x_t=x_T) - t1 = time.time() - - if ddim_use_x0_pred: - sample = intermediates['pred_x0'][-1] - - x_sample = model.decode_first_stage(sample) - - try: - x_sample_noquant = model.decode_first_stage(sample, force_not_quantize=True) - log["sample_noquant"] = x_sample_noquant - log["sample_diff"] = torch.abs(x_sample_noquant - x_sample) - except Exception: - pass - - log["sample"] = x_sample - log["time"] = t1 - t0 - - return log diff --git a/modules/ldsr/sd_hijack_autoencoder.py b/modules/ldsr/sd_hijack_autoencoder.py deleted file mode 100644 index 4fd7d67fe..000000000 --- a/modules/ldsr/sd_hijack_autoencoder.py +++ /dev/null @@ -1,290 +0,0 @@ -# The content of this file comes from the ldm/models/autoencoder.py file of the compvis/stable-diffusion repo -# The VQModel & VQModelInterface were subsequently removed from ldm/models/autoencoder.py when we moved to the stability-ai/stablediffusion repo -# As the LDSR upscaler relies on VQModel & VQModelInterface, the hijack aims to put them back into the ldm.models.autoencoder -from contextlib import contextmanager -import numpy as np -import torch -import pytorch_lightning as pl -import torch.nn.functional as F -from torch.optim.lr_scheduler import LambdaLR -from taming.modules.vqvae.quantize import VectorQuantizer2 as VectorQuantizer -from ldm.modules.ema import LitEma -from ldm.modules.diffusionmodules.model import Encoder, Decoder -from ldm.util import instantiate_from_config - -import ldm.models.autoencoder -from packaging import version - -class VQModel(pl.LightningModule): - def __init__(self, - ddconfig, - lossconfig, - n_embed, - embed_dim, - ckpt_path=None, - ignore_keys=None, - image_key="image", - colorize_nlabels=None, - monitor=None, - batch_resize_range=None, - scheduler_config=None, - lr_g_factor=1.0, - remap=None, - sane_index_shape=False, # tell vector quantizer to return indices as bhw - use_ema=False - ): - super().__init__() - self.embed_dim = embed_dim - self.n_embed = n_embed - self.image_key = image_key - self.encoder = Encoder(**ddconfig) - self.decoder = Decoder(**ddconfig) - self.loss = instantiate_from_config(lossconfig) - self.quantize = VectorQuantizer(n_embed, embed_dim, beta=0.25, - remap=remap, - sane_index_shape=sane_index_shape) - self.quant_conv = torch.nn.Conv2d(ddconfig["z_channels"], embed_dim, 1) - self.post_quant_conv = torch.nn.Conv2d(embed_dim, ddconfig["z_channels"], 1) - if colorize_nlabels is not None: - assert type(colorize_nlabels)==int - self.register_buffer("colorize", torch.randn(3, colorize_nlabels, 1, 1)) - if monitor is not None: - self.monitor = monitor - self.batch_resize_range = batch_resize_range - if self.batch_resize_range is not None: - print(f"{self.__class__.__name__}: Using per-batch resizing in range {batch_resize_range}.") - - self.use_ema = use_ema - if self.use_ema: - self.model_ema = LitEma(self) - print(f"Keeping EMAs of {len(list(self.model_ema.buffers()))}.") - - if ckpt_path is not None: - self.init_from_ckpt(ckpt_path, ignore_keys=ignore_keys or []) - self.scheduler_config = scheduler_config - self.lr_g_factor = lr_g_factor - - @contextmanager - def ema_scope(self, context=None): - if self.use_ema: - self.model_ema.store(self.parameters()) - self.model_ema.copy_to(self) - if context is not None: - print(f"{context}: Switched to EMA weights") - try: - yield None - finally: - if self.use_ema: - self.model_ema.restore(self.parameters()) - if context is not None: - print(f"{context}: Restored training weights") - - def init_from_ckpt(self, path, ignore_keys=None): - sd = torch.load(path, map_location="cpu")["state_dict"] - keys = list(sd.keys()) - for k in keys: - for ik in ignore_keys or []: - if k.startswith(ik): - print(f"Deleting key {k} from state_dict.") - del sd[k] - missing, unexpected = self.load_state_dict(sd, strict=False) - print(f"Restored from {path} with {len(missing)} missing and {len(unexpected)} unexpected keys") - if len(missing) > 0: - print(f"Missing Keys: {missing}") - print(f"Unexpected Keys: {unexpected}") - - def on_train_batch_end(self, *args, **kwargs): - if self.use_ema: - self.model_ema(self) - - def encode(self, x): - h = self.encoder(x) - h = self.quant_conv(h) - quant, emb_loss, info = self.quantize(h) - return quant, emb_loss, info - - def encode_to_prequant(self, x): - h = self.encoder(x) - h = self.quant_conv(h) - return h - - def decode(self, quant): - quant = self.post_quant_conv(quant) - dec = self.decoder(quant) - return dec - - def decode_code(self, code_b): - quant_b = self.quantize.embed_code(code_b) - dec = self.decode(quant_b) - return dec - - def forward(self, input, return_pred_indices=False): - quant, diff, (_,_,ind) = self.encode(input) - dec = self.decode(quant) - if return_pred_indices: - return dec, diff, ind - return dec, diff - - def get_input(self, batch, k): - x = batch[k] - if len(x.shape) == 3: - x = x[..., None] - x = x.permute(0, 3, 1, 2).to(memory_format=torch.contiguous_format).float() - if self.batch_resize_range is not None: - lower_size = self.batch_resize_range[0] - upper_size = self.batch_resize_range[1] - if self.global_step <= 4: - # do the first few batches with max size to avoid later oom - new_resize = upper_size - else: - new_resize = np.random.choice(np.arange(lower_size, upper_size+16, 16)) - if new_resize != x.shape[2]: - x = F.interpolate(x, size=new_resize, mode="bicubic") - x = x.detach() - return x - - def training_step(self, batch, batch_idx, optimizer_idx): - # https://github.com/pytorch/pytorch/issues/37142 - # try not to fool the heuristics - x = self.get_input(batch, self.image_key) - xrec, qloss, ind = self(x, return_pred_indices=True) - - if optimizer_idx == 0: - # autoencode - aeloss, log_dict_ae = self.loss(qloss, x, xrec, optimizer_idx, self.global_step, - last_layer=self.get_last_layer(), split="train", - predicted_indices=ind) - - self.log_dict(log_dict_ae, prog_bar=False, logger=True, on_step=True, on_epoch=True) - return aeloss - - if optimizer_idx == 1: - # discriminator - discloss, log_dict_disc = self.loss(qloss, x, xrec, optimizer_idx, self.global_step, - last_layer=self.get_last_layer(), split="train") - self.log_dict(log_dict_disc, prog_bar=False, logger=True, on_step=True, on_epoch=True) - return discloss - - def validation_step(self, batch, batch_idx): - log_dict = self._validation_step(batch, batch_idx) - with self.ema_scope(): - self._validation_step(batch, batch_idx, suffix="_ema") - return log_dict - - def _validation_step(self, batch, batch_idx, suffix=""): - x = self.get_input(batch, self.image_key) - xrec, qloss, ind = self(x, return_pred_indices=True) - aeloss, log_dict_ae = self.loss(qloss, x, xrec, 0, - self.global_step, - last_layer=self.get_last_layer(), - split="val"+suffix, - predicted_indices=ind - ) - - discloss, log_dict_disc = self.loss(qloss, x, xrec, 1, - self.global_step, - last_layer=self.get_last_layer(), - split="val"+suffix, - predicted_indices=ind - ) - rec_loss = log_dict_ae[f"val{suffix}/rec_loss"] - self.log(f"val{suffix}/rec_loss", rec_loss, - prog_bar=True, logger=True, on_step=False, on_epoch=True, sync_dist=True) - self.log(f"val{suffix}/aeloss", aeloss, - prog_bar=True, logger=True, on_step=False, on_epoch=True, sync_dist=True) - if version.parse(pl.__version__) >= version.parse('1.4.0'): - del log_dict_ae[f"val{suffix}/rec_loss"] - self.log_dict(log_dict_ae) - self.log_dict(log_dict_disc) - return self.log_dict - - def configure_optimizers(self): - lr_d = self.learning_rate - lr_g = self.lr_g_factor*self.learning_rate - print("lr_d", lr_d) - print("lr_g", lr_g) - opt_ae = torch.optim.Adam(list(self.encoder.parameters())+ - list(self.decoder.parameters())+ - list(self.quantize.parameters())+ - list(self.quant_conv.parameters())+ - list(self.post_quant_conv.parameters()), - lr=lr_g, betas=(0.5, 0.9)) - opt_disc = torch.optim.Adam(self.loss.discriminator.parameters(), - lr=lr_d, betas=(0.5, 0.9)) - - if self.scheduler_config is not None: - scheduler = instantiate_from_config(self.scheduler_config) - - print("Setting up LambdaLR scheduler...") - scheduler = [ - { - 'scheduler': LambdaLR(opt_ae, lr_lambda=scheduler.schedule), - 'interval': 'step', - 'frequency': 1 - }, - { - 'scheduler': LambdaLR(opt_disc, lr_lambda=scheduler.schedule), - 'interval': 'step', - 'frequency': 1 - }, - ] - return [opt_ae, opt_disc], scheduler - return [opt_ae, opt_disc], [] - - def get_last_layer(self): - return self.decoder.conv_out.weight - - def log_images(self, batch, only_inputs=False, plot_ema=False, **kwargs): - log = {} - x = self.get_input(batch, self.image_key) - x = x.to(self.device) - if only_inputs: - log["inputs"] = x - return log - xrec, _ = self(x) - if x.shape[1] > 3: - # colorize with random projection - assert xrec.shape[1] > 3 - x = self.to_rgb(x) - xrec = self.to_rgb(xrec) - log["inputs"] = x - log["reconstructions"] = xrec - if plot_ema: - with self.ema_scope(): - xrec_ema, _ = self(x) - if x.shape[1] > 3: - xrec_ema = self.to_rgb(xrec_ema) - log["reconstructions_ema"] = xrec_ema - return log - - def to_rgb(self, x): - assert self.image_key == "segmentation" - if not hasattr(self, "colorize"): - self.register_buffer("colorize", torch.randn(3, x.shape[1], 1, 1).to(x)) - x = F.conv2d(x, weight=self.colorize) - x = 2.*(x-x.min())/(x.max()-x.min()) - 1. - return x - - -class VQModelInterface(VQModel): - def __init__(self, embed_dim, *args, **kwargs): - super().__init__(*args, embed_dim=embed_dim, **kwargs) - self.embed_dim = embed_dim - - def encode(self, x): - h = self.encoder(x) - h = self.quant_conv(h) - return h - - def decode(self, h, force_not_quantize=False): - # also go through quantization layer - if not force_not_quantize: - quant, emb_loss, info = self.quantize(h) - else: - quant = h - quant = self.post_quant_conv(quant) - dec = self.decoder(quant) - return dec - -ldm.models.autoencoder.VQModel = VQModel -ldm.models.autoencoder.VQModelInterface = VQModelInterface diff --git a/modules/ldsr/sd_hijack_ddpm_v1.py b/modules/ldsr/sd_hijack_ddpm_v1.py deleted file mode 100644 index 2cf506990..000000000 --- a/modules/ldsr/sd_hijack_ddpm_v1.py +++ /dev/null @@ -1,1439 +0,0 @@ -# This script is copied from the compvis/stable-diffusion repo (aka the SD V1 repo) -# Original filename: ldm/models/diffusion/ddpm.py -# The purpose to reinstate the old DDPM logic which works with VQ, whereas the V2 one doesn't -# Some models such as LDSR require VQ to work correctly -# The classes are suffixed with "V1" and added back to the "ldm.models.diffusion.ddpm" module - -import torch -import torch.nn as nn -import numpy as np -import pytorch_lightning as pl -from torch.optim.lr_scheduler import LambdaLR -from einops import rearrange, repeat -from contextlib import contextmanager -from functools import partial -from tqdm import tqdm -from torchvision.utils import make_grid -from pytorch_lightning.utilities.distributed import rank_zero_only -from ldm.util import log_txt_as_img, exists, default, ismap, isimage, mean_flat, count_params, instantiate_from_config -from ldm.modules.ema import LitEma -from ldm.modules.distributions.distributions import normal_kl, DiagonalGaussianDistribution -from ldm.models.autoencoder import VQModelInterface, IdentityFirstStage, AutoencoderKL -from ldm.modules.diffusionmodules.util import make_beta_schedule, extract_into_tensor, noise_like -from ldm.models.diffusion.ddim import DDIMSampler -import ldm.models.diffusion.ddpm - -__conditioning_keys__ = {'concat': 'c_concat', - 'crossattn': 'c_crossattn', - 'adm': 'y'} - - -def disabled_train(self, mode=True): - """Overwrite model.train with this function to make sure train/eval mode - does not change anymore.""" - return self - - -def uniform_on_device(r1, r2, shape, device): - return (r1 - r2) * torch.rand(*shape, device=device) + r2 - - -class DDPMV1(pl.LightningModule): - # classic DDPM with Gaussian diffusion, in image space - def __init__(self, - unet_config, - timesteps=1000, - beta_schedule="linear", - loss_type="l2", - ckpt_path=None, - ignore_keys=None, - load_only_unet=False, - monitor="val/loss", - use_ema=True, - first_stage_key="image", - image_size=256, - channels=3, - log_every_t=100, - clip_denoised=True, - linear_start=1e-4, - linear_end=2e-2, - cosine_s=8e-3, - given_betas=None, - original_elbo_weight=0., - v_posterior=0., # weight for choosing posterior variance as sigma = (1-v) * beta_tilde + v * beta - l_simple_weight=1., - conditioning_key=None, - parameterization="eps", # all assuming fixed variance schedules - scheduler_config=None, - use_positional_encodings=False, - learn_logvar=False, - logvar_init=0., - ): - super().__init__() - assert parameterization in ["eps", "x0"], 'currently only supporting "eps" and "x0"' - self.parameterization = parameterization - print(f"{self.__class__.__name__}: Running in {self.parameterization}-prediction mode") - self.cond_stage_model = None - self.clip_denoised = clip_denoised - self.log_every_t = log_every_t - self.first_stage_key = first_stage_key - self.image_size = image_size # try conv? - self.channels = channels - self.use_positional_encodings = use_positional_encodings - self.model = DiffusionWrapperV1(unet_config, conditioning_key) - count_params(self.model, verbose=True) - self.use_ema = use_ema - if self.use_ema: - self.model_ema = LitEma(self.model) - print(f"Keeping EMAs of {len(list(self.model_ema.buffers()))}.") - - self.use_scheduler = scheduler_config is not None - if self.use_scheduler: - self.scheduler_config = scheduler_config - - self.v_posterior = v_posterior - self.original_elbo_weight = original_elbo_weight - self.l_simple_weight = l_simple_weight - - if monitor is not None: - self.monitor = monitor - if ckpt_path is not None: - self.init_from_ckpt(ckpt_path, ignore_keys=ignore_keys or [], only_model=load_only_unet) - - self.register_schedule(given_betas=given_betas, beta_schedule=beta_schedule, timesteps=timesteps, - linear_start=linear_start, linear_end=linear_end, cosine_s=cosine_s) - - self.loss_type = loss_type - - self.learn_logvar = learn_logvar - self.logvar = torch.full(fill_value=logvar_init, size=(self.num_timesteps,)) - if self.learn_logvar: - self.logvar = nn.Parameter(self.logvar, requires_grad=True) - - - def register_schedule(self, given_betas=None, beta_schedule="linear", timesteps=1000, - linear_start=1e-4, linear_end=2e-2, cosine_s=8e-3): - if exists(given_betas): - betas = given_betas - else: - betas = make_beta_schedule(beta_schedule, timesteps, linear_start=linear_start, linear_end=linear_end, - cosine_s=cosine_s) - alphas = 1. - betas - alphas_cumprod = np.cumprod(alphas, axis=0) - alphas_cumprod_prev = np.append(1., alphas_cumprod[:-1]) - - timesteps, = betas.shape - self.num_timesteps = int(timesteps) - self.linear_start = linear_start - self.linear_end = linear_end - assert alphas_cumprod.shape[0] == self.num_timesteps, 'alphas have to be defined for each timestep' - - to_torch = partial(torch.tensor, dtype=torch.float32) - - self.register_buffer('betas', to_torch(betas)) - self.register_buffer('alphas_cumprod', to_torch(alphas_cumprod)) - self.register_buffer('alphas_cumprod_prev', to_torch(alphas_cumprod_prev)) - - # calculations for diffusion q(x_t | x_{t-1}) and others - self.register_buffer('sqrt_alphas_cumprod', to_torch(np.sqrt(alphas_cumprod))) - self.register_buffer('sqrt_one_minus_alphas_cumprod', to_torch(np.sqrt(1. - alphas_cumprod))) - self.register_buffer('log_one_minus_alphas_cumprod', to_torch(np.log(1. - alphas_cumprod))) - self.register_buffer('sqrt_recip_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod))) - self.register_buffer('sqrt_recipm1_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod - 1))) - - # calculations for posterior q(x_{t-1} | x_t, x_0) - posterior_variance = (1 - self.v_posterior) * betas * (1. - alphas_cumprod_prev) / ( - 1. - alphas_cumprod) + self.v_posterior * betas - # above: equal to 1. / (1. / (1. - alpha_cumprod_tm1) + alpha_t / beta_t) - self.register_buffer('posterior_variance', to_torch(posterior_variance)) - # below: log calculation clipped because the posterior variance is 0 at the beginning of the diffusion chain - self.register_buffer('posterior_log_variance_clipped', to_torch(np.log(np.maximum(posterior_variance, 1e-20)))) - self.register_buffer('posterior_mean_coef1', to_torch( - betas * np.sqrt(alphas_cumprod_prev) / (1. - alphas_cumprod))) - self.register_buffer('posterior_mean_coef2', to_torch( - (1. - alphas_cumprod_prev) * np.sqrt(alphas) / (1. - alphas_cumprod))) - - if self.parameterization == "eps": - lvlb_weights = self.betas ** 2 / ( - 2 * self.posterior_variance * to_torch(alphas) * (1 - self.alphas_cumprod)) - elif self.parameterization == "x0": - lvlb_weights = 0.5 * np.sqrt(torch.Tensor(alphas_cumprod)) / (2. * 1 - torch.Tensor(alphas_cumprod)) - else: - raise NotImplementedError("mu not supported") - lvlb_weights[0] = lvlb_weights[1] - self.register_buffer('lvlb_weights', lvlb_weights, persistent=False) - assert not torch.isnan(self.lvlb_weights).all() - - @contextmanager - def ema_scope(self, context=None): - if self.use_ema: - self.model_ema.store(self.model.parameters()) - self.model_ema.copy_to(self.model) - if context is not None: - print(f"{context}: Switched to EMA weights") - try: - yield None - finally: - if self.use_ema: - self.model_ema.restore(self.model.parameters()) - if context is not None: - print(f"{context}: Restored training weights") - - def init_from_ckpt(self, path, ignore_keys=None, only_model=False): - sd = torch.load(path, map_location="cpu") - if "state_dict" in list(sd.keys()): - sd = sd["state_dict"] - keys = list(sd.keys()) - for k in keys: - for ik in ignore_keys or []: - if k.startswith(ik): - print("Deleting key {} from state_dict.".format(k)) - del sd[k] - missing, unexpected = self.load_state_dict(sd, strict=False) if not only_model else self.model.load_state_dict( - sd, strict=False) - print(f"Restored from {path} with {len(missing)} missing and {len(unexpected)} unexpected keys") - if len(missing) > 0: - print(f"Missing Keys: {missing}") - if len(unexpected) > 0: - print(f"Unexpected Keys: {unexpected}") - - def q_mean_variance(self, x_start, t): - """ - Get the distribution q(x_t | x_0). - :param x_start: the [N x C x ...] tensor of noiseless inputs. - :param t: the number of diffusion steps (minus 1). Here, 0 means one step. - :return: A tuple (mean, variance, log_variance), all of x_start's shape. - """ - mean = (extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start) - variance = extract_into_tensor(1.0 - self.alphas_cumprod, t, x_start.shape) - log_variance = extract_into_tensor(self.log_one_minus_alphas_cumprod, t, x_start.shape) - return mean, variance, log_variance - - def predict_start_from_noise(self, x_t, t, noise): - return ( - extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t - - extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * noise - ) - - def q_posterior(self, x_start, x_t, t): - posterior_mean = ( - extract_into_tensor(self.posterior_mean_coef1, t, x_t.shape) * x_start + - extract_into_tensor(self.posterior_mean_coef2, t, x_t.shape) * x_t - ) - posterior_variance = extract_into_tensor(self.posterior_variance, t, x_t.shape) - posterior_log_variance_clipped = extract_into_tensor(self.posterior_log_variance_clipped, t, x_t.shape) - return posterior_mean, posterior_variance, posterior_log_variance_clipped - - def p_mean_variance(self, x, t, clip_denoised: bool): - model_out = self.model(x, t) - if self.parameterization == "eps": - x_recon = self.predict_start_from_noise(x, t=t, noise=model_out) - elif self.parameterization == "x0": - x_recon = model_out - if clip_denoised: - x_recon.clamp_(-1., 1.) - - model_mean, posterior_variance, posterior_log_variance = self.q_posterior(x_start=x_recon, x_t=x, t=t) - return model_mean, posterior_variance, posterior_log_variance - - @torch.no_grad() - def p_sample(self, x, t, clip_denoised=True, repeat_noise=False): - b, *_, device = *x.shape, x.device - model_mean, _, model_log_variance = self.p_mean_variance(x=x, t=t, clip_denoised=clip_denoised) - noise = noise_like(x.shape, device, repeat_noise) - # no noise when t == 0 - nonzero_mask = (1 - (t == 0).float()).reshape(b, *((1,) * (len(x.shape) - 1))) - return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise - - @torch.no_grad() - def p_sample_loop(self, shape, return_intermediates=False): - device = self.betas.device - b = shape[0] - img = torch.randn(shape, device=device) - intermediates = [img] - for i in tqdm(reversed(range(0, self.num_timesteps)), desc='Sampling t', total=self.num_timesteps): - img = self.p_sample(img, torch.full((b,), i, device=device, dtype=torch.long), - clip_denoised=self.clip_denoised) - if i % self.log_every_t == 0 or i == self.num_timesteps - 1: - intermediates.append(img) - if return_intermediates: - return img, intermediates - return img - - @torch.no_grad() - def sample(self, batch_size=16, return_intermediates=False): - image_size = self.image_size - channels = self.channels - return self.p_sample_loop((batch_size, channels, image_size, image_size), - return_intermediates=return_intermediates) - - def q_sample(self, x_start, t, noise=None): - noise = default(noise, lambda: torch.randn_like(x_start)) - return (extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start + - extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise) - - def get_loss(self, pred, target, mean=True): - if self.loss_type == 'l1': - loss = (target - pred).abs() - if mean: - loss = loss.mean() - elif self.loss_type == 'l2': - if mean: - loss = torch.nn.functional.mse_loss(target, pred) - else: - loss = torch.nn.functional.mse_loss(target, pred, reduction='none') - else: - raise NotImplementedError("unknown loss type '{loss_type}'") - - return loss - - def p_losses(self, x_start, t, noise=None): - noise = default(noise, lambda: torch.randn_like(x_start)) - x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise) - model_out = self.model(x_noisy, t) - - loss_dict = {} - if self.parameterization == "eps": - target = noise - elif self.parameterization == "x0": - target = x_start - else: - raise NotImplementedError(f"Paramterization {self.parameterization} not yet supported") - - loss = self.get_loss(model_out, target, mean=False).mean(dim=[1, 2, 3]) - - log_prefix = 'train' if self.training else 'val' - - loss_dict.update({f'{log_prefix}/loss_simple': loss.mean()}) - loss_simple = loss.mean() * self.l_simple_weight - - loss_vlb = (self.lvlb_weights[t] * loss).mean() - loss_dict.update({f'{log_prefix}/loss_vlb': loss_vlb}) - - loss = loss_simple + self.original_elbo_weight * loss_vlb - - loss_dict.update({f'{log_prefix}/loss': loss}) - - return loss, loss_dict - - def forward(self, x, *args, **kwargs): - # b, c, h, w, device, img_size, = *x.shape, x.device, self.image_size - # assert h == img_size and w == img_size, f'height and width of image must be {img_size}' - t = torch.randint(0, self.num_timesteps, (x.shape[0],), device=self.device).long() - return self.p_losses(x, t, *args, **kwargs) - - def get_input(self, batch, k): - x = batch[k] - if len(x.shape) == 3: - x = x[..., None] - x = rearrange(x, 'b h w c -> b c h w') - x = x.to(memory_format=torch.contiguous_format).float() - return x - - def shared_step(self, batch): - x = self.get_input(batch, self.first_stage_key) - loss, loss_dict = self(x) - return loss, loss_dict - - def training_step(self, batch, batch_idx): - loss, loss_dict = self.shared_step(batch) - - self.log_dict(loss_dict, prog_bar=True, - logger=True, on_step=True, on_epoch=True) - - self.log("global_step", self.global_step, - prog_bar=True, logger=True, on_step=True, on_epoch=False) - - if self.use_scheduler: - lr = self.optimizers().param_groups[0]['lr'] - self.log('lr_abs', lr, prog_bar=True, logger=True, on_step=True, on_epoch=False) - - return loss - - @torch.no_grad() - def validation_step(self, batch, batch_idx): - _, loss_dict_no_ema = self.shared_step(batch) - with self.ema_scope(): - _, loss_dict_ema = self.shared_step(batch) - loss_dict_ema = {key + '_ema': loss_dict_ema[key] for key in loss_dict_ema} - self.log_dict(loss_dict_no_ema, prog_bar=False, logger=True, on_step=False, on_epoch=True) - self.log_dict(loss_dict_ema, prog_bar=False, logger=True, on_step=False, on_epoch=True) - - def on_train_batch_end(self, *args, **kwargs): - if self.use_ema: - self.model_ema(self.model) - - def _get_rows_from_list(self, samples): - n_imgs_per_row = len(samples) - denoise_grid = rearrange(samples, 'n b c h w -> b n c h w') - denoise_grid = rearrange(denoise_grid, 'b n c h w -> (b n) c h w') - denoise_grid = make_grid(denoise_grid, nrow=n_imgs_per_row) - return denoise_grid - - @torch.no_grad() - def log_images(self, batch, N=8, n_row=2, sample=True, return_keys=None, **kwargs): - log = {} - x = self.get_input(batch, self.first_stage_key) - N = min(x.shape[0], N) - n_row = min(x.shape[0], n_row) - x = x.to(self.device)[:N] - log["inputs"] = x - - # get diffusion row - diffusion_row = [] - x_start = x[:n_row] - - for t in range(self.num_timesteps): - if t % self.log_every_t == 0 or t == self.num_timesteps - 1: - t = repeat(torch.tensor([t]), '1 -> b', b=n_row) - t = t.to(self.device).long() - noise = torch.randn_like(x_start) - x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise) - diffusion_row.append(x_noisy) - - log["diffusion_row"] = self._get_rows_from_list(diffusion_row) - - if sample: - # get denoise row - with self.ema_scope("Plotting"): - samples, denoise_row = self.sample(batch_size=N, return_intermediates=True) - - log["samples"] = samples - log["denoise_row"] = self._get_rows_from_list(denoise_row) - - if return_keys: - if np.intersect1d(list(log.keys()), return_keys).shape[0] == 0: - return log - else: - return {key: log[key] for key in return_keys} - return log - - def configure_optimizers(self): - lr = self.learning_rate - params = list(self.model.parameters()) - if self.learn_logvar: - params = params + [self.logvar] - opt = torch.optim.AdamW(params, lr=lr) - return opt - - -class LatentDiffusionV1(DDPMV1): - """main class""" - def __init__(self, - first_stage_config, - cond_stage_config, - num_timesteps_cond=None, - cond_stage_key="image", - cond_stage_trainable=False, - concat_mode=True, - cond_stage_forward=None, - conditioning_key=None, - scale_factor=1.0, - scale_by_std=False, - *args, **kwargs): - self.num_timesteps_cond = default(num_timesteps_cond, 1) - self.scale_by_std = scale_by_std - assert self.num_timesteps_cond <= kwargs['timesteps'] - # for backwards compatibility after implementation of DiffusionWrapper - if conditioning_key is None: - conditioning_key = 'concat' if concat_mode else 'crossattn' - if cond_stage_config == '__is_unconditional__': - conditioning_key = None - ckpt_path = kwargs.pop("ckpt_path", None) - ignore_keys = kwargs.pop("ignore_keys", []) - super().__init__(*args, conditioning_key=conditioning_key, **kwargs) - self.concat_mode = concat_mode - self.cond_stage_trainable = cond_stage_trainable - self.cond_stage_key = cond_stage_key - try: - self.num_downs = len(first_stage_config.params.ddconfig.ch_mult) - 1 - except Exception: - self.num_downs = 0 - if not scale_by_std: - self.scale_factor = scale_factor - else: - self.register_buffer('scale_factor', torch.tensor(scale_factor)) - self.instantiate_first_stage(first_stage_config) - self.instantiate_cond_stage(cond_stage_config) - self.cond_stage_forward = cond_stage_forward - self.clip_denoised = False - self.bbox_tokenizer = None - - self.restarted_from_ckpt = False - if ckpt_path is not None: - self.init_from_ckpt(ckpt_path, ignore_keys) - self.restarted_from_ckpt = True - - def make_cond_schedule(self, ): - self.cond_ids = torch.full(size=(self.num_timesteps,), fill_value=self.num_timesteps - 1, dtype=torch.long) - ids = torch.round(torch.linspace(0, self.num_timesteps - 1, self.num_timesteps_cond)).long() - self.cond_ids[:self.num_timesteps_cond] = ids - - @rank_zero_only - @torch.no_grad() - def on_train_batch_start(self, batch, batch_idx, dataloader_idx): - # only for very first batch - if self.scale_by_std and self.current_epoch == 0 and self.global_step == 0 and batch_idx == 0 and not self.restarted_from_ckpt: - assert self.scale_factor == 1., 'rather not use custom rescaling and std-rescaling simultaneously' - # set rescale weight to 1./std of encodings - print("### USING STD-RESCALING ###") - x = super().get_input(batch, self.first_stage_key) - x = x.to(self.device) - encoder_posterior = self.encode_first_stage(x) - z = self.get_first_stage_encoding(encoder_posterior).detach() - del self.scale_factor - self.register_buffer('scale_factor', 1. / z.flatten().std()) - print(f"setting self.scale_factor to {self.scale_factor}") - print("### USING STD-RESCALING ###") - - def register_schedule(self, - given_betas=None, beta_schedule="linear", timesteps=1000, - linear_start=1e-4, linear_end=2e-2, cosine_s=8e-3): - super().register_schedule(given_betas, beta_schedule, timesteps, linear_start, linear_end, cosine_s) - - self.shorten_cond_schedule = self.num_timesteps_cond > 1 - if self.shorten_cond_schedule: - self.make_cond_schedule() - - def instantiate_first_stage(self, config): - model = instantiate_from_config(config) - self.first_stage_model = model.eval() - self.first_stage_model.train = disabled_train - for param in self.first_stage_model.parameters(): - param.requires_grad = False - - def instantiate_cond_stage(self, config): - if not self.cond_stage_trainable: - if config == "__is_first_stage__": - print("Using first stage also as cond stage.") - self.cond_stage_model = self.first_stage_model - elif config == "__is_unconditional__": - print(f"Training {self.__class__.__name__} as an unconditional model.") - self.cond_stage_model = None - # self.be_unconditional = True - else: - model = instantiate_from_config(config) - self.cond_stage_model = model.eval() - self.cond_stage_model.train = disabled_train - for param in self.cond_stage_model.parameters(): - param.requires_grad = False - else: - assert config != '__is_first_stage__' - assert config != '__is_unconditional__' - model = instantiate_from_config(config) - self.cond_stage_model = model - - def _get_denoise_row_from_list(self, samples, desc='', force_no_decoder_quantization=False): - denoise_row = [] - for zd in tqdm(samples, desc=desc): - denoise_row.append(self.decode_first_stage(zd.to(self.device), - force_not_quantize=force_no_decoder_quantization)) - n_imgs_per_row = len(denoise_row) - denoise_row = torch.stack(denoise_row) # n_log_step, n_row, C, H, W - denoise_grid = rearrange(denoise_row, 'n b c h w -> b n c h w') - denoise_grid = rearrange(denoise_grid, 'b n c h w -> (b n) c h w') - denoise_grid = make_grid(denoise_grid, nrow=n_imgs_per_row) - return denoise_grid - - def get_first_stage_encoding(self, encoder_posterior): - if isinstance(encoder_posterior, DiagonalGaussianDistribution): - z = encoder_posterior.sample() - elif isinstance(encoder_posterior, torch.Tensor): - z = encoder_posterior - else: - raise NotImplementedError(f"encoder_posterior of type '{type(encoder_posterior)}' not yet implemented") - return self.scale_factor * z - - def get_learned_conditioning(self, c): - if self.cond_stage_forward is None: - if hasattr(self.cond_stage_model, 'encode') and callable(self.cond_stage_model.encode): - c = self.cond_stage_model.encode(c) - if isinstance(c, DiagonalGaussianDistribution): - c = c.mode() - else: - c = self.cond_stage_model(c) - else: - assert hasattr(self.cond_stage_model, self.cond_stage_forward) - c = getattr(self.cond_stage_model, self.cond_stage_forward)(c) - return c - - def meshgrid(self, h, w): - y = torch.arange(0, h).view(h, 1, 1).repeat(1, w, 1) - x = torch.arange(0, w).view(1, w, 1).repeat(h, 1, 1) - - arr = torch.cat([y, x], dim=-1) - return arr - - def delta_border(self, h, w): - """ - :param h: height - :param w: width - :return: normalized distance to image border, - wtith min distance = 0 at border and max dist = 0.5 at image center - """ - lower_right_corner = torch.tensor([h - 1, w - 1]).view(1, 1, 2) - arr = self.meshgrid(h, w) / lower_right_corner - dist_left_up = torch.min(arr, dim=-1, keepdims=True)[0] - dist_right_down = torch.min(1 - arr, dim=-1, keepdims=True)[0] - edge_dist = torch.min(torch.cat([dist_left_up, dist_right_down], dim=-1), dim=-1)[0] - return edge_dist - - def get_weighting(self, h, w, Ly, Lx, device): - weighting = self.delta_border(h, w) - weighting = torch.clip(weighting, self.split_input_params["clip_min_weight"], - self.split_input_params["clip_max_weight"], ) - weighting = weighting.view(1, h * w, 1).repeat(1, 1, Ly * Lx).to(device) - - if self.split_input_params["tie_braker"]: - L_weighting = self.delta_border(Ly, Lx) - L_weighting = torch.clip(L_weighting, - self.split_input_params["clip_min_tie_weight"], - self.split_input_params["clip_max_tie_weight"]) - - L_weighting = L_weighting.view(1, 1, Ly * Lx).to(device) - weighting = weighting * L_weighting - return weighting - - def get_fold_unfold(self, x, kernel_size, stride, uf=1, df=1): # todo load once not every time, shorten code - """ - :param x: img of size (bs, c, h, w) - :return: n img crops of size (n, bs, c, kernel_size[0], kernel_size[1]) - """ - bs, nc, h, w = x.shape - - # number of crops in image - Ly = (h - kernel_size[0]) // stride[0] + 1 - Lx = (w - kernel_size[1]) // stride[1] + 1 - - if uf == 1 and df == 1: - fold_params = dict(kernel_size=kernel_size, dilation=1, padding=0, stride=stride) - unfold = torch.nn.Unfold(**fold_params) - - fold = torch.nn.Fold(output_size=x.shape[2:], **fold_params) - - weighting = self.get_weighting(kernel_size[0], kernel_size[1], Ly, Lx, x.device).to(x.dtype) - normalization = fold(weighting).view(1, 1, h, w) # normalizes the overlap - weighting = weighting.view((1, 1, kernel_size[0], kernel_size[1], Ly * Lx)) - - elif uf > 1 and df == 1: - fold_params = dict(kernel_size=kernel_size, dilation=1, padding=0, stride=stride) - unfold = torch.nn.Unfold(**fold_params) - - fold_params2 = dict(kernel_size=(kernel_size[0] * uf, kernel_size[0] * uf), - dilation=1, padding=0, - stride=(stride[0] * uf, stride[1] * uf)) - fold = torch.nn.Fold(output_size=(x.shape[2] * uf, x.shape[3] * uf), **fold_params2) - - weighting = self.get_weighting(kernel_size[0] * uf, kernel_size[1] * uf, Ly, Lx, x.device).to(x.dtype) - normalization = fold(weighting).view(1, 1, h * uf, w * uf) # normalizes the overlap - weighting = weighting.view((1, 1, kernel_size[0] * uf, kernel_size[1] * uf, Ly * Lx)) - - elif df > 1 and uf == 1: - fold_params = dict(kernel_size=kernel_size, dilation=1, padding=0, stride=stride) - unfold = torch.nn.Unfold(**fold_params) - - fold_params2 = dict(kernel_size=(kernel_size[0] // df, kernel_size[0] // df), - dilation=1, padding=0, - stride=(stride[0] // df, stride[1] // df)) - fold = torch.nn.Fold(output_size=(x.shape[2] // df, x.shape[3] // df), **fold_params2) - - weighting = self.get_weighting(kernel_size[0] // df, kernel_size[1] // df, Ly, Lx, x.device).to(x.dtype) - normalization = fold(weighting).view(1, 1, h // df, w // df) # normalizes the overlap - weighting = weighting.view((1, 1, kernel_size[0] // df, kernel_size[1] // df, Ly * Lx)) - - else: - raise NotImplementedError - - return fold, unfold, normalization, weighting - - @torch.no_grad() - def get_input(self, batch, k, return_first_stage_outputs=False, force_c_encode=False, - cond_key=None, return_original_cond=False, bs=None): - x = super().get_input(batch, k) - if bs is not None: - x = x[:bs] - x = x.to(self.device) - encoder_posterior = self.encode_first_stage(x) - z = self.get_first_stage_encoding(encoder_posterior).detach() - - if self.model.conditioning_key is not None: - if cond_key is None: - cond_key = self.cond_stage_key - if cond_key != self.first_stage_key: - if cond_key in ['caption', 'coordinates_bbox']: - xc = batch[cond_key] - elif cond_key == 'class_label': - xc = batch - else: - xc = super().get_input(batch, cond_key).to(self.device) - else: - xc = x - if not self.cond_stage_trainable or force_c_encode: - if isinstance(xc, dict) or isinstance(xc, list): - # import pudb; pudb.set_trace() - c = self.get_learned_conditioning(xc) - else: - c = self.get_learned_conditioning(xc.to(self.device)) - else: - c = xc - if bs is not None: - c = c[:bs] - - if self.use_positional_encodings: - pos_x, pos_y = self.compute_latent_shifts(batch) - ckey = __conditioning_keys__[self.model.conditioning_key] - c = {ckey: c, 'pos_x': pos_x, 'pos_y': pos_y} - - else: - c = None - xc = None - if self.use_positional_encodings: - pos_x, pos_y = self.compute_latent_shifts(batch) - c = {'pos_x': pos_x, 'pos_y': pos_y} - out = [z, c] - if return_first_stage_outputs: - xrec = self.decode_first_stage(z) - out.extend([x, xrec]) - if return_original_cond: - out.append(xc) - return out - - @torch.no_grad() - def decode_first_stage(self, z, predict_cids=False, force_not_quantize=False): - if predict_cids: - if z.dim() == 4: - z = torch.argmax(z.exp(), dim=1).long() - z = self.first_stage_model.quantize.get_codebook_entry(z, shape=None) - z = rearrange(z, 'b h w c -> b c h w').contiguous() - - z = 1. / self.scale_factor * z - - if hasattr(self, "split_input_params"): - if self.split_input_params["patch_distributed_vq"]: - ks = self.split_input_params["ks"] # eg. (128, 128) - stride = self.split_input_params["stride"] # eg. (64, 64) - uf = self.split_input_params["vqf"] - bs, nc, h, w = z.shape - if ks[0] > h or ks[1] > w: - ks = (min(ks[0], h), min(ks[1], w)) - print("reducing Kernel") - - if stride[0] > h or stride[1] > w: - stride = (min(stride[0], h), min(stride[1], w)) - print("reducing stride") - - fold, unfold, normalization, weighting = self.get_fold_unfold(z, ks, stride, uf=uf) - - z = unfold(z) # (bn, nc * prod(**ks), L) - # 1. Reshape to img shape - z = z.view((z.shape[0], -1, ks[0], ks[1], z.shape[-1])) # (bn, nc, ks[0], ks[1], L ) - - # 2. apply model loop over last dim - if isinstance(self.first_stage_model, VQModelInterface): - output_list = [self.first_stage_model.decode(z[:, :, :, :, i], - force_not_quantize=predict_cids or force_not_quantize) - for i in range(z.shape[-1])] - else: - - output_list = [self.first_stage_model.decode(z[:, :, :, :, i]) - for i in range(z.shape[-1])] - - o = torch.stack(output_list, axis=-1) # # (bn, nc, ks[0], ks[1], L) - o = o * weighting - # Reverse 1. reshape to img shape - o = o.view((o.shape[0], -1, o.shape[-1])) # (bn, nc * ks[0] * ks[1], L) - # stitch crops together - decoded = fold(o) - decoded = decoded / normalization # norm is shape (1, 1, h, w) - return decoded - else: - if isinstance(self.first_stage_model, VQModelInterface): - return self.first_stage_model.decode(z, force_not_quantize=predict_cids or force_not_quantize) - else: - return self.first_stage_model.decode(z) - - else: - if isinstance(self.first_stage_model, VQModelInterface): - return self.first_stage_model.decode(z, force_not_quantize=predict_cids or force_not_quantize) - else: - return self.first_stage_model.decode(z) - - # same as above but without decorator - def differentiable_decode_first_stage(self, z, predict_cids=False, force_not_quantize=False): - if predict_cids: - if z.dim() == 4: - z = torch.argmax(z.exp(), dim=1).long() - z = self.first_stage_model.quantize.get_codebook_entry(z, shape=None) - z = rearrange(z, 'b h w c -> b c h w').contiguous() - - z = 1. / self.scale_factor * z - - if hasattr(self, "split_input_params"): - if self.split_input_params["patch_distributed_vq"]: - ks = self.split_input_params["ks"] # eg. (128, 128) - stride = self.split_input_params["stride"] # eg. (64, 64) - uf = self.split_input_params["vqf"] - bs, nc, h, w = z.shape - if ks[0] > h or ks[1] > w: - ks = (min(ks[0], h), min(ks[1], w)) - print("reducing Kernel") - - if stride[0] > h or stride[1] > w: - stride = (min(stride[0], h), min(stride[1], w)) - print("reducing stride") - - fold, unfold, normalization, weighting = self.get_fold_unfold(z, ks, stride, uf=uf) - - z = unfold(z) # (bn, nc * prod(**ks), L) - # 1. Reshape to img shape - z = z.view((z.shape[0], -1, ks[0], ks[1], z.shape[-1])) # (bn, nc, ks[0], ks[1], L ) - - # 2. apply model loop over last dim - if isinstance(self.first_stage_model, VQModelInterface): - output_list = [self.first_stage_model.decode(z[:, :, :, :, i], - force_not_quantize=predict_cids or force_not_quantize) - for i in range(z.shape[-1])] - else: - - output_list = [self.first_stage_model.decode(z[:, :, :, :, i]) - for i in range(z.shape[-1])] - - o = torch.stack(output_list, axis=-1) # # (bn, nc, ks[0], ks[1], L) - o = o * weighting - # Reverse 1. reshape to img shape - o = o.view((o.shape[0], -1, o.shape[-1])) # (bn, nc * ks[0] * ks[1], L) - # stitch crops together - decoded = fold(o) - decoded = decoded / normalization # norm is shape (1, 1, h, w) - return decoded - else: - if isinstance(self.first_stage_model, VQModelInterface): - return self.first_stage_model.decode(z, force_not_quantize=predict_cids or force_not_quantize) - else: - return self.first_stage_model.decode(z) - - else: - if isinstance(self.first_stage_model, VQModelInterface): - return self.first_stage_model.decode(z, force_not_quantize=predict_cids or force_not_quantize) - else: - return self.first_stage_model.decode(z) - - @torch.no_grad() - def encode_first_stage(self, x): - if hasattr(self, "split_input_params"): - if self.split_input_params["patch_distributed_vq"]: - ks = self.split_input_params["ks"] # eg. (128, 128) - stride = self.split_input_params["stride"] # eg. (64, 64) - df = self.split_input_params["vqf"] - self.split_input_params['original_image_size'] = x.shape[-2:] - bs, nc, h, w = x.shape - if ks[0] > h or ks[1] > w: - ks = (min(ks[0], h), min(ks[1], w)) - print("reducing Kernel") - - if stride[0] > h or stride[1] > w: - stride = (min(stride[0], h), min(stride[1], w)) - print("reducing stride") - - fold, unfold, normalization, weighting = self.get_fold_unfold(x, ks, stride, df=df) - z = unfold(x) # (bn, nc * prod(**ks), L) - # Reshape to img shape - z = z.view((z.shape[0], -1, ks[0], ks[1], z.shape[-1])) # (bn, nc, ks[0], ks[1], L ) - - output_list = [self.first_stage_model.encode(z[:, :, :, :, i]) - for i in range(z.shape[-1])] - - o = torch.stack(output_list, axis=-1) - o = o * weighting - - # Reverse reshape to img shape - o = o.view((o.shape[0], -1, o.shape[-1])) # (bn, nc * ks[0] * ks[1], L) - # stitch crops together - decoded = fold(o) - decoded = decoded / normalization - return decoded - - else: - return self.first_stage_model.encode(x) - else: - return self.first_stage_model.encode(x) - - def shared_step(self, batch, **kwargs): - x, c = self.get_input(batch, self.first_stage_key) - loss = self(x, c) - return loss - - def forward(self, x, c, *args, **kwargs): - t = torch.randint(0, self.num_timesteps, (x.shape[0],), device=self.device).long() - if self.model.conditioning_key is not None: - assert c is not None - if self.cond_stage_trainable: - c = self.get_learned_conditioning(c) - if self.shorten_cond_schedule: - tc = self.cond_ids[t].to(self.device) - c = self.q_sample(x_start=c, t=tc, noise=torch.randn_like(c.float())) - return self.p_losses(x, c, t, *args, **kwargs) - - def apply_model(self, x_noisy, t, cond, return_ids=False): - - if isinstance(cond, dict): - # hybrid case, cond is exptected to be a dict - pass - else: - if not isinstance(cond, list): - cond = [cond] - key = 'c_concat' if self.model.conditioning_key == 'concat' else 'c_crossattn' - cond = {key: cond} - - if hasattr(self, "split_input_params"): - assert len(cond) == 1 # todo can only deal with one conditioning atm - assert not return_ids - ks = self.split_input_params["ks"] # eg. (128, 128) - stride = self.split_input_params["stride"] # eg. (64, 64) - - h, w = x_noisy.shape[-2:] - - fold, unfold, normalization, weighting = self.get_fold_unfold(x_noisy, ks, stride) - - z = unfold(x_noisy) # (bn, nc * prod(**ks), L) - # Reshape to img shape - z = z.view((z.shape[0], -1, ks[0], ks[1], z.shape[-1])) # (bn, nc, ks[0], ks[1], L ) - z_list = [z[:, :, :, :, i] for i in range(z.shape[-1])] - - if self.cond_stage_key in ["image", "LR_image", "segmentation", - 'bbox_img'] and self.model.conditioning_key: # todo check for completeness - c_key = next(iter(cond.keys())) # get key - c = next(iter(cond.values())) # get value - assert (len(c) == 1) # todo extend to list with more than one elem - c = c[0] # get element - - c = unfold(c) - c = c.view((c.shape[0], -1, ks[0], ks[1], c.shape[-1])) # (bn, nc, ks[0], ks[1], L ) - - cond_list = [{c_key: [c[:, :, :, :, i]]} for i in range(c.shape[-1])] - - elif self.cond_stage_key == 'coordinates_bbox': - assert 'original_image_size' in self.split_input_params, 'BoudingBoxRescaling is missing original_image_size' - - # assuming padding of unfold is always 0 and its dilation is always 1 - n_patches_per_row = int((w - ks[0]) / stride[0] + 1) - full_img_h, full_img_w = self.split_input_params['original_image_size'] - # as we are operating on latents, we need the factor from the original image size to the - # spatial latent size to properly rescale the crops for regenerating the bbox annotations - num_downs = self.first_stage_model.encoder.num_resolutions - 1 - rescale_latent = 2 ** (num_downs) - - # get top left postions of patches as conforming for the bbbox tokenizer, therefore we - # need to rescale the tl patch coordinates to be in between (0,1) - tl_patch_coordinates = [(rescale_latent * stride[0] * (patch_nr % n_patches_per_row) / full_img_w, - rescale_latent * stride[1] * (patch_nr // n_patches_per_row) / full_img_h) - for patch_nr in range(z.shape[-1])] - - # patch_limits are tl_coord, width and height coordinates as (x_tl, y_tl, h, w) - patch_limits = [(x_tl, y_tl, - rescale_latent * ks[0] / full_img_w, - rescale_latent * ks[1] / full_img_h) for x_tl, y_tl in tl_patch_coordinates] - # patch_values = [(np.arange(x_tl,min(x_tl+ks, 1.)),np.arange(y_tl,min(y_tl+ks, 1.))) for x_tl, y_tl in tl_patch_coordinates] - - # tokenize crop coordinates for the bounding boxes of the respective patches - patch_limits_tknzd = [torch.LongTensor(self.bbox_tokenizer._crop_encoder(bbox))[None].to(self.device) - for bbox in patch_limits] # list of length l with tensors of shape (1, 2) - print(patch_limits_tknzd[0].shape) - # cut tknzd crop position from conditioning - assert isinstance(cond, dict), 'cond must be dict to be fed into model' - cut_cond = cond['c_crossattn'][0][..., :-2].to(self.device) - print(cut_cond.shape) - - adapted_cond = torch.stack([torch.cat([cut_cond, p], dim=1) for p in patch_limits_tknzd]) - adapted_cond = rearrange(adapted_cond, 'l b n -> (l b) n') - print(adapted_cond.shape) - adapted_cond = self.get_learned_conditioning(adapted_cond) - print(adapted_cond.shape) - adapted_cond = rearrange(adapted_cond, '(l b) n d -> l b n d', l=z.shape[-1]) - print(adapted_cond.shape) - - cond_list = [{'c_crossattn': [e]} for e in adapted_cond] - - else: - cond_list = [cond for i in range(z.shape[-1])] - - # apply model by loop over crops - output_list = [self.model(z_list[i], t, **cond_list[i]) for i in range(z.shape[-1])] - assert not isinstance(output_list[0], - tuple) # todo cant deal with multiple model outputs check this never happens - - o = torch.stack(output_list, axis=-1) - o = o * weighting - # Reverse reshape to img shape - o = o.view((o.shape[0], -1, o.shape[-1])) # (bn, nc * ks[0] * ks[1], L) - # stitch crops together - x_recon = fold(o) / normalization - - else: - x_recon = self.model(x_noisy, t, **cond) - - if isinstance(x_recon, tuple) and not return_ids: - return x_recon[0] - else: - return x_recon - - def _predict_eps_from_xstart(self, x_t, t, pred_xstart): - return (extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t - pred_xstart) / \ - extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) - - def _prior_bpd(self, x_start): - """ - Get the prior KL term for the variational lower-bound, measured in - bits-per-dim. - This term can't be optimized, as it only depends on the encoder. - :param x_start: the [N x C x ...] tensor of inputs. - :return: a batch of [N] KL values (in bits), one per batch element. - """ - batch_size = x_start.shape[0] - t = torch.tensor([self.num_timesteps - 1] * batch_size, device=x_start.device) - qt_mean, _, qt_log_variance = self.q_mean_variance(x_start, t) - kl_prior = normal_kl(mean1=qt_mean, logvar1=qt_log_variance, mean2=0.0, logvar2=0.0) - return mean_flat(kl_prior) / np.log(2.0) - - def p_losses(self, x_start, cond, t, noise=None): - noise = default(noise, lambda: torch.randn_like(x_start)) - x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise) - model_output = self.apply_model(x_noisy, t, cond) - - loss_dict = {} - prefix = 'train' if self.training else 'val' - - if self.parameterization == "x0": - target = x_start - elif self.parameterization == "eps": - target = noise - else: - raise NotImplementedError - - loss_simple = self.get_loss(model_output, target, mean=False).mean([1, 2, 3]) - loss_dict.update({f'{prefix}/loss_simple': loss_simple.mean()}) - - logvar_t = self.logvar[t].to(self.device) - loss = loss_simple / torch.exp(logvar_t) + logvar_t - # loss = loss_simple / torch.exp(self.logvar) + self.logvar - if self.learn_logvar: - loss_dict.update({f'{prefix}/loss_gamma': loss.mean()}) - loss_dict.update({'logvar': self.logvar.data.mean()}) - - loss = self.l_simple_weight * loss.mean() - - loss_vlb = self.get_loss(model_output, target, mean=False).mean(dim=(1, 2, 3)) - loss_vlb = (self.lvlb_weights[t] * loss_vlb).mean() - loss_dict.update({f'{prefix}/loss_vlb': loss_vlb}) - loss += (self.original_elbo_weight * loss_vlb) - loss_dict.update({f'{prefix}/loss': loss}) - - return loss, loss_dict - - def p_mean_variance(self, x, c, t, clip_denoised: bool, return_codebook_ids=False, quantize_denoised=False, - return_x0=False, score_corrector=None, corrector_kwargs=None): - t_in = t - model_out = self.apply_model(x, t_in, c, return_ids=return_codebook_ids) - - if score_corrector is not None: - assert self.parameterization == "eps" - model_out = score_corrector.modify_score(self, model_out, x, t, c, **corrector_kwargs) - - if return_codebook_ids: - model_out, logits = model_out - - if self.parameterization == "eps": - x_recon = self.predict_start_from_noise(x, t=t, noise=model_out) - elif self.parameterization == "x0": - x_recon = model_out - else: - raise NotImplementedError - - if clip_denoised: - x_recon.clamp_(-1., 1.) - if quantize_denoised: - x_recon, _, [_, _, indices] = self.first_stage_model.quantize(x_recon) - model_mean, posterior_variance, posterior_log_variance = self.q_posterior(x_start=x_recon, x_t=x, t=t) - if return_codebook_ids: - return model_mean, posterior_variance, posterior_log_variance, logits - elif return_x0: - return model_mean, posterior_variance, posterior_log_variance, x_recon - else: - return model_mean, posterior_variance, posterior_log_variance - - @torch.no_grad() - def p_sample(self, x, c, t, clip_denoised=False, repeat_noise=False, - return_codebook_ids=False, quantize_denoised=False, return_x0=False, - temperature=1., noise_dropout=0., score_corrector=None, corrector_kwargs=None): - b, *_, device = *x.shape, x.device - outputs = self.p_mean_variance(x=x, c=c, t=t, clip_denoised=clip_denoised, - return_codebook_ids=return_codebook_ids, - quantize_denoised=quantize_denoised, - return_x0=return_x0, - score_corrector=score_corrector, corrector_kwargs=corrector_kwargs) - if return_codebook_ids: - raise DeprecationWarning("Support dropped.") - model_mean, _, model_log_variance, logits = outputs - elif return_x0: - model_mean, _, model_log_variance, x0 = outputs - else: - model_mean, _, model_log_variance = outputs - - noise = noise_like(x.shape, device, repeat_noise) * temperature - if noise_dropout > 0.: - noise = torch.nn.functional.dropout(noise, p=noise_dropout) - # no noise when t == 0 - nonzero_mask = (1 - (t == 0).float()).reshape(b, *((1,) * (len(x.shape) - 1))) - - if return_codebook_ids: - return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise, logits.argmax(dim=1) - if return_x0: - return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise, x0 - else: - return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise - - @torch.no_grad() - def progressive_denoising(self, cond, shape, verbose=True, callback=None, quantize_denoised=False, - img_callback=None, mask=None, x0=None, temperature=1., noise_dropout=0., - score_corrector=None, corrector_kwargs=None, batch_size=None, x_T=None, start_T=None, - log_every_t=None): - if not log_every_t: - log_every_t = self.log_every_t - timesteps = self.num_timesteps - if batch_size is not None: - b = batch_size if batch_size is not None else shape[0] - shape = [batch_size] + list(shape) - else: - b = batch_size = shape[0] - if x_T is None: - img = torch.randn(shape, device=self.device) - else: - img = x_T - intermediates = [] - if cond is not None: - if isinstance(cond, dict): - cond = {key: cond[key][:batch_size] if not isinstance(cond[key], list) else - [x[:batch_size] for x in cond[key]] for key in cond} - else: - cond = [c[:batch_size] for c in cond] if isinstance(cond, list) else cond[:batch_size] - - if start_T is not None: - timesteps = min(timesteps, start_T) - iterator = tqdm(reversed(range(0, timesteps)), desc='Progressive Generation', - total=timesteps) if verbose else reversed( - range(0, timesteps)) - if type(temperature) == float: - temperature = [temperature] * timesteps - - for i in iterator: - ts = torch.full((b,), i, device=self.device, dtype=torch.long) - if self.shorten_cond_schedule: - assert self.model.conditioning_key != 'hybrid' - tc = self.cond_ids[ts].to(cond.device) - cond = self.q_sample(x_start=cond, t=tc, noise=torch.randn_like(cond)) - - img, x0_partial = self.p_sample(img, cond, ts, - clip_denoised=self.clip_denoised, - quantize_denoised=quantize_denoised, return_x0=True, - temperature=temperature[i], noise_dropout=noise_dropout, - score_corrector=score_corrector, corrector_kwargs=corrector_kwargs) - if mask is not None: - assert x0 is not None - img_orig = self.q_sample(x0, ts) - img = img_orig * mask + (1. - mask) * img - - if i % log_every_t == 0 or i == timesteps - 1: - intermediates.append(x0_partial) - if callback: - callback(i) - if img_callback: - img_callback(img, i) - return img, intermediates - - @torch.no_grad() - def p_sample_loop(self, cond, shape, return_intermediates=False, - x_T=None, verbose=True, callback=None, timesteps=None, quantize_denoised=False, - mask=None, x0=None, img_callback=None, start_T=None, - log_every_t=None): - - if not log_every_t: - log_every_t = self.log_every_t - device = self.betas.device - b = shape[0] - if x_T is None: - img = torch.randn(shape, device=device) - else: - img = x_T - - intermediates = [img] - if timesteps is None: - timesteps = self.num_timesteps - - if start_T is not None: - timesteps = min(timesteps, start_T) - iterator = tqdm(reversed(range(0, timesteps)), desc='Sampling t', total=timesteps) if verbose else reversed( - range(0, timesteps)) - - if mask is not None: - assert x0 is not None - assert x0.shape[2:3] == mask.shape[2:3] # spatial size has to match - - for i in iterator: - ts = torch.full((b,), i, device=device, dtype=torch.long) - if self.shorten_cond_schedule: - assert self.model.conditioning_key != 'hybrid' - tc = self.cond_ids[ts].to(cond.device) - cond = self.q_sample(x_start=cond, t=tc, noise=torch.randn_like(cond)) - - img = self.p_sample(img, cond, ts, - clip_denoised=self.clip_denoised, - quantize_denoised=quantize_denoised) - if mask is not None: - img_orig = self.q_sample(x0, ts) - img = img_orig * mask + (1. - mask) * img - - if i % log_every_t == 0 or i == timesteps - 1: - intermediates.append(img) - if callback: - callback(i) - if img_callback: - img_callback(img, i) - - if return_intermediates: - return img, intermediates - return img - - @torch.no_grad() - def sample(self, cond, batch_size=16, return_intermediates=False, x_T=None, - verbose=True, timesteps=None, quantize_denoised=False, - mask=None, x0=None, shape=None,**kwargs): - if shape is None: - shape = (batch_size, self.channels, self.image_size, self.image_size) - if cond is not None: - if isinstance(cond, dict): - cond = {key: cond[key][:batch_size] if not isinstance(cond[key], list) else - [x[:batch_size] for x in cond[key]] for key in cond} - else: - cond = [c[:batch_size] for c in cond] if isinstance(cond, list) else cond[:batch_size] - return self.p_sample_loop(cond, - shape, - return_intermediates=return_intermediates, x_T=x_T, - verbose=verbose, timesteps=timesteps, quantize_denoised=quantize_denoised, - mask=mask, x0=x0) - - @torch.no_grad() - def sample_log(self,cond,batch_size,ddim, ddim_steps,**kwargs): - - if ddim: - ddim_sampler = DDIMSampler(self) - shape = (self.channels, self.image_size, self.image_size) - samples, intermediates =ddim_sampler.sample(ddim_steps,batch_size, - shape,cond,verbose=False,**kwargs) - - else: - samples, intermediates = self.sample(cond=cond, batch_size=batch_size, - return_intermediates=True,**kwargs) - - return samples, intermediates - - - @torch.no_grad() - def log_images(self, batch, N=8, n_row=4, sample=True, ddim_steps=200, ddim_eta=1., return_keys=None, - quantize_denoised=True, inpaint=True, plot_denoise_rows=False, plot_progressive_rows=True, - plot_diffusion_rows=True, **kwargs): - - use_ddim = ddim_steps is not None - - log = {} - z, c, x, xrec, xc = self.get_input(batch, self.first_stage_key, - return_first_stage_outputs=True, - force_c_encode=True, - return_original_cond=True, - bs=N) - N = min(x.shape[0], N) - n_row = min(x.shape[0], n_row) - log["inputs"] = x - log["reconstruction"] = xrec - if self.model.conditioning_key is not None: - if hasattr(self.cond_stage_model, "decode"): - xc = self.cond_stage_model.decode(c) - log["conditioning"] = xc - elif self.cond_stage_key in ["caption"]: - xc = log_txt_as_img((x.shape[2], x.shape[3]), batch["caption"]) - log["conditioning"] = xc - elif self.cond_stage_key == 'class_label': - xc = log_txt_as_img((x.shape[2], x.shape[3]), batch["human_label"]) - log['conditioning'] = xc - elif isimage(xc): - log["conditioning"] = xc - if ismap(xc): - log["original_conditioning"] = self.to_rgb(xc) - - if plot_diffusion_rows: - # get diffusion row - diffusion_row = [] - z_start = z[:n_row] - for t in range(self.num_timesteps): - if t % self.log_every_t == 0 or t == self.num_timesteps - 1: - t = repeat(torch.tensor([t]), '1 -> b', b=n_row) - t = t.to(self.device).long() - noise = torch.randn_like(z_start) - z_noisy = self.q_sample(x_start=z_start, t=t, noise=noise) - diffusion_row.append(self.decode_first_stage(z_noisy)) - - diffusion_row = torch.stack(diffusion_row) # n_log_step, n_row, C, H, W - diffusion_grid = rearrange(diffusion_row, 'n b c h w -> b n c h w') - diffusion_grid = rearrange(diffusion_grid, 'b n c h w -> (b n) c h w') - diffusion_grid = make_grid(diffusion_grid, nrow=diffusion_row.shape[0]) - log["diffusion_row"] = diffusion_grid - - if sample: - # get denoise row - with self.ema_scope("Plotting"): - samples, z_denoise_row = self.sample_log(cond=c,batch_size=N,ddim=use_ddim, - ddim_steps=ddim_steps,eta=ddim_eta) - # samples, z_denoise_row = self.sample(cond=c, batch_size=N, return_intermediates=True) - x_samples = self.decode_first_stage(samples) - log["samples"] = x_samples - if plot_denoise_rows: - denoise_grid = self._get_denoise_row_from_list(z_denoise_row) - log["denoise_row"] = denoise_grid - - if quantize_denoised and not isinstance(self.first_stage_model, AutoencoderKL) and not isinstance( - self.first_stage_model, IdentityFirstStage): - # also display when quantizing x0 while sampling - with self.ema_scope("Plotting Quantized Denoised"): - samples, z_denoise_row = self.sample_log(cond=c,batch_size=N,ddim=use_ddim, - ddim_steps=ddim_steps,eta=ddim_eta, - quantize_denoised=True) - # samples, z_denoise_row = self.sample(cond=c, batch_size=N, return_intermediates=True, - # quantize_denoised=True) - x_samples = self.decode_first_stage(samples.to(self.device)) - log["samples_x0_quantized"] = x_samples - - if inpaint: - # make a simple center square - h, w = z.shape[2], z.shape[3] - mask = torch.ones(N, h, w).to(self.device) - # zeros will be filled in - mask[:, h // 4:3 * h // 4, w // 4:3 * w // 4] = 0. - mask = mask[:, None, ...] - with self.ema_scope("Plotting Inpaint"): - - samples, _ = self.sample_log(cond=c,batch_size=N,ddim=use_ddim, eta=ddim_eta, - ddim_steps=ddim_steps, x0=z[:N], mask=mask) - x_samples = self.decode_first_stage(samples.to(self.device)) - log["samples_inpainting"] = x_samples - log["mask"] = mask - - # outpaint - with self.ema_scope("Plotting Outpaint"): - samples, _ = self.sample_log(cond=c, batch_size=N, ddim=use_ddim,eta=ddim_eta, - ddim_steps=ddim_steps, x0=z[:N], mask=mask) - x_samples = self.decode_first_stage(samples.to(self.device)) - log["samples_outpainting"] = x_samples - - if plot_progressive_rows: - with self.ema_scope("Plotting Progressives"): - img, progressives = self.progressive_denoising(c, - shape=(self.channels, self.image_size, self.image_size), - batch_size=N) - prog_row = self._get_denoise_row_from_list(progressives, desc="Progressive Generation") - log["progressive_row"] = prog_row - - if return_keys: - if np.intersect1d(list(log.keys()), return_keys).shape[0] == 0: - return log - else: - return {key: log[key] for key in return_keys} - return log - - def configure_optimizers(self): - lr = self.learning_rate - params = list(self.model.parameters()) - if self.cond_stage_trainable: - print(f"{self.__class__.__name__}: Also optimizing conditioner params!") - params = params + list(self.cond_stage_model.parameters()) - if self.learn_logvar: - print('Diffusion model optimizing logvar') - params.append(self.logvar) - opt = torch.optim.AdamW(params, lr=lr) - if self.use_scheduler: - assert 'target' in self.scheduler_config - scheduler = instantiate_from_config(self.scheduler_config) - - print("Setting up LambdaLR scheduler...") - scheduler = [ - { - 'scheduler': LambdaLR(opt, lr_lambda=scheduler.schedule), - 'interval': 'step', - 'frequency': 1 - }] - return [opt], scheduler - return opt - - @torch.no_grad() - def to_rgb(self, x): - x = x.float() - if not hasattr(self, "colorize"): - self.colorize = torch.randn(3, x.shape[1], 1, 1).to(x) - x = nn.functional.conv2d(x, weight=self.colorize) - x = 2. * (x - x.min()) / (x.max() - x.min()) - 1. - return x - - -class DiffusionWrapperV1(pl.LightningModule): - def __init__(self, diff_model_config, conditioning_key): - super().__init__() - self.diffusion_model = instantiate_from_config(diff_model_config) - self.conditioning_key = conditioning_key - assert self.conditioning_key in [None, 'concat', 'crossattn', 'hybrid', 'adm'] - - def forward(self, x, t, c_concat: list = None, c_crossattn: list = None): - if self.conditioning_key is None: - out = self.diffusion_model(x, t) - elif self.conditioning_key == 'concat': - xc = torch.cat([x] + c_concat, dim=1) - out = self.diffusion_model(xc, t) - elif self.conditioning_key == 'crossattn': - cc = torch.cat(c_crossattn, 1) - out = self.diffusion_model(x, t, context=cc) - elif self.conditioning_key == 'hybrid': - xc = torch.cat([x] + c_concat, dim=1) - cc = torch.cat(c_crossattn, 1) - out = self.diffusion_model(xc, t, context=cc) - elif self.conditioning_key == 'adm': - cc = c_crossattn[0] - out = self.diffusion_model(x, t, y=cc) - else: - raise NotImplementedError - - return out - - -class Layout2ImgDiffusionV1(LatentDiffusionV1): - def __init__(self, cond_stage_key, *args, **kwargs): - assert cond_stage_key == 'coordinates_bbox', 'Layout2ImgDiffusion only for cond_stage_key="coordinates_bbox"' - super().__init__(*args, cond_stage_key=cond_stage_key, **kwargs) - - def log_images(self, batch, N=8, *args, **kwargs): - logs = super().log_images(*args, batch=batch, N=N, **kwargs) - - key = 'train' if self.training else 'validation' - dset = self.trainer.datamodule.datasets[key] - mapper = dset.conditional_builders[self.cond_stage_key] - - bbox_imgs = [] - map_fn = lambda catno: dset.get_textual_label(dset.get_category_id(catno)) - for tknzd_bbox in batch[self.cond_stage_key][:N]: - bboximg = mapper.plot(tknzd_bbox.detach().cpu(), map_fn, (256, 256)) - bbox_imgs.append(bboximg) - - cond_img = torch.stack(bbox_imgs, dim=0) - logs['bbox_image'] = cond_img - return logs - -ldm.models.diffusion.ddpm.DDPMV1 = DDPMV1 -ldm.models.diffusion.ddpm.LatentDiffusionV1 = LatentDiffusionV1 -ldm.models.diffusion.ddpm.DiffusionWrapperV1 = DiffusionWrapperV1 -ldm.models.diffusion.ddpm.Layout2ImgDiffusionV1 = Layout2ImgDiffusionV1 diff --git a/modules/loader.py b/modules/loader.py index d3a33c52e..e1c4cb5a2 100644 --- a/modules/loader.py +++ b/modules/loader.py @@ -16,6 +16,7 @@ logging.getLogger("DeepSpeed").disabled = True np = None try: + os.environ.setdefault('NEP50_DISABLE_WARNING', '1') import numpy as np # pylint: disable=W0611,C0411 import numpy.random # pylint: disable=W0611,C0411 # this causes failure if numpy version changed def obj2sctype(obj): @@ -24,16 +25,22 @@ try: np.obj2sctype = obj2sctype # noqa: NPY201 np.bool8 = np.bool np.float_ = np.float64 # noqa: NPY201 + def dummy_npwarn_decorator_factory(): + def npwarn_decorator(x): + return x + return npwarn_decorator + np._no_nep50_warning = getattr(np, '_no_nep50_warning', dummy_npwarn_decorator_factory) # pylint: disable=protected-access except Exception as e: errors.log.error(f'Loader: numpy=={np.__version__ if np is not None else None} {e}') errors.log.error('Please restart the app to fix this issue') sys.exit(1) timer.startup.record("numpy") +scipy = None try: import scipy # pylint: disable=W0611,C0411 except Exception as e: - errors.log.error(f'Loader: scipy=={np.__version__ if np is not None else None} {e}') + errors.log.error(f'Loader: scipy=={scipy.__version__ if scipy is not None else None} {e}') errors.log.error('Please restart the app to fix this issue') sys.exit(1) timer.startup.record("scipy") @@ -50,8 +57,15 @@ except Exception: urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning) warnings.filterwarnings(action="ignore", category=UserWarning, module="torchvision") -import torchvision # pylint: disable=W0611,C0411 -import pytorch_lightning # pytorch_lightning should be imported after torch, but it re-enables warnings on import so import once to disable them # pylint: disable=W0611,C0411 +torchvision = None +try: + import torchvision # pylint: disable=W0611,C0411 + import pytorch_lightning # pytorch_lightning should be imported after torch, but it re-enables warnings on import so import once to disable them # pylint: disable=W0611,C0411 +except Exception as e: + errors.log.error(f'Loader: torchvision=={torchvision.__version__ if "torchvision" in sys.modules else None} {e}') + if '_no_nep' in str(e): + errors.log.error('Loaded versions of packaged are not compatible') + errors.log.error('Please restart the app to fix this issue') logging.getLogger("xformers").addFilter(lambda record: 'A matching Triton is not available' not in record.getMessage()) logging.getLogger("pytorch_lightning").disabled = True warnings.filterwarnings(action="ignore", category=DeprecationWarning) @@ -87,10 +101,13 @@ timer.startup.record("transformers") import accelerate # pylint: disable=W0611,C0411 timer.startup.record("accelerate") -import onnxruntime # pylint: disable=W0611,C0411 -onnxruntime.set_default_logger_severity(4) -onnxruntime.set_default_logger_verbosity(1) -onnxruntime.disable_telemetry_events() +try: + import onnxruntime # pylint: disable=W0611,C0411 + onnxruntime.set_default_logger_severity(4) + onnxruntime.set_default_logger_verbosity(1) + onnxruntime.disable_telemetry_events() +except Exception as e: + errors.log.warning(f'Torch onnxruntime: {e}') timer.startup.record("onnx") from fastapi import FastAPI # pylint: disable=W0611,C0411 @@ -183,5 +200,22 @@ diffusers.utils.deprecation_utils.deprecate = deprecate_warn diffusers.utils.deprecate = deprecate_warn +def patch_torch_version(): + if not hasattr(torch, '__version_backup__'): + torch.__version_backup__ = torch.__version__ + # Convert string version to tuple format to solve TypeError caused by BnB + version_parts = torch.__version__.split('+')[0].split('.') + torch.__version_tuple__ = tuple(int(x) for x in version_parts[:3]) + # Support both string and tuple for version check + class VersionString(str): + def __ge__(self, other): + if isinstance(other, tuple): + self_tuple = tuple(int(x) for x in self.split('+')[0].split('.')[:len(other)]) + return self_tuple >= other + return super().__ge__(other) + torch.__version__ = VersionString(torch.__version__) + + +patch_torch_version() errors.log.info(f'Torch: torch=={torch.__version__} torchvision=={torchvision.__version__}') errors.log.info(f'Packages: diffusers=={diffusers.__version__} transformers=={transformers.__version__} accelerate=={accelerate.__version__} gradio=={gradio.__version__} pydantic=={pydantic.__version__} numpy=={np.__version__}') diff --git a/modules/lora/lora_apply.py b/modules/lora/lora_apply.py index 2fcea174c..0343d7b14 100644 --- a/modules/lora/lora_apply.py +++ b/modules/lora/lora_apply.py @@ -158,6 +158,7 @@ def network_add_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn.G use_quantized_matmul=shared.opts.sdnq_use_quantized_matmul, use_quantized_matmul_conv=shared.opts.sdnq_use_quantized_matmul_conv, dequantize_fp32=shared.opts.sdnq_dequantize_fp32, + non_blocking=False, quantization_device=devices.device, return_device=device, param_name=getattr(self, 'network_layer_name', None), diff --git a/modules/lora/lora_overrides.py b/modules/lora/lora_overrides.py index f18dfd9b1..9b51b62e2 100644 --- a/modules/lora/lora_overrides.py +++ b/modules/lora/lora_overrides.py @@ -33,6 +33,10 @@ force_models_diffusers = [ # forced always 'hunyuandit', 'auraflow', 'lumina2', + 'qwen', + 'bria', + 'flite', + 'cosmos', # video models 'hunyuanvideo', 'cogvideo', diff --git a/modules/memstats.py b/modules/memstats.py index 47469fd77..fdc94df37 100644 --- a/modules/memstats.py +++ b/modules/memstats.py @@ -7,7 +7,10 @@ from modules import shared, errors fail_once = False +ram = {} +gpu = {} mem = {} +process = None docker_limit = None runpod_limit = None @@ -40,40 +43,62 @@ def get_runpod_limit(): return runpod_limit -def memory_stats(): - global fail_once # pylint: disable=global-statement - mem.clear() +def ram_stats(): + global process, fail_once # pylint: disable=global-statement try: - process = psutil.Process(os.getpid()) + if process is None: + process = psutil.Process(os.getpid()) res = process.memory_info() - ram_total = 100 * res.rss / process.memory_percent() - ram_total = min(ram_total, get_docker_limit(), get_runpod_limit()) - ram = { 'used': gb(res.rss), 'total': gb(ram_total) } - mem.update({ 'ram': ram }) + if 'total' not in ram: + process = psutil.Process(os.getpid()) + ram_total = 100 * res.rss / process.memory_percent() + ram_total = min(ram_total, get_docker_limit(), get_runpod_limit()) + ram['total'] = gb(ram_total) + ram['used'] = gb(res.rss) + ram['free'] = round(ram['total'] - ram['used']) except Exception as e: + ram['total'] = 0 + ram['used'] = 0 + ram['error'] = str(e) if not fail_once: - shared.log.error(f'Memory stats: {e}') - errors.display(e, 'Memory stats') + shared.log.error(f'RAM stats: {e}') + errors.display(e, 'RAM stats') fail_once = True - mem.update({ 'ram': { 'error': str(e) } }) + return ram + + +def gpu_stats(): + global fail_once # pylint: disable=global-statement try: free, total = torch.cuda.mem_get_info() - gpu = { 'used': gb(total - free), 'total': gb(total) } + gpu['used'] = gb(total - free) + gpu['total'] = gb(total) stats = dict(torch.cuda.memory_stats()) if stats.get('num_ooms', 0) > 0: shared.state.oom = True - mem.update({ - 'gpu': gpu, - '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), - 'job': shared.state.job, - }) - mem['swap'] = round(mem['active'] - mem['gpu']['used'], 2) if mem['active'] > mem['gpu']['used'] else 0 - return mem + gpu['active'] = gb(stats.get('active_bytes.all.current', 0)) + gpu['peak'] = gb(stats.get('active_bytes.all.peak', 0)) + gpu['retries'] = stats.get('num_alloc_retries', 0) + gpu['oom'] = stats.get('num_ooms', 0) + except Exception as e: + gpu['total'] = 0 + gpu['used'] = 0 + gpu['error'] = str(e) + if not fail_once: + shared.log.error(f'GPU stats: {e}') + # errors.display(e, 'GPU stats') + fail_once = True + return gpu + + +def memory_stats(): + mem['ram'] = ram_stats() + mem['gpu'] = gpu_stats() + mem['job'] = shared.state.job + try: + mem['gpu']['swap'] = round(mem['gpu']['active'] - mem['gpu']['used']) if mem['gpu']['active'] > mem['gpu']['used'] else 0 except Exception: - pass + mem['gpu']['swap'] = 0 return mem @@ -84,18 +109,6 @@ def reset_stats(): pass -def ram_stats(): - try: - process = psutil.Process(os.getpid()) - res = process.memory_info() - ram_total = 100 * res.rss / process.memory_percent() - ram_total = min(ram_total, get_docker_limit(), get_runpod_limit()) - ram = { 'used': gb(res.rss), 'total': gb(ram_total) } - return ram - except Exception: - return { 'used': 0, 'total': 0 } - - class Object: pattern = r"'(.*?)'" diff --git a/modules/mit_nunchaku.py b/modules/mit_nunchaku.py index 798de661e..b929a1d8b 100644 --- a/modules/mit_nunchaku.py +++ b/modules/mit_nunchaku.py @@ -51,8 +51,9 @@ def install_nunchaku(): suffix = 'x86_64' if arch == 'linux' else 'win_amd64' url = os.environ.get('NUNCHAKU_COMMAND', None) if url is None: + arch = f'{arch}_' if arch == 'linux' else '' url = f'https://huggingface.co/mit-han-lab/nunchaku/resolve/main/nunchaku-{ver}' - url += f'+torch{torch_ver}-cp{python_ver}-cp{python_ver}-{arch}_{suffix}.whl' + url += f'+torch{torch_ver}-cp{python_ver}-cp{python_ver}-{arch}{suffix}.whl' cmd = f'install --upgrade {url}' # pip install https://huggingface.co/mit-han-lab/nunchaku/resolve/main/nunchaku-0.2.0+torch2.6-cp311-cp311-linux_x86_64.whl log.debug(f'Nunchaku: install="{url}"') diff --git a/modules/model_quant.py b/modules/model_quant.py index 66c538f90..686703b9b 100644 --- a/modules/model_quant.py +++ b/modules/model_quant.py @@ -101,8 +101,27 @@ def create_quanto_config(kwargs = None, allow: bool = True, module: str = 'Model return kwargs -def create_sdnq_config(kwargs = None, allow: bool = True, module: str = 'Model', weights_dtype: str = None, modules_to_not_convert: list = []): +def get_sdnq_devices(): from modules import devices, shared + if shared.opts.device_map == "gpu": + quantization_device = devices.device + return_device = devices.device + elif shared.opts.device_map == "cpu": + quantization_device = devices.cpu + return_device = devices.cpu + elif shared.opts.diffusers_offload_mode in {"none", "model"}: + quantization_device = devices.device if shared.opts.sdnq_quantize_with_gpu else devices.cpu + return_device = devices.device + elif shared.opts.sdnq_quantize_with_gpu: + quantization_device = devices.device + return_device = devices.device if shared.opts.diffusers_to_gpu else devices.cpu + else: + quantization_device = None + return_device = None + return quantization_device, return_device + +def create_sdnq_config(kwargs = None, allow: bool = True, module: str = 'Model', weights_dtype: str = None, modules_to_not_convert: list = [], modules_dtype_dict: dict = {}): + from modules import shared if allow and (shared.opts.sdnq_quantize_mode in {'pre', 'auto'}) and (module == 'any' or module in shared.opts.sdnq_quantize_weights): from modules.sdnq import SDNQQuantizer, SDNQConfig diffusers.quantizers.auto.AUTO_QUANTIZER_MAPPING["sdnq"] = SDNQQuantizer @@ -118,18 +137,7 @@ def create_sdnq_config(kwargs = None, allow: bool = True, module: str = 'Model', if weights_dtype is None or weights_dtype == 'none': return kwargs - if shared.opts.device_map == "gpu": - quantization_device = devices.device - return_device = devices.device - elif shared.opts.diffusers_offload_mode in {"none", "model"}: - quantization_device = devices.device if shared.opts.sdnq_quantize_with_gpu else devices.cpu - return_device = devices.device - elif shared.opts.sdnq_quantize_with_gpu: - quantization_device = devices.device - return_device = devices.cpu - else: - quantization_device = None - return_device = None + quantization_device, return_device = get_sdnq_devices() sdnq_config = SDNQConfig( weights_dtype=weights_dtype, @@ -138,11 +146,13 @@ def create_sdnq_config(kwargs = None, allow: bool = True, module: str = 'Model', use_quantized_matmul=shared.opts.sdnq_use_quantized_matmul, use_quantized_matmul_conv=shared.opts.sdnq_use_quantized_matmul_conv, dequantize_fp32=shared.opts.sdnq_dequantize_fp32, + non_blocking=shared.opts.diffusers_offload_nonblocking, quantization_device=quantization_device, return_device=return_device, modules_to_not_convert=modules_to_not_convert, + modules_dtype_dict=modules_dtype_dict, ) - log.debug(f'Quantization: module="{module}" type=sdnq dtype={weights_dtype} matmul={shared.opts.sdnq_use_quantized_matmul} group_size={shared.opts.sdnq_quantize_weights_group_size} quant_conv={shared.opts.sdnq_quantize_conv_layers} matmul_conv={shared.opts.sdnq_use_quantized_matmul_conv} dequantize_fp32={shared.opts.sdnq_dequantize_fp32} quantize_with_gpu={shared.opts.sdnq_quantize_with_gpu} quantization_device={quantization_device} return_device={return_device}') + log.debug(f'Quantization: module="{module}" type=sdnq dtype={weights_dtype} matmul={shared.opts.sdnq_use_quantized_matmul} group_size={shared.opts.sdnq_quantize_weights_group_size} quant_conv={shared.opts.sdnq_quantize_conv_layers} matmul_conv={shared.opts.sdnq_use_quantized_matmul_conv} dequantize_fp32={shared.opts.sdnq_dequantize_fp32} quantize_with_gpu={shared.opts.sdnq_quantize_with_gpu} quantization_device={quantization_device} return_device={return_device} device_map={shared.opts.device_map} offload_mode={shared.opts.diffusers_offload_mode} non_blocking={shared.opts.diffusers_offload_nonblocking}') if kwargs is None: return sdnq_config else: @@ -169,10 +179,10 @@ def check_nunchaku(module: str = ''): return True -def create_config(kwargs = None, allow: bool = True, module: str = 'Model', modules_to_not_convert = []): +def create_config(kwargs = None, allow: bool = True, module: str = 'Model', modules_to_not_convert = [], modules_dtype_dict = {}): if kwargs is None: kwargs = {} - kwargs = create_sdnq_config(kwargs, allow=allow, module=module, modules_to_not_convert=modules_to_not_convert) + kwargs = create_sdnq_config(kwargs, allow=allow, module=module, modules_to_not_convert=modules_to_not_convert, modules_dtype_dict=modules_dtype_dict) if kwargs is not None and 'quantization_config' in kwargs: if debug: log.trace(f'Quantization: type=sdnq config={kwargs.get("quantization_config", None)}') @@ -370,9 +380,9 @@ def apply_layerwise(sd_model, quiet:bool=False): log.error(f'Quantization: type=layerwise {e}') -def sdnq_quantize_model(model, op=None, sd_model=None, do_gc: bool = True, weights_dtype: str = None, modules_to_not_convert: list = []): +def sdnq_quantize_model(model, op=None, sd_model=None, do_gc: bool = True, weights_dtype: str = None, modules_to_not_convert: list = [], modules_dtype_dict: dict = {}): global quant_last_model_name, quant_last_model_device # pylint: disable=global-statement - from modules import devices, shared + from modules import devices, shared, timer from modules.sdnq import apply_sdnq_to_module if weights_dtype is None: @@ -386,15 +396,7 @@ def sdnq_quantize_model(model, op=None, sd_model=None, do_gc: bool = True, weigh if debug: log.trace(f'Quantization: type=SDNQ op={op} cls={model.__class__} dtype={weights_dtype} mode{shared.opts.diffusers_offload_mode}') - if shared.opts.diffusers_offload_mode in {"none", "model"}: - quantization_device = devices.device if shared.opts.sdnq_quantize_with_gpu else devices.cpu - return_device = devices.device - elif shared.opts.sdnq_quantize_with_gpu: - quantization_device = devices.device - return_device = getattr(model, "device", devices.cpu) - else: - quantization_device = None - return_device = None + quantization_device, return_device = get_sdnq_devices() if getattr(model, "_keep_in_fp32_modules", None) is not None: modules_to_not_convert.extend(model._keep_in_fp32_modules) # pylint: disable=protected-access @@ -402,12 +404,18 @@ def sdnq_quantize_model(model, op=None, sd_model=None, do_gc: bool = True, weigh modules_to_not_convert.extend(model._skip_layerwise_casting_patterns) # pylint: disable=protected-access if model.__class__.__name__ == "ChromaTransformer2DModel": modules_to_not_convert.append("distilled_guidance_layer") + if model.__class__.__name__ == "QwenImageTransformer2DModel": + if "minimum_6bit" not in modules_dtype_dict.keys(): + modules_dtype_dict["minimum_6bit"] = ["img_mod", "pos_embed", "time_text_embed", "img_in", "txt_in", "norm_out"] + else: + modules_dtype_dict["minimum_6bit"].extend(["img_mod", "pos_embed", "time_text_embed", "img_in", "txt_in", "norm_out"]) model.eval() backup_embeddings = None if hasattr(model, "get_input_embeddings"): backup_embeddings = copy.deepcopy(model.get_input_embeddings()) + t0 = time.time() model = apply_sdnq_to_module( model, weights_dtype=weights_dtype, @@ -417,11 +425,15 @@ def sdnq_quantize_model(model, op=None, sd_model=None, do_gc: bool = True, weigh use_quantized_matmul=shared.opts.sdnq_use_quantized_matmul, use_quantized_matmul_conv=shared.opts.sdnq_use_quantized_matmul_conv, dequantize_fp32=shared.opts.sdnq_dequantize_fp32, + non_blocking=shared.opts.diffusers_offload_nonblocking, quantization_device=quantization_device, return_device=return_device, - param_name=op, modules_to_not_convert=modules_to_not_convert, + modules_dtype_dict=modules_dtype_dict, + op=op, ) + t1 = time.time() + timer.load.add('sdnq', t1 - t0) model.quantization_method = 'SDNQ' if hasattr(model, "set_input_embeddings") and backup_embeddings is not None: @@ -443,7 +455,7 @@ def sdnq_quantize_model(model, op=None, sd_model=None, do_gc: bool = True, weigh quant_last_model_name = None quant_last_model_device = None model.to(devices.device) - elif shared.opts.diffusers_offload_mode != "none": + elif (shared.opts.diffusers_offload_mode != "none") and (not shared.opts.diffusers_to_gpu): model = model.to(devices.cpu) if do_gc: devices.torch_gc(force=True, reason='sdnq') @@ -605,7 +617,7 @@ def torchao_quantization(sd_model): return sd_model -def get_dit_args(load_config:dict={}, module:str=None, device_map:bool=False, allow_quant:bool=True, modules_to_not_convert: list = []): +def get_dit_args(load_config:dict={}, module:str=None, device_map:bool=False, allow_quant:bool=True, modules_to_not_convert: list = [], modules_dtype_dict: dict = {}): from modules import shared, devices config = load_config.copy() if 'torch_dtype' not in config: @@ -628,7 +640,7 @@ def get_dit_args(load_config:dict={}, module:str=None, device_map:bool=False, al elif shared.opts.device_map == 'gpu': config['device_map'] = devices.device if allow_quant: - quant_args = create_config(module=module, modules_to_not_convert=modules_to_not_convert) + quant_args = create_config(module=module, modules_to_not_convert=modules_to_not_convert, modules_dtype_dict=modules_dtype_dict) else: quant_args = {} return config, quant_args @@ -637,11 +649,15 @@ def get_dit_args(load_config:dict={}, module:str=None, device_map:bool=False, al def do_post_load_quant(sd_model, allow=True): from modules import shared if shared.opts.sdnq_quantize_weights and (shared.opts.sdnq_quantize_mode == 'post' or (allow and shared.opts.sdnq_quantize_mode == 'auto')): + shared.log.debug('Load model: post_quant=sdnq') sd_model = sdnq_quantize_weights(sd_model) if len(shared.opts.optimum_quanto_weights) > 0: + shared.log.debug('Load model: post_quant=quanto') sd_model = optimum_quanto_weights(sd_model) if shared.opts.torchao_quantization and (shared.opts.torchao_quantization_mode == 'post' or (allow and shared.opts.torchao_quantization_mode == 'auto')): + shared.log.debug('Load model: post_quant=torchao') sd_model = torchao_quantization(sd_model) if shared.opts.layerwise_quantization: + shared.log.debug('Load model: post_quant=layerwise') apply_layerwise(sd_model) return sd_model diff --git a/modules/model_te.py b/modules/model_te.py index b47d73675..82a6a40a0 100644 --- a/modules/model_te.py +++ b/modules/model_te.py @@ -15,10 +15,14 @@ def load_t5(name=None, cache_dir=None): global loaded_te # pylint: disable=global-statement if name is None: return None + cache_dir = cache_dir or shared.opts.hfcache_dir from modules import modelloader modelloader.hf_login() repo_id = 'stabilityai/stable-diffusion-3-medium-diffusers' - fn = te_dict.get(name) if name in te_dict else None + if os.path.exists(name): + fn = name + else: + fn = te_dict.get(name) if name in te_dict else None if fn is not None and name.lower().endswith('gguf'): from modules import ggml @@ -46,12 +50,13 @@ def load_t5(name=None, cache_dir=None): except Exception: shared.log.error(f"T5: Failed to cast text encoder to {devices.dtype}, set dtype to {t5.dtype}") raise + del state_dict elif fn is not None: with open(os.path.join('configs', 'flux', 'text_encoder_2', 'config.json'), encoding='utf8') as f: t5_config = transformers.T5Config(**json.load(f)) state_dict = load_file(fn) - t5 = transformers.T5EncoderModel.from_pretrained(None, state_dict=state_dict, config=t5_config) + t5 = transformers.T5EncoderModel.from_pretrained(None, state_dict=state_dict, config=t5_config, torch_dtype=devices.dtype) elif 'fp16' in name.lower(): t5 = transformers.T5EncoderModel.from_pretrained(repo_id, subfolder='text_encoder_3', cache_dir=cache_dir, torch_dtype=devices.dtype) @@ -141,6 +146,7 @@ def load_vit_l(): te = transformers.CLIPTextModel.from_pretrained(pretrained_model_name_or_path=None, state_dict=state_dict, config=config) te = te.to(dtype=devices.dtype) loaded_te = shared.opts.sd_text_encoder + del state_dict return te @@ -151,6 +157,7 @@ def load_vit_g(): te = transformers.CLIPTextModelWithProjection.from_pretrained(pretrained_model_name_or_path=None, state_dict=state_dict, config=config) te = te.to(dtype=devices.dtype) loaded_te = shared.opts.sd_text_encoder + del state_dict return te diff --git a/modules/modeldata.py b/modules/modeldata.py index 48d21028d..0cdbcb19d 100644 --- a/modules/modeldata.py +++ b/modules/modeldata.py @@ -58,6 +58,8 @@ def get_model_type(pipe): model_type = 'pixartalpha' elif "Bria" in name: model_type = 'bria' + elif 'Qwen' in name: + model_type = 'qwen' # video models elif "CogVideo" in name: model_type = 'cogvideo' diff --git a/modules/modelloader.py b/modules/modelloader.py index 5ad1d540b..c4a1d3e00 100644 --- a/modules/modelloader.py +++ b/modules/modelloader.py @@ -1,14 +1,11 @@ import io import os import time -import json import shutil import importlib import contextlib from typing import Dict from urllib.parse import urlparse -from PIL import Image -import rich.progress as p import huggingface_hub as hf from installer import install, log from modules import shared, errors, files_cache @@ -48,193 +45,6 @@ def hf_login(token=None): return True -def download_civit_meta(model_path: str, model_id): - fn = os.path.splitext(model_path)[0] + '.json' - url = f'https://civitai.com/api/v1/models/{model_id}' - r = shared.req(url) - if r.status_code == 200: - try: - shared.writefile(r.json(), filename=fn, mode='w', silent=True) - msg = f'CivitAI download: id={model_id} url={url} file="{fn}"' - shared.log.info(msg) - return msg - except Exception as e: - msg = f'CivitAI download error: id={model_id} url={url} file="{fn}" {e}' - errors.display(e, 'CivitAI download error') - shared.log.error(msg) - return msg - 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}"' - r = shared.req(preview_url, stream=True) - total_size = int(r.headers.get('content-length', 0)) - block_size = 16384 # 16KB blocks - 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 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) - pbar.update(task, advance=block_size) - if written < 1024: # min threshold - os.remove(preview_file) - raise ValueError(f'removed invalid download: bytes={written}') - if is_video: - img = save_video_frame(preview_file) - else: - img = Image.open(preview_file) - except Exception as e: - # 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() - if img is None: - return res - shared.log.info(f'{res} size={total_size} image={img.size}') - img.close() - return res - - -download_pbar = None - -def download_civit_model_thread(model_name: str, model_url: str, model_path: str = "", model_type: str = "Model", token: str = None): - import hashlib - sha256 = hashlib.sha256() - sha256.update(model_url.encode('utf-8')) - temp_file = sha256.hexdigest()[:8] + '.tmp' - - headers = {} - starting_pos = 0 - if os.path.isfile(temp_file): - starting_pos = os.path.getsize(temp_file) - headers['Range'] = f'bytes={starting_pos}-' - if token is None: - token = shared.opts.civitai_token - if token is not None and len(token) > 0: - headers['Authorization'] = f'Bearer {token}' - - r = shared.req(model_url, headers=headers, stream=True) - total_size = int(r.headers.get('content-length', 0)) - if model_name is None or len(model_name) == 0: - cn = r.headers.get('content-disposition', '') - model_name = cn.split('filename=')[-1].strip('"') - - if model_type == 'LoRA': - model_file = os.path.join(shared.opts.lora_dir, model_path, model_name) - temp_file = os.path.join(shared.opts.lora_dir, model_path, temp_file) - elif model_type == 'Embedding': - model_file = os.path.join(shared.opts.embeddings_dir, model_path, model_name) - temp_file = os.path.join(shared.opts.embeddings_dir, model_path, temp_file) - elif model_type == 'VAE': - model_file = os.path.join(shared.opts.vae_dir, model_path, model_name) - temp_file = os.path.join(shared.opts.vae_dir, model_path, temp_file) - else: - model_file = os.path.join(shared.opts.ckpt_dir, model_path, model_name) - temp_file = os.path.join(shared.opts.ckpt_dir, model_path, temp_file) - - res = f'Model download: name="{model_name}" url="{model_url}" path="{model_path}" temp="{temp_file}"' - if os.path.isfile(model_file): - res += ' already exists' - shared.log.warning(res) - return res - - res += f' size={round((starting_pos + total_size)/1024/1024, 2)}Mb' - shared.log.info(res) - shared.state.begin('CivitAI') - block_size = 16384 # 16KB blocks - written = starting_pos - global download_pbar # pylint: disable=global-statement - if download_pbar is None: - download_pbar = p.Progress(p.TextColumn('[cyan]{task.description}'), p.DownloadColumn(), p.BarColumn(), p.TaskProgressColumn(), p.TimeRemainingColumn(), p.TimeElapsedColumn(), p.TransferSpeedColumn(), p.TextColumn('[cyan]{task.fields[name]}'), console=shared.console) - with download_pbar: - task = download_pbar.add_task(description="Download starting", total=starting_pos+total_size, name=model_name) - try: - with open(temp_file, 'ab') as f: - for data in r.iter_content(block_size): - if written == 0: - try: # check if response is JSON message instead of bytes - shared.log.error(f'Model download: response={json.loads(data.decode("utf-8"))}') - raise ValueError('response: type=json expected=bytes') - except Exception: # this is good - pass - written = written + len(data) - f.write(data) - download_pbar.update(task, description="Download", completed=written) - if written < 1024: # min threshold - os.remove(temp_file) - raise ValueError(f'removed invalid download: bytes={written}') - """ - if preview is not None: - preview_file = os.path.splitext(model_file)[0] + '.jpg' - preview.save(preview_file) - res += f' preview={preview_file}' - """ - except Exception as e: - shared.log.error(f'{res} {e}') - finally: - download_pbar.stop_task(task) - download_pbar.remove_task(task) - if starting_pos+total_size != written: - shared.log.warning(f'{res} written={round(written/1024/1024)}Mb incomplete download') - elif os.path.exists(temp_file): - shared.log.debug(f'Model download complete: temp="{temp_file}" path="{model_file}"') - os.rename(temp_file, model_file) - shared.state.end() - if os.path.exists(model_file): - return model_file - else: - return None - - -def download_civit_model(model_url: str, model_name: str, model_path: str, model_type: str, token: str = None): - import threading - if model_name is None or len(model_name) == 0: - err = 'Model download: no target model name provided' - shared.log.error(err) - return err - thread = threading.Thread(target=download_civit_model_thread, args=(model_name, model_url, model_path, model_type, token)) - thread.start() - return f'Model download: name={model_name} url={model_url} path={model_path}' - - def download_diffusers_model(hub_id: str, cache_dir: str = None, download_config: Dict[str, str] = None, token = None, variant = None, revision = None, mirror = None, custom_pipeline = None): if hub_id is None or len(hub_id) == 0: return None @@ -438,6 +248,7 @@ def load_civitai(model: str, url: str): return name # already downloaded else: shared.log.debug(f'Reference download start: model="{name}"') + from modules.civitai.download_civitai import download_civit_model_thread download_civit_model_thread(model_name=model, model_url=url, model_path='', model_type='safetensors', token=shared.opts.civitai_token) shared.log.debug(f'Reference download complete: model="{name}"') sd_models.list_models() diff --git a/modules/models_hf.py b/modules/models_hf.py new file mode 100644 index 000000000..23e979821 --- /dev/null +++ b/modules/models_hf.py @@ -0,0 +1,43 @@ +import os +import gradio as gr +from modules.shared import log, opts + + +def hf_init(): + os.environ.setdefault('HF_HUB_DISABLE_EXPERIMENTAL_WARNING', '1') + os.environ.setdefault('HF_HUB_DISABLE_SYMLINKS_WARNING', '1') + os.environ.setdefault('HF_HUB_DISABLE_IMPLICIT_TOKEN', '1') + os.environ.setdefault('HUGGINGFACE_HUB_VERBOSITY', 'warning') + os.environ.setdefault('HF_ENABLE_PARALLEL_LOADING', 'true' if opts.sd_parallel_load else 'false') + + +def hf_search(keyword): + hf_init() + import huggingface_hub as hf + hf_api = hf.HfApi() + models = hf_api.list_models(model_name=keyword, full=True, library="diffusers", limit=50, sort="downloads", direction=-1) + data = [] + for model in models: + tags = [t for t in model.tags if not t.startswith('diffusers') and not t.startswith('license') and not t.startswith('arxiv') and len(t) > 2] + data.append([model.id, model.pipeline_tag, tags, model.downloads, model.lastModified, f'https://huggingface.co/{model.id}']) + return data + + +def hf_select(evt: gr.SelectData, data): + return data[evt.index[0]][0] + + +def hf_download_model(hub_id: str, token, variant, revision, mirror, custom_pipeline): + hf_init() + from modules.modelloader import download_diffusers_model + download_diffusers_model(hub_id, cache_dir=opts.diffusers_dir, token=token, variant=variant, revision=revision, mirror=mirror, custom_pipeline=custom_pipeline) + from modules.sd_models import list_models # pylint: disable=W0621 + list_models() + log.info(f'Diffuser model downloaded: model="{hub_id}"') + return f'Diffuser model downloaded: model="{hub_id}"' + + +def hf_update_token(token): + log.debug('Huggingface update token') + opts.huggingface_token = token + opts.save() diff --git a/modules/modelstats.py b/modules/modelstats.py index 8ef02ca35..3657e523a 100644 --- a/modules/modelstats.py +++ b/modules/modelstats.py @@ -4,6 +4,30 @@ import torch from modules import shared, sd_models +def walk(folder: str): + files = [] + for root, _, filenames in os.walk(folder): + for filename in filenames: + files.append(os.path.join(root, filename)) + return files + + +def stat(fn: str): + if fn is None or len(fn) == 0 or not os.path.exists(fn): + return 0, datetime.fromtimestamp(0) + fs_stat = os.stat(fn, follow_symlinks=False) + mtime = datetime.fromtimestamp(fs_stat.st_mtime).replace(microsecond=0) + if os.path.islink(fn): + size = 0 + elif os.path.isfile(fn): + size = round(fs_stat.st_size) + elif os.path.isdir(fn): + size = round(sum(stat(fn)[0] for fn in walk(fn))) + else: + size = 0 + return size, mtime + + class Module(): name: str = '' cls: str = None @@ -11,6 +35,7 @@ class Module(): dtype: str = None params: int = 0 modules: int = 0 + quant: str = None config: dict = None def __init__(self, name, module): @@ -25,6 +50,7 @@ class Module(): self.dtype = getattr(module, 'dtype', None) self.params = sum(p.numel() for p in module.parameters(recurse=True)) self.modules = len(list(module.modules())) + self.quant = getattr(module, 'quantization_method', None) def __repr__(self): s = f'name="{self.name}" cls={self.cls} config={self.config is not None}' @@ -58,17 +84,15 @@ class Model(): self.name = self.info.name or self.name self.hash = self.info.shorthash or '' self.meta = self.info.metadata or {} - if os.path.exists(self.info.filename): - stat = os.stat(self.info.filename) - self.mtime = datetime.fromtimestamp(stat.st_mtime).replace(microsecond=0) - if os.path.isfile(self.info.filename): - self.size = round(stat.st_size) + self.size, self.mtime = stat(self.info.filename) def __repr__(self): return f'model="{self.name}" type={self.type} class={self.cls} size={self.size} mtime="{self.mtime}" modules={self.modules}' def analyze(): + if not shared.sd_loaded: + return None model = Model(shared.opts.sd_model_checkpoint) if model.cls == '': return model diff --git a/modules/onnx_impl/__init__.py b/modules/onnx_impl/__init__.py index a8e04b691..5a009a741 100644 --- a/modules/onnx_impl/__init__.py +++ b/modules/onnx_impl/__init__.py @@ -4,6 +4,7 @@ import torch import diffusers import onnxruntime as ort + initialized = False diff --git a/modules/options_handler.py b/modules/options_handler.py index ac974f6de..5dd7f5435 100644 --- a/modules/options_handler.py +++ b/modules/options_handler.py @@ -9,6 +9,7 @@ from installer import log if TYPE_CHECKING: from modules.options import OptionInfo + cmd_opts = cmd_args.parse_args() compatibility_opts = ['clip_skip', 'uni_pc_lower_order_final', 'uni_pc_order'] diff --git a/modules/paths.py b/modules/paths.py index 5e0515ebd..f5d0f190c 100644 --- a/modules/paths.py +++ b/modules/paths.py @@ -44,23 +44,6 @@ if os.environ.get('SD_PATH_DEBUG', None) is not None: log.debug(f'Paths: script-path="{script_path}" data-dir="{data_path}" models-dir="{models_path}" config="{config_path}"') -def register_paths(): - log.debug('Register paths') - sys.path.insert(0, script_path) - # sd_path = os.path.join(script_path, 'repositories') - path_dirs = [ - # (os.path.join(sd_path, 'codeformer'), 'inference_codeformer.py', 'CodeFormer', []), - ] - for d, must_exist, what, _options in path_dirs: - must_exist_path = os.path.abspath(os.path.join(script_path, d, must_exist)) - if not os.path.exists(must_exist_path): - log.error(f'Required path not found: path={must_exist_path} item={what}') - else: - d = os.path.abspath(d) - sys.path.append(d) - paths[what] = d - - def create_path(folder): if folder is None or folder == '': return @@ -103,6 +86,7 @@ def create_paths(opts): create_path(fix_path('temp_dir')) create_path(fix_path('ckpt_dir')) create_path(fix_path('diffusers_dir')) + create_path(fix_path('hfcache_dir')) create_path(fix_path('vae_dir')) create_path(fix_path('unet_dir')) create_path(fix_path('te_dir')) @@ -139,3 +123,14 @@ class Prioritize: def __exit__(self, exc_type, exc_val, exc_tb): sys.path = self.path self.path = None + + +def check_cache(opts): + prev_default = os.environ.get("SD_HFCACHEDIR", None) or os.path.join(os.path.expanduser('~'), '.cache', 'huggingface', 'hub') + from modules.modelstats import stat + if opts.hfcache_dir != prev_default: + size, _mtime = stat(prev_default) + if size//1024//1024 > 0: + log.warning(f'Cache location changed: previous="{prev_default}" size={size//1024//1024} MB') + size, _mtime = stat(opts.hfcache_dir) + log.debug(f'Huggingface cache: path="{opts.hfcache_dir}" size={size//1024//1024} MB') diff --git a/modules/paths_internal.py b/modules/paths_internal.py index 3a408329d..a9dabdd0f 100644 --- a/modules/paths_internal.py +++ b/modules/paths_internal.py @@ -1,30 +1,3 @@ # no longer used, all paths are defined in paths.py from modules.paths import modules_path, script_path, sd_configs_path, sd_default_config, sd_model_file, default_sd_model_file, data_path, models_path, extensions_dir, extensions_builtin_dir # pylint: disable=unused-import - -""" -import argparse -import os - -modules_path = os.path.dirname(os.path.realpath(__file__)) -script_path = os.path.dirname(modules_path) -sd_configs_path = os.path.join(script_path, "configs") -sd_default_config = os.path.join(sd_configs_path, "v1-inference.yaml") - -# Parse the --data-dir flag first so we can use it as a base for our other argument default values -parser_pre = argparse.ArgumentParser(add_help=False) -parser_pre.add_argument("--ckpt", type=str, default=os.environ.get("SD_MODEL", None), help="Path to model checkpoint to load immediately, default: %(default)s") -parser_pre.add_argument("--data-dir", type=str, default=os.environ.get("SD_DATADIR", ''), help="Base path where all user data is stored, default: %(default)s") -parser_pre.add_argument("--models-dir", type=str, default=os.environ.get("SD_MODELSDIR", 'models'), help="Base path where all models are stored, default: %(default)s",) -cmd_opts_pre = parser_pre.parse_known_args()[0] - -# parser_pre.add_argument("--config", type=str, default=os.environ.get("SD_CONFIG", os.path.join(data_path, 'config.json')), help="Use specific server configuration file, default: %(default)s") - -data_path = cmd_opts_pre.data_dir -models_path = cmd_opts_pre.models_dir if os.path.isabs(cmd_opts_pre.models_dir) else os.path.join(data_path, cmd_opts_pre.models_dir) -extensions_dir = os.path.join(data_path, "extensions") -extensions_builtin_dir = "extensions-builtin" - -sd_model_file = cmd_opts_pre.ckpt or os.path.join(script_path, 'model.ckpt') # not used -default_sd_model_file = sd_model_file # not used -""" diff --git a/modules/postprocess/ldsr_model.py b/modules/postprocess/ldsr_model.py deleted file mode 100644 index 8abb7081a..000000000 --- a/modules/postprocess/ldsr_model.py +++ /dev/null @@ -1,75 +0,0 @@ -import os -import sys -import traceback -from modules.upscaler import Upscaler, UpscalerData -from modules import shared, script_callbacks - - -class Dummy: - pass - -cls = Upscaler if not shared.native else Dummy - -class UpscalerLDSR(cls): - def __init__(self, user_path): - self.name = "LDSR" - self.user_path = user_path - self.model_url = "https://heibox.uni-heidelberg.de/f/578df07c8fc04ffbadf3/?dl=1" - self.yaml_url = "https://heibox.uni-heidelberg.de/f/31a76b13ea27482981b4/?dl=1" - super().__init__() - scaler_data = UpscalerData("LDSR", None, self) - self.scalers = [scaler_data] - - def load_model(self, path: str): - from modules.ldsr.ldsr_model_arch import LDSR - import modules.ldsr.sd_hijack_autoencoder # pylint: disable=unused-import - import modules.ldsr.sd_hijack_ddpm_v1 # pylint: disable=unused-import - # Remove incorrect project.yaml file if too big - yaml_path = os.path.join(self.model_path, "project.yaml") - old_model_path = os.path.join(self.model_path, "model.pth") - new_model_path = os.path.join(self.model_path, "model.ckpt") - - local_model_paths = self.find_models(ext_filter=[".ckpt", ".safetensors"]) - local_ckpt_path = next(iter([local_model for local_model in local_model_paths if local_model.endswith("model.ckpt")]), None) - local_safetensors_path = next(iter([local_model for local_model in local_model_paths if local_model.endswith("model.safetensors")]), None) - local_yaml_path = next(iter([local_model for local_model in local_model_paths if local_model.endswith("project.yaml")]), None) - - if os.path.exists(yaml_path): - statinfo = os.stat(yaml_path) - if statinfo.st_size >= 10485760: - print("Removing invalid LDSR YAML file.") - os.remove(yaml_path) - - if os.path.exists(old_model_path): - print("Renaming model from model.pth to model.ckpt") - os.rename(old_model_path, new_model_path) - - from modules.modelloader import load_file_from_url - if local_safetensors_path is not None and os.path.exists(local_safetensors_path): - model = local_safetensors_path - else: - model = local_ckpt_path if local_ckpt_path is not None else load_file_from_url(url=self.model_url, model_dir=self.model_download_path, file_name="model.ckpt", progress=True) - - yaml = local_yaml_path if local_yaml_path is not None else load_file_from_url(url=self.yaml_url, model_dir=self.model_download_path, file_name="project.yaml", progress=True) - - try: - return LDSR(model, yaml) - except Exception: - print("Error importing LDSR:", file=sys.stderr) - print(traceback.format_exc(), file=sys.stderr) - return None - - def do_upscale(self, img, selected_model): - ldsr = self.load_model(selected_model) - if ldsr is None: - print("NO LDSR!") - return img - ddim_steps = shared.opts.ldsr_steps - return ldsr.super_resolution(img, ddim_steps, self.scale) - - -def on_ui_settings(): - import gradio as gr - shared.opts.add_option("ldsr_steps", shared.OptionInfo(100, "LDSR processing steps", gr.Slider, {"minimum": 1, "maximum": 200, "step": 1}, section=('postprocessing', "Postprocessing"))) - -script_callbacks.on_ui_settings(on_ui_settings) diff --git a/modules/processing.py b/modules/processing.py index 268661647..2a5087d37 100644 --- a/modules/processing.py +++ b/modules/processing.py @@ -119,6 +119,9 @@ def process_images(p: StableDiffusionProcessing) -> Processed: if not hasattr(p.sd_model, 'sd_checkpoint_info'): shared.log.error('Processing: incomplete model') return None + if p.abort: + shared.log.debug('Processing: aborted') + return None if p.scripts is not None and isinstance(p.scripts, scripts_manager.ScriptRunner): p.scripts.before_process(p) stored_opts = {} @@ -158,11 +161,12 @@ def process_images(p: StableDiffusionProcessing) -> Processed: shared.prompt_styles.apply_styles_to_extra(p) shared.prompt_styles.extract_comments(p) + vae_scale_factor = sd_vae.get_vae_scale_factor() if p.width is not None: - p.width = 8 * int(p.width / 8) + p.width = vae_scale_factor * int(p.width / vae_scale_factor) if p.height is not None: - p.height = 8 * int(p.height / 8) + p.height = vae_scale_factor * int(p.height / vae_scale_factor) script_callbacks.before_process_callback(p) timer.process.record('pre') diff --git a/modules/processing_args.py b/modules/processing_args.py index 24ea8e6e5..c34cdc2b9 100644 --- a/modules/processing_args.py +++ b/modules/processing_args.py @@ -7,7 +7,7 @@ import inspect import torch import numpy as np from PIL import Image -from modules import shared, errors, sd_models, processing, processing_vae, processing_helpers, sd_hijack_hypertile, prompt_parser_diffusers, timer, extra_networks +from modules import shared, errors, sd_models, processing, processing_vae, processing_helpers, sd_hijack_hypertile, prompt_parser_diffusers, timer, extra_networks, sd_vae from modules.processing_callbacks import diffusers_callback_legacy, diffusers_callback, set_callbacks_p from modules.processing_helpers import resize_hires, fix_prompts, calculate_base_steps, calculate_hires_steps, calculate_refiner_steps, get_generator, set_latents, apply_circular # pylint: disable=unused-import from modules.api import helpers @@ -19,6 +19,7 @@ disable_pbar = os.environ.get('SD_DISABLE_PBAR', None) is not None def task_specific_kwargs(p, model): + vae_scale_factor = sd_vae.get_vae_scale_factor(model) task_args = {} is_img2img_model = bool('Zero123' in shared.sd_model.__class__.__name__) if len(getattr(p, 'init_images', [])) > 0: @@ -30,8 +31,8 @@ def task_specific_kwargs(p, model): p.ops.append('txt2img') if hasattr(p, 'width') and hasattr(p, 'height'): task_args = { - 'width': 8 * math.ceil(p.width / 8), - 'height': 8 * math.ceil(p.height / 8), + 'width': vae_scale_factor * math.ceil(p.width / vae_scale_factor), + 'height': vae_scale_factor * math.ceil(p.height / vae_scale_factor), } elif (sd_models.get_diffusers_task(model) == sd_models.DiffusersTaskType.IMAGE_2_IMAGE or is_img2img_model) and len(getattr(p, 'init_images', [])) > 0: if shared.sd_model_type == 'sdxl' and hasattr(model, 'register_to_config'): @@ -50,19 +51,18 @@ def task_specific_kwargs(p, model): } if model.__class__.__name__ == 'FluxImg2ImgPipeline' or model.__class__.__name__ == 'FluxKontextPipeline': # needs explicit width/height if torch.is_tensor(p.init_images[0]): - p.width, p.height = p.init_images[0].shape[-1] * 16, p.init_images[0].shape[-2] * 16 + p.width, p.height = p.init_images[0].shape[-1] * vae_scale_factor, p.init_images[0].shape[-2] * vae_scale_factor else: - p.width, p.height = 8 * math.ceil(p.init_images[0].width / 8), 8 * math.ceil(p.init_images[0].height / 8) + p.width, p.height = 8 * math.ceil(p.init_images[0].width / vae_scale_factor), 8 * math.ceil(p.init_images[0].height / vae_scale_factor) if model.__class__.__name__ == 'FluxKontextPipeline': aspect_ratio = p.width / p.height - vae_scale_factor = 16 max_area = max(p.width, p.height)**2 p.width, p.height = round((max_area * aspect_ratio) ** 0.5), round((max_area / aspect_ratio) ** 0.5) p.width, p.height = p.width // vae_scale_factor * vae_scale_factor, p.height // vae_scale_factor * vae_scale_factor task_args['max_area'] = max_area task_args['width'], task_args['height'] = p.width, p.height elif model.__class__.__name__ == 'OmniGenPipeline' or model.__class__.__name__ == 'OmniGen2Pipeline': - p.width, p.height = 16 * math.ceil(p.init_images[0].width / 16), 16 * math.ceil(p.init_images[0].height / 16) + p.width, p.height = vae_scale_factor * math.ceil(p.init_images[0].width / vae_scale_factor), vae_scale_factor * math.ceil(p.init_images[0].height / vae_scale_factor) task_args = { 'width': p.width, 'height': p.height, @@ -71,8 +71,8 @@ def task_specific_kwargs(p, model): elif sd_models.get_diffusers_task(model) == sd_models.DiffusersTaskType.INSTRUCT and len(getattr(p, 'init_images', [])) > 0: p.ops.append('instruct') task_args = { - 'width': 8 * math.ceil(p.width / 8) if hasattr(p, 'width') else None, - 'height': 8 * math.ceil(p.height / 8) if hasattr(p, 'height') else None, + 'width': vae_scale_factor * math.ceil(p.width / vae_scale_factor) if hasattr(p, 'width') else None, + 'height': vae_scale_factor * math.ceil(p.height / vae_scale_factor) if hasattr(p, 'height') else None, 'image': p.init_images, 'strength': p.denoising_strength, } @@ -99,6 +99,8 @@ def task_specific_kwargs(p, model): 'height': height, 'width': width, } + + # model specific args if model.__class__.__name__ == 'LatentConsistencyModelPipeline' and hasattr(p, 'init_images') and len(p.init_images) > 0: p.ops.append('lcm') init_latents = [processing_vae.vae_encode(image, model=shared.sd_model, vae_type=p.vae_type).squeeze(dim=0) for image in p.init_images] @@ -376,18 +378,21 @@ def set_pipeline_args(p, model, prompts:list, negative_prompts:list, prompts_2:t # handle missing resolution if args.get('image', None) is not None and ('width' not in args or 'height' not in args): if 'width' in possible and 'height' in possible: + vae_scale_factor = sd_vae.get_vae_scale_factor(model) if isinstance(args['image'], torch.Tensor) or isinstance(args['image'], np.ndarray): - args['width'] = 8 * args['image'].shape[-1] - args['height'] = 8 * args['image'].shape[-2] + args['width'] = vae_scale_factor * args['image'].shape[-1] + args['height'] = vae_scale_factor * args['image'].shape[-2] elif isinstance(args['image'], Image.Image): args['width'] = args['image'].width args['height'] = args['image'].height elif isinstance(args['image'][0], torch.Tensor) or isinstance(args['image'][0], np.ndarray): - args['width'] = 8 * args['image'][0].shape[-1] - args['height'] = 8 * args['image'][0].shape[-2] + args['width'] = vae_scale_factor * args['image'][0].shape[-1] + args['height'] = vae_scale_factor * args['image'][0].shape[-2] else: - args['width'] = 8 * math.ceil(args['image'][0].width / 8) - args['height'] = 8 * math.ceil(args['image'][0].height / 8) + args['width'] = vae_scale_factor * math.ceil(args['image'][0].width / vae_scale_factor) + args['height'] = vae_scale_factor * math.ceil(args['image'][0].height / vae_scale_factor) + if 'max_area' in possible and 'width' in args and 'height' in args and 'max_area' not in args: + args['max_area'] = args['width'] * args['height'] # handle implicit controlnet if 'control_image' in possible and 'control_image' not in args and 'image' in args: diff --git a/modules/processing_class.py b/modules/processing_class.py index 4a9d93b90..8b02e38aa 100644 --- a/modules/processing_class.py +++ b/modules/processing_class.py @@ -6,7 +6,7 @@ from typing import Any, Dict, List from dataclasses import dataclass, field import numpy as np from PIL import Image, ImageOps -from modules import shared, images, scripts_manager, masking, sd_models, processing_helpers +from modules import shared, images, scripts_manager, masking, sd_models, sd_vae, processing_helpers debug = shared.log.trace if os.environ.get('SD_PROCESS_DEBUG', None) is not None else lambda *args, **kwargs: None @@ -15,6 +15,7 @@ debug = shared.log.trace if os.environ.get('SD_PROCESS_DEBUG', None) is not None @dataclass(repr=False) class StableDiffusionProcessing: def __init__(self, + sd_model_checkpoint: str = None, # # used only to set sd_model sd_model=None, # pylint: disable=unused-argument # local instance of sd_model # base params prompt: str = "", @@ -355,6 +356,17 @@ class StableDiffusionProcessing: self.prompt_attention_masks = [] self.negative_prompt_attention_mask = [] self.xyz = xyz + self.abort = False + + # set model + if sd_model_checkpoint is not None and len(sd_model_checkpoint) > 0: + from modules import sd_checkpoint + if sd_checkpoint.select_checkpoint(op='model', sd_model_checkpoint=sd_model_checkpoint) is None: + shared.log.error(f'Processing: model="{sd_model_checkpoint}" not found') + self.abort = True + else: + shared.opts.sd_model_checkpoint = sd_model_checkpoint + sd_models.reload_model_weights() def __str__(self): return f'{self.__class__.__name__}: {self.__dict__}' @@ -449,10 +461,11 @@ class StableDiffusionProcessingImg2Img(StableDiffusionProcessing): def init(self, all_prompts=None, all_seeds=None, all_subseeds=None): if self.init_images is not None and len(self.init_images) > 0: + vae_scale_factor = sd_vae.get_vae_scale_factor() if self.width is None or self.width == 0: - self.width = int(8 * (self.init_images[0].width * self.scale_by // 8)) + self.width = int(vae_scale_factor * (self.init_images[0].width * self.scale_by // vae_scale_factor)) if self.height is None or self.height == 0: - self.height = int(8 * (self.init_images[0].height * self.scale_by // 8)) + self.height = int(vae_scale_factor * (self.init_images[0].height * self.scale_by // vae_scale_factor)) if getattr(self, 'image_mask', None) is not None: shared.sd_model = sd_models.set_diffuser_pipe(self.sd_model, sd_models.DiffusersTaskType.INPAINTING) elif getattr(self, 'init_images', None) is not None: @@ -554,7 +567,8 @@ class StableDiffusionProcessingControl(StableDiffusionProcessingImg2Img): self.hr_force = force self.hr_upscaler = upscaler if use_scale: - self.hr_upscale_to_x, self.hr_upscale_to_y = 8 * int(self.width * scale / 8), 8 * int(self.height * scale / 8) + vae_scale_factor = sd_vae.get_vae_scale_factor() + self.hr_upscale_to_x, self.hr_upscale_to_y = vae_scale_factor * int(self.width * scale / vae_scale_factor), vae_scale_factor * int(self.height * scale / vae_scale_factor) else: self.hr_upscale_to_x, self.hr_upscale_to_y = self.hr_resize_x, self.hr_resize_y diff --git a/modules/processing_diffusers.py b/modules/processing_diffusers.py index 41ed7834c..8343559c4 100644 --- a/modules/processing_diffusers.py +++ b/modules/processing_diffusers.py @@ -451,8 +451,15 @@ def update_pipeline(sd_model, p: processing.StableDiffusionProcessing): def validate_pipeline(p: processing.StableDiffusionProcessing): - is_video_model = ('video' in shared.sd_model_type.lower()) or ('video' in shared.sd_model.__class__.__name__.lower()) - is_video_pipeline = 'video' in p.__class__.__name__.lower() + from modules.video_models.models_def import models as video_models + models_cls = [] + for family in video_models: + for m in video_models[family]: + if m.repo_cls is not None: + models_cls.append(m.repo_cls.__name__) + is_video_model = shared.sd_model.__class__.__name__ in models_cls + override_video_pipelines = ['WanPipeline'] + is_video_pipeline = ('video' in p.__class__.__name__.lower()) or (shared.sd_model.__class__.__name__ in override_video_pipelines) if is_video_model and not is_video_pipeline: shared.log.error(f'Mismatch: type={shared.sd_model_type} cls={shared.sd_model.__class__.__name__} request={p.__class__.__name__} video model with non-video pipeline') return False diff --git a/modules/processing_helpers.py b/modules/processing_helpers.py index 4777463ba..8729d2465 100644 --- a/modules/processing_helpers.py +++ b/modules/processing_helpers.py @@ -8,7 +8,7 @@ import numpy as np import cv2 from PIL import Image from blendmodes.blend import blendLayers, BlendType -from modules import shared, devices, images, sd_models, sd_samplers, sd_hijack_hypertile, processing_vae, timer +from modules import shared, devices, images, sd_models, sd_samplers, sd_vae, sd_hijack_hypertile, processing_vae, timer debug = shared.log.trace if os.environ.get('SD_PROCESS_DEBUG', None) is not None else lambda *args, **kwargs: None @@ -282,7 +282,8 @@ def resize_init_images(p): if getattr(p, 'image', None) is not None and getattr(p, 'init_images', None) is None: p.init_images = [p.image] if getattr(p, 'init_images', None) is not None and len(p.init_images) > 0: - tgt_width, tgt_height = 8 * math.ceil(p.init_images[0].width / 8), 8 * math.ceil(p.init_images[0].height / 8) + vae_scale_factor = sd_vae.get_vae_scale_factor() + tgt_width, tgt_height = vae_scale_factor * math.ceil(p.init_images[0].width / vae_scale_factor), vae_scale_factor * math.ceil(p.init_images[0].height / vae_scale_factor) if p.init_images[0].size != (tgt_width, tgt_height): shared.log.debug(f'Resizing init images: original={p.init_images[0].width}x{p.init_images[0].height} target={tgt_width}x{tgt_height}') p.init_images = [images.resize_image(1, image, tgt_width, tgt_height, upscaler_name=None) for image in p.init_images] diff --git a/modules/processing_vae.py b/modules/processing_vae.py index 6637e0e5b..62273045b 100644 --- a/modules/processing_vae.py +++ b/modules/processing_vae.py @@ -168,7 +168,8 @@ def full_vae_decode(latents, model): if debug: log_debug(f'VAE memory: {shared.mem_mon.read()}') vae_name = os.path.splitext(os.path.basename(sd_vae.loaded_vae_file))[0] if sd_vae.loaded_vae_file is not None else "default" - shared.log.debug(f'Decode: vae="{vae_name}" upcast={upcast} slicing={getattr(model.vae, "use_slicing", None)} tiling={getattr(model.vae, "use_tiling", None)} latents={list(latents.shape)}:{latents.device}:{latents.dtype} time={t1-t0:.3f}') + vae_scale_factor = sd_vae.get_vae_scale_factor(model) + shared.log.debug(f'Decode: vae="{vae_name}" scale={vae_scale_factor} upcast={upcast} slicing={getattr(model.vae, "use_slicing", None)} tiling={getattr(model.vae, "use_tiling", None)} latents={list(latents.shape)}:{latents.device}:{latents.dtype} time={t1-t0:.3f}') return decoded diff --git a/modules/prompt_parser.py b/modules/prompt_parser.py index 92389c1fd..57587e3b5 100644 --- a/modules/prompt_parser.py +++ b/modules/prompt_parser.py @@ -14,7 +14,8 @@ from typing import List import lark import torch from compel import Compel -from modules.shared import opts, log, native +from modules.shared import opts, log + # a prompt like this: "fantasy landscape with a [mountain:lake:0.25] and [an oak:a christmas tree:0.75][ in foreground::0.6][ in background:0.25] [shoddy:masterful:0.5]" # will be represented with prompt_schedule like this (assuming steps=100): @@ -24,6 +25,7 @@ from modules.shared import opts, log, native # [75, 'fantasy landscape with a lake and an oak in background masterful'] # [100, 'fantasy landscape with a lake and a christmas tree in background masterful'] + round_bracket_multiplier = 1.1 square_bracket_multiplier = 1.0 / 1.1 re_AND = re.compile(r"\bAND\b") @@ -334,7 +336,7 @@ def parse_prompt_attention(text): whitespace = '' else: re_attention = re_attention_v2 - if native and opts.sd_textencder_linebreak: + if opts.sd_textencder_linebreak: text = text.replace('\n', ' BREAK ') else: text = text.replace('\n', ' ') diff --git a/modules/prompt_parser_xhinker.py b/modules/prompt_parser_xhinker.py index fed488c2a..377a9f3cc 100644 --- a/modules/prompt_parser_xhinker.py +++ b/modules/prompt_parser_xhinker.py @@ -439,13 +439,13 @@ def get_weighted_text_embeddings_sdxl( , pad_last_block=pad_last_block ) - prompt_token_groups_2, prompt_weight_groups_2 = group_tokens_and_weights( + prompt_token_groups_2, _prompt_weight_groups_2 = group_tokens_and_weights( prompt_tokens_2.copy() , prompt_weights_2.copy() , pad_last_block=pad_last_block ) - neg_prompt_token_groups_2, neg_prompt_weight_groups_2 = group_tokens_and_weights( + neg_prompt_token_groups_2, _neg_prompt_weight_groups_2 = group_tokens_and_weights( neg_prompt_tokens_2.copy() , neg_prompt_weights_2.copy() , pad_last_block=pad_last_block @@ -609,7 +609,6 @@ def get_weighted_text_embeddings_sdxl_refiner( , generator = torch.Generator(text2img_pipe.device).manual_seed(2) ).images[0] """ - import math eos = 49407 # pipe.tokenizer.eos_token_id # tokenizer 2 @@ -1148,13 +1147,13 @@ def get_weighted_text_embeddings_sd3( , pad_last_block=pad_last_block ) - prompt_token_groups_2, prompt_weight_groups_2 = group_tokens_and_weights( + prompt_token_groups_2, _prompt_weight_groups_2 = group_tokens_and_weights( prompt_tokens_2.copy() , prompt_weights_2.copy() , pad_last_block=pad_last_block ) - neg_prompt_token_groups_2, neg_prompt_weight_groups_2 = group_tokens_and_weights( + neg_prompt_token_groups_2, _neg_prompt_weight_groups_2 = group_tokens_and_weights( neg_prompt_tokens_2.copy() , neg_prompt_weights_2.copy() , pad_last_block=pad_last_block @@ -1374,7 +1373,7 @@ def get_weighted_text_embeddings_flux1( pipe.tokenizer_2, prompt2 ) - prompt_token_groups, prompt_weight_groups = group_tokens_and_weights( + prompt_token_groups, _prompt_weight_groups = group_tokens_and_weights( prompt_tokens.copy() , prompt_weights.copy() , pad_last_block=True diff --git a/modules/schedulers/scheduler_flashflow.py b/modules/schedulers/scheduler_flashflow.py new file mode 100644 index 000000000..122f8ed74 --- /dev/null +++ b/modules/schedulers/scheduler_flashflow.py @@ -0,0 +1,428 @@ +# Copyright 2024 Stability AI, Katherine Crowson and The HuggingFace 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 math +from dataclasses import dataclass +from typing import List, Optional, Tuple, Union + +import numpy as np +import torch +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.schedulers.scheduling_utils import SchedulerMixin +from diffusers.utils import BaseOutput, is_scipy_available, logging +from diffusers.utils.torch_utils import randn_tensor + +if is_scipy_available(): + import scipy.stats + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +@dataclass +class FlashFlowMatchEulerDiscreteSchedulerOutput(BaseOutput): + """ + Output class for the scheduler's `step` function output. + + Args: + prev_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images): + Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the + denoising loop. + """ + + prev_sample: torch.FloatTensor + + +class FlashFlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin): + """ + Euler scheduler. + + This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic + methods the library implements for all schedulers such as loading and saving. + + Args: + num_train_timesteps (`int`, defaults to 1000): + The number of diffusion steps to train the model. + timestep_spacing (`str`, defaults to `"linspace"`): + The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and + Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information. + shift (`float`, defaults to 1.0): + The shift value for the timestep schedule. + """ + + _compatibles = [] + order = 1 + + @register_to_config + def __init__( + self, + num_train_timesteps: int = 1000, + shift: float = 1.0, + use_dynamic_shifting=False, + base_shift: Optional[float] = 0.5, + max_shift: Optional[float] = 1.15, + base_image_seq_len: Optional[int] = 256, + max_image_seq_len: Optional[int] = 4096, + invert_sigmas: bool = False, + use_karras_sigmas: Optional[bool] = False, + use_exponential_sigmas: Optional[bool] = False, + use_beta_sigmas: Optional[bool] = False, + ): + if self.config.use_beta_sigmas and not is_scipy_available(): + raise ImportError("Make sure to install scipy if you want to use beta sigmas.") + if sum([self.config.use_beta_sigmas, self.config.use_exponential_sigmas, self.config.use_karras_sigmas]) > 1: + raise ValueError( + "Only one of `config.use_beta_sigmas`, `config.use_exponential_sigmas`, `config.use_karras_sigmas` can be used." + ) + timesteps = np.linspace(1, num_train_timesteps, num_train_timesteps, dtype=np.float32)[::-1].copy() + timesteps = torch.from_numpy(timesteps).to(dtype=torch.float32) + + sigmas = timesteps / num_train_timesteps + if not use_dynamic_shifting: + # when use_dynamic_shifting is True, we apply the timestep shifting on the fly based on the image resolution + sigmas = shift * sigmas / (1 + (shift - 1) * sigmas) + + self.timesteps = sigmas * num_train_timesteps + + self._step_index = None + self._begin_index = None + + self.sigmas = sigmas.to("cpu") # to avoid too much CPU/GPU communication + self.sigma_min = self.sigmas[-1].item() + self.sigma_max = self.sigmas[0].item() + + @property + def step_index(self): + """ + The index counter for current timestep. It will increase 1 after each scheduler step. + """ + return self._step_index + + @property + def begin_index(self): + """ + The index for the first timestep. It should be set from pipeline with `set_begin_index` method. + """ + return self._begin_index + + # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_begin_index + def set_begin_index(self, begin_index: int = 0): + """ + Sets the begin index for the scheduler. This function should be run from pipeline before the inference. + + Args: + begin_index (`int`): + The begin index for the scheduler. + """ + self._begin_index = begin_index + + def scale_noise( + self, + sample: torch.FloatTensor, + timestep: Union[float, torch.FloatTensor], + noise: Optional[torch.FloatTensor] = None, + ) -> torch.FloatTensor: + """ + Forward process in flow-matching + + Args: + sample (`torch.FloatTensor`): + The input sample. + timestep (`int`, *optional*): + The current timestep in the diffusion chain. + + Returns: + `torch.FloatTensor`: + A scaled input sample. + """ + # Make sure sigmas and timesteps have the same device and dtype as original_samples + sigmas = self.sigmas.to(device=sample.device, dtype=sample.dtype) + + if sample.device.type == "mps" and torch.is_floating_point(timestep): + # mps does not support float64 + schedule_timesteps = self.timesteps.to(sample.device, dtype=torch.float32) + timestep = timestep.to(sample.device, dtype=torch.float32) + else: + schedule_timesteps = self.timesteps.to(sample.device) + timestep = timestep.to(sample.device) + + # self.begin_index is None when scheduler is used for training, or pipeline does not implement set_begin_index + if self.begin_index is None: + step_indices = [self.index_for_timestep(t, schedule_timesteps) for t in timestep] + elif self.step_index is not None: + # add_noise is called after first denoising step (for inpainting) + step_indices = [self.step_index] * timestep.shape[0] + else: + # add noise is called before first denoising step to create initial latent(img2img) + step_indices = [self.begin_index] * timestep.shape[0] + + sigma = sigmas[step_indices].flatten() + while len(sigma.shape) < len(sample.shape): + sigma = sigma.unsqueeze(-1) + + sample = sigma * noise + (1.0 - sigma) * sample + + return sample + + def _sigma_to_t(self, sigma): + return sigma * self.config.num_train_timesteps + + def time_shift(self, mu: float, sigma: float, t: torch.Tensor): + return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma) + + def set_timesteps( + self, + num_inference_steps: int = None, + device: Union[str, torch.device] = None, + sigmas: Optional[List[float]] = None, + mu: Optional[float] = None, + ): + """ + Sets the discrete timesteps used for the diffusion chain (to be run before inference). + + Args: + num_inference_steps (`int`): + The number of diffusion steps used when generating samples with a pre-trained model. + device (`str` or `torch.device`, *optional*): + The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. + """ + if self.config.use_dynamic_shifting and mu is None: + raise ValueError(" you have a pass a value for `mu` when `use_dynamic_shifting` is set to be `True`") + + if sigmas is None: + timesteps = np.linspace( + self._sigma_to_t(self.sigma_max), self._sigma_to_t(self.sigma_min), num_inference_steps + ) + + sigmas = timesteps / self.config.num_train_timesteps + else: + sigmas = np.array(sigmas).astype(np.float32) + num_inference_steps = len(sigmas) + self.num_inference_steps = num_inference_steps + + if self.config.use_dynamic_shifting: + sigmas = self.time_shift(mu, 1.0, sigmas) + else: + sigmas = self.config.shift * sigmas / (1 + (self.config.shift - 1) * sigmas) + + if self.config.use_karras_sigmas: + sigmas = self._convert_to_karras(in_sigmas=sigmas, num_inference_steps=num_inference_steps) + + elif self.config.use_exponential_sigmas: + sigmas = self._convert_to_exponential(in_sigmas=sigmas, num_inference_steps=num_inference_steps) + + elif self.config.use_beta_sigmas: + sigmas = self._convert_to_beta(in_sigmas=sigmas, num_inference_steps=num_inference_steps) + + sigmas = torch.from_numpy(sigmas).to(dtype=torch.float32, device=device) + timesteps = sigmas * self.config.num_train_timesteps + + if self.config.invert_sigmas: + sigmas = 1.0 - sigmas + timesteps = sigmas * self.config.num_train_timesteps + sigmas = torch.cat([sigmas, torch.ones(1, device=sigmas.device)]) + else: + sigmas = torch.cat([sigmas, torch.zeros(1, device=sigmas.device)]) + + self.timesteps = timesteps.to(device=device) + self.sigmas = sigmas + self._step_index = None + self._begin_index = None + + def index_for_timestep(self, timestep, schedule_timesteps=None): + if schedule_timesteps is None: + schedule_timesteps = self.timesteps + + indices = (schedule_timesteps == timestep).nonzero() + + # The sigma index that is taken for the **very** first `step` + # is always the second index (or the last index if there is only 1) + # This way we can ensure we don't accidentally skip a sigma in + # case we start in the middle of the denoising schedule (e.g. for image-to-image) + pos = 1 if len(indices) > 1 else 0 + + return indices[pos].item() + + def _init_step_index(self, timestep): + if self.begin_index is None: + if isinstance(timestep, torch.Tensor): + timestep = timestep.to(self.timesteps.device) + self._step_index = self.index_for_timestep(timestep) + else: + self._step_index = self._begin_index + + def step( + self, + model_output: torch.FloatTensor, + timestep: Union[float, torch.FloatTensor], + sample: torch.FloatTensor, + s_churn: float = 0.0, + s_tmin: float = 0.0, + s_tmax: float = float("inf"), + s_noise: float = 1.0, + generator: Optional[torch.Generator] = None, + return_dict: bool = True, + ) -> Union[FlashFlowMatchEulerDiscreteSchedulerOutput, Tuple]: + """ + Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion + process from the learned model outputs (most often the predicted noise). + + Args: + model_output (`torch.FloatTensor`): + The direct output from learned diffusion model. + timestep (`float`): + The current discrete timestep in the diffusion chain. + sample (`torch.FloatTensor`): + A current instance of a sample created by the diffusion process. + s_churn (`float`): + s_tmin (`float`): + s_tmax (`float`): + s_noise (`float`, defaults to 1.0): + Scaling factor for noise added to the sample. + generator (`torch.Generator`, *optional*): + A random number generator. + return_dict (`bool`): + Whether or not to return a [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] or + tuple. + + Returns: + [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] or `tuple`: + If return_dict is `True`, [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] is + returned, otherwise a tuple is returned where the first element is the sample tensor. + """ + + if ( + isinstance(timestep, int) + or isinstance(timestep, torch.IntTensor) + or isinstance(timestep, torch.LongTensor) + ): + raise ValueError( + ( + "Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to" + " `EulerDiscreteScheduler.step()` is not supported. Make sure to pass" + " one of the `scheduler.timesteps` as a timestep." + ), + ) + + if self.step_index is None: + self._init_step_index(timestep) + + # Upcast to avoid precision issues when computing prev_sample + + sigma = self.sigmas[self.step_index] + + # Upcast to avoid precision issues when computing prev_sample + sample = sample.to(torch.float32) + + denoised = sample - model_output * sigma + + if self.step_index < self.num_inference_steps - 1: + sigma_next = self.sigmas[self.step_index + 1] + noise = randn_tensor( + model_output.shape, + generator=generator, + device=model_output.device, + dtype=denoised.dtype, + ) + sample = sigma_next * noise + (1.0 - sigma_next) * denoised + + self._step_index += 1 + sample = sample.to(model_output.dtype) + + if not return_dict: + return (sample,) + + return FlashFlowMatchEulerDiscreteSchedulerOutput(prev_sample=sample) + + # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_karras + def _convert_to_karras(self, in_sigmas: torch.Tensor, num_inference_steps) -> torch.Tensor: + """Constructs the noise schedule of Karras et al. (2022).""" + + # Hack to make sure that other schedulers which copy this function don't break + # TODO: Add this logic to the other schedulers + if hasattr(self.config, "sigma_min"): + sigma_min = self.config.sigma_min + else: + sigma_min = None + + if hasattr(self.config, "sigma_max"): + sigma_max = self.config.sigma_max + else: + sigma_max = None + + sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item() + sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item() + + rho = 7.0 # 7.0 is the value used in the paper + ramp = np.linspace(0, 1, num_inference_steps) + min_inv_rho = sigma_min ** (1 / rho) + max_inv_rho = sigma_max ** (1 / rho) + sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho + return sigmas + + # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_exponential + def _convert_to_exponential(self, in_sigmas: torch.Tensor, num_inference_steps: int) -> torch.Tensor: + """Constructs an exponential noise schedule.""" + + # Hack to make sure that other schedulers which copy this function don't break + # TODO: Add this logic to the other schedulers + if hasattr(self.config, "sigma_min"): + sigma_min = self.config.sigma_min + else: + sigma_min = None + + if hasattr(self.config, "sigma_max"): + sigma_max = self.config.sigma_max + else: + sigma_max = None + + sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item() + sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item() + + sigmas = np.exp(np.linspace(math.log(sigma_max), math.log(sigma_min), num_inference_steps)) + return sigmas + + # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_beta + def _convert_to_beta( + self, in_sigmas: torch.Tensor, num_inference_steps: int, alpha: float = 0.6, beta: float = 0.6 + ) -> torch.Tensor: + """From "Beta Sampling is All You Need" [arXiv:2407.12173] (Lee et. al, 2024)""" + + # Hack to make sure that other schedulers which copy this function don't break + # TODO: Add this logic to the other schedulers + if hasattr(self.config, "sigma_min"): + sigma_min = self.config.sigma_min + else: + sigma_min = None + + if hasattr(self.config, "sigma_max"): + sigma_max = self.config.sigma_max + else: + sigma_max = None + + sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item() + sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item() + + sigmas = np.array( + [ + sigma_min + (ppf * (sigma_max - sigma_min)) + for ppf in [ + scipy.stats.beta.ppf(timestep, alpha, beta) + for timestep in 1 - np.linspace(0, 1, num_inference_steps) + ] + ] + ) + return sigmas + + def __len__(self): + return self.config.num_train_timesteps diff --git a/modules/sd_checkpoint.py b/modules/sd_checkpoint.py index 1c9bfe117..6afc5b953 100644 --- a/modules/sd_checkpoint.py +++ b/modules/sd_checkpoint.py @@ -127,7 +127,7 @@ def list_models(): global checkpoints_list # pylint: disable=global-statement checkpoints_list.clear() checkpoint_aliases.clear() - ext_filter = [".safetensors"] if shared.opts.sd_disable_ckpt or shared.native else [".ckpt", ".safetensors"] + ext_filter = [".safetensors"] model_list = list(modelloader.load_models(model_path=model_path, model_url=None, command_path=shared.opts.ckpt_dir, ext_filter=ext_filter, download_name=None, ext_blacklist=[".vae.ckpt", ".vae.safetensors"])) safetensors_list = [] for filename in sorted(model_list, key=str.lower): @@ -136,21 +136,16 @@ def list_models(): if checkpoint_info.name is not None: checkpoint_info.register() diffusers_list = [] - if shared.native: - for repo in modelloader.load_diffusers_models(clear=True): - checkpoint_info = CheckpointInfo(repo['name'], sha=repo['hash']) - diffusers_list.append(checkpoint_info) - if checkpoint_info.name is not None: - checkpoint_info.register() + for repo in modelloader.load_diffusers_models(clear=True): + checkpoint_info = CheckpointInfo(repo['name'], sha=repo['hash']) + diffusers_list.append(checkpoint_info) + if checkpoint_info.name is not None: + checkpoint_info.register() if shared.cmd_opts.ckpt is not None: - if not os.path.exists(shared.cmd_opts.ckpt) and not shared.native: - if shared.cmd_opts.ckpt.lower() != "none": - shared.log.warning(f'Load model: path="{shared.cmd_opts.ckpt}" not found') - else: - checkpoint_info = CheckpointInfo(shared.cmd_opts.ckpt) - if checkpoint_info.name is not None: - checkpoint_info.register() - shared.opts.data['sd_model_checkpoint'] = checkpoint_info.title + checkpoint_info = CheckpointInfo(shared.cmd_opts.ckpt) + if checkpoint_info.name is not None: + checkpoint_info.register() + shared.opts.data['sd_model_checkpoint'] = checkpoint_info.title elif shared.cmd_opts.ckpt != shared.default_sd_model_file and shared.cmd_opts.ckpt is not None: shared.log.warning(f'Load model: path="{shared.cmd_opts.ckpt}" not found') shared.log.info(f'Available Models: safetensors="{shared.opts.ckpt_dir}":{len(safetensors_list)} diffusers="{shared.opts.diffusers_dir}":{len(diffusers_list)} items={len(checkpoints_list)} time={time.time()-t0:.2f}') @@ -158,20 +153,42 @@ def list_models(): def update_model_hashes(): - txt = [] + def update_model_hashes_table(rows): + html = """ + + + + + + {tbody} + +
NameTypeHash
+ """ + tbody = '' + for row in rows: + try: + tbody += f""" + + {row.name} + {row.type} + {row.shorthash} + + """ + except Exception as e: + shared.log.error(f'Model list: row={row} {e}') + return html.format(tbody=tbody) + lst = [ckpt for ckpt in checkpoints_list.values() if ckpt.hash is None] for ckpt in lst: ckpt.hash = model_hash(ckpt.filename) lst = [ckpt for ckpt in checkpoints_list.values() if ckpt.sha256 is None or ckpt.shorthash is None] shared.log.info(f'Models list: hash missing={len(lst)} total={len(checkpoints_list)}') + updated = [] for ckpt in lst: ckpt.sha256 = hashes.sha256(ckpt.filename, f"checkpoint/{ckpt.name}") ckpt.shorthash = ckpt.sha256[0:10] if ckpt.sha256 is not None else None - if ckpt.sha256 is not None: - txt.append(f'Hash: {ckpt.title} {ckpt.shorthash}') - txt.append(f'Updated hashes for {len(lst)} out of {len(checkpoints_list)} models') - txt = '
'.join(txt) - return txt + updated.append(ckpt) + yield update_model_hashes_table(updated) def remove_hash(s): @@ -184,6 +201,7 @@ def get_closet_checkpoint_match(s: str) -> CheckpointInfo: checkpoint_info = CheckpointInfo(model_name) # create a virutal model info checkpoint_info.type = 'huggingface' return checkpoint_info + if s.startswith('huggingface/'): model_name = s.replace('huggingface/', '') checkpoint_info = CheckpointInfo(model_name) # create a virutal model info @@ -226,7 +244,8 @@ def get_closet_checkpoint_match(s: str) -> CheckpointInfo: # civitai search if shared.opts.sd_checkpoint_autodownload and s.startswith("https://civitai.com/api/download/models"): - fn = modelloader.download_civit_model_thread(model_name=None, model_url=s, model_path='', model_type='Model', token=None) + from modules.civitai.download_civitai import download_civit_model_thread + fn = download_civit_model_thread(model_name=None, model_url=s, model_path='', model_type='Model', token=None) if fn is not None: checkpoint_info = CheckpointInfo(fn) return checkpoint_info @@ -250,8 +269,8 @@ def model_hash(filename): return 'NOHASH' -def select_checkpoint(op='model'): - model_checkpoint = shared.opts.data.get('sd_model_refiner', None) if op == 'refiner' else shared.opts.data.get('sd_model_checkpoint', None) +def select_checkpoint(op='model', sd_model_checkpoint=None): + model_checkpoint = sd_model_checkpoint or (shared.opts.data.get('sd_model_refiner', None) if op == 'refiner' else shared.opts.data.get('sd_model_checkpoint', None)) if model_checkpoint is None or model_checkpoint == 'None' or len(model_checkpoint) < 3: return None checkpoint_info = get_closet_checkpoint_match(model_checkpoint) diff --git a/modules/sd_detect.py b/modules/sd_detect.py index b16478580..52d32f113 100644 --- a/modules/sd_detect.py +++ b/modules/sd_detect.py @@ -84,7 +84,7 @@ def guess_by_name(fn, current_guess): return 'Stable Diffusion 3' elif 'hidream' in fn.lower(): return 'HiDream' - elif 'chroma' in fn.lower(): + elif 'chroma' in fn.lower() and 'xl' not in fn.lower(): return 'Chroma' elif 'flux' in fn.lower() or 'flex.1' in fn.lower(): size = round(os.path.getsize(fn) / 1024 / 1024) if os.path.isfile(fn) else 0 @@ -101,13 +101,25 @@ def guess_by_name(fn, current_guess): return 'WanAI' elif 'bria' in fn.lower(): return 'Bria' + elif 'qwen' in fn.lower(): + return 'Qwen' + elif 'kandinsky-2-1' in fn.lower(): + return 'Kandinsky 2.1' + elif 'kandinsky-2-2' in fn.lower(): + return 'Kandinsky 2.2' + elif 'kandinsky-3' in fn.lower(): + return 'Kandinsky 3.0' return current_guess def guess_by_diffusers(fn, current_guess): + exclude_by_name = ['ostris/Flex.2-preview'] # pipeline may be misleading index = os.path.join(fn, 'model_index.json') if os.path.exists(index) and os.path.isfile(index): index = shared.readfile(index, silent=True) + name = index.get('_name_or_path', None) + if name is not None and name in exclude_by_name: + return current_guess, None cls = index.get('_class_name', None) if cls is not None: pipeline = getattr(diffusers, cls, None) diff --git a/modules/sd_hijack_te.py b/modules/sd_hijack_te.py index 2d0f165ff..960d87430 100644 --- a/modules/sd_hijack_te.py +++ b/modules/sd_hijack_te.py @@ -8,10 +8,12 @@ def hijack_encode_prompt(*args, **kwargs): t0 = time.time() if 'max_sequence_length' in kwargs: kwargs['max_sequence_length'] = max(kwargs['max_sequence_length'], os.environ.get('HIDREAM_MAX_SEQUENCE_LENGTH', 256)) + # if hasattr(shared.sd_model, 'text_encoder') and shared.sd_model.text_encoder is not None: + # sd_models.move_model(shared.sd_model.text_encoder, devices.device) try: res = shared.sd_model.orig_encode_prompt(*args, **kwargs) except Exception as e: - shared.log.error(f'Eencode prompt: {e}') + shared.log.error(f'Encode prompt: {e}') errors.display(e, 'Encode prompt') res = None t1 = time.time() diff --git a/modules/sd_models.py b/modules/sd_models.py index 1ba384732..94824218a 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -10,8 +10,7 @@ import diffusers.loaders.single_file_utils import torch import huggingface_hub as hf from installer import log -from modules import paths, shared, shared_state, shared_items, modelloader, devices, script_callbacks, sd_vae, sd_unet, errors, sd_models_compile, sd_hijack_accelerate, sd_detect, model_quant, sd_hijack_te -from modules.timer import Timer, process as process_timer +from modules import timer, paths, shared, shared_items, modelloader, devices, script_callbacks, sd_vae, sd_unet, errors, sd_models_compile, sd_hijack_accelerate, sd_detect, model_quant, sd_hijack_te from modules.memstats import memory_stats from modules.modeldata import model_data from modules.sd_checkpoint import CheckpointInfo, select_checkpoint, list_models, checkpoints_list, checkpoint_titles, get_closet_checkpoint_match, model_hash, update_model_hashes, setup_model, write_metadata, read_metadata_from_safetensors # pylint: disable=unused-import @@ -30,6 +29,7 @@ debug_load = os.environ.get('SD_LOAD_DEBUG', None) debug_process = log.trace if os.environ.get('SD_PROCESS_DEBUG', None) is not None else lambda *args, **kwargs: None diffusers_version = int(diffusers.__version__.split('.')[1]) checkpoint_tiles = checkpoint_titles # legacy compatibility +allow_post_quant = None pipe_switch_task_exclude = [ 'AnimateDiffPipeline', 'AnimateDiffSDXLPipeline', 'FluxControlPipeline', @@ -73,25 +73,39 @@ def copy_diffuser_options(new_pipe, orig_pipe): set_accelerate(new_pipe) +def set_huggingface_options(op: str, model_type: str): + if model_type is not None: # overrides + pass + if shared.opts.diffusers_to_gpu: # and model_type.startswith('Stable Diffusion'): + shared.log.debug(f'Setting {op}: component=accelerate direct={shared.opts.diffusers_to_gpu}') + sd_hijack_accelerate.hijack_accelerate() + else: + sd_hijack_accelerate.restore_accelerate() + if shared.opts.sd_parallel_load: + shared.log.debug(f'Setting {op}: component=huggingface parallel={shared.opts.sd_parallel_load}') + + def set_vae_options(sd_model, vae=None, op:str='model', quiet:bool=False): + ops = {} if hasattr(sd_model, "vae"): if vae is not None: sd_model.vae = vae - shared.log.quiet(quiet, f'Setting {op}: component=VAE name="{sd_vae.loaded_vae_file}"') + ops['name'] = f"{sd_vae.loaded_vae_file}" if shared.opts.diffusers_vae_upcast != 'default': sd_model.vae.config.force_upcast = True if shared.opts.diffusers_vae_upcast == 'true' else False - shared.log.quiet(quiet, f'Setting {op}: component=VAE upcast={sd_model.vae.config.force_upcast}') + ops['upcast'] = sd_model.vae.config.force_upcast if shared.opts.no_half_vae and op not in {'decode', 'encode'}: devices.dtype_vae = torch.float32 sd_model.vae.to(devices.dtype_vae) - shared.log.quiet(quiet, f'Setting {op}: component=VAE no-half=True') - if hasattr(sd_model, "enable_vae_slicing"): + ops['no-half'] = True + if hasattr(sd_model, "enable_vae_slicing") and hasattr(sd_model, "disable_vae_slicing"): + ops['slicing'] = shared.opts.diffusers_vae_slicing if shared.opts.diffusers_vae_slicing: - shared.log.quiet(quiet, f'Setting {op}: component=VAE slicing=True') sd_model.enable_vae_slicing() else: sd_model.disable_vae_slicing() if hasattr(sd_model, "enable_vae_tiling") and hasattr(sd_model, "disable_vae_tiling"): + ops['tiling'] = shared.opts.diffusers_vae_tiling if shared.opts.diffusers_vae_tiling: if hasattr(sd_model, 'vae') and hasattr(sd_model.vae, 'config') and hasattr(sd_model.vae.config, 'sample_size') and isinstance(sd_model.vae.config.sample_size, int): if getattr(sd_model.vae, "tile_sample_min_size_backup", None) is None: @@ -108,15 +122,16 @@ def set_vae_options(sd_model, vae=None, op:str='model', quiet:bool=False): sd_model.vae.tile_overlap_factor = float(shared.opts.diffusers_vae_tile_overlap) else: sd_model.vae.tile_overlap_factor = getattr(sd_model.vae, "tile_overlap_factor_backup", sd_model.vae.tile_overlap_factor) - shared.log.quiet(quiet, f'Setting {op}: component=VAE tiling=True tile={sd_model.vae.tile_sample_min_size} overlap={sd_model.vae.tile_overlap_factor}') - else: - shared.log.quiet(quiet, f'Setting {op}: component=VAE tiling=True') + ops['tile'] = sd_model.vae.tile_sample_min_size + ops['overlap'] = sd_model.vae.tile_overlap_factor sd_model.enable_vae_tiling() else: sd_model.disable_vae_tiling() if hasattr(sd_model, "vqvae"): - shared.log.quiet(quiet, f'Setting {op}: component=VQVAE upcast=True') + ops['upcast'] = True sd_model.vqvae.to(torch.float32) # vqvae is producing nans in fp16 + if not quiet and len(ops) > 0: + shared.log.quiet(quiet, f'Setting {op}: component=vae {ops}') def set_diffuser_options(sd_model, vae=None, op:str='model', offload:bool=True, quiet:bool=False): @@ -220,15 +235,17 @@ def move_model(model, device=None, force=False): pass # ignore model move if quantization is enabled elif 'already been set to the correct devices' in str(e0): pass # ignore errors on pre-quant models + elif 'Casting a quantized model to' in str(e0): + pass # ignore errors on quantized models else: raise e0 t1 = time.time() except Exception as e1: t1 = time.time() shared.log.error(f'Model move: device={device} {e1}') - if 'move' not in process_timer.records: - process_timer.records['move'] = 0 - process_timer.records['move'] += t1 - t0 + if 'move' not in timer.process.records: + timer.process.records['move'] = 0 + timer.process.records['move'] += t1 - t0 if os.environ.get('SD_MOVE_DEBUG', None) is not None or (t1-t0) > 2: shared.log.debug(f'Model move: device={device} class={model.__class__.__name__} accelerate={getattr(model, "has_accelerate", False)} fn={fn} time={t1-t0:.2f}') # pylint: disable=protected-access devices.torch_gc() @@ -271,7 +288,7 @@ def load_diffuser_initial(diffusers_load_config, op='model'): def load_diffuser_force(model_type, checkpoint_info, diffusers_load_config, op='model'): sd_model = None - allow_post_quant = True + global allow_post_quant # pylint: disable=global-statement unload_model_weights(op=op) shared.sd_model = None try: @@ -312,21 +329,23 @@ def load_diffuser_force(model_type, checkpoint_info, diffusers_load_config, op=' allow_post_quant = True elif model_type in ['FLUX']: from pipelines.model_flux import load_flux - sd_model, allow_post_quant = load_flux(checkpoint_info, diffusers_load_config) + sd_model = load_flux(checkpoint_info, diffusers_load_config) + allow_post_quant = False elif model_type in ['FLEX']: from pipelines.model_flex import load_flex sd_model = load_flex(checkpoint_info, diffusers_load_config) allow_post_quant = False elif model_type in ['Chroma']: from pipelines.model_chroma import load_chroma - sd_model, allow_post_quant = load_chroma(checkpoint_info, diffusers_load_config) + sd_model = load_chroma(checkpoint_info, diffusers_load_config) + allow_post_quant = False elif model_type in ['Lumina 2']: from pipelines.model_lumina import load_lumina2 sd_model = load_lumina2(checkpoint_info, diffusers_load_config) allow_post_quant = False elif model_type in ['Stable Diffusion 3']: from pipelines.model_sd3 import load_sd3 - sd_model = load_sd3(checkpoint_info, cache_dir=shared.opts.diffusers_dir, config=diffusers_load_config.get('config', None)) + sd_model = load_sd3(checkpoint_info, diffusers_load_config) allow_post_quant = False elif model_type in ['CogView 3']: # forced pipeline from pipelines.model_cogview import load_cogview3 @@ -368,12 +387,32 @@ def load_diffuser_force(model_type, checkpoint_info, diffusers_load_config, op=' from pipelines.model_bria import load_bria sd_model = load_bria(checkpoint_info, diffusers_load_config) allow_post_quant = False + elif model_type in ['Qwen']: + from pipelines.model_qwen import load_qwen + sd_model = load_qwen(checkpoint_info, diffusers_load_config) + allow_post_quant = False + elif model_type in ['HunyuanDiT']: + from pipelines.model_hunyuandit import load_hunyuandit + sd_model = load_hunyuandit(checkpoint_info, diffusers_load_config) + allow_post_quant = False + elif model_type in ['Kandinsky 2.1']: + from pipelines.model_kandinsky import load_kandinsky21 + sd_model = load_kandinsky21(checkpoint_info, diffusers_load_config) + allow_post_quant = True + elif model_type in ['Kandinsky 2.2']: + from pipelines.model_kandinsky import load_kandinsky22 + sd_model = load_kandinsky22(checkpoint_info, diffusers_load_config) + allow_post_quant = False + elif model_type in ['Kandinsky 3.0']: + from pipelines.model_kandinsky import load_kandinsky3 + sd_model = load_kandinsky3(checkpoint_info, diffusers_load_config) + allow_post_quant = False except Exception as e: shared.log.error(f'Load {op}: path="{checkpoint_info.path}" {e}') if debug_load: errors.display(e, 'Load') return None, True - return sd_model, allow_post_quant + return sd_model def load_diffuser_folder(model_type, pipeline, checkpoint_info, diffusers_load_config, op='model'): @@ -461,7 +500,7 @@ def load_diffuser_file(model_type, pipeline, checkpoint_info, diffusers_load_con diffusers_load_config['config'] = model_config if model_type.startswith('Stable Diffusion 3'): from pipelines.model_sd3 import load_sd3 - sd_model = load_sd3(checkpoint_info=checkpoint_info, cache_dir=shared.opts.diffusers_dir, config=diffusers_load_config.get('config', None)) + sd_model = load_sd3(checkpoint_info, diffusers_load_config) elif hasattr(pipeline, 'from_single_file'): diffusers.loaders.single_file_utils.CHECKPOINT_KEY_NAMES["clip"] = "cond_stage_model.transformer.text_model.embeddings.position_embedding.weight" # patch for diffusers==0.28.0 diffusers_load_config['use_safetensors'] = True @@ -471,11 +510,6 @@ def load_diffuser_file(model_type, pipeline, checkpoint_info, diffusers_load_con if shared.opts.disable_accelerate: from diffusers.utils import import_utils import_utils._accelerate_available = False # pylint: disable=protected-access - if shared.opts.diffusers_to_gpu and model_type.startswith('Stable Diffusion'): - shared.log.debug(f'Setting {op}: component=accelerate direct={shared.opts.diffusers_to_gpu}') - sd_hijack_accelerate.hijack_accelerate() - else: - sd_hijack_accelerate.restore_accelerate() sd_model = pipeline.from_single_file(checkpoint_info.path, **diffusers_load_config) # sd_model = patch_diffuser_config(sd_model, checkpoint_info.path) elif hasattr(pipeline, 'from_ckpt'): @@ -535,11 +569,11 @@ def set_defaults(sd_model, checkpoint_info): sd_model.set_progress_bar_config(bar_format='Progress {rate_fmt}{postfix} {bar} {percentage:3.0f}% {n_fmt}/{total_fmt} {elapsed} {remaining}', ncols=80, colour='#327fba') -def load_diffuser(checkpoint_info=None, timer=None, op='model', revision=None): # pylint: disable=unused-argument - if timer is None: - timer = Timer() +def load_diffuser(checkpoint_info=None, op='model', revision=None): # pylint: disable=unused-argument + global allow_post_quant # pylint: disable=global-statement + allow_post_quant = True # assume default logging.getLogger("diffusers").setLevel(logging.ERROR) - timer.record("diffusers") + timer.load.record("diffusers") diffusers_load_config = { "low_cpu_mem_usage": True, "torch_dtype": devices.dtype, @@ -566,7 +600,6 @@ def load_diffuser(checkpoint_info=None, timer=None, op='model', revision=None): return sd_model = None - allow_post_quant = True try: # initial load only if sd_model is None: @@ -581,6 +614,7 @@ def load_diffuser(checkpoint_info=None, timer=None, op='model', revision=None): # detect pipeline pipeline, model_type = sd_detect.detect_pipeline(checkpoint_info.path, op) + set_huggingface_options(op, model_type) # preload vae so it can be used as param vae = None @@ -594,11 +628,11 @@ def load_diffuser(checkpoint_info=None, timer=None, op='model', revision=None): vae = sd_vae.load_vae_diffusers(checkpoint_info.path, vae_file, vae_source) if vae is not None: diffusers_load_config["vae"] = vae - timer.record("vae") + timer.load.record("vae") # load with custom loader if sd_model is None: - sd_model, allow_post_quant = load_diffuser_force(model_type, checkpoint_info, diffusers_load_config, op) + sd_model = load_diffuser_force(model_type, checkpoint_info, diffusers_load_config, op) if sd_model is not None and not sd_model: shared.log.error(f'Load {op}: type="{model_type}" pipeline="{pipeline}" not loaded') return @@ -628,7 +662,7 @@ def load_diffuser(checkpoint_info=None, timer=None, op='model', revision=None): add_noise_pred_to_diffusers_callback(sd_model) - timer.record("load") + timer.load.record("load") if op == 'refiner': model_data.sd_refiner = sd_model @@ -636,7 +670,7 @@ def load_diffuser(checkpoint_info=None, timer=None, op='model', revision=None): model_data.sd_model = sd_model reload_text_encoder(initial=True) # must be before embeddings - timer.record("te") + timer.load.record("te") if debug_load: shared.log.trace(f'Model components: {list(get_signature(sd_model).values())}') @@ -645,7 +679,7 @@ def load_diffuser(checkpoint_info=None, timer=None, op='model', revision=None): sd_model.embedding_db = textual_inversion.EmbeddingDatabase() sd_model.embedding_db.add_embedding_dir(shared.opts.embeddings_dir) sd_model.embedding_db.load_textual_inversion_embeddings(force_reload=True) - timer.record("embeddings") + timer.load.record("embeddings") from modules import prompt_parser_diffusers prompt_parser_diffusers.insert_parser_highjack(sd_model.__class__.__name__) @@ -653,7 +687,7 @@ def load_diffuser(checkpoint_info=None, timer=None, op='model', revision=None): set_diffuser_options(sd_model, vae, op, offload=False) sd_model = model_quant.do_post_load_quant(sd_model, allow=allow_post_quant) # run this before move model so it can be compressed in CPU - timer.record("options") + timer.load.record("options") set_diffuser_offload(sd_model, op) @@ -665,14 +699,14 @@ def load_diffuser(checkpoint_info=None, timer=None, op='model', revision=None): move_model(sd_model, devices.cpu) else: move_model(sd_model, devices.device) - timer.record("move") + timer.load.record("move") if shared.opts.ipex_optimize: sd_model = sd_models_compile.ipex_optimize(sd_model) if ('Model' in shared.opts.cuda_compile and shared.opts.cuda_compile_backend != 'none'): sd_model = sd_models_compile.compile_diffusers(sd_model) - timer.record("compile") + timer.load.record("compile") except Exception as e: shared.log.error(f"Load {op}: {e}") @@ -687,7 +721,7 @@ def load_diffuser(checkpoint_info=None, timer=None, op='model', revision=None): from modules import modelstats modelstats.analyze() - shared.log.info(f"Load {op}: family={shared.sd_model_type} time={timer.dct()} native={get_native(sd_model)} memory={memory_stats()}") + shared.log.info(f"Load {op}: family={shared.sd_model_type} time={timer.load.dct()} native={get_native(sd_model)} memory={memory_stats()}") class DiffusersTaskType(Enum): @@ -837,6 +871,7 @@ def backup_pipe_components(pipe): 'feature_extractor': getattr(pipe, "feature_extractor", None), 'mask_processor': getattr(pipe, "mask_processor", None), 'restore_pipeline': getattr(pipe, "restore_pipeline", None), + 'task_args': getattr(pipe, "task_args", None), } @@ -860,6 +895,8 @@ def restore_pipe_components(pipe, components): pipe.mask_processor = components['mask_processor'] if components['restore_pipeline'] is not None: pipe.restore_pipeline = components['restore_pipeline'] + if components['task_args'] is not None: + pipe.task_args = components['task_args'] if pipe.__class__.__name__ in ['FluxPipeline', 'StableDiffusion3Pipeline']: pipe.register_modules(image_encoder = components['image_encoder']) @@ -906,7 +943,7 @@ def set_diffuser_pipe(pipe, new_pipe_type): elif new_pipe_type == DiffusersTaskType.INPAINTING: new_pipe = diffusers.AutoPipelineForInpainting.from_pipe(pipe) else: - shared.log.error(f'Pipeline class change failed: type={new_pipe_type} pipeline={cls}') + shared.log.warning(f'Pipeline class change failed: type={new_pipe_type} pipeline={cls}') return pipe except Exception as e: # pylint: disable=unused-variable shared.log.warning(f'Pipeline class change failed: type={new_pipe_type} pipeline={cls} {e}') @@ -1045,6 +1082,8 @@ def reload_text_encoder(initial=False): from modules.model_te import set_t5 shared.log.debug(f'Load module: type=t5 path="{shared.opts.sd_text_encoder}" module="text_encoder_3"') set_t5(pipe=shared.sd_model, module='text_encoder_3', t5=shared.opts.sd_text_encoder, cache_dir=shared.opts.diffusers_dir) + clear_caches() + apply_balanced_offload(shared.sd_model) def reload_model_weights(sd_model=None, info=None, op='model', force=False, revision=None): @@ -1053,7 +1092,6 @@ def reload_model_weights(sd_model=None, info=None, op='model', force=False, revi unload_model_weights(op=op) return None orig_state = copy.deepcopy(shared.state) - shared.state = shared_state.State() shared.state.begin('Load') if sd_model is None: sd_model = model_data.sd_model if op == 'model' or op == 'dict' else model_data.sd_refiner @@ -1067,12 +1105,12 @@ def reload_model_weights(sd_model=None, info=None, op='model', force=False, revi move_model(sd_model, devices.cpu) unload_model_weights(op=op) sd_model = None - timer = Timer() + timer.load = timer.Timer() # TODO model load: implement model in-memory caching - timer.record("config") + timer.load.record("config") if sd_model is None or force: sd_model = None - load_diffuser(checkpoint_info, timer=timer, op=op, revision=revision) + load_diffuser(checkpoint_info, op=op, revision=revision) shared.state.end() shared.state = orig_state if op == 'model': diff --git a/modules/sd_offload.py b/modules/sd_offload.py index 689c49e8c..4a253a428 100644 --- a/modules/sd_offload.py +++ b/modules/sd_offload.py @@ -16,7 +16,7 @@ debug_move = log.trace if debug else lambda *args, **kwargs: None offload_warn = ['sc', 'sd3', 'f1', 'h1', 'hunyuandit', 'auraflow', 'omnigen', 'omnigen2', 'cogview4', 'cosmos', 'chroma'] offload_post = ['h1'] offload_hook_instance = None -balanced_offload_exclude = ['CogView4Pipeline'] +balanced_offload_exclude = ['CogView4Pipeline', 'MeissonicPipeline'] accelerate_dtype_byte_size = None @@ -351,6 +351,16 @@ def apply_balanced_offload_to_module(module, op="apply"): devices.torch_gc(fast=True, force=True, reason='offload') +def report_model_stats(module_name, module): + try: + size = offload_hook_instance.offload_map.get(module_name, 0) + quant = getattr(module, "quantization_method", None) + params = sum(p.numel() for p in module.parameters(recurse=True)) + shared.log.debug(f'Module: name={module_name} cls={module.__class__.__name__} size={size:.3f} params={params} quant={quant}') + except Exception as e: + shared.log.error(f'Module stats: name={module_name} {e}') + + def apply_balanced_offload(sd_model=None, exclude=[]): global offload_hook_instance # pylint: disable=global-statement if shared.opts.diffusers_offload_mode != "balanced": @@ -382,6 +392,7 @@ def apply_balanced_offload(sd_model=None, exclude=[]): module.module_name = module_name module.offload_dir = os.path.join(shared.opts.accelerate_offload_path, checkpoint_name, module_name) apply_balanced_offload_to_module(module, op='apply') + report_model_stats(module_name, module) set_accelerate(sd_model) t = time.time() - t0 process_timer.add('offload', t) diff --git a/modules/sd_samplers.py b/modules/sd_samplers.py index c7502cb4b..644e9dd16 100644 --- a/modules/sd_samplers.py +++ b/modules/sd_samplers.py @@ -78,7 +78,7 @@ def create_sampler(name, model): if model is not None: if getattr(model, "default_scheduler", None) is None: model.default_scheduler = copy.deepcopy(model.scheduler) - requires_flow = ('FlowMatch' in model.default_scheduler.__class__.__name__) or (getattr(model.scheduler.config, 'prediction_type', None) == 'flow_prediction') + requires_flow = ('FlowMatch' in model.default_scheduler.__class__.__name__) or (getattr(model.default_scheduler.config, 'prediction_type', None) == 'flow_prediction') else: requires_flow = False diff --git a/modules/sd_samplers_diffusers.py b/modules/sd_samplers_diffusers.py index b052042bc..6f06eceaa 100644 --- a/modules/sd_samplers_diffusers.py +++ b/modules/sd_samplers_diffusers.py @@ -51,6 +51,7 @@ try: from modules.schedulers.scheduler_bdia import BDIA_DDIMScheduler # pylint: disable=ungrouped-imports from modules.schedulers.scheduler_ufogen import UFOGenScheduler # pylint: disable=ungrouped-imports from modules.schedulers.scheduler_unipc_flowmatch import FlowUniPCMultistepScheduler # pylint: disable=ungrouped-imports + from modules.schedulers.scheduler_flashflow import FlashFlowMatchEulerDiscreteScheduler # pylint: disable=ungrouped-imports from modules.perflow import PeRFlowScheduler # pylint: disable=ungrouped-imports except Exception as e: shared.log.error(f'Sampler import: version={diffusers.__version__} error: {e}') @@ -104,6 +105,7 @@ config = { 'VDM Solver': { 'clip_sample_range': 2.0, }, 'TCD': { 'set_alpha_to_one': True, 'rescale_betas_zero_snr': False, 'beta_schedule': 'scaled_linear' }, 'TDD': { }, + 'Flash FlowMatch': { 'shift': 1, 'use_dynamic_shifting': False, 'use_karras_sigmas': False, 'use_exponential_sigmas': False, 'use_beta_sigmas': False }, 'PeRFlow': { 'prediction_type': 'ddim_eps' }, 'UFOGen': { }, 'BDIA DDIM': { 'clip_sample': False, 'set_alpha_to_one': True, 'steps_offset': 0, 'clip_sample_range': 1.0, 'sample_max_value': 1.0, 'timestep_spacing': 'leading', 'rescale_betas_zero_snr': False, 'thresholding': False, 'gamma': 1.0 }, @@ -153,6 +155,7 @@ samplers_data_diffusers = [ SamplerData('Heun', lambda model: DiffusionSampler('Heun', HeunDiscreteScheduler, model), [], {}), SamplerData('Heun FlowMatch', lambda model: DiffusionSampler('Heun FlowMatch', FlowMatchHeunDiscreteScheduler, model), [], {}), + SamplerData('Flash FlowMatch', lambda model: DiffusionSampler('Flash FlowMatch', FlashFlowMatchEulerDiscreteScheduler, model), [], {}), SamplerData('DEIS', lambda model: DiffusionSampler('DEIS', DEISMultistepScheduler, model), [], {}), SamplerData('SA Solver', lambda model: DiffusionSampler('SA Solver', SASolverScheduler, model), [], {}), @@ -317,6 +320,12 @@ class DiffusionSampler: self.sampler = None return + # monkey-patch to allow sdxl pipeline to execute flowmatch samplers + if not hasattr(sampler, 'scale_model_input'): + sampler.scale_model_input = lambda x, _y: x + if not hasattr(sampler, 'init_noise_sigma'): + sampler.init_noise_sigma = 1.0 + self.sampler = sampler # shared.log.debug_log(f'Sampler: class="{self.sampler.__class__.__name__}" config={self.sampler.config}') diff --git a/modules/sd_unet.py b/modules/sd_unet.py index 13b6302fb..2bd1741f4 100644 --- a/modules/sd_unet.py +++ b/modules/sd_unet.py @@ -37,7 +37,7 @@ def load_unet(model): if prior_text_encoder is not None: model.prior_pipe.text_encoder = None # Prevent OOM model.prior_pipe.text_encoder = prior_text_encoder.to(devices.device, dtype=devices.dtype) - elif any([m in model.__class__.__name__ for m in dit_models]): # noqa: C419 # pylint: disable=use-a-generator + elif any([m in model.__class__.__name__ for m in dit_models]) or hasattr(model, 'transformer'): # noqa: C419 # pylint: disable=use-a-generator loaded_unet = shared.opts.sd_unet sd_models.load_diffuser() # TODO model load: force-reloading entire model as loading transformers only leads to massive memory usage else: diff --git a/modules/sd_vae.py b/modules/sd_vae.py index b6c3982e9..7986b2568 100644 --- a/modules/sd_vae.py +++ b/modules/sd_vae.py @@ -12,6 +12,35 @@ checkpoint_info = None vae_path = os.path.abspath(os.path.join(paths.models_path, 'VAE')) debug = os.environ.get('SD_LOAD_DEBUG', None) is not None unspecified = object() +vae_scale_override = { + 'WanPipeline': 16, +} + + +def get_vae_scale_factor(model=None): + patch_size = 1 + if model is None: + model = shared.sd_model + if model is None: + vae_scale_factor = 8 + elif model.__class__.__name__ in vae_scale_override: + vae_scale_factor = vae_scale_override[model.__class__.__name__] + elif hasattr(model, 'vae_scale_factor_spatial'): + vae_scale_factor = model.vae_scale_factor_spatial + elif hasattr(model, 'vae_scale_factor'): + vae_scale_factor = model.vae_scale_factor + elif hasattr(model, 'pipe') and hasattr(model.pipe, 'vae_scale_factor'): + vae_scale_factor = model.pipe.vae_scale_factor + elif hasattr(model, 'config') and hasattr(model.config, 'vae_scale_factor'): + vae_scale_factor = model.config.vae_scale_factor + else: + # shared.log.warning(f'VAE: cls={model.__class__.__name__ if model else "None"} scale=unknown') + vae_scale_factor = 8 + if hasattr(model, 'patch_size'): + patch_size = model.patch_size + if debug: + shared.log.trace(f'VAE: cls={model.__class__.__name__ if model else "None"} scale={vae_scale_factor} patch={patch_size}') + return vae_scale_factor * patch_size def load_vae_dict(filename): diff --git a/modules/sd_vae_taesd.py b/modules/sd_vae_taesd.py index 7fc9c35f7..aed995f1b 100644 --- a/modules/sd_vae_taesd.py +++ b/modules/sd_vae_taesd.py @@ -8,7 +8,7 @@ import os import threading from PIL import Image import torch -from modules import devices, paths +from modules import devices, paths, shared TAESD_MODELS = { @@ -36,22 +36,20 @@ prev_cls = '' prev_type = '' prev_model = '' lock = threading.Lock() -supported = ['sd', 'sdxl', 'sd3', 'f1', 'h1', 'lumina2', 'hunyuanvideo', 'wanai', 'mochivideo', 'pixartsigma', 'pixartalpha', 'hunyuandit', 'omnigen'] +supported = ['sd', 'sdxl', 'sd3', 'f1', 'h1', 'lumina2', 'hunyuanvideo', 'wanai', 'mochivideo', 'pixartsigma', 'pixartalpha', 'hunyuandit', 'omnigen', 'qwen'] def warn_once(msg, variant=None): - from modules import shared variant = variant or shared.opts.taesd_variant global prev_warnings # pylint: disable=global-statement if not prev_warnings: prev_warnings = True - shared.log.error(f'Decode: type="taesd" variant="{variant}": {msg}') + shared.log.warning(f'Decode: type="taesd" variant="{variant}": {msg}') return Image.new('RGB', (8, 8), color = (0, 0, 0)) def get_model(model_type = 'decoder', variant = None): global prev_cls, prev_type, prev_model # pylint: disable=global-statement - from modules import shared model_cls = shared.sd_model_type if model_cls is None or model_cls == 'none': return None, variant @@ -61,7 +59,7 @@ def get_model(model_type = 'decoder', variant = None): model_cls = 'sdxl' elif model_cls in {'h1', 'lumina2', 'chroma'}: model_cls = 'f1' - elif model_cls in {'wanai'}: + elif model_cls in {'wanai', 'qwen'}: variant = variant or 'TAE WanVideo' elif model_cls not in supported: warn_once(f'cls={shared.sd_model.__class__.__name__} type={model_cls} unsuppported', variant=variant) diff --git a/modules/sdnq/__init__.py b/modules/sdnq/__init__.py index 985e99a2a..2654a06e7 100644 --- a/modules/sdnq/__init__.py +++ b/modules/sdnq/__init__.py @@ -3,6 +3,7 @@ from typing import Any, Dict, List, Tuple, Optional, Union from dataclasses import dataclass from enum import Enum + import torch from diffusers.quantizers.base import DiffusersQuantizer from diffusers.quantizers.quantization_config import QuantizationConfigMixin @@ -14,8 +15,38 @@ from .dequantizer import dequantizer_dict from .forward import get_forward_func +class QuantizationMethod(str, Enum): + SDNQ = "sdnq" + + +def get_scale_asymmetric(weight: torch.FloatTensor, reduction_axes: Union[int, List[int]], weights_dtype: str) -> Tuple[torch.FloatTensor, torch.FloatTensor]: + zero_point = torch.amin(weight, dim=reduction_axes, keepdims=True) + scale = torch.amax(weight, dim=reduction_axes, keepdims=True).sub_(zero_point).div_(dtype_dict[weights_dtype]["max"] - dtype_dict[weights_dtype]["min"]) + if dtype_dict[weights_dtype]["min"] != 0: + zero_point.sub_(torch.mul(scale, dtype_dict[weights_dtype]["min"])) + return scale, zero_point + + +def get_scale_symmetric(weight: torch.FloatTensor, reduction_axes: Union[int, List[int]], weights_dtype: str) -> torch.FloatTensor: + return torch.amax(weight.abs(), dim=reduction_axes, keepdims=True).div_(dtype_dict[weights_dtype]["max"]) + + +def quantize_weight(weight: torch.FloatTensor, reduction_axes: Union[int, List[int]], weights_dtype: str) -> Tuple[torch.Tensor, torch.FloatTensor, torch.FloatTensor]: + if dtype_dict[weights_dtype]["is_unsigned"]: + scale, zero_point = get_scale_asymmetric(weight, reduction_axes, weights_dtype) + quantized_weight = torch.sub(weight, zero_point).div_(scale) + else: + scale = get_scale_symmetric(weight, reduction_axes, weights_dtype) + quantized_weight = torch.div(weight, scale) + zero_point = None + if dtype_dict[weights_dtype]["is_integer"]: + quantized_weight.round_() + quantized_weight = quantized_weight.clamp_(dtype_dict[weights_dtype]["min"], dtype_dict[weights_dtype]["max"]).to(dtype_dict[weights_dtype]["torch_dtype"]) + return quantized_weight, scale, zero_point + + @devices.inference_context() -def sdnq_quantize_layer(layer, weights_dtype="int8", torch_dtype=None, group_size=0, quant_conv=False, use_quantized_matmul=False, use_quantized_matmul_conv=False, dequantize_fp32=False, quantization_device=None, return_device=None, param_name=None): # pylint: disable=unused-argument +def sdnq_quantize_layer(layer, weights_dtype="int8", torch_dtype=None, group_size=0, quant_conv=False, use_quantized_matmul=False, use_quantized_matmul_conv=False, dequantize_fp32=False, non_blocking=False, quantization_device=None, return_device=None, param_name=None): # pylint: disable=unused-argument layer_class_name = layer.__class__.__name__ if layer_class_name in allowed_types: is_conv_type = False @@ -117,7 +148,7 @@ def sdnq_quantize_layer(layer, weights_dtype="int8", torch_dtype=None, group_siz if return_device is None: return_device = layer.weight.device if quantization_device is not None: - layer.weight.data = layer.weight.to(quantization_device) + layer.weight.data = layer.weight.to(quantization_device, non_blocking=non_blocking) if layer.weight.dtype != torch.float32: layer.weight.data = layer.weight.to(dtype=torch.float32) @@ -129,8 +160,7 @@ def sdnq_quantize_layer(layer, weights_dtype="int8", torch_dtype=None, group_siz if use_quantized_matmul: scale = scale.transpose(0,1) - if dtype_dict[weights_dtype]["num_bits"] == 8: - layer.weight.data = layer.weight.transpose(0,1) + layer.weight.data = layer.weight.transpose(0,1) if not dtype_dict[weights_dtype]["is_integer"]: stride = layer.weight.stride() if stride[0] > stride[1] and stride[1] == 1: @@ -148,23 +178,40 @@ def sdnq_quantize_layer(layer, weights_dtype="int8", torch_dtype=None, group_siz weights_dtype=weights_dtype, use_quantized_matmul=use_quantized_matmul, ) - layer.weight.data = layer.sdnq_dequantizer.pack_weight(layer.weight).to(return_device) - layer.sdnq_dequantizer = layer.sdnq_dequantizer.to(return_device) + layer.weight.data = layer.sdnq_dequantizer.pack_weight(layer.weight).to(return_device, non_blocking=non_blocking) + layer.sdnq_dequantizer = layer.sdnq_dequantizer.to(return_device, non_blocking=non_blocking) layer.forward = get_forward_func(layer_class_name, use_quantized_matmul, dtype_dict[weights_dtype]["is_integer"], use_tensorwise_fp8_matmul) layer.forward = layer.forward.__get__(layer, layer.__class__) - #devices.torch_gc(force=False, reason=f"SDNQ param_name: {param_name}") return layer -def apply_sdnq_to_module(model, weights_dtype="int8", torch_dtype=None, group_size=0, quant_conv=False, use_quantized_matmul=False, use_quantized_matmul_conv=False, dequantize_fp32=False, quantization_device=None, return_device=None, param_name=None, modules_to_not_convert: List[str] = []): # pylint: disable=unused-argument +def apply_sdnq_to_module(model, weights_dtype="int8", torch_dtype=None, group_size=0, quant_conv=False, use_quantized_matmul=False, use_quantized_matmul_conv=False, dequantize_fp32=False, non_blocking=False, quantization_device=None, return_device=None, modules_to_not_convert: List[str] = [], modules_dtype_dict: Dict[str, List[str]] = {}, op=None): # pylint: disable=unused-argument has_children = list(model.children()) if not has_children: return model - for module_param_name, module in model.named_children(): - if module_param_name in modules_to_not_convert: + for param_name, module in model.named_children(): + if param_name in modules_to_not_convert: continue if hasattr(module, "weight") and module.weight is not None: + if len(modules_dtype_dict.keys()) > 0: + for key, value in modules_dtype_dict.items(): + if param_name in value: + key = key.lower() + if key in {"8bit", "8bits"}: + if dtype_dict[weights_dtype]["num_bits"] != 8: + weights_dtype = "int8" + elif key.startswith("minimum_"): + minimum_bits_str = key.removeprefix("minimum_").removesuffix("bits").removesuffix("bit") + minimum_bits = int(minimum_bits_str) + if dtype_dict[weights_dtype]["num_bits"] < minimum_bits: + weights_dtype = "int" + minimum_bits_str + if minimum_bits <= 4: + weights_dtype = "u" + weights_dtype + else: + weights_dtype = key + break + module = sdnq_quantize_layer( module, weights_dtype=weights_dtype, @@ -174,9 +221,10 @@ def apply_sdnq_to_module(model, weights_dtype="int8", torch_dtype=None, group_si use_quantized_matmul=use_quantized_matmul, use_quantized_matmul_conv=use_quantized_matmul_conv, dequantize_fp32=dequantize_fp32, + non_blocking=non_blocking, quantization_device=quantization_device, return_device=return_device, - param_name=module_param_name, + param_name=param_name, ) module = apply_sdnq_to_module( module, @@ -187,49 +235,16 @@ def apply_sdnq_to_module(model, weights_dtype="int8", torch_dtype=None, group_si use_quantized_matmul=use_quantized_matmul, use_quantized_matmul_conv=use_quantized_matmul_conv, dequantize_fp32=dequantize_fp32, + non_blocking=non_blocking, quantization_device=quantization_device, return_device=return_device, - param_name=module_param_name, modules_to_not_convert=modules_to_not_convert, + modules_dtype_dict=modules_dtype_dict, + op=op, ) return model -def get_scale_asymmetric(weight: torch.FloatTensor, reduction_axes: Union[int, List[int]], weights_dtype: str) -> Tuple[torch.FloatTensor, torch.FloatTensor]: - zero_point = torch.amin(weight, dim=reduction_axes, keepdims=True) - scale = torch.amax(weight, dim=reduction_axes, keepdims=True).sub_(zero_point).div_(dtype_dict[weights_dtype]["max"] - dtype_dict[weights_dtype]["min"]) - eps = torch.finfo(scale.dtype).eps # prevent divison by 0 - scale = torch.where(torch.abs(scale) < eps, eps, scale) - if dtype_dict[weights_dtype]["min"] != 0: - zero_point.sub_(torch.mul(scale, dtype_dict[weights_dtype]["min"])) - return scale, zero_point - - -def get_scale_symmetric(weight: torch.FloatTensor, reduction_axes: Union[int, List[int]], weights_dtype: str) -> torch.FloatTensor: - scale = torch.amax(weight.abs(), dim=reduction_axes, keepdims=True).div_(dtype_dict[weights_dtype]["max"]) - eps = torch.finfo(scale.dtype).eps # prevent divison by 0 - scale = torch.where(torch.abs(scale) < eps, eps, scale) - return scale - - -def quantize_weight(weight: torch.FloatTensor, reduction_axes: Union[int, List[int]], weights_dtype: str) -> Tuple[torch.Tensor, torch.FloatTensor, torch.FloatTensor]: - if dtype_dict[weights_dtype]["is_unsigned"]: - scale, zero_point = get_scale_asymmetric(weight, reduction_axes, weights_dtype) - quantized_weight = torch.sub(weight, zero_point).div_(scale) - else: - scale = get_scale_symmetric(weight, reduction_axes, weights_dtype) - quantized_weight = torch.div(weight, scale) - zero_point = None - if dtype_dict[weights_dtype]["is_integer"]: - quantized_weight.round_() - quantized_weight = quantized_weight.clamp_(dtype_dict[weights_dtype]["min"], dtype_dict[weights_dtype]["max"]).to(dtype_dict[weights_dtype]["torch_dtype"]) - return quantized_weight, scale, zero_point - - -class QuantizationMethod(str, Enum): - SDNQ = "sdnq" - - class SDNQQuantizer(DiffusersQuantizer): r""" Diffusers Quantizer for SDNQ @@ -284,6 +299,26 @@ class SDNQQuantizer(DiffusersQuantizer): unexpected_keys: List[str], # pylint: disable=unused-argument **kwargs, # pylint: disable=unused-argument ): + weights_dtype = self.quantization_config.weights_dtype + if len(self.quantization_config.modules_dtype_dict.keys()) > 0: + split_param_name = param_name.split(".") + for key, value in self.quantization_config.modules_dtype_dict.items(): + if param_name in value or any(param in split_param_name for param in value): + key = key.lower() + if key in {"8bit", "8bits"}: + if dtype_dict[weights_dtype]["num_bits"] != 8: + weights_dtype = "int8" + elif key.startswith("minimum_"): + minimum_bits_str = key.removeprefix("minimum_").removesuffix("bits").removesuffix("bit") + minimum_bits = int(minimum_bits_str) + if dtype_dict[weights_dtype]["num_bits"] < minimum_bits: + weights_dtype = "int" + minimum_bits_str + if minimum_bits <= 4: + weights_dtype = "u" + weights_dtype + else: + weights_dtype = key + break + if self.quantization_config.return_device is not None: return_device = self.quantization_config.return_device else: @@ -295,19 +330,20 @@ class SDNQQuantizer(DiffusersQuantizer): if param_value.dtype == torch.float32 and devices.same_device(param_value.device, target_device): param_value = param_value.clone() else: - param_value = param_value.to(target_device).to(dtype=torch.float32) + param_value = param_value.to(target_device, non_blocking=self.quantization_config.non_blocking).to(dtype=torch.float32) layer, _ = get_module_from_name(model, param_name) layer.weight = torch.nn.Parameter(param_value, requires_grad=False) layer = sdnq_quantize_layer( layer, - weights_dtype=self.quantization_config.weights_dtype, + weights_dtype=weights_dtype, torch_dtype=self.torch_dtype, group_size=self.quantization_config.group_size, quant_conv=self.quantization_config.quant_conv, use_quantized_matmul=self.quantization_config.use_quantized_matmul, use_quantized_matmul_conv=self.quantization_config.use_quantized_matmul_conv, dequantize_fp32=self.quantization_config.dequantize_fp32, + non_blocking=self.quantization_config.non_blocking, quantization_device=None, return_device=return_device, param_name=param_name, @@ -376,6 +412,12 @@ class SDNQQuantizer(DiffusersQuantizer): """ return expected_keys + def update_param_name(self, param_name: str) -> str: + """ + needed for transformers compatibilty, no-op function + """ + return param_name + @property def is_trainable(self): return False @@ -409,6 +451,8 @@ class SDNQConfig(QuantizationConfigMixin): Same as use_quantized_matmul_conv but for the convolutional layers with UNets like SDXL. dequantize_fp32 (`bool`, *optional*, defaults to `False`): Enabling this option will use FP32 on the dequantization step. + non_blocking (`bool`, *optional*, defaults to `False`): + Enabling this option will use non blocking ops when moving layers between the quantization device and the return device. quantization_device (`torch.device`, *optional*, defaults to `None`): Used to set which device will be used for the quantization calculation on model load. return_device (`torch.device`, *optional*, defaults to `None`): @@ -416,6 +460,8 @@ class SDNQConfig(QuantizationConfigMixin): modules_to_not_convert (`list`, *optional*, default to `None`): The list of modules to not quantize, useful for quantizing models that explicitly require to have some modules left in their original precision (e.g. Whisper encoder, Llava encoder, Mixtral gate layers). + modules_dtype_dict (`dict`, *optional*, default to `None`): + The dict of dtypes and list of modules, useful for quantizing some modules with a different dtype. """ def __init__( # pylint: disable=super-init-not-called @@ -426,9 +472,11 @@ class SDNQConfig(QuantizationConfigMixin): use_quantized_matmul: bool = False, use_quantized_matmul_conv: bool = False, dequantize_fp32: bool = False, + non_blocking: bool = False, quantization_device: Optional[torch.device] = None, return_device: Optional[torch.device] = None, modules_to_not_convert: Optional[List[str]] = None, + modules_dtype_dict: Optional[Dict[str, List[str]]] = None, **kwargs, # pylint: disable=unused-argument ): self.weights_dtype = weights_dtype @@ -438,9 +486,11 @@ class SDNQConfig(QuantizationConfigMixin): self.use_quantized_matmul = use_quantized_matmul self.use_quantized_matmul_conv = use_quantized_matmul_conv self.dequantize_fp32 = dequantize_fp32 + self.non_blocking = non_blocking self.quantization_device = quantization_device self.return_device = return_device self.modules_to_not_convert = modules_to_not_convert + self.modules_dtype_dict = modules_dtype_dict self.post_init() self.is_integer = dtype_dict[self.weights_dtype]["is_integer"] @@ -456,3 +506,6 @@ class SDNQConfig(QuantizationConfigMixin): self.modules_to_not_convert = [] elif not isinstance(self.modules_to_not_convert, list): self.modules_to_not_convert = [self.modules_to_not_convert] + + if self.modules_dtype_dict is None: + self.modules_dtype_dict = {} diff --git a/modules/sdnq/common.py b/modules/sdnq/common.py index 9dab864c5..d3dc3e3b5 100644 --- a/modules/sdnq/common.py +++ b/modules/sdnq/common.py @@ -1,7 +1,8 @@ # pylint: disable=redefined-builtin,no-member,protected-access +import os import torch -from modules import devices +from modules import devices, shared torch_version = float(torch.__version__[:3]) @@ -30,7 +31,9 @@ if hasattr(torch, "float8_e4m3fnuz"): if hasattr(torch, "float8_e5m2fnuz"): dtype_dict["float8_e5m2fnuz"] = {"min": -57344, "max": 57344, "num_bits": 8, "target_dtype": "fp8", "torch_dtype": torch.float8_e5m2fnuz, "storage_dtype": torch.float8_e5m2fnuz, "is_unsigned": False, "is_integer": False} -use_tensorwise_fp8_matmul = True # Direct tensorwise only exist on H100 hardware, sdnq will use software tensorwise with this setting +use_torch_compile = shared.opts.sdnq_dequantize_compile # this setting requires a full restart of the webui to apply +use_tensorwise_fp8_matmul = os.environ.get('SDNQ_USE_TENSORWISE_FP8_MATMUL', "1").lower() not in {"0", "false", "no"} # row-wise FP8 only exist on H100 hardware, sdnq will use software row-wise with tensorwise hardware with this setting + quantized_matmul_dtypes = ("int8", "int7", "int6", "int5", "int4", "int3", "int2", "float8_e4m3fn", "float8_e5m2") if devices.backend in {"cpu", "openvino"}: quantized_matmul_dtypes += ("float8_e4m3fnuz", "float8_e5m2fnuz") @@ -39,3 +42,7 @@ linear_types = ("Linear",) conv_types = ("Conv1d", "Conv2d", "Conv3d") conv_transpose_types = ("ConvTranspose1d", "ConvTranspose2d", "ConvTranspose3d") allowed_types = linear_types + conv_types + conv_transpose_types + +if use_torch_compile: + torch._dynamo.config.cache_size_limit = max(8192, torch._dynamo.config.cache_size_limit) + torch._dynamo.config.accumulated_recompile_limit = max(8192, torch._dynamo.config.accumulated_recompile_limit) diff --git a/modules/sdnq/dequantizer.py b/modules/sdnq/dequantizer.py index db6fb622c..76e830f59 100644 --- a/modules/sdnq/dequantizer.py +++ b/modules/sdnq/dequantizer.py @@ -1,10 +1,9 @@ # pylint: disable=redefined-builtin,no-member,protected-access import torch -from modules import shared -from .common import dtype_dict -from .packed_int import pack_int_symetric, unpack_int_symetric, packed_int_function_dict +from .common import dtype_dict, use_torch_compile +from .packed_int import pack_int_symetric, unpack_int_symetric, pack_int_asymetric, unpack_int_asymetric def dequantize_asymmetric(weight: torch.ByteTensor, scale: torch.FloatTensor, zero_point: torch.FloatTensor, dtype: torch.dtype, result_shape: torch.Size) -> torch.FloatTensor: @@ -15,10 +14,9 @@ def dequantize_asymmetric(weight: torch.ByteTensor, scale: torch.FloatTensor, ze def dequantize_symmetric(weight: torch.CharTensor, scale: torch.FloatTensor, dtype: torch.dtype, result_shape: torch.Size, skip_quantized_matmul: bool = False) -> torch.FloatTensor: + result = weight.to(dtype=scale.dtype).mul_(scale).to(dtype=dtype) if skip_quantized_matmul: - result = weight.transpose(0,1).to(dtype=scale.dtype).mul_(scale.transpose(0,1)).to(dtype=dtype) - else: - result = weight.to(dtype=scale.dtype).mul_(scale).to(dtype=dtype) + result = result.t() if result_shape is not None: result = result.reshape(result_shape) return result @@ -29,14 +27,11 @@ def dequantize_symmetric_with_bias(weight: torch.CharTensor, scale: torch.FloatT def dequantize_packed_int_asymmetric(weight: torch.ByteTensor, scale: torch.FloatTensor, zero_point: torch.FloatTensor, shape: torch.Size, dtype: torch.dtype, result_shape: torch.Size, weights_dtype: str) -> torch.FloatTensor: - return dequantize_asymmetric(packed_int_function_dict[weights_dtype]["unpack"](weight, shape), scale, zero_point, dtype, result_shape) + return dequantize_asymmetric(unpack_int_asymetric(weight, shape, weights_dtype), scale, zero_point, dtype, result_shape) def dequantize_packed_int_symmetric(weight: torch.ByteTensor, scale: torch.FloatTensor, shape: torch.Size, dtype: torch.dtype, result_shape: torch.Size, weights_dtype: str, skip_quantized_matmul: bool = False) -> torch.FloatTensor: - if skip_quantized_matmul: - return dequantize_symmetric(unpack_int_symetric(weight, shape, weights_dtype, dtype=scale.dtype), scale.transpose(0,1), dtype, result_shape) - else: - return dequantize_symmetric(unpack_int_symetric(weight, shape, weights_dtype, dtype=scale.dtype), scale, dtype, result_shape) + return dequantize_symmetric(unpack_int_symetric(weight, shape, weights_dtype, dtype=scale.dtype), scale, dtype, result_shape, skip_quantized_matmul=skip_quantized_matmul) class AsymmetricWeightsDequantizer(torch.nn.Module): @@ -115,7 +110,7 @@ class PackedINTAsymmetricWeightsDequantizer(torch.nn.Module): self.register_buffer("zero_point", zero_point) def pack_weight(self, weight: torch.Tensor) -> torch.Tensor: - return packed_int_function_dict[self.weights_dtype]["pack"](weight.to(dtype=dtype_dict[self.weights_dtype]["torch_dtype"])) + return pack_int_asymetric(weight, self.weights_dtype) def forward(self, weight, **kwargs): # pylint: disable=unused-argument return dequantize_packed_int_asymmetric_compiled(weight, self.scale, self.zero_point, self.quantized_weight_shape, self.result_dtype, self.result_shape, self.weights_dtype) @@ -173,20 +168,11 @@ dequantizer_dict = { } -if shared.opts.sdnq_dequantize_compile: - try: - torch._dynamo.config.cache_size_limit = max(8192, torch._dynamo.config.cache_size_limit) - torch._dynamo.config.accumulated_recompile_limit = max(8192, torch._dynamo.config.accumulated_recompile_limit) - dequantize_asymmetric_compiled = torch.compile(dequantize_asymmetric, fullgraph=True, dynamic=False) - dequantize_symmetric_compiled = torch.compile(dequantize_symmetric, fullgraph=True, dynamic=False) - dequantize_packed_int_asymmetric_compiled = torch.compile(dequantize_packed_int_asymmetric, fullgraph=True, dynamic=False) - dequantize_packed_int_symmetric_compiled = torch.compile(dequantize_packed_int_symmetric, fullgraph=True, dynamic=False) - except Exception as e: - shared.log.warning(f"Quantization: type=sdnq Dequantize using torch.compile is not available: {e}") - dequantize_asymmetric_compiled = dequantize_asymmetric - dequantize_symmetric_compiled = dequantize_symmetric - dequantize_packed_int_asymmetric_compiled = dequantize_packed_int_asymmetric - dequantize_packed_int_symmetric_compiled = dequantize_packed_int_symmetric +if use_torch_compile: + dequantize_asymmetric_compiled = torch.compile(dequantize_asymmetric, fullgraph=True, dynamic=False) + dequantize_symmetric_compiled = torch.compile(dequantize_symmetric, fullgraph=True, dynamic=False) + dequantize_packed_int_asymmetric_compiled = torch.compile(dequantize_packed_int_asymmetric, fullgraph=True, dynamic=False) + dequantize_packed_int_symmetric_compiled = torch.compile(dequantize_packed_int_symmetric, fullgraph=True, dynamic=False) else: dequantize_asymmetric_compiled = dequantize_asymmetric dequantize_symmetric_compiled = dequantize_symmetric diff --git a/modules/sdnq/forward.py b/modules/sdnq/forward.py index 4abc4d7e8..ff2923dfb 100644 --- a/modules/sdnq/forward.py +++ b/modules/sdnq/forward.py @@ -1,410 +1,48 @@ -# pylint: disable=redefined-builtin,no-member,protected-access +# pylint: disable=protected-access -from typing import Callable, List, Tuple, Optional -import torch -from modules import shared +from typing import Callable from .common import conv_types, conv_transpose_types -from .dequantizer import dequantize_symmetric, dequantize_symmetric_with_bias -from .packed_int import unpack_int_symetric def get_forward_func(layer_class_name: str, use_quantized_matmul: bool, is_integer: bool, use_tensorwise_fp8_matmul: bool) -> Callable: # pylint: disable=inconsistent-return-statements if layer_class_name in conv_types: if use_quantized_matmul: if is_integer: + from .layers.conv.conv_int8 import quantized_conv_forward_int8_matmul return quantized_conv_forward_int8_matmul else: if use_tensorwise_fp8_matmul: + from .layers.conv.conv_fp8_tensorwise import quantized_conv_forward_fp8_matmul_tensorwise return quantized_conv_forward_fp8_matmul_tensorwise else: + from .layers.conv.conv_fp8 import quantized_conv_forward_fp8_matmul return quantized_conv_forward_fp8_matmul else: + from .layers.conv.forward import quantized_conv_forward return quantized_conv_forward elif layer_class_name in conv_transpose_types: if layer_class_name.endswith("1d"): + from .layers.conv.forward import quantized_conv_transpose_1d_forward return quantized_conv_transpose_1d_forward elif layer_class_name.endswith("2d"): + from .layers.conv.forward import quantized_conv_transpose_2d_forward return quantized_conv_transpose_2d_forward elif layer_class_name.endswith("3d"): + from .layers.conv.forward import quantized_conv_transpose_3d_forward return quantized_conv_transpose_3d_forward else: if use_quantized_matmul: if is_integer: + from .layers.linear.linear_int8 import quantized_linear_forward_int8_matmul return quantized_linear_forward_int8_matmul else: if use_tensorwise_fp8_matmul: + from .layers.linear.linear_fp8_tensorwise import quantized_linear_forward_fp8_matmul_tensorwise return quantized_linear_forward_fp8_matmul_tensorwise else: + from .layers.linear.linear_fp8 import quantized_linear_forward_fp8_matmul return quantized_linear_forward_fp8_matmul else: + from .layers.linear.forward import quantized_linear_forward return quantized_linear_forward - - -def quantize_fp8_matmul_input(input: torch.FloatTensor) -> Tuple[torch.Tensor, torch.FloatTensor]: - input = input.flatten(0,-2).contiguous() - input_scale = torch.amax(input.abs(), dim=-1, keepdims=True).div_(448) - input = torch.div(input, input_scale).clamp_(-448, 448).to(dtype=torch.float8_e4m3fn) - input_scale = input_scale.to(dtype=torch.float32) - return input, input_scale - - -def quantize_fp8_matmul_input_tensorwise(input: torch.FloatTensor, scale: torch.FloatTensor) -> Tuple[torch.Tensor, torch.FloatTensor]: - input = input.flatten(0,-2).contiguous() - input_scale = torch.amax(input.abs(), dim=-1, keepdims=True).div_(448) - input = torch.div(input, input_scale).clamp_(-448, 448).to(dtype=torch.float8_e4m3fn) - scale = torch.mul(input_scale, scale) - if scale.dtype == torch.float16: # fp16 will overflow - scale = scale.to(dtype=torch.float32) - return input, scale - - -def quantize_int8_matmul_input(input: torch.FloatTensor, scale: torch.FloatTensor) -> Tuple[torch.CharTensor, torch.FloatTensor]: - input = input.flatten(0,-2).contiguous() - input_scale = torch.amax(input.abs(), dim=-1, keepdims=True).div_(127) - input = torch.div(input, input_scale).round_().clamp_(-128, 127).to(dtype=torch.int8) - scale = torch.mul(input_scale, scale) - if scale.dtype == torch.float16: # fp16 will overflow - scale = scale.to(dtype=torch.float32) - return input, scale - - -def fp8_matmul( - input: torch.FloatTensor, - weight: torch.Tensor, - bias: torch.FloatTensor, - scale: torch.FloatTensor, -) -> torch.FloatTensor: - return_dtype = input.dtype - output_shape = list(input.shape) - output_shape[-1] = weight.shape[-1] - input, input_scale = quantize_fp8_matmul_input(input) - return torch._scaled_mm(input, weight, scale_a=input_scale, scale_b=scale, bias=bias, out_dtype=return_dtype).reshape(output_shape) - - -# sm89 doesn't support row wise scale in Windows -def fp8_matmul_tensorwise( - input: torch.FloatTensor, - weight: torch.Tensor, - bias: torch.FloatTensor, - scale: torch.FloatTensor, -) -> torch.FloatTensor: - return_dtype = input.dtype - output_shape = list(input.shape) - output_shape[-1] = weight.shape[-1] - dummy_input_scale = torch.ones(1, device=input.device, dtype=torch.float32) - input, scale = quantize_fp8_matmul_input_tensorwise(input, scale) - if bias is not None: - return dequantize_symmetric_with_bias(torch._scaled_mm(input, weight, scale_a=dummy_input_scale, scale_b=dummy_input_scale, bias=None, out_dtype=scale.dtype), scale, bias, return_dtype, output_shape) - else: - return dequantize_symmetric(torch._scaled_mm(input, weight, scale_a=dummy_input_scale, scale_b=dummy_input_scale, bias=None, out_dtype=scale.dtype), scale, return_dtype, output_shape) - - -def int8_matmul( - input: torch.FloatTensor, - weight: torch.Tensor, - bias: torch.FloatTensor, - scale: torch.FloatTensor, - quantized_weight_shape: torch.Size, - weights_dtype: str, -) -> torch.FloatTensor: - if quantized_weight_shape is not None: - weight = unpack_int_symetric(weight, quantized_weight_shape, weights_dtype, dtype=torch.int8, transpose=True) - return_dtype = input.dtype - output_shape = list(input.shape) - output_shape[-1] = weight.shape[-1] - input, scale = quantize_int8_matmul_input(input, scale) - if bias is not None: - return dequantize_symmetric_with_bias(torch._int_mm(input, weight), scale, bias, return_dtype, output_shape) - else: - return dequantize_symmetric(torch._int_mm(input, weight), scale, return_dtype, output_shape) - - -def process_conv_input(conv_type, input, reversed_padding_repeated_twice, padding_mode, result_shape, stride, padding, dilation): - if conv_type == 1: - batch_size, _, L_in = input.shape - C_out, _, K_l = result_shape - L_out = (L_in + 2 * padding[1] - dilation[1] * (K_l - 1) - 1) // stride[1] + 1 - mm_output_shape = (batch_size, L_out, C_out) - kernel_size = (1, K_l) - if conv_type == 2: - batch_size, _, H_in, W_in = input.shape - C_out, _, K_h, K_w = result_shape - H_out = (H_in + 2 * padding[0] - dilation[0] * (K_h - 1) - 1) // stride[0] + 1 - W_out = (W_in + 2 * padding[1] - dilation[1] * (K_w - 1) - 1) // stride[1] + 1 - mm_output_shape = (batch_size, H_out, W_out, C_out) - kernel_size = (K_h, K_w) - else: - batch_size, _, D_in, H_in, W_in = input.shape - C_out, _, K_d, K_h, K_w = result_shape - D_out = (D_in + 2 * padding[0] - dilation[0] * (K_d - 1) - 1) // stride[0] + 1 - H_out = (H_in + 2 * padding[1] - dilation[1] * (K_h - 1) - 1) // stride[1] + 1 - W_out = (W_in + 2 * padding[2] - dilation[2] * (K_w - 1) - 1) // stride[2] + 1 - mm_output_shape = (batch_size, D_out, H_out, W_out, C_out) - kernel_size = (K_d, K_h, K_w) - - if padding_mode != "zeros": - input = torch.nn.functional.pad(input, reversed_padding_repeated_twice, mode=padding_mode) - padding = (0,) * (conv_type if conv_type != 1 else 2) - elif conv_type == 3: - input = torch.nn.functional.pad(input, reversed_padding_repeated_twice) - - if conv_type == 1: - input = input.unsqueeze(2) - - if conv_type == 3: - K_D_eff = K_d + (K_d - 1) * (dilation[0] - 1) - K_H_eff = K_h + (K_h - 1) * (dilation[0] - 1) - K_W_eff = K_w + (K_w - 1) * (dilation[0] - 1) - input = input.unfold(2, K_D_eff, stride[0]).unfold(3, K_H_eff, stride[1]).unfold(4, K_W_eff, stride[2]) - if dilation[0] > 1: - input = input[..., ::dilation[0], :, :] - if dilation[1] > 1: - input = input[..., ::dilation[1], :] - if dilation[2] > 1: - input = input[..., ::dilation[2]] - input = input.permute(0, 2, 3, 4, 1, 5, 6, 7).reshape(batch_size, D_out * H_out * W_out, -1) - else: - input = torch.nn.functional.unfold(input, kernel_size=kernel_size, padding=padding, stride=stride, dilation=dilation).transpose(1,2) - return input, mm_output_shape - - -def conv_fp8_matmul( - input: torch.FloatTensor, - weight: torch.Tensor, - bias: torch.FloatTensor, - scale: torch.FloatTensor, - result_shape: torch.Size, - reversed_padding_repeated_twice: List[int], - padding_mode: str, conv_type: int, - groups: int, stride: List[int], - padding: List[int], dilation: List[int], -) -> torch.FloatTensor: - return_dtype = input.dtype - input, mm_output_shape = process_conv_input(conv_type, input, reversed_padding_repeated_twice, padding_mode, result_shape, stride, padding, dilation) - input, input_scale = quantize_fp8_matmul_input(input) - - if groups == 1: - result = torch._scaled_mm(input, weight, scale_a=input_scale, scale_b=scale, bias=bias, out_dtype=return_dtype).reshape(mm_output_shape) - else: - scale = scale.reshape(groups, 1, scale.shape[1] // groups) - input_scale = input_scale.reshape(groups, input_scale.shape[0] // groups, 1) - weight = weight.reshape(weight.shape[0], groups, weight.shape[1] // groups).transpose(0,1) - input = input.reshape(input.shape[0], groups, input.shape[1] // groups).transpose(0,1) - result = [] - if bias is not None: - bias = bias.reshape(groups, bias.shape[0] // groups) - for i in range(groups): - result.append(torch._scaled_mm(input[i], weight[i], scale_a=input_scale[i], scale_b=scale[i], bias=bias[i], out_dtype=return_dtype)) - else: - for i in range(groups): - result.append(torch._scaled_mm(input[i], weight[i], scale_a=input_scale[i], scale_b=scale[i], bias=None, out_dtype=return_dtype)) - result = torch.cat(result, dim=-1).reshape(mm_output_shape) - - if conv_type == 1: - result = result.transpose(1,2) - elif conv_type == 2: - result = result.permute(0,3,1,2) - elif conv_type == 3: - result = result.permute(0,4,1,2,3) - return result - - -def conv_fp8_matmul_tensorwise( - input: torch.FloatTensor, - weight: torch.Tensor, - bias: torch.FloatTensor, - scale: torch.FloatTensor, - result_shape: torch.Size, - reversed_padding_repeated_twice: List[int], - padding_mode: str, conv_type: int, - groups: int, stride: List[int], - padding: List[int], dilation: List[int], -) -> torch.FloatTensor: - return_dtype = input.dtype - input, mm_output_shape = process_conv_input(conv_type, input, reversed_padding_repeated_twice, padding_mode, result_shape, stride, padding, dilation) - input, scale = quantize_fp8_matmul_input_tensorwise(input, scale) - dummy_input_scale = torch.ones(1, device=input.device, dtype=torch.float32) - - if groups == 1: - result = torch._scaled_mm(input, weight, scale_a=dummy_input_scale, scale_b=dummy_input_scale, bias=None, out_dtype=scale.dtype) - else: - weight = weight.reshape(weight.shape[0], groups, weight.shape[1] // groups).transpose(0,1) - input = input.reshape(input.shape[0], groups, input.shape[1] // groups).transpose(0,1) - result = [] - for i in range(groups): - result.append(torch._scaled_mm(input[i], weight[i], scale_a=dummy_input_scale, scale_b=dummy_input_scale, bias=None, out_dtype=scale.dtype)) - result = torch.cat(result, dim=-1) - if bias is not None: - dequantize_symmetric_with_bias(result, scale, bias, return_dtype, mm_output_shape) - else: - dequantize_symmetric(result, scale, return_dtype, mm_output_shape) - - if conv_type == 1: - result = result.transpose(1,2) - elif conv_type == 2: - result = result.permute(0,3,1,2) - elif conv_type == 3: - result = result.permute(0,4,1,2,3) - return result - - -def conv_int8_matmul( - input: torch.FloatTensor, - weight: torch.CharTensor, - bias: torch.FloatTensor, - scale: torch.FloatTensor, - result_shape: torch.Size, - quantized_weight_shape: torch.Size, - weights_dtype: str, - reversed_padding_repeated_twice: List[int], - padding_mode: str, conv_type: int, - groups: int, stride: List[int], - padding: List[int], dilation: List[int], -) -> torch.FloatTensor: - return_dtype = input.dtype - input, mm_output_shape = process_conv_input(conv_type, input, reversed_padding_repeated_twice, padding_mode, result_shape, stride, padding, dilation) - input, scale = quantize_int8_matmul_input(input, scale) - if quantized_weight_shape is not None: - weight = unpack_int_symetric(weight, quantized_weight_shape, weights_dtype, dtype=torch.int8, transpose=True) - - if groups == 1: - result = torch._int_mm(input, weight) - else: - weight = weight.reshape(weight.shape[0], groups, weight.shape[1] // groups).transpose(0,1) - input = input.reshape(input.shape[0], groups, input.shape[1] // groups).transpose(0,1) - result = [] - for i in range(groups): - result.append(torch._int_mm(input[i], weight[i])) - result = torch.cat(result, dim=-1) - if bias is not None: - result = dequantize_symmetric_with_bias(result, scale, bias, return_dtype, mm_output_shape) - else: - result = dequantize_symmetric(result, scale, return_dtype, mm_output_shape) - - if conv_type == 1: - result = result.transpose(1,2) - elif conv_type == 2: - result = result.permute(0,3,1,2) - elif conv_type == 3: - result = result.permute(0,4,1,2,3) - return result - - -def quantized_linear_forward_fp8_matmul(self, input: torch.FloatTensor) -> torch.FloatTensor: - if torch.numel(input) / input.shape[-1] < 32: - return torch.nn.functional.linear(input, self.sdnq_dequantizer(self.weight, skip_quantized_matmul=True), self.bias) - return fp8_matmul(input, self.weight, self.bias, self.sdnq_dequantizer.scale) - - -def quantized_linear_forward_fp8_matmul_tensorwise(self, input: torch.FloatTensor) -> torch.FloatTensor: - if torch.numel(input) / input.shape[-1] < 32: - return torch.nn.functional.linear(input, self.sdnq_dequantizer(self.weight, skip_quantized_matmul=True), self.bias) - return fp8_matmul_tensorwise(input, self.weight, self.bias, self.sdnq_dequantizer.scale) - - -def quantized_linear_forward_int8_matmul(self, input: torch.FloatTensor) -> torch.FloatTensor: - if torch.numel(input) / input.shape[-1] < 32: - return torch.nn.functional.linear(input, self.sdnq_dequantizer(self.weight, skip_quantized_matmul=True), self.bias) - return int8_matmul(input, self.weight, self.bias, self.sdnq_dequantizer.scale, getattr(self.sdnq_dequantizer, "quantized_weight_shape", None), self.sdnq_dequantizer.weights_dtype) - - -def quantized_linear_forward(self, input: torch.FloatTensor) -> torch.FloatTensor: - return torch.nn.functional.linear(input, self.sdnq_dequantizer(self.weight), self.bias) - - -def get_conv_args(input_ndim: int, stride, padding, dilation): - if input_ndim == 3: - conv_type = 1 - elif input_ndim == 4: - conv_type = 2 - else: - conv_type = 3 - if isinstance(stride, int): - stride = (stride,) * conv_type - if isinstance(padding, int): - padding = (padding,) * conv_type - if isinstance(dilation, int): - dilation = (dilation,) * conv_type - if conv_type == 1: - stride = (1, stride[0]) - padding = (0, padding[0]) - dilation = (1, dilation[0]) - return conv_type, stride, padding, dilation - - -def quantized_conv_forward_fp8_matmul(self, input) -> torch.FloatTensor: - if torch.numel(input) / input.shape[2] < 32: - return self._conv_forward(input, self.sdnq_dequantizer(self.weight, skip_quantized_matmul=True), self.bias) - conv_type, stride, padding, dilation = get_conv_args(input.ndim, self.stride, self.padding, self.dilation) - return conv_fp8_matmul( - input, self.weight, self.bias, - self.sdnq_dequantizer.scale, - self.sdnq_dequantizer.result_shape, - self._reversed_padding_repeated_twice, - self.padding_mode, conv_type, - self.groups, stride, padding, dilation, - ) - - -def quantized_conv_forward_fp8_matmul_tensorwise(self, input) -> torch.FloatTensor: - if torch.numel(input) / input.shape[2] < 32: - return self._conv_forward(input, self.sdnq_dequantizer(self.weight, skip_quantized_matmul=True), self.bias) - conv_type, stride, padding, dilation = get_conv_args(input.ndim, self.stride, self.padding, self.dilation) - return conv_fp8_matmul_tensorwise( - input, self.weight, self.bias, - self.sdnq_dequantizer.scale, - self.sdnq_dequantizer.result_shape, - self._reversed_padding_repeated_twice, - self.padding_mode, conv_type, - self.groups, stride, padding, dilation, - ) - - -def quantized_conv_forward_int8_matmul(self, input) -> torch.FloatTensor: - if torch.numel(input) / input.shape[2] < 32: - return self._conv_forward(input, self.sdnq_dequantizer(self.weight, skip_quantized_matmul=True), self.bias) - conv_type, stride, padding, dilation = get_conv_args(input.ndim, self.stride, self.padding, self.dilation) - return conv_int8_matmul( - input, self.weight, self.bias, - self.sdnq_dequantizer.scale, - self.sdnq_dequantizer.result_shape, - getattr(self.sdnq_dequantizer, "quantized_weight_shape", None), - self.sdnq_dequantizer.weights_dtype, - self._reversed_padding_repeated_twice, - self.padding_mode, conv_type, - self.groups, stride, padding, dilation, - ) - - -def quantized_conv_forward(self, input) -> torch.FloatTensor: - return self._conv_forward(input, self.sdnq_dequantizer(self.weight), self.bias) - - -def quantized_conv_transpose_1d_forward(self, input: torch.FloatTensor, output_size: Optional[list[int]] = None) -> torch.FloatTensor: - output_padding = self._output_padding(input, output_size, self.stride, self.padding, self.kernel_size, 1, self.dilation) - return torch.nn.functional.conv_transpose1d(input, self.sdnq_dequantizer(self.weight), self.bias, self.stride, self.padding, output_padding, self.groups, self.dilation) - - -def quantized_conv_transpose_2d_forward(self, input: torch.FloatTensor, output_size: Optional[list[int]] = None) -> torch.FloatTensor: - output_padding = self._output_padding(input, output_size, self.stride, self.padding, self.kernel_size, 2, self.dilation) - return torch.nn.functional.conv_transpose2d(input, self.sdnq_dequantizer(self.weight), self.bias, self.stride, self.padding, output_padding, self.groups, self.dilation) - - -def quantized_conv_transpose_3d_forward(self, input: torch.FloatTensor, output_size: Optional[list[int]] = None) -> torch.FloatTensor: - output_padding = self._output_padding(input, output_size, self.stride, self.padding, self.kernel_size, 3, self.dilation) - return torch.nn.functional.conv_transpose3d(input, self.sdnq_dequantizer(self.weight), self.bias, self.stride, self.padding, output_padding, self.groups, self.dilation) - - -if shared.opts.sdnq_dequantize_compile: - try: - torch._dynamo.config.cache_size_limit = max(8192, torch._dynamo.config.cache_size_limit) - torch._dynamo.config.accumulated_recompile_limit = max(8192, torch._dynamo.config.accumulated_recompile_limit) - int8_matmul = torch.compile(int8_matmul, fullgraph=True, dynamic=False) - fp8_matmul = torch.compile(fp8_matmul, fullgraph=True, dynamic=False) - fp8_matmul_tensorwise = torch.compile(fp8_matmul_tensorwise, fullgraph=True, dynamic=False) - conv_int8_matmul = torch.compile(conv_int8_matmul, fullgraph=True, dynamic=False) - conv_fp8_matmul = torch.compile(conv_fp8_matmul, fullgraph=True, dynamic=False) - conv_fp8_matmul_tensorwise = torch.compile(conv_fp8_matmul_tensorwise, fullgraph=True, dynamic=False) - except Exception as e: - shared.log.warning(f"Quantization: type=sdnq MatMul using torch.compile is not available: {e}") diff --git a/modules/sdnq/layers/conv/conv_fp8.py b/modules/sdnq/layers/conv/conv_fp8.py new file mode 100644 index 000000000..f93d39519 --- /dev/null +++ b/modules/sdnq/layers/conv/conv_fp8.py @@ -0,0 +1,68 @@ +# pylint: disable=relative-beyond-top-level,redefined-builtin,protected-access + +from typing import List + +import torch + +from ...common import use_torch_compile # noqa: TID252 +from ..linear.linear_fp8 import quantize_fp8_matmul_input # noqa: TID252 +from .forward import get_conv_args, process_conv_input + + +def conv_fp8_matmul( + input: torch.FloatTensor, + weight: torch.Tensor, + bias: torch.FloatTensor, + scale: torch.FloatTensor, + result_shape: torch.Size, + reversed_padding_repeated_twice: List[int], + padding_mode: str, conv_type: int, + groups: int, stride: List[int], + padding: List[int], dilation: List[int], +) -> torch.FloatTensor: + return_dtype = input.dtype + input, mm_output_shape = process_conv_input(conv_type, input, reversed_padding_repeated_twice, padding_mode, result_shape, stride, padding, dilation) + input, input_scale = quantize_fp8_matmul_input(input) + + if groups == 1: + result = torch._scaled_mm(input, weight, scale_a=input_scale, scale_b=scale, bias=bias, out_dtype=return_dtype).reshape(mm_output_shape) + else: + scale = scale.reshape(groups, 1, scale.shape[1] // groups) + input_scale = input_scale.reshape(groups, input_scale.shape[0] // groups, 1) + weight = weight.reshape(weight.shape[0], groups, weight.shape[1] // groups).transpose(0,1) + input = input.reshape(input.shape[0], groups, input.shape[1] // groups).transpose(0,1) + result = [] + if bias is not None: + bias = bias.reshape(groups, bias.shape[0] // groups) + for i in range(groups): + result.append(torch._scaled_mm(input[i], weight[i], scale_a=input_scale[i], scale_b=scale[i], bias=bias[i], out_dtype=return_dtype)) + else: + for i in range(groups): + result.append(torch._scaled_mm(input[i], weight[i], scale_a=input_scale[i], scale_b=scale[i], bias=None, out_dtype=return_dtype)) + result = torch.cat(result, dim=-1).reshape(mm_output_shape) + + if conv_type == 1: + result = result.transpose(1,2) + elif conv_type == 2: + result = result.permute(0,3,1,2) + elif conv_type == 3: + result = result.permute(0,4,1,2,3) + return result + + +def quantized_conv_forward_fp8_matmul(self, input) -> torch.FloatTensor: + if torch.numel(input) / input.shape[2] < 32: + return self._conv_forward(input, self.sdnq_dequantizer(self.weight, skip_quantized_matmul=True), self.bias) + conv_type, stride, padding, dilation = get_conv_args(input.ndim, self.stride, self.padding, self.dilation) + return conv_fp8_matmul( + input, self.weight, self.bias, + self.sdnq_dequantizer.scale, + self.sdnq_dequantizer.result_shape, + self._reversed_padding_repeated_twice, + self.padding_mode, conv_type, + self.groups, stride, padding, dilation, + ) + + +if use_torch_compile: + conv_fp8_matmul = torch.compile(conv_fp8_matmul, fullgraph=True, dynamic=False) diff --git a/modules/sdnq/layers/conv/conv_fp8_tensorwise.py b/modules/sdnq/layers/conv/conv_fp8_tensorwise.py new file mode 100644 index 000000000..46b53a2d6 --- /dev/null +++ b/modules/sdnq/layers/conv/conv_fp8_tensorwise.py @@ -0,0 +1,67 @@ +# pylint: disable=relative-beyond-top-level,redefined-builtin,protected-access + +from typing import List + +import torch + +from ...common import use_torch_compile # noqa: TID252 +from ...dequantizer import dequantize_symmetric, dequantize_symmetric_with_bias # noqa: TID252 +from ..linear.linear_fp8_tensorwise import quantize_fp8_matmul_input_tensorwise # noqa: TID252 +from .forward import get_conv_args, process_conv_input + + +def conv_fp8_matmul_tensorwise( + input: torch.FloatTensor, + weight: torch.Tensor, + bias: torch.FloatTensor, + scale: torch.FloatTensor, + result_shape: torch.Size, + reversed_padding_repeated_twice: List[int], + padding_mode: str, conv_type: int, + groups: int, stride: List[int], + padding: List[int], dilation: List[int], +) -> torch.FloatTensor: + return_dtype = input.dtype + input, mm_output_shape = process_conv_input(conv_type, input, reversed_padding_repeated_twice, padding_mode, result_shape, stride, padding, dilation) + input, scale = quantize_fp8_matmul_input_tensorwise(input, scale) + dummy_input_scale = torch.ones(1, device=input.device, dtype=torch.float32) + + if groups == 1: + result = torch._scaled_mm(input, weight, scale_a=dummy_input_scale, scale_b=dummy_input_scale, bias=None, out_dtype=scale.dtype) + else: + weight = weight.reshape(weight.shape[0], groups, weight.shape[1] // groups).transpose(0,1) + input = input.reshape(input.shape[0], groups, input.shape[1] // groups).transpose(0,1) + result = [] + for i in range(groups): + result.append(torch._scaled_mm(input[i], weight[i], scale_a=dummy_input_scale, scale_b=dummy_input_scale, bias=None, out_dtype=scale.dtype)) + result = torch.cat(result, dim=-1) + if bias is not None: + dequantize_symmetric_with_bias(result, scale, bias, return_dtype, mm_output_shape) + else: + dequantize_symmetric(result, scale, return_dtype, mm_output_shape) + + if conv_type == 1: + result = result.transpose(1,2) + elif conv_type == 2: + result = result.permute(0,3,1,2) + elif conv_type == 3: + result = result.permute(0,4,1,2,3) + return result + + +def quantized_conv_forward_fp8_matmul_tensorwise(self, input) -> torch.FloatTensor: + if torch.numel(input) / input.shape[2] < 32: + return self._conv_forward(input, self.sdnq_dequantizer(self.weight, skip_quantized_matmul=True), self.bias) + conv_type, stride, padding, dilation = get_conv_args(input.ndim, self.stride, self.padding, self.dilation) + return conv_fp8_matmul_tensorwise( + input, self.weight, self.bias, + self.sdnq_dequantizer.scale, + self.sdnq_dequantizer.result_shape, + self._reversed_padding_repeated_twice, + self.padding_mode, conv_type, + self.groups, stride, padding, dilation, + ) + + +if use_torch_compile: + conv_fp8_matmul_tensorwise = torch.compile(conv_fp8_matmul_tensorwise, fullgraph=True, dynamic=False) diff --git a/modules/sdnq/layers/conv/conv_int8.py b/modules/sdnq/layers/conv/conv_int8.py new file mode 100644 index 000000000..02b553d15 --- /dev/null +++ b/modules/sdnq/layers/conv/conv_int8.py @@ -0,0 +1,73 @@ +# pylint: disable=relative-beyond-top-level,redefined-builtin,protected-access + +from typing import List + +import torch + +from ...common import use_torch_compile # noqa: TID252 +from ...packed_int import unpack_int_symetric # noqa: TID252 +from ...dequantizer import dequantize_symmetric, dequantize_symmetric_with_bias # noqa: TID252 +from ..linear.linear_int8 import quantize_int8_matmul_input # noqa: TID252 +from .forward import get_conv_args, process_conv_input + + +def conv_int8_matmul( + input: torch.FloatTensor, + weight: torch.CharTensor, + bias: torch.FloatTensor, + scale: torch.FloatTensor, + result_shape: torch.Size, + quantized_weight_shape: torch.Size, + weights_dtype: str, + reversed_padding_repeated_twice: List[int], + padding_mode: str, conv_type: int, + groups: int, stride: List[int], + padding: List[int], dilation: List[int], +) -> torch.FloatTensor: + return_dtype = input.dtype + input, mm_output_shape = process_conv_input(conv_type, input, reversed_padding_repeated_twice, padding_mode, result_shape, stride, padding, dilation) + input, scale = quantize_int8_matmul_input(input, scale) + if quantized_weight_shape is not None: + weight = unpack_int_symetric(weight, quantized_weight_shape, weights_dtype, dtype=torch.int8) + + if groups == 1: + result = torch._int_mm(input, weight) + else: + weight = weight.reshape(weight.shape[0], groups, weight.shape[1] // groups).transpose(0,1) + input = input.reshape(input.shape[0], groups, input.shape[1] // groups).transpose(0,1) + result = [] + for i in range(groups): + result.append(torch._int_mm(input[i], weight[i])) + result = torch.cat(result, dim=-1) + if bias is not None: + result = dequantize_symmetric_with_bias(result, scale, bias, return_dtype, mm_output_shape) + else: + result = dequantize_symmetric(result, scale, return_dtype, mm_output_shape) + + if conv_type == 1: + result = result.transpose(1,2) + elif conv_type == 2: + result = result.permute(0,3,1,2) + elif conv_type == 3: + result = result.permute(0,4,1,2,3) + return result + + +def quantized_conv_forward_int8_matmul(self, input) -> torch.FloatTensor: + if torch.numel(input) / input.shape[2] < 32: + return self._conv_forward(input, self.sdnq_dequantizer(self.weight, skip_quantized_matmul=True), self.bias) + conv_type, stride, padding, dilation = get_conv_args(input.ndim, self.stride, self.padding, self.dilation) + return conv_int8_matmul( + input, self.weight, self.bias, + self.sdnq_dequantizer.scale, + self.sdnq_dequantizer.result_shape, + getattr(self.sdnq_dequantizer, "quantized_weight_shape", None), + self.sdnq_dequantizer.weights_dtype, + self._reversed_padding_repeated_twice, + self.padding_mode, conv_type, + self.groups, stride, padding, dilation, + ) + + +if use_torch_compile: + conv_int8_matmul = torch.compile(conv_int8_matmul, fullgraph=True, dynamic=False) diff --git a/modules/sdnq/layers/conv/forward.py b/modules/sdnq/layers/conv/forward.py new file mode 100644 index 000000000..44c90b061 --- /dev/null +++ b/modules/sdnq/layers/conv/forward.py @@ -0,0 +1,93 @@ +# pylint: disable=relative-beyond-top-level,redefined-builtin,protected-access + +from typing import Optional + +import torch + + +def get_conv_args(input_ndim: int, stride, padding, dilation): + if input_ndim == 3: + conv_type = 1 + elif input_ndim == 4: + conv_type = 2 + else: + conv_type = 3 + if isinstance(stride, int): + stride = (stride,) * conv_type + if isinstance(padding, int): + padding = (padding,) * conv_type + if isinstance(dilation, int): + dilation = (dilation,) * conv_type + if conv_type == 1: + stride = (1, stride[0]) + padding = (0, padding[0]) + dilation = (1, dilation[0]) + return conv_type, stride, padding, dilation + + +def process_conv_input(conv_type, input, reversed_padding_repeated_twice, padding_mode, result_shape, stride, padding, dilation): + if conv_type == 1: + batch_size, _, L_in = input.shape + C_out, _, K_l = result_shape + L_out = (L_in + 2 * padding[1] - dilation[1] * (K_l - 1) - 1) // stride[1] + 1 + mm_output_shape = (batch_size, L_out, C_out) + kernel_size = (1, K_l) + if conv_type == 2: + batch_size, _, H_in, W_in = input.shape + C_out, _, K_h, K_w = result_shape + H_out = (H_in + 2 * padding[0] - dilation[0] * (K_h - 1) - 1) // stride[0] + 1 + W_out = (W_in + 2 * padding[1] - dilation[1] * (K_w - 1) - 1) // stride[1] + 1 + mm_output_shape = (batch_size, H_out, W_out, C_out) + kernel_size = (K_h, K_w) + else: + batch_size, _, D_in, H_in, W_in = input.shape + C_out, _, K_d, K_h, K_w = result_shape + D_out = (D_in + 2 * padding[0] - dilation[0] * (K_d - 1) - 1) // stride[0] + 1 + H_out = (H_in + 2 * padding[1] - dilation[1] * (K_h - 1) - 1) // stride[1] + 1 + W_out = (W_in + 2 * padding[2] - dilation[2] * (K_w - 1) - 1) // stride[2] + 1 + mm_output_shape = (batch_size, D_out, H_out, W_out, C_out) + kernel_size = (K_d, K_h, K_w) + + if padding_mode != "zeros": + input = torch.nn.functional.pad(input, reversed_padding_repeated_twice, mode=padding_mode) + padding = (0,) * (conv_type if conv_type != 1 else 2) + elif conv_type == 3: + input = torch.nn.functional.pad(input, reversed_padding_repeated_twice) + + if conv_type == 1: + input = input.unsqueeze(2) + + if conv_type == 3: + K_D_eff = K_d + (K_d - 1) * (dilation[0] - 1) + K_H_eff = K_h + (K_h - 1) * (dilation[0] - 1) + K_W_eff = K_w + (K_w - 1) * (dilation[0] - 1) + input = input.unfold(2, K_D_eff, stride[0]).unfold(3, K_H_eff, stride[1]).unfold(4, K_W_eff, stride[2]) + if dilation[0] > 1: + input = input[..., ::dilation[0], :, :] + if dilation[1] > 1: + input = input[..., ::dilation[1], :] + if dilation[2] > 1: + input = input[..., ::dilation[2]] + input = input.permute(0, 2, 3, 4, 1, 5, 6, 7).reshape(batch_size, D_out * H_out * W_out, -1) + else: + input = torch.nn.functional.unfold(input, kernel_size=kernel_size, padding=padding, stride=stride, dilation=dilation).transpose(1,2) + return input, mm_output_shape + + +def quantized_conv_forward(self, input) -> torch.FloatTensor: + return self._conv_forward(input, self.sdnq_dequantizer(self.weight), self.bias) + + +def quantized_conv_transpose_1d_forward(self, input: torch.FloatTensor, output_size: Optional[list[int]] = None) -> torch.FloatTensor: + output_padding = self._output_padding(input, output_size, self.stride, self.padding, self.kernel_size, 1, self.dilation) + return torch.nn.functional.conv_transpose1d(input, self.sdnq_dequantizer(self.weight), self.bias, self.stride, self.padding, output_padding, self.groups, self.dilation) + + +def quantized_conv_transpose_2d_forward(self, input: torch.FloatTensor, output_size: Optional[list[int]] = None) -> torch.FloatTensor: + output_padding = self._output_padding(input, output_size, self.stride, self.padding, self.kernel_size, 2, self.dilation) + return torch.nn.functional.conv_transpose2d(input, self.sdnq_dequantizer(self.weight), self.bias, self.stride, self.padding, output_padding, self.groups, self.dilation) + + +def quantized_conv_transpose_3d_forward(self, input: torch.FloatTensor, output_size: Optional[list[int]] = None) -> torch.FloatTensor: + output_padding = self._output_padding(input, output_size, self.stride, self.padding, self.kernel_size, 3, self.dilation) + return torch.nn.functional.conv_transpose3d(input, self.sdnq_dequantizer(self.weight), self.bias, self.stride, self.padding, output_padding, self.groups, self.dilation) diff --git a/modules/sdnq/layers/linear/forward.py b/modules/sdnq/layers/linear/forward.py new file mode 100644 index 000000000..20224204c --- /dev/null +++ b/modules/sdnq/layers/linear/forward.py @@ -0,0 +1,7 @@ +# pylint: disable=relative-beyond-top-level,redefined-builtin,protected-access + +import torch + + +def quantized_linear_forward(self, input: torch.FloatTensor) -> torch.FloatTensor: + return torch.nn.functional.linear(input, self.sdnq_dequantizer(self.weight), self.bias) diff --git a/modules/sdnq/layers/linear/linear_fp8.py b/modules/sdnq/layers/linear/linear_fp8.py new file mode 100644 index 000000000..f6da5c2ef --- /dev/null +++ b/modules/sdnq/layers/linear/linear_fp8.py @@ -0,0 +1,37 @@ +# pylint: disable=relative-beyond-top-level,redefined-builtin,protected-access + +from typing import Tuple + +import torch + +from ...common import use_torch_compile # noqa: TID252 + + +def quantize_fp8_matmul_input(input: torch.FloatTensor) -> Tuple[torch.Tensor, torch.FloatTensor]: + input = input.flatten(0,-2).contiguous().to(dtype=torch.float32) + input_scale = torch.amax(input.abs(), dim=-1, keepdims=True).div_(448) + input = torch.div(input, input_scale).clamp_(-448, 448).to(dtype=torch.float8_e4m3fn) + return input, input_scale + + +def fp8_matmul( + input: torch.FloatTensor, + weight: torch.Tensor, + bias: torch.FloatTensor, + scale: torch.FloatTensor, +) -> torch.FloatTensor: + return_dtype = input.dtype + output_shape = list(input.shape) + output_shape[-1] = weight.shape[-1] + input, input_scale = quantize_fp8_matmul_input(input) + return torch._scaled_mm(input, weight, scale_a=input_scale, scale_b=scale, bias=bias, out_dtype=return_dtype).reshape(output_shape) + + +def quantized_linear_forward_fp8_matmul(self, input: torch.FloatTensor) -> torch.FloatTensor: + if torch.numel(input) / input.shape[-1] < 32: + return torch.nn.functional.linear(input, self.sdnq_dequantizer(self.weight, skip_quantized_matmul=True), self.bias) + return fp8_matmul(input, self.weight, self.bias, self.sdnq_dequantizer.scale) + + +if use_torch_compile: + fp8_matmul = torch.compile(fp8_matmul, fullgraph=True, dynamic=False) diff --git a/modules/sdnq/layers/linear/linear_fp8_tensorwise.py b/modules/sdnq/layers/linear/linear_fp8_tensorwise.py new file mode 100644 index 000000000..3fdecaf57 --- /dev/null +++ b/modules/sdnq/layers/linear/linear_fp8_tensorwise.py @@ -0,0 +1,45 @@ +# pylint: disable=relative-beyond-top-level,redefined-builtin,protected-access + +from typing import Tuple + +import torch + +from ...common import use_torch_compile # noqa: TID252 +from ...dequantizer import dequantize_symmetric, dequantize_symmetric_with_bias # noqa: TID252 + + +def quantize_fp8_matmul_input_tensorwise(input: torch.FloatTensor, scale: torch.FloatTensor) -> Tuple[torch.Tensor, torch.FloatTensor]: + input = input.flatten(0,-2).contiguous().to(dtype=scale.dtype) + input_scale = torch.amax(input.abs(), dim=-1, keepdims=True).div_(448) + input = torch.div(input, input_scale).clamp_(-448, 448).to(dtype=torch.float8_e4m3fn) + scale = torch.mul(input_scale, scale) + if scale.dtype == torch.float16: # fp16 will overflow + scale = scale.to(dtype=torch.float32) + return input, scale + + +def fp8_matmul_tensorwise( + input: torch.FloatTensor, + weight: torch.Tensor, + bias: torch.FloatTensor, + scale: torch.FloatTensor, +) -> torch.FloatTensor: + return_dtype = input.dtype + output_shape = list(input.shape) + output_shape[-1] = weight.shape[-1] + dummy_input_scale = torch.ones(1, device=input.device, dtype=torch.float32) + input, scale = quantize_fp8_matmul_input_tensorwise(input, scale) + if bias is not None: + return dequantize_symmetric_with_bias(torch._scaled_mm(input, weight, scale_a=dummy_input_scale, scale_b=dummy_input_scale, bias=None, out_dtype=scale.dtype), scale, bias, return_dtype, output_shape) + else: + return dequantize_symmetric(torch._scaled_mm(input, weight, scale_a=dummy_input_scale, scale_b=dummy_input_scale, bias=None, out_dtype=scale.dtype), scale, return_dtype, output_shape) + + +def quantized_linear_forward_fp8_matmul_tensorwise(self, input: torch.FloatTensor) -> torch.FloatTensor: + if torch.numel(input) / input.shape[-1] < 32: + return torch.nn.functional.linear(input, self.sdnq_dequantizer(self.weight, skip_quantized_matmul=True), self.bias) + return fp8_matmul_tensorwise(input, self.weight, self.bias, self.sdnq_dequantizer.scale) + + +if use_torch_compile: + fp8_matmul_tensorwise = torch.compile(fp8_matmul_tensorwise, fullgraph=True, dynamic=False) diff --git a/modules/sdnq/layers/linear/linear_int8.py b/modules/sdnq/layers/linear/linear_int8.py new file mode 100644 index 000000000..3c0184056 --- /dev/null +++ b/modules/sdnq/layers/linear/linear_int8.py @@ -0,0 +1,49 @@ +# pylint: disable=relative-beyond-top-level,redefined-builtin,protected-access + +from typing import Tuple + +import torch + +from ...common import use_torch_compile # noqa: TID252 +from ...packed_int import unpack_int_symetric # noqa: TID252 +from ...dequantizer import dequantize_symmetric, dequantize_symmetric_with_bias # noqa: TID252 + + +def quantize_int8_matmul_input(input: torch.FloatTensor, scale: torch.FloatTensor) -> Tuple[torch.CharTensor, torch.FloatTensor]: + input = input.flatten(0,-2).contiguous().to(dtype=scale.dtype) + input_scale = torch.amax(input.abs(), dim=-1, keepdims=True).div_(127) + input = torch.div(input, input_scale).round_().clamp_(-128, 127).to(dtype=torch.int8) + scale = torch.mul(input_scale, scale) + if scale.dtype == torch.float16: # fp16 will overflow + scale = scale.to(dtype=torch.float32) + return input, scale + + +def int8_matmul( + input: torch.FloatTensor, + weight: torch.Tensor, + bias: torch.FloatTensor, + scale: torch.FloatTensor, + quantized_weight_shape: torch.Size, + weights_dtype: str, +) -> torch.FloatTensor: + if quantized_weight_shape is not None: + weight = unpack_int_symetric(weight, quantized_weight_shape, weights_dtype, dtype=torch.int8) + return_dtype = input.dtype + output_shape = list(input.shape) + output_shape[-1] = weight.shape[-1] + input, scale = quantize_int8_matmul_input(input, scale) + if bias is not None: + return dequantize_symmetric_with_bias(torch._int_mm(input, weight), scale, bias, return_dtype, output_shape) + else: + return dequantize_symmetric(torch._int_mm(input, weight), scale, return_dtype, output_shape) + + +def quantized_linear_forward_int8_matmul(self, input: torch.FloatTensor) -> torch.FloatTensor: + if torch.numel(input) / input.shape[-1] < 32: + return torch.nn.functional.linear(input, self.sdnq_dequantizer(self.weight, skip_quantized_matmul=True), self.bias) + return int8_matmul(input, self.weight, self.bias, self.sdnq_dequantizer.scale, getattr(self.sdnq_dequantizer, "quantized_weight_shape", None), self.sdnq_dequantizer.weights_dtype) + + +if use_torch_compile: + int8_matmul = torch.compile(int8_matmul, fullgraph=True, dynamic=False) diff --git a/modules/sdnq/packed_int.py b/modules/sdnq/packed_int.py index b20c61818..d946ef4ba 100644 --- a/modules/sdnq/packed_int.py +++ b/modules/sdnq/packed_int.py @@ -1,6 +1,7 @@ # pylint: disable=redefined-builtin,no-member,protected-access from typing import Optional + import torch from .common import dtype_dict @@ -10,13 +11,18 @@ def pack_int_symetric(tensor: torch.CharTensor, weights_dtype: str) -> torch.Byt return packed_int_function_dict[weights_dtype]["pack"](tensor.sub_(dtype_dict[weights_dtype]["min"]).to(dtype=dtype_dict[weights_dtype]["storage_dtype"])) -def unpack_int_symetric(packed_tensor: torch.ByteTensor, shape: torch.Size, weights_dtype: str, dtype: Optional[torch.dtype] = None, transpose: Optional[bool] = False) -> torch.CharTensor: +def pack_int_asymetric(tensor: torch.CharTensor, weights_dtype: str) -> torch.ByteTensor: + return packed_int_function_dict[weights_dtype]["pack"](tensor.to(dtype=dtype_dict[weights_dtype]["storage_dtype"])) + + +def unpack_int_symetric(packed_tensor: torch.ByteTensor, shape: torch.Size, weights_dtype: str, dtype: Optional[torch.dtype] = None) -> torch.CharTensor: if dtype is None: dtype = dtype_dict[weights_dtype]["torch_dtype"] - result = packed_int_function_dict[weights_dtype]["unpack"](packed_tensor, shape).to(dtype=dtype).add_(dtype_dict[weights_dtype]["min"]) - if transpose: - result = result.transpose(0,1) - return result + return packed_int_function_dict[weights_dtype]["unpack"](packed_tensor, shape).to(dtype=dtype).add_(dtype_dict[weights_dtype]["min"]) + + +def unpack_int_asymetric(packed_tensor: torch.ByteTensor, shape: torch.Size, weights_dtype: str) -> torch.CharTensor: + return packed_int_function_dict[weights_dtype]["unpack"](packed_tensor, shape) def pack_uint7(tensor: torch.ByteTensor) -> torch.ByteTensor: diff --git a/modules/shared.py b/modules/shared.py index 6fed4d343..2a366248d 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -24,6 +24,11 @@ import modules.paths as paths from installer import log, print_dict, console, get_version # pylint: disable=unused-import +class Backend(Enum): + ORIGINAL = 1 + DIFFUSERS = 2 + + errors.install([gr]) demo: gr.Blocks = None api = None @@ -57,17 +62,12 @@ restricted_opts = { "outdir_init_images" } resize_modes = ["None", "Fixed", "Crop", "Fill", "Outpaint", "Context aware"] -max_workers = 8 -default_hfcache_dir = os.environ.get("SD_HFCACHEDIR", None) or os.path.join(os.path.expanduser('~'), '.cache', 'huggingface', 'hub') +max_workers = 12 sdnq_quant_modes = ["int8", "float8_e4m3fn", "int7", "int6", "int5", "uint4", "uint3", "uint2", "float8_e5m2", "float8_e4m3fnuz", "float8_e5m2fnuz", "uint8", "uint7", "uint6", "uint5", "int4", "int3", "int2", "uint1"] +default_hfcache_dir = os.environ.get("SD_HFCACHEDIR", None) or os.path.join(paths.models_path, 'huggingface') state = shared_state.State() -class Backend(Enum): - ORIGINAL = 1 - DIFFUSERS = 2 - - # early select backend backend = Backend.DIFFUSERS if not hasattr(cmd_opts, "use_openvino"): @@ -142,10 +142,11 @@ options_templates.update(options_section(('sd', "Model Loading"), { "advanced_sep": OptionInfo("

Advanced Options

", "", gr.HTML), "sd_checkpoint_autoload": OptionInfo(True, "Model auto-load on start"), + "sd_parallel_load": OptionInfo(True, "Model load using multiple threads"), "sd_checkpoint_autodownload": OptionInfo(True, "Model auto-download on demand"), "stream_load": OptionInfo(False, "Model load using streams", gr.Checkbox), "diffusers_to_gpu": OptionInfo(False, "Model load model direct to GPU"), - "diffusers_eval": OptionInfo(True, "Force model eval", gr.Checkbox, {"visible": True }), + "diffusers_eval": OptionInfo(False, "Force model eval", gr.Checkbox, {"visible": True }), "device_map": OptionInfo('default', "Model load device map", gr.Radio, {"choices": ['default', 'gpu', 'cpu'] }), "disable_accelerate": OptionInfo(False, "Disable accelerate", gr.Checkbox, {"visible": False }), "sd_checkpoint_cache": OptionInfo(0, "Cached models", gr.Slider, {"minimum": 0, "maximum": 10, "step": 1, "visible": False }), @@ -248,6 +249,7 @@ options_templates.update(options_section(('text_encoder', "Text Encoder"), { "diffusers_zeros_prompt_pad": OptionInfo(False, "Use zeros for prompt padding", gr.Checkbox), "te_hijack": OptionInfo(True, "Offload after prompt encode", gr.Checkbox), "te_optional_sep": OptionInfo("

Optional

", "", gr.HTML), + "te_shared_t5": OptionInfo(True, "T5: Use shared instance of text encoder"), "te_pooled_embeds": OptionInfo(False, "SDXL: Use weighted pooled embeds"), "te_complex_human_instruction": OptionInfo(True, "Sana: Use complex human instructions"), "te_use_mask": OptionInfo(True, "Lumina: Use mask in transformers"), @@ -264,7 +266,7 @@ options_templates.update(options_section(('cuda', "Compute Settings"), { "diffusers_generator_device": OptionInfo("GPU", "Generator device", gr.Radio, {"choices": ["GPU", "CPU", "Unset"]}), "cross_attention_sep": OptionInfo("

Cross Attention

", "", gr.HTML), - "cross_attention_optimization": OptionInfo(startup_cross_attention, "Attention optimization method", gr.Radio, lambda: {"choices": shared_items.list_crossattention(native)}), + "cross_attention_optimization": OptionInfo(startup_cross_attention, "Attention optimization method", gr.Radio, lambda: {"choices": shared_items.list_crossattention()}), "sdp_options": OptionInfo(startup_sdp_options, "SDP options", gr.CheckboxGroup, {"choices": ['Flash attention', 'Memory attention', 'Math attention', 'Dynamic attention', 'CK Flash attention', 'Sage attention']}), "xformers_options": OptionInfo(['Flash attention'], "xFormers options", gr.CheckboxGroup, {"choices": ['Flash attention'] }), "dynamic_attention_slice_rate": OptionInfo(0.5, "Dynamic Attention slicing rate in GB", gr.Slider, {"minimum": 0.01, "maximum": max(gpu_memory,4), "step": 0.01}), @@ -424,7 +426,6 @@ options_templates.update(options_section(('system-paths', "System Paths"), { "realesrgan_models_path": OptionInfo(os.path.join(paths.models_path, 'RealESRGAN'), "Folder with RealESRGAN models", folder=True), "scunet_models_path": OptionInfo(os.path.join(paths.models_path, 'SCUNet'), "Folder with SCUNet models", folder=True), "swinir_models_path": OptionInfo(os.path.join(paths.models_path, 'SwinIR'), "Folder with SwinIR models", folder=True), - "ldsr_models_path": OptionInfo(os.path.join(paths.models_path, 'LDSR'), "Folder with LDSR models", folder=True), "clip_models_path": OptionInfo(os.path.join(paths.models_path, 'CLIP'), "Folder with CLIP models", folder=True), "other_paths_sep_options": OptionInfo("

Cache folders

", "", gr.HTML), "clean_temp_dir_at_start": OptionInfo(True, "Cleanup temporary folder on startup"), @@ -514,26 +515,44 @@ options_templates.update(options_section(('image-metadata', "Image Metadata"), { })) options_templates.update(options_section(('ui', "User Interface"), { + "themes_sep_ui": OptionInfo("

Theme options

", "", gr.HTML), "theme_type": OptionInfo("Standard", "Theme type", gr.Radio, {"choices": ["Modern", "Standard", "None"]}), "theme_style": OptionInfo("Auto", "Theme mode", gr.Radio, {"choices": ["Auto", "Dark", "Light"]}), "gradio_theme": OptionInfo("black-teal", "UI theme", gr.Dropdown, lambda: {"choices": theme.list_themes()}, refresh=theme.refresh_themes), - "ui_locale": OptionInfo("Auto", "UI locale", gr.Dropdown, lambda: {"choices": theme.list_locales()}), - "subpath": OptionInfo("", "Mount URL subpath"), + + "quicksetting_sep_images": OptionInfo("

Quicksettings

", "", gr.HTML), + "quicksettings_list": OptionInfo(["sd_model_checkpoint"], "Quicksettings list", gr.Dropdown, lambda: {"multiselect":True, "choices": opts.list()}), + + "server_sep_ui": OptionInfo("

Startup & Server Options

", "", gr.HTML), "autolaunch": OptionInfo(False, "Autolaunch browser upon startup"), - "font_size": OptionInfo(14, "Font size", gr.Slider, {"minimum": 8, "maximum": 32, "step": 1}), - "aspect_ratios": OptionInfo("1:1, 4:3, 3:2, 16:9, 16:10, 21:9, 2:3, 3:4, 9:16, 10:16, 9:21", "Allowed aspect ratios"), - "logmonitor_show": OptionInfo(True, "Show log view"), - "logmonitor_refresh_period": OptionInfo(5000, "Log view update period", gr.Slider, {"minimum": 0, "maximum": 30000, "step": 25}), - "ui_request_timeout": OptionInfo(30000, "UI request timeout", gr.Slider, {"minimum": 1000, "maximum": 120000, "step": 10}), "motd": OptionInfo(False, "Show MOTD"), + "subpath": OptionInfo("", "Mount URL subpath"), + "ui_request_timeout": OptionInfo(120000, "UI request timeout", gr.Slider, {"minimum": 1000, "maximum": 300000, "step": 10}), + + "cards_sep_ui": OptionInfo("

Card options

", "", gr.HTML), + "extra_networks_card_size": OptionInfo(140, "UI card size (px)", gr.Slider, {"minimum": 20, "maximum": 2000, "step": 1}), + "extra_networks_card_cover": OptionInfo("sidebar", "UI position", gr.Radio, {"choices": ["cover", "inline", "sidebar"]}), + "extra_networks_card_square": OptionInfo(True, "UI disable variable aspect ratio"), + + "other_sep_ui": OptionInfo("

Other...

", "", gr.HTML), + "ui_locale": OptionInfo("Auto", "UI locale", gr.Dropdown, lambda: {"choices": theme.list_locales()}), + "font_size": OptionInfo(14, "Font size", gr.Slider, {"minimum": 8, "maximum": 32, "step": 1}), + "gpu_monitor": OptionInfo(3000, "GPU monitor interval", gr.Slider, {"minimum": 100, "maximum": 60000, "step": 100}), + "aspect_ratios": OptionInfo("1:1, 4:3, 3:2, 16:9, 16:10, 21:9, 2:3, 3:4, 9:16, 10:16, 9:21", "Allowed aspect ratios"), "compact_view": OptionInfo(False, "Compact view"), "ui_columns": OptionInfo(4, "Gallery view columns", gr.Slider, {"minimum": 1, "maximum": 8, "step": 1}), + + "images_sep_log": OptionInfo("

Log Display

", "", gr.HTML), + "logmonitor_show": OptionInfo(True, "Show log view"), + "logmonitor_refresh_period": OptionInfo(5000, "Log view update period", gr.Slider, {"minimum": 0, "maximum": 30000, "step": 25}), + + "images_sep_ui": OptionInfo("

Outputs & Images

", "", gr.HTML), "return_grid": OptionInfo(True, "Show grid in results"), "return_mask": OptionInfo(False, "Inpainting include greyscale mask in results"), "return_mask_composite": OptionInfo(False, "Inpainting include masked composite in results"), "send_seed": OptionInfo(True, "Send seed when sending prompt or image to other interface", gr.Checkbox, {"visible": False}), "send_size": OptionInfo(False, "Send size when sending prompt or image to another interface", gr.Checkbox, {"visible": False}), - "quicksettings_list": OptionInfo(["sd_model_checkpoint"], "Quicksettings list", gr.Dropdown, lambda: {"multiselect":True, "choices": opts.list()}), + })) options_templates.update(options_section(('live-preview', "Live Previews"), { @@ -642,11 +661,8 @@ options_templates.update(options_section(('extra_networks', "Networks"), { "extra_networks": OptionInfo(["All"], "Available networks", gr.Dropdown, lambda: {"multiselect":True, "choices": ['All'] + [en.title for en in extra_networks]}), "extra_networks_sort": OptionInfo("Default", "Sort order", gr.Dropdown, {"choices": ['Default', 'Name [A-Z]', 'Name [Z-A]', 'Date [Newest]', 'Date [Oldest]', 'Size [Largest]', 'Size [Smallest]']}), "extra_networks_view": OptionInfo("gallery", "UI view", gr.Radio, {"choices": ["gallery", "list"]}), - "extra_networks_card_cover": OptionInfo("sidebar", "UI position", gr.Radio, {"choices": ["cover", "inline", "sidebar"]}), - "extra_networks_height": OptionInfo(0, "UI height (%)", gr.Slider, {"minimum": 0, "maximum": 100, "step": 1}), # set in ui_javascript "extra_networks_sidebar_width": OptionInfo(35, "UI sidebar width (%)", gr.Slider, {"minimum": 10, "maximum": 80, "step": 1}), - "extra_networks_card_size": OptionInfo(140, "UI card size (px)", gr.Slider, {"minimum": 20, "maximum": 2000, "step": 1}), - "extra_networks_card_square": OptionInfo(True, "UI disable variable aspect ratio"), + "extra_networks_height": OptionInfo(0, "UI height (%)", gr.Slider, {"minimum": 0, "maximum": 100, "step": 1}), # set in ui_javascript "extra_networks_fetch": OptionInfo(True, "UI fetch network info on mouse-over"), "extra_network_skip_indexing": OptionInfo(False, "Build info on first access", gr.Checkbox), @@ -680,7 +696,7 @@ options_templates.update(options_section(('extra_networks', "Networks"), { "wildcards_enabled": OptionInfo(True, "Enable file wildcards support"), })) -options_templates.update(options_section((None, "Hidden options"), { +options_templates.update(options_section(('hidden_options', "Hidden options"), { # internal options "diffusers_version": OptionInfo("", "Diffusers version", gr.Textbox, {"visible": False}), "disabled_extensions": OptionInfo([], "Disable these extensions", gr.Textbox, {"visible": False}), diff --git a/modules/shared_items.py b/modules/shared_items.py index 483186905..260ad2f2a 100644 --- a/modules/shared_items.py +++ b/modules/shared_items.py @@ -42,9 +42,10 @@ pipelines = { 'UniDiffuser': getattr(diffusers, 'UniDiffuserPipeline', None), 'Amused': getattr(diffusers, 'AmusedPipeline', None), 'HiDream': getattr(diffusers, 'HiDreamImagePipeline', None), - 'OmniGenPipeline': getattr(diffusers, 'OmniGenPipeline', None), + 'OmniGen': getattr(diffusers, 'OmniGenPipeline', None), 'Cosmos': getattr(diffusers, 'Cosmos2TextToImagePipeline', None), 'WanAI': getattr(diffusers, 'WanPipeline', None), + 'Qwen': getattr(diffusers, 'QwenImagePipeline', None), # dynamically imported and redefined later 'Meissonic': getattr(diffusers, 'DiffusionPipeline', None), @@ -105,26 +106,15 @@ def refresh_te_list(): modules.model_te.refresh_te_list() -def list_crossattention(native:bool=True): - if native: - return [ - "Disabled", - "Scaled-Dot-Product", - "xFormers", - "Batch matrix-matrix", - "Split attention", - "Dynamic Attention BMM" - ] - else: - return [ - "Disabled", - "Scaled-Dot-Product", - "xFormers", - "Doggettx's", - "InvokeAI's", - "Sub-quadratic", - "Split attention" - ] +def list_crossattention(): + return [ + "Disabled", + "Scaled-Dot-Product", + "xFormers", + "Batch matrix-matrix", + "Split attention", + "Dynamic Attention BMM" + ] def get_pipelines(): from installer import log diff --git a/modules/shared_legacy.py b/modules/shared_legacy.py index 032220c23..d83096675 100644 --- a/modules/shared_legacy.py +++ b/modules/shared_legacy.py @@ -9,8 +9,9 @@ class LegacyOption(OptionInfo): super().__init__(*args, **kwargs) -legacy_options = options_section((None, "Legacy options"), { - "interrogate_clip_skip_categories": LegacyOption(["artists", "movements", "flavors"], "CLiP: skip categories", gr.CheckboxGroup, lambda: {"choices": []}, visible=False), +legacy_options = options_section(('legacy_options', "Legacy options"), { + "ldsr_models_path": LegacyOption(os.path.join(paths.models_path, 'LDSR'), "LDSR Path", gr.Textbox, { "visible": False}), + "interrogate_clip_skip_categories": LegacyOption(["artists", "movements", "flavors"], "CLiP: skip categories", gr.CheckboxGroup, {"choices": [], "visible":False}), "lora_legacy": LegacyOption(False, "LoRA load using legacy method", gr.Checkbox, {"visible": False}), "lora_preferred_name": LegacyOption("filename", "LoRA preferred name", gr.Radio, {"choices": ["filename", "alias"], "visible": False}), "img2img_extra_noise": LegacyOption(0.0, "Extra noise multiplier for img2img", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01, "visible": False}), diff --git a/modules/shared_state.py b/modules/shared_state.py index c7c1cb472..616e92cce 100644 --- a/modules/shared_state.py +++ b/modules/shared_state.py @@ -50,6 +50,9 @@ class State: server_start = time.time() oom = False + def __init__(self): + log.debug(f'State initialized: id={id(self)}') + def __str__(self) -> str: status = ' ' status += 'skipped ' if self.skipped else '' @@ -208,7 +211,8 @@ class State: log.trace(f'State end: {self}') self.time_end = time.time() self.history('end') - self.job = "" + self.id = '' + self.job = '' self.job_count = 0 self.job_no = 0 self.frame_count = 0 diff --git a/modules/textual_inversion.py b/modules/textual_inversion.py index f6b1c558d..4d7b76a77 100644 --- a/modules/textual_inversion.py +++ b/modules/textual_inversion.py @@ -13,7 +13,7 @@ supported_models = ['ldm', 'sd', 'sdxl'] def list_embeddings(*dirs): - is_ext = extension_filter(['.SAFETENSORS', '.PT' ] + ( ['.PNG', '.WEBP', '.JXL', '.AVIF', '.BIN' ] if not shared.native else [] )) + is_ext = extension_filter(['.SAFETENSORS', '.PT' ]) is_not_preview = lambda fp: not next(iter(os.path.splitext(fp))).upper().endswith('.PREVIEW') # pylint: disable=unnecessary-lambda-assignment return list(filter(lambda fp: is_ext(fp) and is_not_preview(fp) and os.stat(fp).st_size > 0, directory_files(*dirs))) diff --git a/modules/timer.py b/modules/timer.py index 59c6a1de3..d94735902 100644 --- a/modules/timer.py +++ b/modules/timer.py @@ -72,3 +72,4 @@ startup = Timer() process = Timer() launch = Timer() init = Timer() +load = Timer() diff --git a/modules/ui_common.py b/modules/ui_common.py index 33a76e850..292adcdef 100644 --- a/modules/ui_common.py +++ b/modules/ui_common.py @@ -19,10 +19,14 @@ def gr_show(visible=True): def update_generation_info(generation_info, html_info, img_index): try: - generation_info = json.loads(generation_info) - if img_index < 0 or img_index >= len(generation_info["infotexts"]): - return html_info, generation_info - info = generation_info["infotexts"][img_index] + generation_json = json.loads(generation_info) + if len(generation_json["infotexts"]) == 0: + return html_info, 'no infotexts found' + if img_index == -1: + img_index = 0 + if img_index >= len(generation_json["infotexts"]): + return html_info, 'error fetching infotext' + info = generation_json["infotexts"][img_index] html_info_formatted = infotext_to_html(info) return html_info, html_info_formatted except Exception: @@ -67,7 +71,7 @@ def delete_files(js_data, files, all_files, index): for _image_index, filedata in enumerate(files, start_index): try: fn = filedata['name'] - if os.path.isfile(fn): + if os.path.exists(fn) and os.path.isfile(fn): deleted.append(fn) os.remove(fn) if fn in all_files: @@ -75,11 +79,11 @@ def delete_files(js_data, files, all_files, index): shared.log.info(f'Delete: image="{fn}"') base, _ext = os.path.splitext(fn) desc = f'{base}.txt' - if os.path.exists(desc): + if os.path.exists(desc) and os.path.isfile(desc): os.remove(desc) shared.log.info(f'Delete: text="{fn}"') except Exception as e: - shared.log.error(f'Delete: image="{fn}" {e}') + shared.log.error(f'Delete: file="{fn}" {e}') deleted = ', '.join(deleted) if len(deleted) > 0 else 'none' return all_files, plaintext_to_html(f"Deleted: {deleted}", ['performance']) @@ -247,6 +251,8 @@ def create_output_panel(tabname, preview=True, prompt=None, height=None, transfe ) if prompt is not None: ui_sections.create_interrogate_button(tab=tabname, inputs=result_gallery, outputs=prompt, what='output') + button_image_fit = gr.Button(ui_symbols.resize, elem_id=f"{tabname}_image_fit", elem_classes=['image-fit']) + button_image_fit.click(fn=None, _js="cycleImageFit", inputs=[], outputs=[]) with gr.Column(elem_id=f"{tabname}_footer", elem_classes="gallery_footer"): dummy_component = gr.Label(visible=False) diff --git a/modules/ui_control.py b/modules/ui_control.py index 8d048da97..cfcd191ed 100644 --- a/modules/ui_control.py +++ b/modules/ui_control.py @@ -144,9 +144,9 @@ def create_ui(_blocks: gr.Blocks=None): with gr.Accordion(open=False, label="Size", elem_id="control_size", elem_classes=["small-accordion"]): with gr.Tabs(): - with gr.Tab('Before'): + with gr.Tab('Initial'): resize_mode_before, resize_name_before, resize_context_before, width_before, height_before, scale_by_before, selected_scale_tab_before = ui_sections.create_resize_inputs('control_before', [], accordion=False, latent=True, prefix='before') - with gr.Tab('After'): + with gr.Tab('Post'): resize_mode_after, resize_name_after, resize_context_after, width_after, height_after, scale_by_after, selected_scale_tab_after = ui_sections.create_resize_inputs('control_after', [], accordion=False, latent=False, prefix='after') with gr.Tab('Mask'): resize_mode_mask, resize_name_mask, resize_context_mask, width_mask, height_mask, scale_by_mask, selected_scale_tab_mask = ui_sections.create_resize_inputs('control_mask', [], accordion=False, latent=False, prefix='mask') @@ -215,7 +215,7 @@ def create_ui(_blocks: gr.Blocks=None): gr.HTML('Output

') with gr.Tabs(elem_classes=['control-tabs'], elem_id='control-tab-output') as output_tabs: with gr.Tab('Gallery', id='out-gallery'): - output_gallery, _output_gen_info, _output_html_info, _output_html_info_formatted, output_html_log = ui_common.create_output_panel("control", preview=True, prompt=prompt, height=gr_height) + output_gallery, _output_gen_info, _output_html_info, _output_html_info_formatted, output_html_log = ui_common.create_output_panel("control", preview=False, prompt=prompt, height=gr_height) with gr.Tab('Image', id='out-image'): output_image = gr.Image(label="Output", show_label=False, type="pil", interactive=False, tool="editor", height=gr_height, elem_id='control_output_image', elem_classes=['control-image']) with gr.Tab('Video', id='out-video'): diff --git a/modules/ui_extensions.py b/modules/ui_extensions.py index 543f4b481..1db079b5d 100644 --- a/modules/ui_extensions.py +++ b/modules/ui_extensions.py @@ -279,6 +279,7 @@ def search_extensions(search_text, sort_column): def create_html(search_text, sort_column): # shared.log.debug(f'Extensions manager: refresh list search="{search_text}" sort="{sort_column}"') code = """ +
@@ -421,7 +422,7 @@ def create_html(search_text, sort_column): """ - code += "
{version_code} {install_code}
" + code += "
" shared.log.debug(f'Extension list: processed={stats["processed"]} installed={stats["installed"]} enabled={stats["enabled"]} disabled={stats["installed"] - stats["enabled"]} visible={stats["processed"] - stats["hidden"]} hidden={stats["hidden"]}') return code @@ -438,9 +439,10 @@ def create_ui(): uninstall_extension_button = gr.Button(elem_id="uninstall_extension_button", visible=False) update_extension_button = gr.Button(elem_id="update_extension_button", visible=False) with gr.Column(scale=4): - search_text = gr.Textbox(label="Search") - with gr.Column(scale=1): - sort_column = gr.Dropdown(value="default", label="Sort by", choices=list(sort_ordering.keys()), multiselect=False) + with gr.Row(): + search_text = gr.Textbox(label="Search") + with gr.Row(): + sort_column = gr.Dropdown(value="default", label="Sort by", choices=list(sort_ordering.keys()), multiselect=False) with gr.Column(scale=1): refresh_extensions_button = gr.Button(value="Refresh extension list", variant="primary") check = gr.Button(value="Update all installed", variant="primary") diff --git a/modules/ui_extra_networks.py b/modules/ui_extra_networks.py index 7787ea1fc..06b3428fb 100644 --- a/modules/ui_extra_networks.py +++ b/modules/ui_extra_networks.py @@ -8,7 +8,6 @@ import html import base64 import urllib.parse import threading -from datetime import datetime from types import SimpleNamespace from pathlib import Path from html.parser import HTMLParser @@ -16,9 +15,7 @@ from collections import OrderedDict import gradio as gr from PIL import Image from starlette.responses import FileResponse, JSONResponse -from modules import paths, shared, files_cache, errors, infotext -from modules.ui_components import ToolButton -import modules.ui_symbols as symbols +from modules import paths, shared, files_cache, errors, infotext, ui_symbols, ui_components, modelstats allowed_dirs = [] @@ -27,7 +24,7 @@ extra_pages = shared.extra_networks debug = shared.log.trace if os.environ.get('SD_EN_DEBUG', None) is not None else lambda *args, **kwargs: None debug('Trace: EN') card_full = ''' -
+
{title}
@@ -41,7 +38,7 @@ card_full = '''
''' card_list = ''' -
+
🛈 
{title}  @@ -70,7 +67,7 @@ def init_api(): return FileResponse(filename, headers={"Accept-Ranges": "bytes"}) def get_metadata(page: str = "", item: str = ""): - page = next(iter([x for x in shared.extra_networks if x.name == page]), None) + page = next(iter([x for x in shared.extra_networks if x.name.lower() == page.lower()]), None) if page is None: return JSONResponse({ 'metadata': 'none' }) metadata = page.metadata.get(item, 'none') @@ -80,10 +77,10 @@ def init_api(): return JSONResponse({"metadata": metadata}) def get_info(page: str = "", item: str = ""): - page = next(iter([x for x in get_pages() if x.name == page]), None) + page = next(iter([x for x in get_pages() if x.name.lower() == page.lower()]), None) if page is None: return JSONResponse({ 'info': 'none' }) - item = next(iter([x for x in page.items if x['name'] == item]), None) + item = next(iter([x for x in page.items if x['name'].lower() == item.lower()]), None) if item is None: return JSONResponse({ 'info': 'none' }) info = page.find_info(item.get('filename', None) or item.get('name', None)) @@ -93,10 +90,10 @@ def init_api(): return JSONResponse({"info": info}) def get_desc(page: str = "", item: str = ""): - page = next(iter([x for x in get_pages() if x.name == page]), None) + page = next(iter([x for x in get_pages() if x.name.lower() == page.lower()]), None) if page is None: return JSONResponse({ 'description': 'none' }) - item = next(iter([x for x in page.items if x['name'] == item]), None) + item = next(iter([x for x in page.items if x['name'].lower() == item.lower()]), None) if item is None: return JSONResponse({ 'description': 'none' }) desc = page.find_description(item.get('filename', None) or item.get('name', None)) @@ -105,10 +102,21 @@ def init_api(): # shared.log.debug(f"Networks desc: page='{page.name}' item={item['name']} len={len(desc)}") return JSONResponse({"description": desc}) - shared.api.add_api_route("/sd_extra_networks/thumb", fetch_file, methods=["GET"]) - shared.api.add_api_route("/sd_extra_networks/metadata", get_metadata, methods=["GET"]) - shared.api.add_api_route("/sd_extra_networks/info", get_info, methods=["GET"]) - shared.api.add_api_route("/sd_extra_networks/description", get_desc, methods=["GET"]) + def get_network(page: str = "", item: str = ""): + page = next(iter([x for x in get_pages() if x.name.lower() == page.lower()]), None) + if page is None: + return JSONResponse({ 'page': 'none' }) + item = next(iter([x for x in page.items if (x['alias'].lower() == item.lower() or x['name'].lower() == item.lower())]), None) + if item is None: + return JSONResponse({ 'item': 'none' }) + return JSONResponse(item) + + + shared.api.add_api_route("/sdapi/v1/network", get_network, methods=["GET"]) + shared.api.add_api_route("/sdapi/v1/network/thumb", fetch_file, methods=["GET"]) + shared.api.add_api_route("/sdapi/v1/network/metadata", get_metadata, methods=["GET"]) + shared.api.add_api_route("/sdapi/v1/network/info", get_info, methods=["GET"]) + shared.api.add_api_route("/sdapi/v1/network/desc", get_desc, methods=["GET"]) class ExtraNetworksPage: @@ -131,6 +139,9 @@ class ExtraNetworksPage: self.view = shared.opts.extra_networks_view self.card = card_full if shared.opts.extra_networks_view == 'gallery' else card_list + def __str__(self): + return f'Page(title="{self.title}" name="{self.name}" items={len(self.items)})' + def refresh(self): pass @@ -159,7 +170,7 @@ class ExtraNetworksPage: def link_preview(self, filename): quoted_filename = urllib.parse.quote(filename.replace('\\', '/')) mtime = os.path.getmtime(filename) if os.path.exists(filename) else 0 - preview = f"./sd_extra_networks/thumb?filename={quoted_filename}&mtime={mtime}" + preview = f"/sdapi/v1/network/thumb?filename={quoted_filename}&mtime={mtime}" return preview def create_thumb(self): @@ -311,12 +322,14 @@ class ExtraNetworksPage: return '#{:02x}{:02x}{:02x}'.format(r, g, b) # pylint: disable=consider-using-f-string try: + onclick = f'cardClicked({item.get("prompt", None)})' args = { - "tabname": tabname, + # "tabname": tabname, "page": self.name, "name": item.get('name', ''), "title": os.path.basename(item["name"].replace('_', ' ')), "filename": item.get('filename', ''), + "short": os.path.splitext(os.path.basename(item.get('filename', '')))[0], "tags": '|'.join([item.get('tags')] if isinstance(item.get('tags', {}), str) else list(item.get('tags', {}).keys())), "preview": html.escape(item.get('preview', None) or self.link_preview('html/card-no-preview.png')), "width": 'var(--card-size)', @@ -325,16 +338,16 @@ class ExtraNetworksPage: "prompt": item.get("prompt", None), "search": item.get("search_term", ""), "description": item.get("description") or "", - "card_click": item.get("onclick", '"' + html.escape(f'return cardClicked({item.get("prompt", None)}, {"true" if self.allow_negative_prompt else "false"})') + '"'), + "card_click": item.get("onclick", '"' + html.escape(onclick) + '"'), "mtime": item.get("mtime", 0), "size": item.get("size", 0), "version": item.get("version", ''), "color": random_bright_color(), "reference": "reference" if 'Reference' in item.get('name', '') else "", } - alias = item.get("alias", None) - if alias is not None: - args['title'] += f'\nAlias: {alias}' + # alias = item.get("alias", None) + # if alias is not None: + # args['title'] += f'\nAlias: {alias}' return self.card.format(**args) except Exception as e: shared.log.error(f'Networks: item error: page={tabname} item={item["name"]} {e}') @@ -490,6 +503,8 @@ def register_pages(): register_page(ExtraNetworksPageStyles()) from modules.ui_extra_networks_lora import ExtraNetworksPageLora register_page(ExtraNetworksPageLora()) + from modules.ui_extra_networks_wildcards import ExtraNetworksPageWildcards + register_page(ExtraNetworksPageWildcards()) if shared.opts.latent_history > 0: from modules.ui_extra_networks_history import ExtraNetworksPageHistory register_page(ExtraNetworksPageHistory()) @@ -502,7 +517,7 @@ def get_pages(title=None): visible = shared.opts.extra_networks pages = [] if 'All' in visible or visible == []: # default en sort order - visible = ['Model', 'Lora', 'Style', 'Embedding', 'VAE', 'History', 'Hypernetwork'] + visible = ['Model', 'Lora', 'Style', 'Wildcards', 'Embedding', 'VAE', 'History', 'Hypernetwork'] titles = [page.title for page in shared.extra_networks] if title is None: @@ -596,7 +611,7 @@ def create_ui(container, button_parent, tabname, skip_indexing = False): return is_visible, gr.update(visible=is_visible), gr.update(variant=("secondary-down" if is_visible else "secondary")) with ui.details: - details_close = ToolButton(symbols.close, elem_id=f"{tabname}_extra_details_close", elem_classes=['extra-details-close']) + details_close = ui_components.ToolButton(ui_symbols.close, elem_id=f"{tabname}_extra_details_close", elem_classes=['extra-details-close']) details_close.click(fn=lambda: gr.update(visible=False), inputs=[], outputs=[ui.details]) with gr.Row(): with gr.Column(scale=1): @@ -649,14 +664,14 @@ def create_ui(container, button_parent, tabname, skip_indexing = False): model_visible = page in ['Model'] return [gr.update(visible=scan_visible), gr.update(visible=save_visible), gr.update(visible=model_visible)] - ui.button_refresh = ToolButton(symbols.refresh, elem_id=f"{tabname}_extra_refresh") - ui.button_scan = ToolButton(symbols.scan, elem_id=f"{tabname}_extra_scan", visible=True) - ui.button_quicksave = ToolButton(symbols.book, elem_id=f"{tabname}_extra_quicksave", visible=False) - ui.button_save = ToolButton(symbols.book, elem_id=f"{tabname}_extra_save", visible=False) - ui.button_sort = ToolButton(symbols.sort, elem_id=f"{tabname}_extra_sort", visible=True) - ui.button_view = ToolButton(symbols.view, elem_id=f"{tabname}_extra_view", visible=True) - ui.button_close = ToolButton(symbols.close, elem_id=f"{tabname}_extra_close", visible=True) - ui.button_model = ToolButton(symbols.refine, elem_id=f"{tabname}_extra_model", visible=True) + ui.button_refresh = ui_components.ToolButton(ui_symbols.refresh, elem_id=f"{tabname}_extra_refresh") + ui.button_scan = ui_components.ToolButton(ui_symbols.scan, elem_id=f"{tabname}_extra_scan", visible=True) + ui.button_quicksave = ui_components.ToolButton(ui_symbols.book, elem_id=f"{tabname}_extra_quicksave", visible=False) + ui.button_save = ui_components.ToolButton(ui_symbols.book, elem_id=f"{tabname}_extra_save", visible=False) + ui.button_sort = ui_components.ToolButton(ui_symbols.sort, elem_id=f"{tabname}_extra_sort", visible=True) + ui.button_view = ui_components.ToolButton(ui_symbols.view, elem_id=f"{tabname}_extra_view", visible=True) + ui.button_close = ui_components.ToolButton(ui_symbols.close, elem_id=f"{tabname}_extra_close", visible=True) + ui.button_model = ui_components.ToolButton(ui_symbols.refine, elem_id=f"{tabname}_extra_model", visible=True) ui.search = gr.Textbox('', show_label=False, elem_id=f"{tabname}_extra_search", placeholder="Search...", elem_classes="textbox", lines=2, container=False) ui.description = gr.Textbox('', show_label=False, elem_id=f"{tabname}_description", elem_classes=["textbox", "extra-description"], lines=2, interactive=False, container=False) @@ -722,15 +737,21 @@ def create_ui(container, button_parent, tabname, skip_indexing = False): shared.log.debug(f'Network save desc: item="{ui.last_item.name}" filename="{fn}"') return desc - def fn_delete_desc(desc): + def fn_delete_network(desc): if ui.last_item is None: return desc - fn = os.path.splitext(ui.last_item.filename)[0] + '.txt' - if os.path.exists(fn): - shared.log.debug(f'Network delete desc: item="{ui.last_item.name}" filename="{fn}"') + basename = os.path.splitext(ui.last_item.filename)[0] + extensions = ['.safetensors', '.ckpt', '.txt', '.json', '.thumb.jpg', '.jpg', '.jpeg', '.png', '.webp', '.tiff', '.jp2', '.jxl'] + candidates = [] + for ext in extensions: + fn = basename + ext + if os.path.exists(fn) and os.path.isfile(fn): + candidates.append(fn) + msg = f'Network delete: item="{ui.last_item.name}" files={candidates}' + shared.log.debug(msg) + for fn in candidates: os.remove(fn) - return '' - return desc + return msg def fn_save_info(info): fn = os.path.splitext(ui.last_item.filename)[0] + '.json' @@ -738,16 +759,6 @@ def create_ui(container, button_parent, tabname, skip_indexing = False): shared.log.debug(f'Network save info: item="{ui.last_item.name}" filename="{fn}"') return info - def fn_delete_info(info): - if ui.last_item is None: - return info - fn = os.path.splitext(ui.last_item.filename)[0] + '.json' - if os.path.exists(fn): - shared.log.debug(f'Network delete info: item="{ui.last_item.name}" filename="{fn}"') - os.remove(fn) - return '' - return info - def fn_save_style(info, description, prompt, negative, extra, wildcards): if not isinstance(info, dict) or isinstance(info, list): shared.log.warning(f'Network save style skip: item="{ui.last_item.name}" not a dict: {type(info)}') @@ -774,9 +785,9 @@ def create_ui(container, button_parent, tabname, skip_indexing = False): btn_save_img.click(fn=fn_save_img, _js='closeDetailsEN', inputs=[img], outputs=[img]) btn_delete_img.click(fn=fn_delete_img, _js='closeDetailsEN', inputs=[img], outputs=[img]) btn_save_desc.click(fn=fn_save_desc, _js='closeDetailsEN', inputs=[desc], outputs=[desc]) - btn_delete_desc.click(fn=fn_delete_desc, _js='closeDetailsEN', inputs=[desc], outputs=[desc]) + btn_delete_desc.click(fn=fn_delete_network, _js='closeDetailsEN', inputs=[desc], outputs=[desc]) btn_save_info.click(fn=fn_save_info, _js='closeDetailsEN', inputs=[info], outputs=[info]) - btn_delete_info.click(fn=fn_delete_info, _js='closeDetailsEN', inputs=[info], outputs=[info]) + btn_delete_info.click(fn=fn_delete_network, _js='closeDetailsEN', inputs=[info], outputs=[desc]) btn_save_style.click(fn=fn_save_style, _js='closeDetailsEN', inputs=[info, description, prompt, negative, extra, wildcards], outputs=[info]) btn_delete_style.click(fn=fn_delete_style, _js='closeDetailsEN', inputs=[info], outputs=[info]) @@ -786,7 +797,7 @@ def create_ui(container, button_parent, tabname, skip_indexing = False): is_valid = (item is not None) and hasattr(item, 'name') and hasattr(item, 'filename') if is_valid: - stat = os.stat(item.filename) if os.path.exists(item.filename) else None + stat_size, stat_mtime = modelstats.stat(item.filename) desc = item.description fullinfo = shared.readfile(os.path.splitext(item.filename)[0] + '.json', silent=True) if 'modelVersions' in fullinfo: # sanitize massive objects @@ -876,8 +887,8 @@ def create_ui(container, button_parent, tabname, skip_indexing = False): Alias{getattr(item, 'alias', 'N/A')} Filename{item.filename} Hash{getattr(item, 'hash', 'N/A')} - Size{round(stat.st_size/1024/1024, 2) if stat is not None else 'N/A'} MB - Last modified{datetime.fromtimestamp(stat.st_mtime) if stat is not None else 'N/A'} + Size{round(stat_size/1024/1024, 2)} MB + Last modified{stat_mtime} Source URL{url} {lora} @@ -932,9 +943,9 @@ def create_ui(container, button_parent, tabname, skip_indexing = False): return pages def ui_scan_click(title): - from modules import ui_models - if ui_models.search_metadata_civit is not None: - ui_models.search_metadata_civit(True, title) + from modules.civitai.metadata_civitai import civit_search_metadata + for _generator in civit_search_metadata(title): # need to read generator output so python does not optimize function away + pass return ui_refresh_click(title) def ui_save_click(): diff --git a/modules/ui_extra_networks_checkpoints.py b/modules/ui_extra_networks_checkpoints.py index 3eab7fd41..520492761 100644 --- a/modules/ui_extra_networks_checkpoints.py +++ b/modules/ui_extra_networks_checkpoints.py @@ -2,7 +2,7 @@ import os import html import json import concurrent -from modules import shared, ui_extra_networks, sd_models +from modules import shared, ui_extra_networks, sd_models, modelstats reference_dir = os.path.join('models', 'Reference') @@ -26,16 +26,18 @@ class ExtraNetworksPageCheckpoints(ui_extra_networks.ExtraNetworksPage): else: continue preview = v.get('preview', v['path']) + preview_file = self.find_preview_file(os.path.join(reference_dir, preview)) + _size, mtime = modelstats.stat(preview_file) yield { "type": 'Model', "name": os.path.join(reference_dir, k), "title": os.path.join(reference_dir, k), "filename": url, "preview": self.find_preview(os.path.join(reference_dir, preview)), - "local_preview": self.find_preview_file(os.path.join(reference_dir, preview)), - "onclick": '"' + html.escape(f"""return selectReference({json.dumps(url)})""") + '"', + "local_preview": preview_file, + "onclick": '"' + html.escape(f"selectReference({json.dumps(url)})") + '"', "hash": None, - "mtime": 0, + "mtime": mtime, "size": 0, "info": {}, "metadata": {}, @@ -46,7 +48,7 @@ class ExtraNetworksPageCheckpoints(ui_extra_networks.ExtraNetworksPage): record = None try: checkpoint: sd_models.CheckpointInfo = sd_models.checkpoints_list.get(name) - exists = os.path.exists(checkpoint.filename) + size, mtime = modelstats.stat(checkpoint.filename) record = { "type": 'Model', "name": checkpoint.name, @@ -54,9 +56,9 @@ class ExtraNetworksPageCheckpoints(ui_extra_networks.ExtraNetworksPage): "filename": checkpoint.filename, "hash": checkpoint.shorthash, "metadata": checkpoint.metadata, - "onclick": '"' + html.escape(f"""return selectCheckpoint({json.dumps(name)})""") + '"', - "mtime": os.path.getmtime(checkpoint.filename) if exists else 0, - "size": os.path.getsize(checkpoint.filename) if exists else 0, + "onclick": '"' + html.escape(f"selectCheckpoint({json.dumps(name)})") + '"', + "mtime": mtime, + "size": size, } record["info"] = self.find_info(checkpoint.filename) record["description"] = self.find_description(checkpoint.filename, record["info"]) diff --git a/modules/ui_extra_networks_lora.py b/modules/ui_extra_networks_lora.py index 194f16b41..8ab95ce8d 100644 --- a/modules/ui_extra_networks_lora.py +++ b/modules/ui_extra_networks_lora.py @@ -1,7 +1,7 @@ import os import json import concurrent -from modules import shared, ui_extra_networks +from modules import shared, ui_extra_networks, modelstats from modules.lora import lora_load @@ -85,21 +85,23 @@ class ExtraNetworksPageLora(ui_extra_networks.ExtraNetworksPage): try: # path, _ext = os.path.splitext(l.filename) name = os.path.splitext(os.path.relpath(l.filename, shared.cmd_opts.lora_dir))[0] + size, mtime = modelstats.stat(l.filename) + info = self.find_info(l.filename) item = { "type": 'Lora', "name": name, + "alias": os.path.splitext(os.path.basename(l.filename))[0], "filename": l.filename, "hash": l.shorthash, "prompt": json.dumps(f" "), "metadata": json.dumps(l.metadata, indent=4) if l.metadata else None, - "mtime": os.path.getmtime(l.filename), - "size": os.path.getsize(l.filename), + "mtime": mtime, + "size": size, "version": l.sd_version, + "info": info, + "description": self.find_description(l.filename, info), + "tags": self.get_tags(l, info), } - info = self.find_info(l.filename) - item["info"] = info - item["description"] = self.find_description(l.filename, info) # use existing info instead of double-read - item["tags"] = self.get_tags(l, info) return item except Exception as e: shared.log.error(f'Networks: type=lora file="{name}" {e}') diff --git a/modules/ui_extra_networks_styles.py b/modules/ui_extra_networks_styles.py index 29ac541a8..84d9f7b95 100644 --- a/modules/ui_extra_networks_styles.py +++ b/modules/ui_extra_networks_styles.py @@ -1,6 +1,7 @@ import os import html import json +from datetime import datetime from modules import shared, extra_networks, ui_extra_networks, styles @@ -80,6 +81,7 @@ class ExtraNetworksPageStyles(ui_extra_networks.ExtraNetworksPage): "type": 'Style', "name": name, "title": k, + "alias": os.path.splitext(os.path.basename(style.filename))[0], "filename": style.filename, "preview": style.preview if getattr(style, 'preview', None) is not None and style.preview.startswith('data:') else None, "description": style.description if getattr(style, 'description', None) is not None and len(style.description) > 0 else txt, @@ -89,7 +91,7 @@ class ExtraNetworksPageStyles(ui_extra_networks.ExtraNetworksPage): "wildcards": getattr(style, 'wildcards', ''), "local_preview": f"{fn}.{shared.opts.samples_format}", "onclick": '"' + html.escape(f"""return selectStyle({json.dumps(name)})""") + '"', - "mtime": getattr(style, 'mtime', 0), + "mtime": getattr(style, 'mtime', datetime.fromtimestamp(0)), "size": os.path.getsize(style.filename), } except Exception as e: diff --git a/modules/ui_extra_networks_textual_inversion.py b/modules/ui_extra_networks_textual_inversion.py index 55b5b01d1..6857447f2 100644 --- a/modules/ui_extra_networks_textual_inversion.py +++ b/modules/ui_extra_networks_textual_inversion.py @@ -1,6 +1,6 @@ import json import os -from modules import shared, sd_models, ui_extra_networks, files_cache +from modules import shared, sd_models, ui_extra_networks, files_cache, modelstats from modules.textual_inversion import Embedding @@ -23,17 +23,20 @@ class ExtraNetworksPageTextualInversion(ui_extra_networks.ExtraNetworksPage): if embedding.tag is not None: tags[embedding.tag]=1 name = os.path.splitext(embedding.basename)[0] + size, mtime = modelstats.stat(embedding.filename) + info = self.find_info(embedding.filename) record = { "type": 'Embedding', "name": name, "filename": embedding.filename, + "alias": os.path.splitext(os.path.basename(embedding.filename))[0], "prompt": json.dumps(f" {os.path.splitext(embedding.name)[0]}"), "tags": tags, - "mtime": os.path.getmtime(embedding.filename), - "size": os.path.getsize(embedding.filename), + "mtime": mtime, + "size": size, + "info": info, + "description": self.find_description(embedding.filename, info), } - record["info"] = self.find_info(embedding.filename) - record["description"] = self.find_description(embedding.filename, record["info"]) except Exception as e: shared.log.debug(f'Networks error: type=embedding file="{embedding.filename}" {e}') return record diff --git a/modules/ui_extra_networks_vae.py b/modules/ui_extra_networks_vae.py index de18b5d26..0db733a3c 100644 --- a/modules/ui_extra_networks_vae.py +++ b/modules/ui_extra_networks_vae.py @@ -1,7 +1,7 @@ import html import json import os -from modules import shared, ui_extra_networks, sd_vae, hashes +from modules import shared, ui_extra_networks, sd_vae, hashes, modelstats class ExtraNetworksPageVAEs(ui_extra_networks.ExtraNetworksPage): @@ -14,9 +14,12 @@ class ExtraNetworksPageVAEs(ui_extra_networks.ExtraNetworksPage): def list_items(self): for name, filename in sd_vae.vae_dict.items(): try: + size, mtime = modelstats.stat(filename) + info = self.find_info(filename) record = { "type": 'VAE', "name": name, + "alias": os.path.splitext(os.path.basename(filename))[0], "title": name, "filename": filename, "hash": hashes.sha256_from_cache(filename, f"vae/{filename}"), @@ -24,11 +27,11 @@ class ExtraNetworksPageVAEs(ui_extra_networks.ExtraNetworksPage): "local_preview": f"{os.path.splitext(filename)[0]}.{shared.opts.samples_format}", "metadata": {}, "onclick": '"' + html.escape(f"""return selectVAE({json.dumps(name)})""") + '"', - "mtime": os.path.getmtime(filename), - "size": os.path.getsize(filename), + "mtime": mtime, + "size": size, + "info": info, + "description": self.find_description(filename, info), } - record["info"] = self.find_info(filename) - record["description"] = self.find_description(filename, record["info"]) yield record except Exception as e: shared.log.debug(f'Networks error: type=vae file="{filename}" {e}') diff --git a/modules/ui_extra_networks_wildcards.py b/modules/ui_extra_networks_wildcards.py new file mode 100644 index 000000000..e60af679e --- /dev/null +++ b/modules/ui_extra_networks_wildcards.py @@ -0,0 +1,50 @@ +import os +import json +from modules import shared, ui_extra_networks, modelstats, files_cache + + +wildcards_list = [] + + +class ExtraNetworksPageWildcards(ui_extra_networks.ExtraNetworksPage): + def __init__(self): + super().__init__('Wildcards') + + def parents(self, file): + folder = os.path.dirname(file) + if folder != shared.opts.wildcards_dir and folder not in wildcards_list: + wildcards_list.append(folder) + self.parents(folder) + + def refresh(self): + wildcards_list.clear() + files = files_cache.list_files(shared.opts.wildcards_dir, ext_filter=[".txt"], recursive=True) + for file in files: + wildcards_list.append(file) + self.parents(file) + + def list_items(self): + self.refresh() + for filename in wildcards_list: + relname = os.path.relpath(filename, shared.opts.wildcards_dir) + name = os.path.splitext(relname)[0] + size, mtime = modelstats.stat(filename) + try: + record = { + "type": 'Wildcard', + "name": name, + "filename": filename, + "preview": self.find_preview(filename), + "local_preview": f"{os.path.splitext(filename)[0]}.{shared.opts.samples_format}", + "prompt": json.dumps(f" __{name}__"), + "mtime": mtime, + "size": size, + "description": '', + "info": {}, + } + yield record + except Exception as e: + shared.log.debug(f'Networks error: type=wildcard file="{filename}" {e}') + + def allowed_directories_for_previews(self): + return [v for v in [shared.opts.wildcards_dir] if v is not None] diff --git a/modules/ui_gallery.py b/modules/ui_gallery.py index 3b0363359..07d2dec26 100644 --- a/modules/ui_gallery.py +++ b/modules/ui_gallery.py @@ -1,17 +1,17 @@ import os -from datetime import datetime from urllib.parse import unquote import gradio as gr from PIL import Image -from modules import shared, ui_symbols, ui_common, images, video +from modules import shared, ui_symbols, ui_common, images, video, modelstats from modules.ui_components import ToolButton + def read_media(fn): fn = unquote(fn).replace('%3A', ':') if not os.path.isfile(fn): shared.log.error(f'Gallery not found: file="{fn}"') return [[], None, '', '', f'Media not found: {fn}'] - stat = os.stat(fn) + stat_size, stat_mtime = modelstats.stat(fn) if fn.lower().endswith('.mp4'): frames, fps, duration, w, h, codec, _frame = video.get_video_params(fn) geninfo = '' @@ -21,8 +21,8 @@ def read_media(fn): | Frames {frames:,} | FPS {fps:.2f} | Duration {duration:.2f} - | Size {stat.st_size:,} - | Modified {datetime.fromtimestamp(stat.st_mtime)}


+ | Size {stat_size:,} + | Modified {stat_mtime}


''' return [gr.update(visible=False, value=[]), gr.update(visible=True, value=fn), geninfo, geninfo, log] else: @@ -33,8 +33,8 @@ def read_media(fn):

Image {image.width} x {image.height} | Format {image.format} | Mode {image.mode} - | Size {stat.st_size:,} - | Modified {datetime.fromtimestamp(stat.st_mtime)}


+ | Size {stat_size:,} + | Modified {stat_mtime}


''' return [gr.update(visible=True, value=[image]), gr.update(visible=False), geninfo, geninfo, log] diff --git a/modules/ui_models.py b/modules/ui_models.py index 1a16d1b44..71b96ba93 100644 --- a/modules/ui_models.py +++ b/modules/ui_models.py @@ -1,21 +1,13 @@ import os -import re -import time -import json import inspect -from datetime import datetime import gradio as gr -from modules import errors, sd_models, sd_vae, extras, sd_samplers, ui_symbols, hashes +from modules import errors, sd_models, sd_vae, extras, sd_samplers, ui_symbols, modelstats from modules.ui_components import ToolButton from modules.ui_common import create_refresh_button from modules.call_queue import wrap_gradio_gpu_call -from modules.shared import opts, log, req, readfile, max_workers, native -from modules.merging import merge_methods -from modules.merging.merge_utils import BETA_METHODS, TRIPLE_METHODS, interpolate -from modules.merging.merge_presets import BLOCK_WEIGHTS_PRESETS, SDXL_BLOCK_WEIGHTS_PRESETS +from modules.shared import opts, log -search_metadata_civit = None extra_ui = [] @@ -23,50 +15,146 @@ def create_ui(): dummy_component = gr.Label(visible=False) with gr.Row(elem_id="models_tab"): with gr.Column(elem_id='models_output_container', scale=1): - # models_output = gr.Textbox(elem_id="models_output", value="", show_label=False) - gr.HTML(elem_id="models_progress", value="") - models_image = gr.Image(elem_id="models_image", show_label=False, interactive=False, type='pil') - models_outcome = gr.HTML(elem_id="models_error", value="") + models_outcome = gr.HTML(elem_id="models_outcome", value="") models_file = gr.File(label='', visible=False) with gr.Column(elem_id='models_input_container', scale=3): - with gr.Tab(label="Current"): - def analyze(): - from modules import modelstats - model = modelstats.analyze() - desc = f"Model: {model.name}
Type: {model.type}
Class: {model.cls}
Size: {model.size} bytes
Modified: {model.mtime}
" - meta = model.meta - components = [] - for m in model.modules: + with gr.Tab(label="Current", elem_id="models_current_tab"): + def create_modules_table(rows: list): + html = """ + + + + + + {tbody} + +
ModuleClassDeviceDtypeQuantParamsModulesConfig
+ """ + tbody = '' + for row in rows: try: - component = (m.name, m.cls, str(m.device), str(m.dtype), m.params, m.modules, str(m.config)) - components.append(component) + config = str(row.config) except Exception: - component = (m.name, m.cls, str(m.device), str(m.dtype), m.params, m.modules, '') - components.append(component) - return [desc, components, meta] + config = '{}' + try: + tbody += f""" + + {row.name} + {row.cls} + {row.device} + {row.dtype} + {row.quant} + {row.params} + {row.modules} +
{config}
+ + """ + except Exception as e: + log.error(f'Model list: row={vars(row)} {e}') + return html.format(tbody=tbody) + + def analyze(): + model = modelstats.analyze() + if model is None: + return ["Model not loaded", {}] + meta = model.meta + html = create_modules_table(model.modules) + return [html, meta] with gr.Row(): - gr.HTML('

 Analyze currently loaded model

') + gr.HTML('

Analyze currently loaded model

') with gr.Row(): model_analyze = gr.Button(value="Analyze", variant='primary') with gr.Row(): model_desc = gr.HTML(value="", elem_id="model_desc") - with gr.Row(): - module_headers = ['Module', 'Class', 'Device', 'DType', 'Params', 'Modules', 'Config'] - module_types = ['str', 'str', 'str', 'str', 'number', 'number', 'str'] - model_modules = gr.DataFrame(value=None, label=None, show_label=False, interactive=False, wrap=True, headers=module_headers, datatype=module_types, type='array') with gr.Row(): model_meta = gr.JSON(label="Metadata", value={}, elem_id="model_meta") - model_analyze.click(fn=analyze, inputs=[], outputs=[model_desc, model_modules, model_meta]) + model_analyze.click(fn=analyze, inputs=[], outputs=[model_desc, model_meta]) - with gr.Tab(label="Loader"): + with gr.Tab(label="List", elem_id="models_list_tab"): + def create_models_table(rows: list): + from modules import sd_detect + html = """ + + + + + + {tbody} + +
NameTypeDetectPipelineHashSizeMTime
+ """ + tbody = '' + for row in rows: + try: + f = row.filename + stat_size, stat_mtime = modelstats.stat(f) + if os.path.isfile(f): + typ = os.path.splitext(f)[1][1:] + size = f"{round(stat_size / 1024 / 1024 / 1024, 3)} gb" + elif os.path.isdir(f): + typ = 'diffusers' + size = 'folder' + else: + typ = 'unknown' + size = 'unknown' + guess = 'Stable Diffusion XL' if 'XL' in f.upper() else 'Stable Diffusion' # set default guess + guess = sd_detect.guess_by_size(f, guess) + guess = sd_detect.guess_by_name(f, guess) + guess, pipeline = sd_detect.guess_by_diffusers(f, guess) + guess = sd_detect.guess_variant(f, guess) + pipeline = sd_detect.shared_items.get_pipelines().get(guess, None) if pipeline is None else pipeline + tbody += f""" + + {row.model_name} + {typ} + {guess} + {pipeline.__name__ if pipeline else '(unknown)'} + {row.shorthash} + {size} + {stat_mtime} + + """ + except Exception as e: + log.error(f'Model list: row={vars(row)} {e}') + return html.format(tbody=tbody) + + with gr.Row(): + gr.HTML('

List all locally available models


') + with gr.Row(): + model_list_btn = gr.Button(value="List models", variant='primary') + model_checkhash_btn = gr.Button(value="Calculate missing hashes", variant='secondary') + with gr.Row(): + model_table = gr.HTML(value='', elem_id="model_list_table") + + model_checkhash_btn.click(fn=sd_models.update_model_hashes, inputs=[], outputs=[model_table]) + model_list_btn.click(fn=lambda: create_models_table(sd_models.checkpoints_list.values()), inputs=[], outputs=[model_table]) + + with gr.Tab(label="Metadata", elem_id="models_metadata_tab"): + from modules.civitai.metadata_civitai import civit_search_metadata, civit_update_metadata + with gr.Row(): + gr.HTML('

Fetch model preview metadata


') + with gr.Row(): + civit_previews_btn = gr.Button(value="Scan missing", variant='primary') + civit_update_btn = gr.Button(value="Update all", variant='primary') + with gr.Row(): + civit_metadata = gr.HTML(value='', elem_id="civit_metadata") + civit_previews_btn.click(fn=civit_search_metadata, inputs=[], outputs=[civit_metadata]) + civit_update_btn.click(fn=civit_update_metadata, inputs=[], outputs=[civit_metadata]) + + + with gr.Tab(label="Loader", elem_id="models_loader_tab"): from modules import ui_models_load ui_models_load.create_ui(models_outcome, models_file) - with gr.Tab(label="Merge"): + with gr.Tab(label="Merge", elem_id="models_merge_tab"): + from modules.merging import merge_methods + from modules.merging.merge_utils import BETA_METHODS, TRIPLE_METHODS, interpolate + from modules.merging.merge_presets import BLOCK_WEIGHTS_PRESETS, SDXL_BLOCK_WEIGHTS_PRESETS + def sd_model_choices(): return ['None'] + sd_models.checkpoint_titles() @@ -306,7 +394,7 @@ def create_ui(): ] ) - with gr.Tab(label="Modules"): + with gr.Tab(label="Replace", elem_id="models_replace_tab"): with gr.Row(): gr.HTML('

 Replace model components

') with gr.Row(): @@ -378,80 +466,65 @@ def create_ui(): outputs=[models_outcome] ) - with gr.Tab(label="Validate"): - model_headers = ['name', 'type', 'filename', 'hash', 'added', 'size', 'metadata'] - model_data = [] + with gr.Tab(label="CivitAI", elem_id="models_civitai_tab"): + from modules.civitai.search_civitai import search_civitai, create_model_cards, base_models - with gr.Row(): - gr.HTML('

 List all models

') - with gr.Row(): - model_list_btn = gr.Button(value="List model details", variant='primary') - model_checkhash_btn = gr.Button(value="Calculate hash for all models", variant='primary') - model_checkhash_btn.click(fn=sd_models.update_model_hashes, inputs=[], outputs=[models_outcome]) - with gr.Row(): - model_table = gr.DataFrame( - value=None, - headers=model_headers, - label='Model data', - show_label=True, - interactive=False, - wrap=True, - ) + def civitai_search(civit_search_text, civit_search_tag, civit_nsfw, civit_type, civit_base, civit_token): + results = search_civitai(query=civit_search_text, tag=civit_search_tag, nsfw=civit_nsfw, types=civit_type, base=civit_base, token=civit_token) + html = create_model_cards(results) + return html - def list_models(): - total_size = 0 - model_data.clear() - txt = '' - for m in sd_models.checkpoints_list.values(): - try: - stat = os.stat(m.filename) - m_name = m.name.replace('.ckpt', '').replace('.safetensors', '') - m_type = 'ckpt' if m.name.endswith('.ckpt') else 'safe' - m_meta = len(json.dumps(m.metadata)) - 2 - m_size = round(stat.st_size / 1024 / 1024 / 1024, 3) - m_time = datetime.fromtimestamp(stat.st_mtime) - model_data.append([m_name, m_type, m.filename, m.shorthash, m_time, m_size, m_meta]) - total_size += stat.st_size - except Exception as e: - txt += f"Error: {m.name} {e}
" - txt += f"Model list enumerated {len(sd_models.checkpoints_list.keys())} models in {round(total_size / 1024 / 1024 / 1024, 3)} GB
" - return model_data, txt - - model_list_btn.click(fn=list_models, inputs=[], outputs=[model_table, models_outcome]) - - with gr.Tab(label="Huggingface"): - data = [] - os.environ.setdefault('HF_HUB_DISABLE_EXPERIMENTAL_WARNING', '1') - os.environ.setdefault('HF_HUB_DISABLE_SYMLINKS_WARNING', '1') - os.environ.setdefault('HF_HUB_DISABLE_IMPLICIT_TOKEN', '1') - os.environ.setdefault('HUGGINGFACE_HUB_VERBOSITY', 'warning') - - def hf_search(keyword): - import huggingface_hub as hf - hf_api = hf.HfApi() - models = hf_api.list_models(model_name=keyword, full=True, library="diffusers", limit=50, sort="downloads", direction=-1) - data.clear() - for model in models: - tags = [t for t in model.tags if not t.startswith('diffusers') and not t.startswith('license') and not t.startswith('arxiv') and len(t) > 2] - data.append([model.id, model.pipeline_tag, tags, model.downloads, model.lastModified, f'https://huggingface.co/{model.id}']) - return data - - def hf_select(evt: gr.SelectData, data): - return data[evt.index[0]][0] - - def hf_download_model(hub_id: str, token, variant, revision, mirror, custom_pipeline): - from modules.modelloader import download_diffusers_model - download_diffusers_model(hub_id, cache_dir=opts.diffusers_dir, token=token, variant=variant, revision=revision, mirror=mirror, custom_pipeline=custom_pipeline) - from modules.sd_models import list_models # pylint: disable=W0621 - list_models() - log.info(f'Diffuser model downloaded: model="{hub_id}"') - return f'Diffuser model downloaded: model="{hub_id}"' - - def hf_update_token(token): - log.debug('Huggingface update token') - opts.huggingface_token = token + def civitai_update_token(token): + log.debug('CivitAI update token') + opts.civitai_token = token opts.save() + def civitai_download(model_urls, model_names, model_types, model_path, civit_token, model_output): + from modules.civitai.download_civitai import download_civit_model + for model_url, model_name, model_type in zip(model_urls, model_names, model_types): + msg = f"

Initiating download

{model_name} | {model_type} | {model_url}

" + yield msg + model_output + download_civit_model(model_url, model_name, model_path, model_type, civit_token) + yield model_output + + with gr.Row(): + gr.HTML('

Search & Download

') + with gr.Row(elem_id='civitai_search_row'): + civit_search_text = gr.Textbox(label='', placeholder='keyword', elem_id="civit_search_text") + civit_search_tag = gr.Textbox(label='', placeholder='tag', elem_id="civit_search_text") + civit_search_text_btn = ToolButton(value=ui_symbols.search, interactive=True) + with gr.Accordion(label='Advanced', open=False, elem_id="civitai_search_options"): + civit_download_btn = gr.Button(value="Download model", variant='primary', elem_id="civitai_download_btn", visible=False) + with gr.Row(): + civit_token = gr.Textbox(opts.civitai_token, label='CivitAI token', placeholder='optional access token for private or gated models', elem_id="civitai_token") + with gr.Row(): + civit_nsfw = gr.Checkbox(label='NSFW allowed', value=True) + with gr.Row(): + civit_type = gr.Textbox(label='Target model type', placeholder='Checkpoint, LORA, ...', value='') + with gr.Row(): + # civit_base = gr.Textbox(label='Base model', placeholder='SDXL, ...') + civit_base = gr.Dropdown(choices=base_models, label='Base model', value='') + with gr.Row(): + civit_folder = gr.Textbox(label='Download folder', placeholder='optional folder for downloads') + with gr.Row(): + civitai_models_output = gr.HTML('', elem_id="civitai_models_output") + # sort, period, limit + _dummy = gr.Label(visible=False) # dummy component to get argspec later + civit_inputs = [civit_search_text, civit_search_tag, civit_nsfw, civit_type, civit_base, civit_token] + civit_search_text_btn.click(fn=civitai_search, inputs=civit_inputs, outputs=[civitai_models_output]) + civit_search_text.submit(fn=civitai_search, inputs=civit_inputs, outputs=[civitai_models_output]) + civit_search_tag.submit(fn=civitai_search, inputs=civit_inputs, outputs=[civitai_models_output]) + civit_token.change(fn=civitai_update_token, inputs=[civit_token], outputs=[]) + civit_download_btn.click( + fn=civitai_download, + _js="downloadCivitModel", + inputs=[_dummy, _dummy, _dummy, civit_folder, civit_token, civitai_models_output], + outputs=[civitai_models_output], + show_progress=True, + ) + + with gr.Tab(label="Huggingface", elem_id="models_huggingface_tab"): + from modules.models_hf import hf_search, hf_select, hf_download_model, hf_update_token with gr.Column(scale=6): with gr.Row(): gr.HTML('

 Download model from huggingface

') @@ -459,17 +532,16 @@ def create_ui(): hf_search_text = gr.Textbox('', label='Search models', placeholder='search huggingface models') hf_search_btn = ToolButton(value=ui_symbols.search) with gr.Row(): - with gr.Column(scale=2): - with gr.Row(): - hf_selected = gr.Textbox('', label='Select model', placeholder='select model from search results or enter model name manually') - with gr.Column(scale=1): - with gr.Row(): - hf_variant = gr.Textbox('', label='Specify model variant', placeholder='') - hf_revision = gr.Textbox('', label='Specify model revision', placeholder='') - with gr.Row(): - hf_token = gr.Textbox(opts.huggingface_token, label='Huggingface token', placeholder='optional access token for private or gated models') - hf_mirror = gr.Textbox('', label='Huggingface mirror', placeholder='optional mirror site for downloads') - hf_custom_pipeline = gr.Textbox('', label='Custom pipeline', placeholder='optional pipeline for downloads') + hf_selected = gr.Textbox('', label='Select model', placeholder='select model from search results or enter model name manually') + with gr.Accordion(label='Advanced', open=False, elem_id="hf_search_options"): + with gr.Row(): + hf_token = gr.Textbox(opts.huggingface_token, label='Huggingface token', placeholder='optional access token for private or gated models', elem_id="hf_token") + with gr.Row(): + hf_variant = gr.Textbox('', label='Specify model variant', placeholder='') + hf_revision = gr.Textbox('', label='Specify model revision', placeholder='') + with gr.Row(): + hf_mirror = gr.Textbox('', label='Huggingface mirror', placeholder='optional mirror site for downloads') + hf_custom_pipeline = gr.Textbox('', label='Custom pipeline', placeholder='optional pipeline for downloads') with gr.Column(scale=1): gr.HTML('
') hf_download_model_btn = gr.Button(value="Download model", variant='primary') @@ -485,375 +557,8 @@ def create_ui(): hf_download_model_btn.click(fn=hf_download_model, inputs=[hf_selected, hf_token, hf_variant, hf_revision, hf_mirror, hf_custom_pipeline], outputs=[models_outcome]) hf_token.change(fn=hf_update_token, inputs=[hf_token], outputs=[]) - with gr.Tab(label="CivitAI"): - data = [] - - def civit_search_model(name, tag, model_type): - # types = 'LORA' if model_type == 'LoRA' else 'Checkpoint' - url = 'https://civitai.com/api/v1/models?limit=25&Sort=Newest' - if model_type == 'Model': - url += '&types=Checkpoint' - elif model_type == 'LoRA': - url += '&types=LORA&types=DoRA&types=LoCon' - elif model_type == 'Embedding': - url += '&types=TextualInversion' - elif model_type == 'VAE': - url += '&types=VAE' - if name is not None and len(name) > 0: - url += f'&query={name}' - if tag is not None and len(tag) > 0: - url += f'&tag={tag}' - r = req(url) - log.debug(f'CivitAI search: type={model_type} name="{name}" tag={tag or "none"} url="{url}" status={r.status_code}') - if r.status_code != 200: - log.warning(f'CivitAI search: name="{name}" tag={tag} status={r.status_code}') - return [], gr.update(visible=False, value=[]), gr.update(visible=False, value=None), gr.update(visible=False, value=None) - try: - body = r.json() - except Exception as e: - log.error(f'CivitAI search: name="{name}" tag={tag} {e}') - return [], gr.update(visible=False, value=[]), gr.update(visible=False, value=None), gr.update(visible=False, value=None) - nonlocal data - data = body.get('items', []) - data1 = [] - for model in data: - found = 0 - if model_type == 'LoRA' and model['type'].lower() in ['lora', 'locon', 'dora', 'lycoris']: - found += 1 - elif model_type == 'Embedding' and model['type'].lower() in ['textualinversion', 'embedding']: - found += 1 - elif model_type == 'Model' and model['type'].lower() in ['checkpoint']: - found += 1 - elif model_type == 'VAE' and model['type'].lower() in ['vae']: - found += 1 - elif model_type == 'Other': - found += 1 - if found > 0: - data1.append([ - model['id'], - model['name'], - ', '.join(model['tags']), - model['stats']['downloadCount'], - model['stats']['rating'] - ]) - res = f'Search result: name={name} tag={tag or "none"} type={model_type} models={len(data1)}' - return res, gr.update(visible=len(data1) > 0, value=data1 if len(data1) > 0 else []), gr.update(visible=False, value=None), gr.update(visible=False, value=None) - - def civit_select1(evt: gr.SelectData, in_data): - model_id = in_data[evt.index[0]][0] - data2 = [] - preview_img = None - for model in data: - if model['id'] == model_id: - for d in model['modelVersions']: - try: - if d.get('images') is not None and len(d['images']) > 0 and len(d['images'][0]['url']) > 0: - preview_img = d['images'][0]['url'] - data2.append([d.get('id', None), d.get('modelId', None) or model_id, d.get('name', None), d.get('baseModel', None), d.get('createdAt', None) or d.get('publishedAt', None)]) - except Exception as e: - log.error(f'CivitAI select: model="{in_data[evt.index[0]]}" {e}') - log.error(f'CivitAI version data={type(d)}: {d}') - log.debug(f'CivitAI select: model="{in_data[evt.index[0]]}" versions={len(data2)}') - return data2, None, preview_img - - def civit_select2(evt: gr.SelectData, in_data): - variant_id = in_data[evt.index[0]][0] - model_id = in_data[evt.index[0]][1] - data3 = [] - for model in data: - if model['id'] == model_id: - for variant in model['modelVersions']: - if variant['id'] == variant_id: - for f in variant['files']: - try: - if os.path.splitext(f['name'])[1].lower() in ['.safetensors', '.ckpt', '.pt', '.pth', '.bin']: - data3.append([f['name'], round(f['sizeKB']), json.dumps(f['metadata']), f['downloadUrl']]) - except Exception: - pass - log.debug(f'CivitAI select: model="{in_data[evt.index[0]]}" files={len(data3)}') - return data3 - - def civit_select3(evt: gr.SelectData, in_data): - log.debug(f'CivitAI select: variant={in_data[evt.index[0]]}') - return in_data[evt.index[0]][3], in_data[evt.index[0]][0], gr.update(interactive=True) - - def civit_download_model(model_url: str, model_name: str, model_path: str, model_type: str, token: str = None): - if model_url is None or len(model_url) == 0: - return 'No model selected' - try: - from modules.modelloader import download_civit_model - res = download_civit_model(model_url, model_name, model_path, model_type, token=token) - except Exception as e: - res = f"CivitAI model downloaded error: model={model_url} {e}" - log.error(res) - return res - from modules.sd_models import list_models # pylint: disable=W0621 - list_models() - return res - - def atomic_civit_search_metadata(item, res, rehash): - from modules.modelloader import download_civit_preview, download_civit_meta - if item is None: - return - meta = os.path.splitext(item['filename'])[0] + '.json' - has_meta = os.path.isfile(meta) and os.stat(meta).st_size > 0 - if ('card-no-preview.png' in item['preview'] or not has_meta) and os.path.isfile(item['filename']): - sha = item.get('hash', None) - found = False - if sha is not None and len(sha) > 0: - r = req(f'https://civitai.com/api/v1/model-versions/by-hash/{sha}') - log.debug(f'CivitAI search: name="{item["name"]}" hash={sha} status={r.status_code}') - if r.status_code == 200: - d = r.json() - res.append(download_civit_meta(item['filename'], d['modelId'])) - if d.get('images') is not None: - for i in d['images']: - preview_url = i['url'] - img_res = download_civit_preview(item['filename'], preview_url) - res.append(img_res) - if 'error' not in img_res: - found = True - break - if not found and rehash and os.stat(item['filename']).st_size < (1024 * 1024 * 1024): - sha = hashes.calculate_sha256(item['filename'], quiet=True)[:10] - r = req(f'https://civitai.com/api/v1/model-versions/by-hash/{sha}') - log.debug(f'CivitAI search: name="{item["name"]}" hash={sha} status={r.status_code}') - if r.status_code == 200: - d = r.json() - res.append(download_civit_meta(item['filename'], d['modelId'])) - if d.get('images') is not None: - for i in d['images']: - preview_url = i['url'] - img_res = download_civit_preview(item['filename'], preview_url) - res.append(img_res) - if 'error' not in img_res: - found = True - break - - def civit_search_metadata(rehash, title): - log.debug(f'CivitAI search metadata: type={title if type(title) == str else "all"}') - from modules.ui_extra_networks import get_pages - res = [] - scanned, skipped = 0, 0 - t0 = time.time() - candidates = [] - re_skip = [r.strip() for r in opts.extra_networks_scan_skip.split(',') if len(r.strip()) > 0] - log.debug(f'CivitAI search metadata: skip={re_skip}') - for page in get_pages(): - if type(title) == str: - if page.title != title: - continue - if page.name == 'style': - continue - for item in page.list_items(): - if item is None: - continue - if any(re.search(re_str, item.get('name', '') + item.get('filename', '')) for re_str in re_skip): - skipped += 1 - continue - scanned += 1 - candidates.append(item) - # atomic_civit_search_metadata(item, res, rehash) - import concurrent - with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor: - for fn in candidates: - executor.submit(atomic_civit_search_metadata, fn, res, rehash) - atomic_civit_search_metadata(None, res, rehash) - t1 = time.time() - log.debug(f'CivitAI search metadata: scanned={scanned} skipped={skipped} time={t1-t0:.2f}') - txt = '
'.join([r for r in res if len(r) > 0]) - return txt - - global search_metadata_civit # pylint: disable=global-statement - search_metadata_civit = civit_search_metadata - - def civitai_update_token(token): - log.debug('CivitAI update token') - opts.civitai_token = token - opts.save() - - with gr.Row(): - gr.HTML('

 CivitAI fetch metadata

') - gr.HTML('Fetches preview and metadata information for all models with missing information
Models with existing previews and information are not updated
') - with gr.Row(): - civit_previews_btn = gr.Button(value="Start", variant='primary') - with gr.Row(): - civit_previews_rehash = gr.Checkbox(value=True, label="Check alternative hash") - - with gr.Row(): - gr.HTML('

Search for models

') - with gr.Row(): - with gr.Column(scale=1): - civit_model_type = gr.Dropdown(label='CivitAI model type', choices=['Model', 'LoRA', 'Embedding', 'VAE', 'Other'], value='Model') - with gr.Column(scale=15): - with gr.Row(): - civit_search_text = gr.Textbox('', label='Search models', placeholder='keyword') - civit_search_tag = gr.Textbox('', label='', placeholder='tags') - civit_search_btn = ToolButton(value=ui_symbols.search, interactive=True) - with gr.Row(): - civit_search_res = gr.HTML('') - with gr.Row(): - gr.HTML('

 CivitAI download model

') - with gr.Row(): - civit_download_model_btn = gr.Button(value="Download", variant='primary') - gr.HTML('Select a model, model version and and model variant from the search results to download or enter model URL manually
') - with gr.Row(): - civit_token = gr.Textbox(opts.civitai_token, label='CivitAI token', placeholder='optional access token for private or gated models') - civit_token.change(fn=civitai_update_token, inputs=[civit_token], outputs=[]) - with gr.Row(): - civit_name = gr.Textbox('', label='Model name', placeholder='select model from search results', visible=True) - civit_selected = gr.Textbox('', label='Model URL', placeholder='select model from search results', visible=True) - civit_path = gr.Textbox('', label='Download path', placeholder='optional subfolder path where to save model', visible=True) - with gr.Row(): - gr.HTML('

Search results

') - with gr.Row(): - civit_headers1 = ['ID', 'Name', 'Tags', 'Downloads', 'Rating'] - civit_types1 = ['number', 'str', 'str', 'number', 'number'] - civit_results1 = gr.DataFrame(value=None, label=None, show_label=False, interactive=False, wrap=True, headers=civit_headers1, datatype=civit_types1, type='array', visible=False) - with gr.Row(): - with gr.Column(): - civit_headers2 = ['ID', 'ModelID', 'Name', 'Base', 'Created', 'Preview'] - civit_types2 = ['number', 'number', 'str', 'str', 'date', 'str'] - civit_results2 = gr.DataFrame(value=None, label='Model versions', show_label=True, interactive=False, wrap=True, headers=civit_headers2, datatype=civit_types2, type='array', visible=False) - with gr.Column(): - civit_headers3 = ['Name', 'Size', 'Metadata', 'URL'] - civit_types3 = ['str', 'number', 'str', 'str'] - civit_results3 = gr.DataFrame(value=None, label='Model variants', show_label=True, interactive=False, wrap=True, headers=civit_headers3, datatype=civit_types3, type='array', visible=False) - - def is_visible(component): - visible = len(component) > 0 if component is not None else False - return gr.update(visible=visible) - - civit_search_text.submit(fn=civit_search_model, inputs=[civit_search_text, civit_search_tag, civit_model_type], outputs=[civit_search_res, civit_results1, civit_results2, civit_results3]) - civit_search_tag.submit(fn=civit_search_model, inputs=[civit_search_text, civit_search_tag, civit_model_type], outputs=[civit_search_res, civit_results1, civit_results2, civit_results3]) - civit_search_btn.click(fn=civit_search_model, inputs=[civit_search_text, civit_search_tag, civit_model_type], outputs=[civit_search_res, civit_results1, civit_results2, civit_results3]) - civit_results1.select(fn=civit_select1, inputs=[civit_results1], outputs=[civit_results2, civit_results3, models_image]) - civit_results2.select(fn=civit_select2, inputs=[civit_results2], outputs=[civit_results3]) - civit_results3.select(fn=civit_select3, inputs=[civit_results3], outputs=[civit_selected, civit_name, civit_search_btn]) - civit_results1.change(fn=is_visible, inputs=[civit_results1], outputs=[civit_results1]) - civit_results2.change(fn=is_visible, inputs=[civit_results2], outputs=[civit_results2]) - civit_results3.change(fn=is_visible, inputs=[civit_results3], outputs=[civit_results3]) - civit_download_model_btn.click(fn=civit_download_model, inputs=[civit_selected, civit_name, civit_path, civit_model_type, civit_token], outputs=[models_outcome]) - civit_previews_btn.click(fn=civit_search_metadata, inputs=[civit_previews_rehash, civit_previews_rehash], outputs=[models_outcome]) - - with gr.Tab(label="Update"): - with gr.Row(): - gr.HTML('

 Scan CivitAI for information on latest available model versions

') - with gr.Row(): - civit_update_btn = gr.Button(value="Update", variant='primary') - with gr.Row(): - gr.HTML('

Update scan results

') - with gr.Row(): - civit_headers4 = ['ID', 'File', 'Name', 'Versions', 'Current', 'Latest', 'Update'] - civit_types4 = ['number', 'str', 'str', 'number', 'str', 'str', 'str'] - civit_widths4 = ['10%', '25%', '25%', '5%', '10%', '10%', '15%'] - civit_results4 = gr.DataFrame(value=None, label=None, show_label=False, interactive=False, wrap=True, row_count=20, headers=civit_headers4, datatype=civit_types4, type='array', column_widths=civit_widths4) - with gr.Row(): - gr.HTML('

Select model from the list and download update if available

') - with gr.Row(): - civit_update_download_btn = gr.Button(value="Download", variant='primary', visible=False) - - class CivitModel: - def __init__(self, name, fn, sha = None, meta = {}): - self.name = name - self.id = meta.get('id', 0) - self.fn = fn - self.sha = sha - self.meta = meta - self.versions = 0 - self.vername = '' - self.latest = '' - self.latest_hashes = [] - self.latest_name = '' - self.url = None - self.status = 'Not found' - def array(self): - return [self.id, self.fn, self.name, self.versions, self.vername, self.latest, self.status] - - selected_model: CivitModel = None - update_data = [] - - def civit_update_metadata(): - nonlocal update_data - log.debug('CivitAI update metadata: models') - from modules import ui_extra_networks, modelloader - res = [] - pages = ui_extra_networks.get_pages('Model') - if len(pages) == 0: - return 'CivitAI update metadata: no models found' - page: ui_extra_networks.ExtraNetworksPage = pages[0] - table_data = [] - update_data.clear() - all_hashes = [(item.get('hash', None) or 'XXXXXXXX').upper()[:8] for item in page.list_items()] - for item in page.list_items(): - model = CivitModel(name=item['name'], fn=item['filename'], sha=item.get('hash', None), meta=item.get('metadata', {})) - if model.sha is None or len(model.sha) == 0: - res.append(f'CivitAI skip search: name="{model.name}" hash=None') - else: - r = req(f'https://civitai.com/api/v1/model-versions/by-hash/{model.sha}') - res.append(f'CivitAI search: name="{model.name}" hash={model.sha} status={r.status_code}') - if r.status_code == 200: - d = r.json() - model.id = d['modelId'] - modelloader.download_civit_meta(model.fn, model.id) - fn = os.path.splitext(item['filename'])[0] + '.json' - model.meta = readfile(fn, silent=True) - model.name = model.meta.get('name', model.name) - model.versions = len(model.meta.get('modelVersions', [])) - versions = model.meta.get('modelVersions', []) - if len(versions) > 0: - model.latest = versions[0].get('name', '') - model.latest_hashes.clear() - for v in versions[0].get('files', []): - for h in v.get('hashes', {}).values(): - model.latest_hashes.append(h[:8].upper()) - for ver in versions: - for f in ver.get('files', []): - for h in f.get('hashes', {}).values(): - if h[:8].upper() == model.sha[:8].upper(): - model.vername = ver.get('name', '') - model.url = f.get('downloadUrl', None) - model.latest_name = f.get('name', '') - if model.vername == model.latest: - model.status = 'Latest' - elif any(map(lambda v: v in model.latest_hashes, all_hashes)): # pylint: disable=cell-var-from-loop # noqa: C417 - model.status = 'Downloaded' - else: - model.status = 'Available' - break - log.debug(res[-1]) - update_data.append(model) - table_data.append(model.array()) - yield gr.update(value=table_data), '
'.join([r for r in res if len(r) > 0]) - return '
'.join([r for r in res if len(r) > 0]) - - def civit_update_select(evt: gr.SelectData, in_data): - nonlocal selected_model, update_data - try: - selected_model = next([m for m in update_data if m.fn == in_data[evt.index[0]][1]]) - except Exception: - selected_model = None - if selected_model is None or selected_model.url is None or selected_model.status != 'Available': - return [gr.update(value='Model update not available'), gr.update(visible=False)] - else: - return [gr.update(), gr.update(visible=True)] - - def civit_update_download(): - if selected_model is None or selected_model.url is None or selected_model.status != 'Available': - return 'Model update not available' - if selected_model.latest_name is None or len(selected_model.latest_name) == 0: - model_name = f'{selected_model.name} {selected_model.latest}.safetensors' - else: - model_name = selected_model.latest_name - return civit_download_model(selected_model.url, model_name, model_path='', model_type='Model') - - civit_update_btn.click(fn=civit_update_metadata, inputs=[], outputs=[civit_results4, models_outcome]) - civit_results4.select(fn=civit_update_select, inputs=[civit_results4], outputs=[models_outcome, civit_update_download_btn]) - civit_update_download_btn.click(fn=civit_update_download, inputs=[], outputs=[models_outcome]) - - if native: - from modules.lora.lora_extract import create_ui as lora_extract_ui - lora_extract_ui() + from modules.lora.lora_extract import create_ui as lora_extract_ui + lora_extract_ui() for ui in extra_ui: if callable(ui): diff --git a/modules/ui_postprocessing.py b/modules/ui_postprocessing.py index 24f787259..b575fa67a 100644 --- a/modules/ui_postprocessing.py +++ b/modules/ui_postprocessing.py @@ -38,11 +38,11 @@ def create_ui(): with gr.Row(elem_id=f"{id_part}_generate_box", elem_classes="generate-box"): submit = gr.Button('Generate', elem_id=f"{id_part}_generate", variant='primary') interrupt = gr.Button('Stop', elem_id=f"{id_part}_interrupt", variant='secondary') - interrupt.click(fn=lambda: shared.state.interrupt(), inputs=[], outputs=[]) + interrupt.click(fn=shared.state.interrupt, inputs=[], outputs=[]) skip = gr.Button('Skip', elem_id=f"{id_part}_skip", variant='secondary') - skip.click(fn=lambda: shared.state.skip(), inputs=[], outputs=[]) + skip.click(fn=shared.state.skip, inputs=[], outputs=[]) pause = gr.Button('Pause', elem_id=f"{id_part}_pause") - pause.click(fn=lambda: shared.state.pause(), _js='checkPaused', inputs=[], outputs=[]) + pause.click(fn=shared.state.pause, _js='checkPaused', inputs=[], outputs=[]) result_images, generation_info, html_info, html_info_formatted, html_log = ui_common.create_output_panel("extras") gr.HTML('File metadata') exif_info = gr.HTML(elem_id="pnginfo_html_info") diff --git a/modules/ui_sections.py b/modules/ui_sections.py index 94fbde348..4f72783bb 100644 --- a/modules/ui_sections.py +++ b/modules/ui_sections.py @@ -134,12 +134,13 @@ def create_video_inputs(tab:str, show_always:bool=False): gr.update(visible=video_type not in ['None', 'GIF', 'PNG'] or show_always), gr.update(visible=video_type not in ['None', 'GIF', 'PNG'] or show_always), ] - with gr.Column(): + with gr.Row(): video_codecs = ['None', 'GIF', 'PNG', 'MP4/MP4V', 'MP4/AVC1', 'MP4/JVT3', 'MKV/H264', 'AVI/DIVX', 'AVI/RGBA', 'MJPEG/MJPG', 'MPG/MPG1', 'AVR/AVR1'] - video_type = gr.Dropdown(label='Save video', choices=video_codecs, value='None', elem_id=f"{tab}_video_type") - with gr.Column(): - video_duration = gr.Slider(label='Duration', minimum=0.25, maximum=300, step=0.25, value=2, visible=show_always, elem_id=f"{tab}_video_duration") - video_loop = gr.Checkbox(label='Loop', value=True, visible=show_always, elem_id=f"{tab}_video_loop") + video_type = gr.Dropdown(label='Video format', choices=video_codecs, value='MP4/MP4V', elem_id=f"{tab}_video_type") + with gr.Row(): + video_duration = gr.Slider(label='Video duration', minimum=0.25, maximum=300, step=0.25, value=2, visible=show_always, elem_id=f"{tab}_video_duration") + video_loop = gr.Checkbox(label='Loop video', value=True, visible=show_always, elem_id=f"{tab}_video_loop") + with gr.Row(): video_pad = gr.Slider(label='Pad frames', minimum=0, maximum=24, step=1, value=1, visible=show_always, elem_id=f"{tab}_video_pad") video_interpolate = gr.Slider(label='Interpolate frames', minimum=0, maximum=24, step=1, value=0, visible=show_always, elem_id=f"{tab}_video_interpolate") video_type.change(fn=video_type_change, inputs=[video_type], outputs=[video_duration, video_loop, video_pad, video_interpolate]) diff --git a/modules/ui_settings.py b/modules/ui_settings.py index 0687e9a6f..f15fd8985 100644 --- a/modules/ui_settings.py +++ b/modules/ui_settings.py @@ -119,6 +119,9 @@ def run_settings(*args): changed = [] for key, value, comp in zip(shared.opts.data_labels.keys(), args, components): if comp == dummy_component or value=='dummy': # or getattr(comp, 'visible', True) is False or key in hidden_list: + # actual = shared.opts.data.get(key, None) # ensure the key is in data + # default = shared.opts.data_labels[key].default + # shared.log.warning(f'Setting skip: key={key} value={value} actual={actual} default={default} comp={comp}') continue if not shared.opts.same_type(value, shared.opts.data_labels[key].default): shared.log.error(f'Setting bad value: {key}={value} expecting={type(shared.opts.data_labels[key].default).__name__}') @@ -191,6 +194,7 @@ def create_ui(): result = gr.HTML(elem_id="settings_result") script_callbacks.ui_settings_callback() # let extensions create settings sections = [] + options_count = len(shared.opts.data_labels) for item in shared.opts.data_labels.values(): # get unique sections from all items if len(item.section) == 2: section_id, section_text = item.section @@ -203,7 +207,7 @@ def create_ui(): if (section_id, section_text) not in sections: sections.append((section_id, section_text)) - shared.log.debug(f'Settings: sections={len(sections)} settings={len(shared.opts.list())}/{len(list(shared.opts.data_labels))}') + shared.log.debug(f'Settings: sections={len(sections)} settings={len(shared.opts.list())}/{len(list(shared.opts.data_labels))} quicksettings={len(quicksettings_list)}') with gr.Tabs(elem_id="settings"): quicksettings_list.clear() for (section_id, section_text) in sections: @@ -228,12 +232,14 @@ def create_ui(): current_items.append(key) components.append(component) create_dirty_indicator(section_id, current_items) + components_count = len(components) + if components_count != options_count: + shared.log.error(f'Settings: count mismatch: options={options_count} components={components_count}') with gr.TabItem("Show all pages", elem_id="settings_show_all_pages"): create_dirty_indicator("show_all_pages", []) request_notifications = gr.Button(value='Request browser notifications', elem_id="request_notifications", visible=False) - with gr.TabItem("Update", id="system_update", elem_id="tab_update"): from modules import update update.create_ui() @@ -245,6 +251,22 @@ def create_ui(): with gr.TabItem("History", id="system_history", elem_id="tab_history"): ui_history.create_ui() + with gr.TabItem("GPU Monitor", id="system_gpu", elem_id="tab_gpu"): + with gr.Row(elem_id='gpu-controls'): + gpu_start = gr.Button(value="Start", elem_id="gpu_start", variant="primary") + gpu_stop = gr.Button(value="Stop", elem_id="gpu_stop", variant="primary") + gpu_start.click(fn=lambda: None, _js='startGPU', inputs=[], outputs=[]) + gpu_stop.click(fn=lambda: None, _js='disableGPU', inputs=[], outputs=[]) + gr.HTML(''' +
+ + + +
+
+
+ ''', elem_id='gpu-container', visible=True) + with gr.TabItem("ONNX", id="onnx_config", elem_id="tab_onnx"): from modules.onnx_impl import ui as ui_onnx ui_onnx.create_ui() diff --git a/modules/ui_symbols.py b/modules/ui_symbols.py index 66fbda681..975cd4e69 100644 --- a/modules/ui_symbols.py +++ b/modules/ui_symbols.py @@ -25,8 +25,7 @@ reuse = '⬅️' search = '🔍' preview = '🖼️' image = '🖌️' -mark_diag = '※' -mark_flag = '⁜' +resize = '⁜' interrogate = '♻' int_clip = '✎' int_blip = '✐' diff --git a/modules/ui_video.py b/modules/ui_video.py index e12f08bc9..bb9cace94 100644 --- a/modules/ui_video.py +++ b/modules/ui_video.py @@ -27,7 +27,7 @@ def create_ui(): with gr.Row(elem_id="video_interface", equal_height=False): with gr.Tabs(elem_classes=['video-tabs'], elem_id='video-tabs'): overrides = ui_common.create_override_inputs('video') - with gr.Tab('Generic', id='video-tab') as video_tab: + with gr.Tab('Core', id='video-tab') as video_tab: from modules.video_models import video_ui video_ui.create_ui(prompt, negative, styles, overrides) with gr.Tab('FramePack', id='framepack-tab') as framepack_tab: diff --git a/modules/video.py b/modules/video.py index 7ae09aaba..8efbad070 100644 --- a/modules/video.py +++ b/modules/video.py @@ -71,7 +71,7 @@ def save_video(p, images, filename = None, video_type: str = 'none', duration: f if filename is None and p is not None: filename = namegen.apply(shared.opts.samples_filename_pattern if shared.opts.samples_filename_pattern and len(shared.opts.samples_filename_pattern) > 0 else "[seq]-[prompt_words]") filename = os.path.join(shared.opts.outdir_video, filename) - filename = namegen.sequence(filename, shared.opts.outdir_video, '') + filename = namegen.sequence(filename) else: if os.path.sep not in filename: filename = os.path.join(shared.opts.outdir_video, filename) diff --git a/modules/video_models/models_def.py b/modules/video_models/models_def.py index e6f570775..c95fe8442 100644 --- a/modules/video_models/models_def.py +++ b/modules/video_models/models_def.py @@ -9,13 +9,16 @@ class Model(): url: str = '' repo: str = None repo_cls: classmethod = None + repo_revision: str = None dit: str = None dit_cls: classmethod = None dit_folder: str = 'transformer' + dit_revision: str = None te: str = None te_cls: classmethod = None te_folder: str = 'text_encoder' te_hijack: bool = True + te_revision: str = None image_hijack: bool = True vae_hijack: bool = True vae_remote: bool = False @@ -194,6 +197,56 @@ models = { repo_cls=diffusers.WanImageToVideoPipeline, te_cls=transformers.T5EncoderModel, dit_cls=diffusers.WanTransformer3DModel), + Model(name='WAN 2.1 VACE 1.3B', + url='https://huggingface.co/Wan-AI/Wan2.1-VACE-1.3B-diffusers', + repo='Wan-AI/Wan2.1-VACE-1.3B-diffusers', + repo_cls=diffusers.WanVACEPipeline, + te_cls=transformers.T5EncoderModel, + dit_cls=diffusers.WanTransformer3DModel), + Model(name='WAN 2.1 VACE 14B', + url='https://huggingface.co/Wan-AI/Wan2.1-VACE-14B-diffusers', + repo='Wan-AI/Wan2.1-VACE-14B-diffusers', + repo_cls=diffusers.WanVACEPipeline, + te_cls=transformers.T5EncoderModel, + dit_cls=diffusers.WanTransformer3DModel), + ], + 'SkyReels V2': [ + Model(name='None'), + Model(name='SkyReels-V2 T2I-DF 1.3B-540P', + url='https://huggingface.co/Skywork/SkyReels-V2-DF-1.3B-540P-Diffusers', + repo='Skywork/SkyReels-V2-DF-1.3B-540P-Diffusers', + repo_cls=diffusers.SkyReelsV2DiffusionForcingPipeline, + repo_revision='refs/pr/1', + te_cls=transformers.UMT5EncoderModel, + dit_cls=diffusers.SkyReelsV2Transformer3DModel), + Model(name='SkyReels-V2 T2I-DF 14B-720P', + url='https://huggingface.co/Skywork/SkyReels-V2-DF-14B-720P-Diffusers', + repo='Skywork/SkyReels-V2-DF-14B-720P-Diffusers', + repo_cls=diffusers.SkyReelsV2DiffusionForcingPipeline, + repo_revision='refs/pr/1', + te_cls=transformers.UMT5EncoderModel, + dit_cls=diffusers.SkyReelsV2Transformer3DModel), + Model(name='SkyReels-V2 I2I-DF 14B-720P', + url='https://huggingface.co/Skywork/SkyReels-V2-DF-14B-720P-Diffusers', + repo='Skywork/SkyReels-V2-DF-14B-720P-Diffusers', + repo_cls=diffusers.SkyReelsV2DiffusionForcingImageToVideoPipeline, + repo_revision='refs/pr/1', + te_cls=transformers.UMT5EncoderModel, + dit_cls=diffusers.SkyReelsV2Transformer3DModel), + Model(name='SkyReels-V2 T2I 14B-720P', + url='https://huggingface.co/Skywork/SkyReels-V2-T2V-14B-720P-Diffusers', + repo='Skywork/SkyReels-V2-T2V-14B-720P-Diffusers', + repo_cls=diffusers.SkyReelsV2Pipeline, + repo_revision='refs/pr/1', + te_cls=transformers.UMT5EncoderModel, + dit_cls=diffusers.SkyReelsV2Transformer3DModel), + Model(name='SkyReels-V2 I2I 14B-720P', + url='https://huggingface.co/Skywork/SkyReels-V2-I2V-14B-720P-Diffusers', + repo='Skywork/SkyReels-V2-I2V-14B-720P-Diffusers', + repo_cls=diffusers.SkyReelsV2ImageToVideoPipeline, + repo_revision='refs/pr/1', + te_cls=transformers.UMT5EncoderModel, + dit_cls=diffusers.SkyReelsV2Transformer3DModel), ], 'Mochi Video': [ Model(name='None'), diff --git a/modules/video_models/video_load.py b/modules/video_models/video_load.py index 3eda8f735..53b4d66a0 100644 --- a/modules/video_models/video_load.py +++ b/modules/video_models/video_load.py @@ -28,6 +28,7 @@ def load_model(selected: models_def.Model): text_encoder = selected.te_cls.from_pretrained( pretrained_model_name_or_path=selected.te or selected.repo, subfolder=selected.te_folder, + revision=selected.te_revision or selected.repo_revision, cache_dir=shared.opts.hfcache_dir, torch_dtype=devices.dtype, **quant_args @@ -44,6 +45,7 @@ def load_model(selected: models_def.Model): transformer = selected.dit_cls.from_pretrained( pretrained_model_name_or_path=selected.dit or selected.repo, subfolder=selected.dit_folder, + revision=selected.dit_revision or selected.repo_revision, torch_dtype=devices.dtype, cache_dir=shared.opts.hfcache_dir, **quant_args @@ -63,6 +65,7 @@ def load_model(selected: models_def.Model): pretrained_model_name_or_path=selected.repo, transformer=transformer, text_encoder=text_encoder, + revision=selected.repo_revision, cache_dir=shared.opts.hfcache_dir, torch_dtype=devices.dtype, **kwargs, diff --git a/modules/video_models/video_overrides.py b/modules/video_models/video_overrides.py index 655256e95..07cf1cd23 100644 --- a/modules/video_models/video_overrides.py +++ b/modules/video_models/video_overrides.py @@ -10,10 +10,13 @@ debug = shared.log.trace if os.environ.get('SD_VIDEO_DEBUG', None) is not None e def load_override(selected: Model): kwargs = {} - if selected.name == 'Allegro T2V': + # Allegro + if 'Allegro T2V' in selected.name: kwargs['vae'] = diffusers.AutoencoderKLAllegro.from_pretrained(selected.repo, subfolder="vae", torch_dtype=torch.float32, cache_dir=shared.opts.hfcache_dir) - if selected.name == 'LTXVideo 0.9.5 I2V': + # LTX + if 'LTXVideo 0.9.5 I2V' in selected.name: kwargs['vae'] = diffusers.AutoencoderKLLTXVideo.from_pretrained(selected.repo, subfolder="vae", torch_dtype=torch.float32, cache_dir=shared.opts.hfcache_dir) + # WAN if 'WAN 2.1 14B' in selected.name: kwargs['vae'] = diffusers.AutoencoderKLWan.from_pretrained(selected.repo, subfolder="vae", torch_dtype=torch.float32, cache_dir=shared.opts.hfcache_dir) debug(f'Video overrides: model="{selected.name}" kwargs={list(kwargs)}') diff --git a/modules/video_models/video_run.py b/modules/video_models/video_run.py index 01e3fd387..8e52ca2c1 100644 --- a/modules/video_models/video_run.py +++ b/modules/video_models/video_run.py @@ -52,6 +52,7 @@ def generate(*args, **kwargs): p.state = ui_state p.do_not_save_grid = True p.do_not_save_samples = not save_frames + p.outpath_samples = shared.opts.outdir_samples or shared.opts.outdir_video if 'I2V' in model: if init_image is None: return video_utils.queue_err('init image not set') diff --git a/modules/video_models/video_utils.py b/modules/video_models/video_utils.py index cbca9a3f5..47faeac13 100644 --- a/modules/video_models/video_utils.py +++ b/modules/video_models/video_utils.py @@ -14,7 +14,7 @@ def queue_err(msg): def get_url(url): - return f'  {url}

' if url else '

' + return f'{url}

' if url else '

' def check_av(): diff --git a/modules/zluda.py b/modules/zluda.py index 8c7802c38..12186ffa7 100644 --- a/modules/zluda.py +++ b/modules/zluda.py @@ -2,7 +2,6 @@ import sys from typing import Union import torch from torch._prims_common import DeviceLikeType -import onnxruntime as ort from modules import shared, devices, zluda_installer from modules.zluda_installer import core, default_agent # pylint: disable=unused-import from modules.onnx_impl.execution_providers import available_execution_providers, ExecutionProvider @@ -42,9 +41,14 @@ def initialize_zluda(): torch.backends.cuda.enable_mem_efficient_sdp = do_nothing # ONNX Runtime is not supported - ort.capi._pybind_state.get_available_providers = lambda: [v for v in available_execution_providers if v != ExecutionProvider.CUDA] # pylint: disable=protected-access - ort.get_available_providers = ort.capi._pybind_state.get_available_providers # pylint: disable=protected-access - if shared.opts.onnx_execution_provider == ExecutionProvider.CUDA: + try: + import onnxruntime as ort + ort.capi._pybind_state.get_available_providers = lambda: [v for v in available_execution_providers if v != ExecutionProvider.CUDA] # pylint: disable=protected-access + ort.get_available_providers = ort.capi._pybind_state.get_available_providers # pylint: disable=protected-access + if shared.opts.onnx_execution_provider == ExecutionProvider.CUDA: + shared.opts.onnx_execution_provider = ExecutionProvider.CPU + except Exception as e: + shared.log.warning(f'ZLUDA ONNX runtime: {e}') shared.opts.onnx_execution_provider = ExecutionProvider.CPU device = devices.get_optimal_device() diff --git a/modules/zluda_installer.py b/modules/zluda_installer.py index 3afeb4680..da0ee290f 100644 --- a/modules/zluda_installer.py +++ b/modules/zluda_installer.py @@ -17,7 +17,7 @@ DLL_MAPPING = { 'cufftw.dll': 'cufftw64_10.dll', 'nvrtc.dll': 'nvrtc64_112_0.dll', } -HIPSDK_TARGETS = ['rocblas.dll', 'rocsolver.dll', 'hipfft.dll',] +HIPSDK_TARGETS = ['rocblas.dll', 'rocsolver.dll', 'rocsparse.dll', 'hipfft.dll',] MIOpen_enabled = False diff --git a/package.json b/package.json index 07b2681ec..1b0b95bca 100644 --- a/package.json +++ b/package.json @@ -1,7 +1,7 @@ { "name": "@vladmandic/sdnext", "version": "dev", - "description": "SD.Next: Opinionated implementation of Stable Diffusion", + "description": "SD.Next: All-in-one WebUI for AI generative image and video creation", "author": "Vladimir Mandic ", "bugs": { "url": "https://github.com/vladmandic/sdnext/issues" @@ -19,7 +19,7 @@ "venv": ". venv/bin/activate", "start": ". venv/bin/activate; python launch.py --debug", "localize": "node cli/localize.js", - "packages": ". venv/bin/activate && pip install --upgrade transformers accelerate huggingface_hub safetensors tokenizers peft compel pytorch_lightning", + "packages": ". venv/bin/activate && pip install --upgrade transformers accelerate huggingface_hub safetensors tokenizers peft compel pytorch_lightning pylint ruff", "eslint": "eslint . javascript/ extensions-builtin/sdnext-modernui/javascript/", "ruff": ". venv/bin/activate && ruff check", "pylint": ". venv/bin/activate && pylint *.py modules/ pipelines/ scripts/ extensions-builtin/ | grep -v '^*'", diff --git a/pipelines/chroma/convert_chroma.py b/pipelines/chroma/convert_chroma.py deleted file mode 100644 index ea49d3bee..000000000 --- a/pipelines/chroma/convert_chroma.py +++ /dev/null @@ -1,95 +0,0 @@ - -import os -import torch -import transformers -import diffusers -import huggingface_hub as hf -from rich import print as rprint -from rich.traceback import install as install_traceback - - -convert = True -test = False -upload = True -input_files = [ - 'chroma-unlocked-v48.safetensors', - 'chroma-unlocked-v48-detail-calibrated.safetensors', - 'chroma-unlocked-v46-flash.safetensors', -] -input_folder = '/mnt/models/UNET' -output_folder = '/mnt/models/Diffusers' -cache_dir = '/mnt/models/huggingface' -hf_token = '' -dtype = torch.bfloat16 -device = torch.device('cuda') - - -rprint('starting chroma conversion') -install_traceback(show_locals=False) -rprint(f'torch={torch.__version__} diffusers={diffusers.__version__} transformers={transformers.__version__}') -for input_file in input_files: - input_basename = os.path.splitext(input_file)[0] - input_model = os.path.join(input_folder, input_file) - output_model = os.path.join(output_folder, input_basename) - - if convert: - rprint(f'load transformer: {input_model}') - transformer = diffusers.ChromaTransformer2DModel.from_single_file( - input_model, - torch_dtype=dtype, - cache_dir=cache_dir, - ).to(device) - - rprint('load text-encoder') - text_encoder = transformers.T5EncoderModel.from_pretrained( - "black-forest-labs/FLUX.1-schnell", - subfolder="text_encoder_2", - torch_dtype=dtype, - cache_dir=cache_dir, - ).to(device) - - rprint('load tokenizer') - tokenizer = transformers.T5Tokenizer.from_pretrained( - "black-forest-labs/FLUX.1-schnell", - subfolder="tokenizer_2", - cache_dir=cache_dir, - ) - - rprint('load pipeline') - pipe = diffusers.ChromaPipeline.from_pretrained( - "black-forest-labs/FLUX.1-dev", - transformer=transformer, - text_encoder=text_encoder, - tokenizer=tokenizer, - torch_dtype=dtype, - cache_dir=cache_dir, - ).to(device) - - - rprint(f'save pipeline: {output_model}') - pipe.save_pretrained( - output_model, - ) - - if test: - rprint('test load') - pipe = diffusers.ChromaPipeline.from_pretrained( - output_model, - torch_dtype=dtype, - cache_dir=cache_dir, - ) - - if upload: - rprint('hf login') - hf.logout() - hf.login(token=hf_token, add_to_git_credential=False, write_permission=True) - rprint('upload model') - pipe.push_to_hub( - input_basename, - private=False, - token=hf_token, - ) - - pipe = None - -rprint('done') diff --git a/pipelines/flux/flux_bnb.py b/pipelines/flux/flux_bnb.py new file mode 100644 index 000000000..777678af1 --- /dev/null +++ b/pipelines/flux/flux_bnb.py @@ -0,0 +1,25 @@ +import diffusers +import transformers +from modules import devices, model_quant + + +def load_flux_bnb(checkpoint_info, diffusers_load_config): # pylint: disable=unused-argument + transformer = None + if isinstance(checkpoint_info, str): + repo_path = checkpoint_info + else: + repo_path = checkpoint_info.path + model_quant.load_bnb('Load model: type=FLUX') + quant = model_quant.get_quant(repo_path) + if quant == 'fp8': + quantization_config = transformers.BitsAndBytesConfig(load_in_8bit=True, bnb_4bit_compute_dtype=devices.dtype) + transformer = diffusers.FluxTransformer2DModel.from_single_file(repo_path, **diffusers_load_config, quantization_config=quantization_config) + elif quant == 'fp4': + quantization_config = transformers.BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=devices.dtype, bnb_4bit_quant_type= 'fp4') + transformer = diffusers.FluxTransformer2DModel.from_single_file(repo_path, **diffusers_load_config, quantization_config=quantization_config) + elif quant == 'nf4': + quantization_config = transformers.BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=devices.dtype, bnb_4bit_quant_type= 'nf4') + transformer = diffusers.FluxTransformer2DModel.from_single_file(repo_path, **diffusers_load_config, quantization_config=quantization_config) + else: + transformer = diffusers.FluxTransformer2DModel.from_single_file(repo_path, **diffusers_load_config) + return transformer diff --git a/pipelines/flux/flux_legacy_loader.py b/pipelines/flux/flux_legacy_loader.py new file mode 100644 index 000000000..6b6f9d294 --- /dev/null +++ b/pipelines/flux/flux_legacy_loader.py @@ -0,0 +1,360 @@ +import os +import json +import torch +import diffusers +import transformers +from safetensors.torch import load_file +from huggingface_hub import hf_hub_download +from modules import shared, errors, devices, sd_models, sd_unet, model_te, model_quant, sd_hijack_te + + +debug = shared.log.trace if os.environ.get('SD_LOAD_DEBUG', None) is not None else lambda *args, **kwargs: None + + +def load_flux_quanto(checkpoint_info): + transformer, text_encoder_2 = None, None + quanto = model_quant.load_quanto('Load model: type=FLUX') + + if isinstance(checkpoint_info, str): + repo_path = checkpoint_info + else: + repo_path = checkpoint_info.path + + try: + quantization_map = os.path.join(repo_path, "transformer", "quantization_map.json") + debug(f'Load model: type=FLUX quantization map="{quantization_map}" repo="{checkpoint_info.name}" component="transformer"') + if not os.path.exists(quantization_map): + repo_id = sd_models.path_to_repo(checkpoint_info) + quantization_map = hf_hub_download(repo_id, subfolder='transformer', filename='quantization_map.json', cache_dir=shared.opts.diffusers_dir) + with open(quantization_map, "r", encoding='utf8') as f: + quantization_map = json.load(f) + state_dict = load_file(os.path.join(repo_path, "transformer", "diffusion_pytorch_model.safetensors")) + dtype = state_dict['context_embedder.bias'].dtype + with torch.device("meta"): + transformer = diffusers.FluxTransformer2DModel.from_config(os.path.join(repo_path, "transformer", "config.json")).to(dtype=dtype) + quanto.requantize(transformer, state_dict, quantization_map, device=torch.device("cpu")) + transformer_dtype = transformer.dtype + if transformer_dtype != devices.dtype: + try: + transformer = transformer.to(dtype=devices.dtype) + except Exception: + shared.log.error(f"Load model: type=FLUX Failed to cast transformer to {devices.dtype}, set dtype to {transformer_dtype}") + except Exception as e: + shared.log.error(f"Load model: type=FLUX failed to load Quanto transformer: {e}") + if debug: + errors.display(e, 'FLUX Quanto:') + + try: + quantization_map = os.path.join(repo_path, "text_encoder_2", "quantization_map.json") + debug(f'Load model: type=FLUX quantization map="{quantization_map}" repo="{checkpoint_info.name}" component="text_encoder_2"') + if not os.path.exists(quantization_map): + repo_id = sd_models.path_to_repo(checkpoint_info) + quantization_map = hf_hub_download(repo_id, subfolder='text_encoder_2', filename='quantization_map.json', cache_dir=shared.opts.diffusers_dir) + with open(quantization_map, "r", encoding='utf8') as f: + quantization_map = json.load(f) + with open(os.path.join(repo_path, "text_encoder_2", "config.json"), encoding='utf8') as f: + t5_config = transformers.T5Config(**json.load(f)) + state_dict = load_file(os.path.join(repo_path, "text_encoder_2", "model.safetensors")) + dtype = state_dict['encoder.block.0.layer.0.SelfAttention.relative_attention_bias.weight'].dtype + with torch.device("meta"): + text_encoder_2 = transformers.T5EncoderModel(t5_config).to(dtype=dtype) + quanto.requantize(text_encoder_2, state_dict, quantization_map, device=torch.device("cpu")) + text_encoder_2_dtype = text_encoder_2.dtype + if text_encoder_2_dtype != devices.dtype: + try: + text_encoder_2 = text_encoder_2.to(dtype=devices.dtype) + except Exception: + shared.log.error(f"Load model: type=FLUX Failed to cast text encoder to {devices.dtype}, set dtype to {text_encoder_2_dtype}") + except Exception as e: + shared.log.error(f"Load model: type=FLUX failed to load Quanto text encoder: {e}") + if debug: + errors.display(e, 'FLUX Quanto:') + + return transformer, text_encoder_2 + + +def load_flux_bnb(checkpoint_info, diffusers_load_config): # pylint: disable=unused-argument + transformer, text_encoder_2 = None, None + if isinstance(checkpoint_info, str): + repo_path = checkpoint_info + else: + repo_path = checkpoint_info.path + model_quant.load_bnb('Load model: type=FLUX') + quant = model_quant.get_quant(repo_path) + try: + if quant == 'fp8': + quantization_config = transformers.BitsAndBytesConfig(load_in_8bit=True, bnb_4bit_compute_dtype=devices.dtype) + debug(f'Quantization: {quantization_config}') + transformer = diffusers.FluxTransformer2DModel.from_single_file(repo_path, **diffusers_load_config, quantization_config=quantization_config) + elif quant == 'fp4': + quantization_config = transformers.BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=devices.dtype, bnb_4bit_quant_type= 'fp4') + debug(f'Quantization: {quantization_config}') + transformer = diffusers.FluxTransformer2DModel.from_single_file(repo_path, **diffusers_load_config, quantization_config=quantization_config) + elif quant == 'nf4': + quantization_config = transformers.BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=devices.dtype, bnb_4bit_quant_type= 'nf4') + debug(f'Quantization: {quantization_config}') + transformer = diffusers.FluxTransformer2DModel.from_single_file(repo_path, **diffusers_load_config, quantization_config=quantization_config) + else: + transformer = diffusers.FluxTransformer2DModel.from_single_file(repo_path, **diffusers_load_config) + except Exception as e: + shared.log.error(f"Load model: type=FLUX failed to load BnB transformer: {e}") + transformer, text_encoder_2 = None, None + if debug: + errors.display(e, 'FLUX:') + return transformer, text_encoder_2 + + +def load_quants(kwargs, repo_id, cache_dir, allow_quant): # pylint: disable=unused-argument + try: + diffusers_load_config = { + "torch_dtype": devices.dtype, + "cache_dir": cache_dir, + } + if 'transformer' not in kwargs and model_quant.check_nunchaku('Model'): + import nunchaku + nunchaku_precision = nunchaku.utils.get_precision() + nunchaku_repo = None + if 'flux.1-kontext' in repo_id.lower(): + nunchaku_repo = f"mit-han-lab/nunchaku-flux.1-kontext-dev/svdq-{nunchaku_precision}_r32-flux.1-kontext-dev.safetensors" + elif 'flux.1-dev' in repo_id.lower(): + nunchaku_repo = f"mit-han-lab/nunchaku-flux.1-dev/svdq-{nunchaku_precision}_r32-flux.1-dev.safetensors" + elif 'flux.1-schnell' in repo_id.lower(): + nunchaku_repo = f"mit-han-lab/nunchaku-flux.1-schnell/svdq-{nunchaku_precision}_r32-flux.1-schnell.safetensors" + elif 'flux.1-fill' in repo_id.lower(): + nunchaku_repo = f"mit-han-lab/svdq-fp4-flux.1-fill-dev/svdq-{nunchaku_precision}_r32-flux.1-schnell.safetensors" + elif 'flux.1-depth' in repo_id.lower(): + nunchaku_repo = f"mit-han-lab/svdq-int4-flux.1-depth-dev/svdq-{nunchaku_precision}_r32-flux.1-schnell.safetensors" + elif 'shuttle' in repo_id.lower(): + nunchaku_repo = f"mit-han-lab/nunchaku-shuttle-jaguar/svdq-{nunchaku_precision}_r32-shuttle-jaguar.safetensors" + else: + shared.log.error(f'Load module: quant=Nunchaku module=transformer repo="{repo_id}" unsupported') + if nunchaku_repo is not None: + shared.log.debug(f'Load module: quant=Nunchaku module=transformer repo="{nunchaku_repo}" precision={nunchaku_precision} offload={shared.opts.nunchaku_offload} attention={shared.opts.nunchaku_attention}') + kwargs['transformer'] = nunchaku.NunchakuFluxTransformer2dModel.from_pretrained(nunchaku_repo, offload=shared.opts.nunchaku_offload, torch_dtype=devices.dtype) + kwargs['transformer'].quantization_method = 'SVDQuant' + if shared.opts.nunchaku_attention: + kwargs['transformer'].set_attention_impl("nunchaku-fp16") + if 'transformer' not in kwargs and model_quant.check_quant('Model'): + load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='Model', device_map=True) + kwargs['transformer'] = diffusers.FluxTransformer2DModel.from_pretrained(repo_id, subfolder="transformer", **load_args, **quant_args) + if 'text_encoder_2' not in kwargs and model_quant.check_nunchaku('TE'): + import nunchaku + nunchaku_precision = nunchaku.utils.get_precision() + nunchaku_repo = 'mit-han-lab/nunchaku-t5/awq-int4-flux.1-t5xxl.safetensors' + shared.log.debug(f'Load module: quant=Nunchaku module=t5 repo="{nunchaku_repo}" precision={nunchaku_precision}') + kwargs['text_encoder_2'] = nunchaku.NunchakuT5EncoderModel.from_pretrained(nunchaku_repo, torch_dtype=devices.dtype) + kwargs['text_encoder_2'].quantization_method = 'SVDQuant' + if 'text_encoder_2' not in kwargs and model_quant.check_quant('TE'): + load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='TE', device_map=True) + kwargs['text_encoder_2'] = transformers.T5EncoderModel.from_pretrained(repo_id, subfolder="text_encoder_2", **load_args, **quant_args) + except Exception as e: + shared.log.error(f'Quantization: {e}') + errors.display(e, 'Quantization:') + return kwargs + + +def load_transformer(file_path): # triggered by opts.sd_unet change + if file_path is None or not os.path.exists(file_path): + return None + transformer = None + quant = model_quant.get_quant(file_path) + diffusers_load_config = { + "torch_dtype": devices.dtype, + "cache_dir": shared.opts.hfcache_dir, + } + if quant is not None and quant != 'none': + shared.log.info(f'Load module: type=UNet/Transformer file="{file_path}" offload={shared.opts.diffusers_offload_mode} prequant={quant} dtype={devices.dtype}') + if 'gguf' in file_path.lower(): + from modules import ggml + _transformer = ggml.load_gguf(file_path, cls=diffusers.FluxTransformer2DModel, compute_dtype=devices.dtype) + if _transformer is not None: + transformer = _transformer + elif quant == "fp8": + _transformer = model_quant.load_fp8_model_layerwise(file_path, diffusers.FluxTransformer2DModel.from_single_file, diffusers_load_config) + if _transformer is not None: + transformer = _transformer + elif quant in {'qint8', 'qint4'}: + _transformer, _text_encoder_2 = load_flux_quanto(file_path) + if _transformer is not None: + transformer = _transformer + elif quant in {'fp8', 'fp4', 'nf4'}: + _transformer, _text_encoder_2 = load_flux_bnb(file_path, diffusers_load_config) + if _transformer is not None: + transformer = _transformer + elif 'nf4' in quant: + from pipelines.flux.flux_nf4 import load_flux_nf4 + _transformer, _text_encoder_2 = load_flux_nf4(file_path, prequantized=True) + if _transformer is not None: + transformer = _transformer + else: + quant_args = model_quant.create_bnb_config({}) + if quant_args: + shared.log.info(f'Load module: type=Flux transformer file="{file_path}" offload={shared.opts.diffusers_offload_mode} quant=bnb dtype={devices.dtype}') + from pipelines.flux.flux_nf4 import load_flux_nf4 + transformer, _text_encoder_2 = load_flux_nf4(file_path, prequantized=False) + if transformer is not None: + return transformer + load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='Model', device_map=True) + shared.log.debug(f'Load model: type=Flux transformer file="{file_path}" offload={shared.opts.diffusers_offload_mode} args={load_args}') + transformer = diffusers.FluxTransformer2DModel.from_single_file(file_path, **load_args, **quant_args) + if transformer is None: + shared.log.error('Failed to load UNet model') + shared.opts.sd_unet = 'Default' + return transformer + + +def load_flux(checkpoint_info, diffusers_load_config): # triggered by opts.sd_checkpoint change + repo_id = sd_models.path_to_repo(checkpoint_info) + sd_models.hf_auth_check(checkpoint_info) + allow_post_quant = False + + prequantized = model_quant.get_quant(checkpoint_info.path) + 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}') + + transformer = None + text_encoder_1 = None + text_encoder_2 = None + vae = None + + # unload current model + sd_models.unload_model_weights() + shared.sd_model = None + devices.torch_gc(force=True, reason='load') + + if shared.opts.teacache_enabled: + from modules import teacache + shared.log.debug(f'Transformers cache: type=teacache patch=forward cls={diffusers.FluxTransformer2DModel.__name__}') + diffusers.FluxTransformer2DModel.forward = teacache.teacache_flux_forward # patch must be done before transformer is loaded + + # load overrides if any + 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 = '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 = 'Default' + if debug: + errors.display(e, 'FLUX UNet:') + 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 + if 'vit-l' in shared.opts.sd_text_encoder.lower(): + text_encoder_1 = load_vit_l() + else: + 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 = 'Default' + if debug: + errors.display(e, 'FLUX T5:') + 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 + # vae = sd_vae.load_vae_diffusers(None, sd_vae.vae_dict[shared.opts.sd_vae], 'override') + vae_file = sd_vae.vae_dict[shared.opts.sd_vae] + if os.path.exists(vae_file): + vae_config = os.path.join('configs', 'flux', 'vae', 'config.json') + 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 = 'Default' + if debug: + errors.display(e, 'FLUX VAE:') + + # load quantized components if any + if prequantized == 'nf4': + try: + from pipelines.flux.flux_nf4 import load_flux_nf4 + _transformer, _text_encoder = load_flux_nf4(checkpoint_info) + if _transformer is not None: + transformer = _transformer + if _text_encoder is not None: + text_encoder_2 = _text_encoder + except Exception as e: + shared.log.error(f"Load model: type=FLUX failed to load NF4 components: {e}") + if debug: + errors.display(e, 'FLUX NF4:') + if prequantized == 'qint8' or prequantized == 'qint4': + try: + _transformer, _text_encoder = load_flux_quanto(checkpoint_info) + if _transformer is not None: + transformer = _transformer + if _text_encoder is not None: + text_encoder_2 = _text_encoder + except Exception as e: + shared.log.error(f"Load model: type=FLUX failed to load Quanto components: {e}") + if debug: + errors.display(e, 'FLUX Quanto:') + + # initialize pipeline with pre-loaded components + kwargs = {} + if transformer is not None: + kwargs['transformer'] = transformer + sd_unet.loaded_unet = shared.opts.sd_unet + if text_encoder_1 is not None: + kwargs['text_encoder'] = text_encoder_1 + model_te.loaded_te = shared.opts.sd_text_encoder + if text_encoder_2 is not None: + kwargs['text_encoder_2'] = text_encoder_2 + model_te.loaded_te = shared.opts.sd_text_encoder + if vae is not None: + kwargs['vae'] = vae + 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 + if 'Fill' in repo_id: + cls = diffusers.FluxFillPipeline + elif 'Canny' in repo_id: + cls = diffusers.FluxControlPipeline + elif 'Depth' in repo_id: + cls = diffusers.FluxControlPipeline + elif 'Kontext' in repo_id: + cls = diffusers.FluxKontextPipeline + from diffusers import pipelines + pipelines.auto_pipeline.AUTO_TEXT2IMAGE_PIPELINES_MAPPING["flux1kontext"] = diffusers.FluxKontextPipeline + pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING["flux1kontext"] = diffusers.FluxKontextPipeline + pipelines.auto_pipeline.AUTO_INPAINT_PIPELINES_MAPPING["flux1kontext"] = diffusers.FluxKontextInpaintPipeline + + else: + cls = diffusers.FluxPipeline + shared.log.debug(f'Load model: type=FLUX cls={cls.__name__} preloaded={list(kwargs)} revision={diffusers_load_config.get("revision", None)}') + for c in kwargs: + if getattr(kwargs[c], 'quantization_method', None) is not None or getattr(kwargs[c], 'gguf', None) is not None: + shared.log.debug(f'Load model: type=FLUX component={c} dtype={kwargs[c].dtype} quant={getattr(kwargs[c], "quantization_method", None) or getattr(kwargs[c], "gguf", None)}') + if kwargs[c].dtype == torch.float32 and devices.dtype != torch.float32: + try: + kwargs[c] = kwargs[c].to(dtype=devices.dtype) + shared.log.warning(f'Load model: type=FLUX component={c} dtype={kwargs[c].dtype} cast dtype={devices.dtype} recast') + except Exception: + pass + + 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) + if fn.endswith('.safetensors') and os.path.isfile(fn): + pipe = cls.from_single_file(fn, cache_dir=shared.opts.diffusers_dir, **kwargs, **diffusers_load_config) + allow_post_quant = True + else: + pipe = cls.from_pretrained(repo_id, cache_dir=shared.opts.diffusers_dir, **kwargs, **diffusers_load_config) + + if shared.opts.teacache_enabled and model_quant.check_nunchaku('Model'): + from nunchaku.caching.diffusers_adapters import apply_cache_on_pipe + apply_cache_on_pipe(pipe, residual_diff_threshold=0.12) + + # release memory + transformer = None + text_encoder_1 = None + text_encoder_2 = None + vae = None + for k in kwargs.keys(): + kwargs[k] = None + sd_hijack_te.init_hijack(pipe) + devices.torch_gc(force=True, reason='load') + return pipe, allow_post_quant diff --git a/pipelines/model_flux_nf4.py b/pipelines/flux/flux_nf4.py similarity index 100% rename from pipelines/model_flux_nf4.py rename to pipelines/flux/flux_nf4.py diff --git a/pipelines/flux/flux_nunchaku.py b/pipelines/flux/flux_nunchaku.py new file mode 100644 index 000000000..e21b93a3b --- /dev/null +++ b/pipelines/flux/flux_nunchaku.py @@ -0,0 +1,29 @@ +from modules import shared, devices + + +def load_flux_nunchaku(repo_id): + import nunchaku + nunchaku_precision = nunchaku.utils.get_precision() + nunchaku_repo = None + transformer = None + if 'flux.1-kontext' in repo_id.lower(): + nunchaku_repo = f"mit-han-lab/nunchaku-flux.1-kontext-dev/svdq-{nunchaku_precision}_r32-flux.1-kontext-dev.safetensors" + elif 'flux.1-dev' in repo_id.lower(): + nunchaku_repo = f"mit-han-lab/nunchaku-flux.1-dev/svdq-{nunchaku_precision}_r32-flux.1-dev.safetensors" + elif 'flux.1-schnell' in repo_id.lower(): + nunchaku_repo = f"mit-han-lab/nunchaku-flux.1-schnell/svdq-{nunchaku_precision}_r32-flux.1-schnell.safetensors" + elif 'flux.1-fill' in repo_id.lower(): + nunchaku_repo = f"mit-han-lab/svdq-fp4-flux.1-fill-dev/svdq-{nunchaku_precision}_r32-flux.1-schnell.safetensors" + elif 'flux.1-depth' in repo_id.lower(): + nunchaku_repo = f"mit-han-lab/svdq-int4-flux.1-depth-dev/svdq-{nunchaku_precision}_r32-flux.1-schnell.safetensors" + elif 'shuttle' in repo_id.lower(): + nunchaku_repo = f"mit-han-lab/nunchaku-shuttle-jaguar/svdq-{nunchaku_precision}_r32-shuttle-jaguar.safetensors" + else: + shared.log.error(f'Load module: quant=Nunchaku module=transformer repo="{repo_id}" unsupported') + if nunchaku_repo is not None: + shared.log.debug(f'Load module: quant=Nunchaku module=transformer repo="{nunchaku_repo}" precision={nunchaku_precision} offload={shared.opts.nunchaku_offload} attention={shared.opts.nunchaku_attention}') + transformer = nunchaku.NunchakuFluxTransformer2dModel.from_pretrained(nunchaku_repo, offload=shared.opts.nunchaku_offload, torch_dtype=devices.dtype) + transformer.quantization_method = 'SVDQuant' + if shared.opts.nunchaku_attention: + transformer.set_attention_impl("nunchaku-fp16") + return transformer diff --git a/pipelines/flux/flux_quanto.py b/pipelines/flux/flux_quanto.py new file mode 100644 index 000000000..11e604b62 --- /dev/null +++ b/pipelines/flux/flux_quanto.py @@ -0,0 +1,73 @@ +import os +import json +import torch +import diffusers +import transformers +from safetensors.torch import load_file +from huggingface_hub import hf_hub_download +from modules import shared, errors, devices, sd_models, model_quant + + +debug = shared.log.trace if os.environ.get('SD_LOAD_DEBUG', None) is not None else lambda *args, **kwargs: None + + +def load_flux_quanto(checkpoint_info): + transformer, text_encoder_2 = None, None + quanto = model_quant.load_quanto('Load model: type=FLUX') + + if isinstance(checkpoint_info, str): + repo_path = checkpoint_info + else: + repo_path = checkpoint_info.path + + try: + quantization_map = os.path.join(repo_path, "transformer", "quantization_map.json") + debug(f'Load model: type=FLUX quantization map="{quantization_map}" repo="{checkpoint_info.name}" component="transformer"') + if not os.path.exists(quantization_map): + repo_id = sd_models.path_to_repo(checkpoint_info) + quantization_map = hf_hub_download(repo_id, subfolder='transformer', filename='quantization_map.json', cache_dir=shared.opts.diffusers_dir) + with open(quantization_map, "r", encoding='utf8') as f: + quantization_map = json.load(f) + state_dict = load_file(os.path.join(repo_path, "transformer", "diffusion_pytorch_model.safetensors")) + dtype = state_dict['context_embedder.bias'].dtype + with torch.device("meta"): + transformer = diffusers.FluxTransformer2DModel.from_config(os.path.join(repo_path, "transformer", "config.json")).to(dtype=dtype) + quanto.requantize(transformer, state_dict, quantization_map, device=torch.device("cpu")) + transformer_dtype = transformer.dtype + if transformer_dtype != devices.dtype: + try: + transformer = transformer.to(dtype=devices.dtype) + except Exception: + shared.log.error(f"Load model: type=FLUX Failed to cast transformer to {devices.dtype}, set dtype to {transformer_dtype}") + except Exception as e: + shared.log.error(f"Load model: type=FLUX failed to load Quanto transformer: {e}") + if debug: + errors.display(e, 'FLUX Quanto:') + + try: + quantization_map = os.path.join(repo_path, "text_encoder_2", "quantization_map.json") + debug(f'Load model: type=FLUX quantization map="{quantization_map}" repo="{checkpoint_info.name}" component="text_encoder_2"') + if not os.path.exists(quantization_map): + repo_id = sd_models.path_to_repo(checkpoint_info) + quantization_map = hf_hub_download(repo_id, subfolder='text_encoder_2', filename='quantization_map.json', cache_dir=shared.opts.diffusers_dir) + with open(quantization_map, "r", encoding='utf8') as f: + quantization_map = json.load(f) + with open(os.path.join(repo_path, "text_encoder_2", "config.json"), encoding='utf8') as f: + t5_config = transformers.T5Config(**json.load(f)) + state_dict = load_file(os.path.join(repo_path, "text_encoder_2", "model.safetensors")) + dtype = state_dict['encoder.block.0.layer.0.SelfAttention.relative_attention_bias.weight'].dtype + with torch.device("meta"): + text_encoder_2 = transformers.T5EncoderModel(t5_config).to(dtype=dtype) + quanto.requantize(text_encoder_2, state_dict, quantization_map, device=torch.device("cpu")) + text_encoder_2_dtype = text_encoder_2.dtype + if text_encoder_2_dtype != devices.dtype: + try: + text_encoder_2 = text_encoder_2.to(dtype=devices.dtype) + except Exception: + shared.log.error(f"Load model: type=FLUX Failed to cast text encoder to {devices.dtype}, set dtype to {text_encoder_2_dtype}") + except Exception as e: + shared.log.error(f"Load model: type=FLUX failed to load Quanto text encoder: {e}") + if debug: + errors.display(e, 'FLUX Quanto:') + + return transformer, text_encoder_2 diff --git a/pipelines/generic.py b/pipelines/generic.py new file mode 100644 index 000000000..8416950aa --- /dev/null +++ b/pipelines/generic.py @@ -0,0 +1,162 @@ +import os +import json +import diffusers +import transformers +from modules import shared, devices, errors, sd_models, model_quant + + +debug = os.environ.get('SD_LOAD_DEBUG', None) is not None + + +def load_transformer(repo_id, cls_name, load_config={}, subfolder="transformer", allow_quant=True, variant=None, dtype=None, modules_to_not_convert=[], modules_dtype_dict={}): + transformer = None + try: + load_args, quant_args = model_quant.get_dit_args(load_config, module='Model', device_map=True, allow_quant=allow_quant, modules_to_not_convert=modules_to_not_convert, modules_dtype_dict=modules_dtype_dict) + quant_type = model_quant.get_quant_type(quant_args) + dtype = dtype or devices.dtype + + local_file = None + if shared.opts.sd_unet is not None and shared.opts.sd_unet != 'Default': + from modules import sd_unet + if shared.opts.sd_unet not in list(sd_unet.unet_dict): + shared.log.error(f'Load module: type=transformer file="{shared.opts.sd_unet}" not found') + elif os.path.exists(sd_unet.unet_dict[shared.opts.sd_unet]): + local_file = sd_unet.unet_dict[shared.opts.sd_unet] + + if local_file is not None and local_file.lower().endswith('.gguf'): + shared.log.debug(f'Load model: transformer="{local_file}" cls={cls_name.__name__} quant="{quant_type}" args={load_args}') + from modules import ggml + ggml.install_gguf() + loader = cls_name.from_single_file if hasattr(cls_name, 'from_single_file') else cls_name.from_pretrained + transformer = loader( + local_file, + quantization_config=diffusers.GGUFQuantizationConfig(compute_dtype=dtype), + cache_dir=shared.opts.hfcache_dir, + **load_args, + ) + transformer = model_quant.do_post_load_quant(transformer, allow=quant_type is not None) + elif local_file is not None and local_file.lower().endswith('.safetensors'): + shared.log.debug(f'Load model: transformer="{local_file}" cls={cls_name.__name__} quant="{quant_type}" args={load_args}') + if dtype is not None: + load_args['torch_dtype'] = dtype + loader = cls_name.from_single_file if hasattr(cls_name, 'from_single_file') else cls_name.from_pretrained + transformer = loader( + local_file, + cache_dir=shared.opts.hfcache_dir, + **load_args, + **quant_args, + ) + else: + shared.log.debug(f'Load model: transformer="{repo_id}" cls={cls_name.__name__} quant="{quant_type}" args={load_args}') + if dtype is not None: + load_args['torch_dtype'] = dtype + if subfolder is not None: + load_args['subfolder'] = subfolder + if variant is not None: + load_args['variant'] = variant + transformer = cls_name.from_pretrained( + repo_id, + cache_dir=shared.opts.hfcache_dir, + **load_args, + **quant_args, + ) + sd_models.allow_post_quant = False # we already handled it + if shared.opts.diffusers_offload_mode != 'none' and transformer is not None: + sd_models.move_model(transformer, devices.cpu) + except Exception as e: + shared.log.error(f'Load model: type=transformer {e}') + if debug: + errors.display(e, 'Load:') + raise + return transformer + + +def load_text_encoder(repo_id, cls_name, load_config={}, subfolder="text_encoder", allow_quant=True, allow_shared=True, variant=None, dtype=None, modules_to_not_convert=[], modules_dtype_dict={}): + text_encoder = None + try: + load_args, quant_args = model_quant.get_dit_args(load_config, module='TE', device_map=True, allow_quant=allow_quant, modules_to_not_convert=modules_to_not_convert, modules_dtype_dict=modules_dtype_dict) + quant_type = model_quant.get_quant_type(quant_args) + dtype = dtype or devices.dtype + + # load from local file if specified + local_file = None + if shared.opts.sd_text_encoder is not None and shared.opts.sd_text_encoder != 'Default': + from modules import model_te + if shared.opts.sd_text_encoder not in list(model_te.te_dict): + shared.log.error(f'Load module: type=te file="{shared.opts.sd_text_encoder}" not found') + elif os.path.exists(model_te.te_dict[shared.opts.sd_text_encoder]): + local_file = model_te.te_dict[shared.opts.sd_text_encoder] + + # load from local file gguf + if local_file is not None and local_file.lower().endswith('.gguf'): + shared.log.debug(f'Load model: text_encoder="{local_file}" cls={cls_name.__name__} quant="{quant_type}"') + """ + from modules import ggml + ggml.install_gguf() + text_encoder = cls_name.from_pretrained( + gguf_file=local_file, + quantization_config=diffusers.GGUFQuantizationConfig(compute_dtype=dtype), + cache_dir=shared.opts.hfcache_dir, + **load_args, + ) + text_encoder = model_quant.do_post_load_quant(text_encoder, allow=quant_type is not None) + """ + text_encoder = model_te.load_t5(local_file) + text_encoder = model_quant.do_post_load_quant(text_encoder, allow=quant_type is not None) + # load from local file safetensors + elif local_file is not None and local_file.lower().endswith('.safetensors'): + shared.log.debug(f'Load model: text_encoder="{local_file}" cls={cls_name.__name__} quant="{quant_type}"') + from modules import model_te + text_encoder = model_te.load_t5(local_file) + text_encoder = model_quant.do_post_load_quant(text_encoder, allow=quant_type is not None) + # use shared t5 if possible + elif cls_name == transformers.T5EncoderModel and allow_shared: + with open(os.path.join('configs', 'flux', 'text_encoder_2', 'config.json'), encoding='utf8') as f: + load_args['config'] = transformers.T5Config(**json.load(f)) + if model_quant.check_nunchaku('TE'): + import nunchaku + repo_id = 'nunchaku-tech/nunchaku-t5/awq-int4-flux.1-t5xxl.safetensors' + cls_name = nunchaku.NunchakuT5EncoderModel + shared.log.debug(f'Load model: text_encoder="{repo_id}" cls={cls_name.__name__} quant="SVDQuant"') + text_encoder = nunchaku.NunchakuT5EncoderModel.from_pretrained( + repo_id, + torch_dtype=dtype, + ) + text_encoder.quantization_method = 'SVDQuant' + elif shared.opts.te_shared_t5: + repo_id = 'Disty0/t5-xxl' + shared.log.debug(f'Load model: text_encoder="{repo_id}" cls={cls_name.__name__} quant="{quant_type}" shared={shared.opts.te_shared_t5}') + if dtype is not None: + load_args['torch_dtype'] = dtype + text_encoder = cls_name.from_pretrained( + repo_id, + cache_dir=shared.opts.hfcache_dir, + **load_args, + **quant_args, + ) + + # load from repo + if text_encoder is None: + shared.log.debug(f'Load model: text_encoder="{repo_id}" cls={cls_name.__name__} quant="{quant_type}" shared={shared.opts.te_shared_t5}') + if dtype is not None: + load_args['torch_dtype'] = dtype + if subfolder is not None: + load_args['subfolder'] = subfolder + if variant is not None: + load_args['variant'] = variant + text_encoder = cls_name.from_pretrained( + repo_id, + cache_dir=shared.opts.hfcache_dir, + **load_args, + **quant_args, + ) + + sd_models.allow_post_quant = False # we already handled it + if shared.opts.diffusers_offload_mode != 'none' and text_encoder is not None: + sd_models.move_model(text_encoder, devices.cpu) + except Exception as e: + shared.log.error(f'Load model: type=te {e}') + if debug: + errors.display(e, 'Load:') + raise + return text_encoder diff --git a/pipelines/meissonic/pipeline.py b/pipelines/meissonic/pipeline.py index 4f1bb05a2..34b894081 100644 --- a/pipelines/meissonic/pipeline.py +++ b/pipelines/meissonic/pipeline.py @@ -11,7 +11,6 @@ # 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 sys from typing import Any, Callable, Dict, List, Optional, Tuple, Union import torch @@ -49,7 +48,7 @@ def _prepare_latent_image_ids(batch_size, height, width, device, dtype): return latent_image_ids.to(device=device, dtype=dtype) -class Pipeline(DiffusionPipeline): +class MeissonicPipeline(DiffusionPipeline): image_processor: VaeImageProcessor vqvae: VQModel tokenizer: CLIPTokenizer @@ -212,27 +211,27 @@ class Pipeline(DiffusionPipeline): width = self.transformer.config.sample_size * self.vae_scale_factor if prompt_embeds is None: - input_ids = self.tokenizer( - prompt, - return_tensors="pt", - padding="max_length", - truncation=True, - max_length=77, #self.tokenizer.model_max_length, - ).input_ids.to(self._execution_device) - # input_ids_t5 = self.tokenizer_t5( - # prompt, - # return_tensors="pt", - # padding="max_length", - # truncation=True, - # max_length=512, - # ).input_ids.to(self._execution_device) + input_ids = self.tokenizer( + prompt, + return_tensors="pt", + padding="max_length", + truncation=True, + max_length=77, #self.tokenizer.model_max_length, + ).input_ids.to(self._execution_device) + # input_ids_t5 = self.tokenizer_t5( + # prompt, + # return_tensors="pt", + # padding="max_length", + # truncation=True, + # max_length=512, + # ).input_ids.to(self._execution_device) - outputs = self.text_encoder(input_ids, return_dict=True, output_hidden_states=True) - # outputs_t5 = self.text_encoder_t5(input_ids_t5, decoder_input_ids = input_ids_t5 ,return_dict=True, output_hidden_states=True) - prompt_embeds = outputs.text_embeds - encoder_hidden_states = outputs.hidden_states[-2] - # encoder_hidden_states = outputs_t5.encoder_hidden_states[-2] + outputs = self.text_encoder(input_ids, return_dict=True, output_hidden_states=True) + # outputs_t5 = self.text_encoder_t5(input_ids_t5, decoder_input_ids = input_ids_t5 ,return_dict=True, output_hidden_states=True) + prompt_embeds = outputs.text_embeds + encoder_hidden_states = outputs.hidden_states[-2] + # encoder_hidden_states = outputs_t5.encoder_hidden_states[-2] prompt_embeds = prompt_embeds.repeat(num_images_per_prompt, 1) encoder_hidden_states = encoder_hidden_states.repeat(num_images_per_prompt, 1, 1) diff --git a/pipelines/meissonic/pipeline_img2img.py b/pipelines/meissonic/pipeline_img2img.py index 13e5c3717..2aaf9d987 100644 --- a/pipelines/meissonic/pipeline_img2img.py +++ b/pipelines/meissonic/pipeline_img2img.py @@ -46,7 +46,7 @@ def _prepare_latent_image_ids(batch_size, height, width, device, dtype): return latent_image_ids.to(device=device, dtype=dtype) -class Img2ImgPipeline(DiffusionPipeline): +class MeissonicImg2ImgPipeline(DiffusionPipeline): image_processor: VaeImageProcessor vqvae: VQModel tokenizer: CLIPTokenizer diff --git a/pipelines/meissonic/pipeline_inpaint.py b/pipelines/meissonic/pipeline_inpaint.py index d405afa53..aa352d9b4 100644 --- a/pipelines/meissonic/pipeline_inpaint.py +++ b/pipelines/meissonic/pipeline_inpaint.py @@ -43,7 +43,7 @@ def _prepare_latent_image_ids(batch_size, height, width, device, dtype): return latent_image_ids.to(device=device, dtype=dtype) -class InpaintPipeline(DiffusionPipeline): +class MeissonicInpaintPipeline(DiffusionPipeline): image_processor: VaeImageProcessor vqvae: VQModel tokenizer: CLIPTokenizer diff --git a/pipelines/meissonic/test.py b/pipelines/meissonic/test.py index 5687cbff0..777f40e22 100644 --- a/pipelines/meissonic/test.py +++ b/pipelines/meissonic/test.py @@ -4,7 +4,7 @@ sys.path.append("./") # import torch # from torchvision import transforms from meissonic.transformer import Transformer2DModel as TransformerMeissonic -from meissonic.pipeline import Pipeline as PipelineMeissonic +from meissonic.pipeline import MeissonicPipeline from meissonic.scheduler import Scheduler as MeissonicScheduler from transformers import CLIPTextModelWithProjection, CLIPTokenizer from diffusers import VQModel @@ -21,7 +21,7 @@ vq_model = VQModel.from_pretrained(model_path, subfolder="vqvae", cache_dir=cach text_encoder = CLIPTextModelWithProjection.from_pretrained("laion/CLIP-ViT-H-14-laion2B-s32B-b79K", cache_dir=cache_dir) tokenizer = CLIPTokenizer.from_pretrained(model_path, subfolder="tokenizer") scheduler = MeissonicScheduler.from_pretrained(model_path, subfolder="scheduler") -pipe = PipelineMeissonic(vq_model, tokenizer=tokenizer, text_encoder=text_encoder, transformer=model, scheduler=scheduler) +pipe = MeissonicPipeline(vq_model, tokenizer=tokenizer, text_encoder=text_encoder, transformer=model, scheduler=scheduler) pipe = pipe.to(device) steps = 64 diff --git a/pipelines/model_auraflow.py b/pipelines/model_auraflow.py index 175ba12bf..9576f5923 100644 --- a/pipelines/model_auraflow.py +++ b/pipelines/model_auraflow.py @@ -1,21 +1,29 @@ -import os -import torch +import transformers import diffusers -from modules import shared, sd_models, devices - - -debug = shared.log.trace if os.environ.get('SD_LOAD_DEBUG', None) is not None else lambda *args, **kwargs: None +from modules import shared, sd_models, devices, model_quant, sd_hijack_te +from pipelines import generic def load_auraflow(checkpoint_info, diffusers_load_config={}): repo_id = sd_models.path_to_repo(checkpoint_info) - if 'torch_dtype' not in diffusers_load_config: - diffusers_load_config['torch_dtype'] = torch.float16 - debug(f'Load model: type=AuraFlow repo="{repo_id}" config={diffusers_load_config}') + sd_models.hf_auth_check(checkpoint_info) + + load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config) + shared.log.debug(f'Load model: type=AuraFlow repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') + + transformer = generic.load_transformer(repo_id, cls_name=diffusers.AuraFlowTransformer2DModel, load_config=diffusers_load_config) + text_encoder = generic.load_text_encoder(repo_id, cls_name=transformers.UMT5EncoderModel, load_config=diffusers_load_config) # auraflow uses EleutherAI/pile-t5-xl + pipe = diffusers.AuraFlowPipeline.from_pretrained( repo_id, - cache_dir = shared.opts.diffusers_dir, - **diffusers_load_config, + transformer=transformer, + text_encoder=text_encoder, + cache_dir=shared.opts.diffusers_dir, + **load_args, ) + + del text_encoder + del transformer + sd_hijack_te.init_hijack(pipe) devices.torch_gc(force=True, reason='load') return pipe diff --git a/pipelines/model_bria.py b/pipelines/model_bria.py index bcd458dc4..900c58a2a 100644 --- a/pipelines/model_bria.py +++ b/pipelines/model_bria.py @@ -1,73 +1,26 @@ import os import sys import transformers +import diffusers from modules import shared, devices, sd_models, model_quant, sd_hijack_te - - -def load_transformer(repo_id, diffusers_load_config={}): - load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='Model', device_map=True) - fn = None - - if shared.opts.sd_unet is not None and shared.opts.sd_unet != 'Default': - from modules import sd_unet - if shared.opts.sd_unet not in list(sd_unet.unet_dict): - shared.log.error(f'Load module: type=Transformer not found: {shared.opts.sd_unet}') - return None - fn = sd_unet.unet_dict[shared.opts.sd_unet] if os.path.exists(sd_unet.unet_dict[shared.opts.sd_unet]) else None - - from pipelines.bria.transformer_bria import BriaTransformer2DModel - - if fn is not None and 'gguf' in fn.lower(): - shared.log.error('Load model: type=Bria format="gguf" unsupported') - transformer = None - elif fn is not None and 'safetensors' in fn.lower(): - shared.log.debug(f'Load model: type=Bria transformer="{fn}" quant="{model_quant.get_quant(repo_id)}" args={load_args}') - transformer = BriaTransformer2DModel.from_single_file( - fn, - cache_dir=shared.opts.hfcache_dir, - **load_args, - ) - else: - shared.log.debug(f'Load model: type=Bria transformer="{repo_id}" quant="{model_quant.get_quant_type(quant_args)}" args={load_args}') - transformer = BriaTransformer2DModel.from_pretrained( - repo_id, - subfolder="transformer", - cache_dir=shared.opts.hfcache_dir, - **load_args, - **quant_args, - ) - if shared.opts.diffusers_offload_mode != 'none' and transformer is not None: - sd_models.move_model(transformer, devices.cpu) - return transformer - - -def load_text_encoder(repo_id, diffusers_load_config={}): - load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='TE', device_map=True) - shared.log.debug(f'Load model: type=Bria te="{repo_id}" quant="{model_quant.get_quant_type(quant_args)}" args={load_args}') - text_encoder = transformers.T5EncoderModel.from_pretrained( - repo_id, - subfolder="text_encoder", - cache_dir=shared.opts.hfcache_dir, - **load_args, - **quant_args, - ) - if shared.opts.diffusers_offload_mode != 'none' and text_encoder is not None: - sd_models.move_model(text_encoder, devices.cpu) - return text_encoder +from pipelines import generic def load_bria(checkpoint_info, diffusers_load_config={}): repo_id = sd_models.path_to_repo(checkpoint_info) sd_models.hf_auth_check(checkpoint_info) - transformer = load_transformer(repo_id, diffusers_load_config) - text_encoder = load_text_encoder(repo_id, diffusers_load_config) - - load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config, module='Model') - shared.log.debug(f'Load model: type=Bria model="{checkpoint_info.name}" repo="{repo_id}" offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') - - from pipelines.bria.bria_pipeline import BriaPipeline sys.path.append(os.path.join(os.path.dirname(__file__), 'bria')) + from pipelines.bria.bria_pipeline import BriaPipeline + from pipelines.bria.transformer_bria import BriaTransformer2DModel + diffusers.BriaPipeline = BriaPipeline + diffusers.BriaTransformer2DModel = BriaTransformer2DModel + + load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config, allow_quant=False) + shared.log.debug(f'Load model: type=Bria repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') + + transformer = generic.load_transformer(repo_id, cls_name=BriaTransformer2DModel, load_config=diffusers_load_config) + text_encoder = generic.load_text_encoder(repo_id, cls_name=transformers.T5EncoderModel, load_config=diffusers_load_config) pipe = BriaPipeline.from_pretrained( repo_id, diff --git a/pipelines/model_chroma.py b/pipelines/model_chroma.py index e976bfa56..b5f60d380 100644 --- a/pipelines/model_chroma.py +++ b/pipelines/model_chroma.py @@ -1,277 +1,31 @@ -import os -import json -import torch import diffusers import transformers -from safetensors.torch import load_file -from huggingface_hub import hf_hub_download -from modules import shared, errors, devices, sd_models, sd_unet, model_te, model_quant, sd_hijack_te +from modules import shared, devices, sd_models, model_quant +from pipelines import generic -debug = shared.log.trace if os.environ.get('SD_LOAD_DEBUG', None) is not None else lambda *args, **kwargs: None - - -def load_chroma_quanto(checkpoint_info): - transformer, text_encoder = None, None - quanto = model_quant.load_quanto('Load model: type=Chroma') - - if isinstance(checkpoint_info, str): - repo_path = checkpoint_info - else: - repo_path = checkpoint_info.path - - try: - quantization_map = os.path.join(repo_path, "transformer", "quantization_map.json") - debug(f'Load model: type=Chroma quantization map="{quantization_map}" repo="{checkpoint_info.name}" component="transformer"') - if not os.path.exists(quantization_map): - repo_id = sd_models.path_to_repo(checkpoint_info) - quantization_map = hf_hub_download(repo_id, subfolder='transformer', filename='quantization_map.json', cache_dir=shared.opts.diffusers_dir) - with open(quantization_map, "r", encoding='utf8') as f: - quantization_map = json.load(f) - state_dict = load_file(os.path.join(repo_path, "transformer", "diffusion_pytorch_model.safetensors")) - dtype = state_dict['context_embedder.bias'].dtype - with torch.device("meta"): - transformer = diffusers.ChromaTransformer2DModel.from_config(os.path.join(repo_path, "transformer", "config.json")).to(dtype=dtype) - quanto.requantize(transformer, state_dict, quantization_map, device=torch.device("cpu")) - transformer_dtype = transformer.dtype - if transformer_dtype != devices.dtype: - try: - transformer = transformer.to(dtype=devices.dtype) - except Exception: - shared.log.error(f"Load model: type=Chroma Failed to cast transformer to {devices.dtype}, set dtype to {transformer_dtype}") - except Exception as e: - shared.log.error(f"Load model: type=Chroma failed to load Quanto transformer: {e}") - if debug: - errors.display(e, 'Chroma Quanto:') - - try: - quantization_map = os.path.join(repo_path, "text_encoder", "quantization_map.json") - debug(f'Load model: type=Chroma quantization map="{quantization_map}" repo="{checkpoint_info.name}" component="text_encoder"') - if not os.path.exists(quantization_map): - repo_id = sd_models.path_to_repo(checkpoint_info) - quantization_map = hf_hub_download(repo_id, subfolder='text_encoder', filename='quantization_map.json', cache_dir=shared.opts.diffusers_dir) - with open(quantization_map, "r", encoding='utf8') as f: - quantization_map = json.load(f) - with open(os.path.join(repo_path, "text_encoder", "config.json"), encoding='utf8') as f: - t5_config = transformers.T5Config(**json.load(f)) - state_dict = load_file(os.path.join(repo_path, "text_encoder", "model.safetensors")) - dtype = state_dict['encoder.block.0.layer.0.SelfAttention.relative_attention_bias.weight'].dtype - with torch.device("meta"): - text_encoder = transformers.T5EncoderModel(t5_config).to(dtype=dtype) - quanto.requantize(text_encoder, state_dict, quantization_map, device=torch.device("cpu")) - text_encoder_dtype = text_encoder.dtype - if text_encoder_dtype != devices.dtype: - try: - text_encoder = text_encoder.to(dtype=devices.dtype) - except Exception: - shared.log.error(f"Load model: type=Chroma Failed to cast text encoder to {devices.dtype}, set dtype to {text_encoder_dtype}") - except Exception as e: - shared.log.error(f"Load model: type=Chroma failed to load Quanto text encoder: {e}") - if debug: - errors.display(e, 'Chroma Quanto:') - - return transformer, text_encoder - - -def load_chroma_bnb(checkpoint_info, diffusers_load_config): # pylint: disable=unused-argument - transformer, text_encoder = None, None - if isinstance(checkpoint_info, str): - repo_path = checkpoint_info - else: - repo_path = checkpoint_info.path - model_quant.load_bnb('Load model: type=Chroma') - quant = model_quant.get_quant(repo_path) - try: - # we ignore the distilled guidance layer because it degrades quality too much - # see: https://github.com/huggingface/diffusers/pull/11698#issuecomment-2969717180 for more details - if quant == 'fp8': - quantization_config = transformers.BitsAndBytesConfig(load_in_8bit=True, llm_int8_skip_modules=["distilled_guidance_layer"], bnb_4bit_compute_dtype=devices.dtype) - debug(f'Quantization: {quantization_config}') - transformer = diffusers.ChromaTransformer2DModel.from_single_file(repo_path, **diffusers_load_config, quantization_config=quantization_config) - elif quant == 'fp4': - quantization_config = transformers.BitsAndBytesConfig(load_in_4bit=True, llm_int8_skip_modules=["distilled_guidance_layer"], bnb_4bit_compute_dtype=devices.dtype, bnb_4bit_quant_type= 'fp4') - debug(f'Quantization: {quantization_config}') - transformer = diffusers.ChromaTransformer2DModel.from_single_file(repo_path, **diffusers_load_config, quantization_config=quantization_config) - elif quant == 'nf4': - quantization_config = transformers.BitsAndBytesConfig(load_in_4bit=True, llm_int8_skip_modules=["distilled_guidance_layer"], bnb_4bit_compute_dtype=devices.dtype, bnb_4bit_quant_type= 'nf4') - debug(f'Quantization: {quantization_config}') - transformer = diffusers.ChromaTransformer2DModel.from_single_file(repo_path, **diffusers_load_config, quantization_config=quantization_config) - else: - transformer = diffusers.ChromaTransformer2DModel.from_single_file(repo_path, **diffusers_load_config) - except Exception as e: - shared.log.error(f"Load model: type=Chroma failed to load BnB transformer: {e}") - transformer, text_encoder = None, None - if debug: - errors.display(e, 'Chroma:') - return transformer, text_encoder - - -def load_quants(kwargs, repo_id, cache_dir, allow_quant): # pylint: disable=unused-argument - try: - diffusers_load_config = { - "torch_dtype": devices.dtype, - "cache_dir": cache_dir, - } - if 'transformer' not in kwargs and model_quant.check_nunchaku('Model'): - shared.log.error(f'Load module: quant=Nunchaku module=transformer repo="{repo_id}" unsupported') - if 'transformer' not in kwargs and model_quant.check_quant('Model'): - load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='Model', device_map=True, modules_to_not_convert=["distilled_guidance_layer"]) - kwargs['transformer'] = diffusers.ChromaTransformer2DModel.from_pretrained(repo_id, subfolder="transformer", **load_args, **quant_args) - if 'text_encoder' not in kwargs and model_quant.check_nunchaku('TE'): - import nunchaku - nunchaku_precision = nunchaku.utils.get_precision() - nunchaku_repo = 'mit-han-lab/nunchaku-t5/awq-int4-flux.1-t5xxl.safetensors' - shared.log.debug(f'Load module: quant=Nunchaku module=t5 repo="{nunchaku_repo}" precision={nunchaku_precision}') - kwargs['text_encoder'] = nunchaku.NunchakuT5EncoderModel.from_pretrained(nunchaku_repo, torch_dtype=devices.dtype) - if 'text_encoder' not in kwargs and model_quant.check_quant('TE'): - load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='TE', device_map=True) - kwargs['text_encoder'] = transformers.T5EncoderModel.from_pretrained(repo_id, subfolder="text_encoder", **load_args, **quant_args) - except Exception as e: - shared.log.error(f'Quantization: {e}') - errors.display(e, 'Quantization:') - return kwargs - - -def load_transformer(file_path): # triggered by opts.sd_unet change - if file_path is None or not os.path.exists(file_path): - return None - transformer = None - quant = model_quant.get_quant(file_path) - diffusers_load_config = { - "torch_dtype": devices.dtype, - "cache_dir": shared.opts.hfcache_dir, - } - if quant is not None and quant != 'none': - shared.log.info(f'Load module: type=UNet/Transformer file="{file_path}" offload={shared.opts.diffusers_offload_mode} prequant={quant} dtype={devices.dtype}') - if 'gguf' in file_path.lower(): - from modules import ggml - _transformer = ggml.load_gguf(file_path, cls=diffusers.ChromaTransformer2DModel, compute_dtype=devices.dtype) - if _transformer is not None: - transformer = _transformer - elif quant in {'qint8', 'qint4'}: - _transformer, _text_encoder = load_chroma_quanto(file_path) - if _transformer is not None: - transformer = _transformer - elif quant in {'fp8', 'fp4', 'nf4'}: - _transformer, _text_encoder = load_chroma_bnb(file_path, diffusers_load_config) - if _transformer is not None: - transformer = _transformer - else: - load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='Model', device_map=True, modules_to_not_convert=["distilled_guidance_layer"]) - shared.log.debug(f'Load model: type=Chroma transformer file="{file_path}" offload={shared.opts.diffusers_offload_mode} args={load_args}') - transformer = diffusers.ChromaTransformer2DModel.from_single_file(file_path, **load_args, **quant_args) - if transformer is None: - shared.log.error('Failed to load UNet model') - shared.opts.sd_unet = 'Default' - return transformer - - -def load_chroma(checkpoint_info, diffusers_load_config): # triggered by opts.sd_checkpoint change - fn = checkpoint_info.path +def load_chroma(checkpoint_info, diffusers_load_config={}): repo_id = sd_models.path_to_repo(checkpoint_info) sd_models.hf_auth_check(checkpoint_info) - allow_post_quant = False - prequantized = model_quant.get_quant(checkpoint_info.path) - shared.log.debug(f'Load model: type=Chroma model="{checkpoint_info.name}" repo={repo_id or "none"} 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=Chroma config={diffusers_load_config}') + load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config, allow_quant=False) + shared.log.debug(f'Load model: type=Chroma repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') - transformer = None - text_encoder = None - vae = None + transformer = generic.load_transformer(repo_id, cls_name=diffusers.ChromaTransformer2DModel, load_config=diffusers_load_config, modules_to_not_convert=["distilled_guidance_layer"]) + text_encoder = generic.load_text_encoder(repo_id, cls_name=transformers.T5EncoderModel, load_config=diffusers_load_config) + + pipe = diffusers.ChromaPipeline.from_pretrained( + repo_id, + transformer=transformer, + text_encoder=text_encoder, + cache_dir=shared.opts.diffusers_dir, + **load_args, + ) + + diffusers.pipelines.auto_pipeline.AUTO_TEXT2IMAGE_PIPELINES_MAPPING["chroma"] = diffusers.ChromaPipeline + diffusers.pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING["chroma"] = diffusers.ChromaImg2ImgPipeline + del text_encoder + del transformer - # unload current model - sd_models.unload_model_weights() - shared.sd_model = None devices.torch_gc(force=True, reason='load') - - if shared.opts.teacache_enabled: - from modules import teacache - shared.log.debug(f'Transformers cache: type=teacache patch=forward cls={diffusers.ChromaTransformer2DModel.__name__}') - diffusers.ChromaTransformer2DModel.forward = teacache.teacache_chroma_forward # patch must be done before transformer is loaded - - # load overrides if any - if shared.opts.sd_unet != 'Default': - try: - debug(f'Load model: type=Chroma unet="{shared.opts.sd_unet}"') - transformer = load_transformer(sd_unet.unet_dict[shared.opts.sd_unet]) - if transformer is 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=Chroma failed to load UNet: {e}") - shared.opts.sd_unet = 'Default' - if debug: - errors.display(e, 'Chroma UNet:') - if shared.opts.sd_text_encoder != 'Default': - try: - debug(f'Load model: type=Chroma te="{shared.opts.sd_text_encoder}"') - from modules.model_te import load_t5 - text_encoder = 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=Chroma failed to load T5: {e}") - shared.opts.sd_text_encoder = 'Default' - if debug: - errors.display(e, 'Chroma T5:') - if shared.opts.sd_vae != 'Default' and shared.opts.sd_vae != 'Automatic': - try: - debug(f'Load model: type=Chroma vae="{shared.opts.sd_vae}"') - from modules import sd_vae - # vae = sd_vae.load_vae_diffusers(None, sd_vae.vae_dict[shared.opts.sd_vae], 'override') - vae_file = sd_vae.vae_dict[shared.opts.sd_vae] - if os.path.exists(vae_file): - vae_config = os.path.join('configs', 'chroma', 'vae', 'config.json') - 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=Chroma failed to load VAE: {e}") - shared.opts.sd_vae = 'Default' - if debug: - errors.display(e, 'Chroma VAE:') - - # initialize pipeline with pre-loaded components - kwargs = {} - if transformer is not None: - kwargs['transformer'] = transformer - sd_unet.loaded_unet = shared.opts.sd_unet - if text_encoder is not None: - kwargs['text_encoder'] = text_encoder - model_te.loaded_te = shared.opts.sd_text_encoder - if vae is not None: - kwargs['vae'] = vae - - # TODO model load: add ChromaFillPipeline, ChromaControlPipeline, ChromaImg2ImgPipeline etc when available - # Chroma will support inpainting *after* its training has finished: https://huggingface.co/lodestones/Chroma/discussions/28#6826dd2ed86f53ff983add5c - cls = diffusers.ChromaPipeline - shared.log.debug(f'Load model: type=Chroma cls={cls.__name__} preloaded={list(kwargs)} revision={diffusers_load_config.get("revision", None)}') - for c in kwargs: - if getattr(kwargs[c], 'quantization_method', None) is not None or getattr(kwargs[c], 'gguf', None) is not None: - shared.log.debug(f'Load model: type=Chroma component={c} dtype={kwargs[c].dtype} quant={getattr(kwargs[c], "quantization_method", None) or getattr(kwargs[c], "gguf", None)}') - if kwargs[c].dtype == torch.float32 and devices.dtype != torch.float32: - try: - kwargs[c] = kwargs[c].to(dtype=devices.dtype) - shared.log.warning(f'Load model: type=Chroma component={c} dtype={kwargs[c].dtype} cast dtype={devices.dtype} recast') - except Exception: - pass - - allow_quant = 'gguf' not in (sd_unet.loaded_unet or '') and (prequantized is None or prequantized == 'none') - 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) - if fn.endswith('.safetensors') and os.path.isfile(fn): - pipe = cls.from_single_file(fn, cache_dir=shared.opts.diffusers_dir, **kwargs, **diffusers_load_config) - allow_post_quant = True - else: - pipe = cls.from_pretrained(repo_id, cache_dir=shared.opts.diffusers_dir, **kwargs, **diffusers_load_config) - - if shared.opts.teacache_enabled and model_quant.check_nunchaku('Model'): - from nunchaku.caching.diffusers_adapters import apply_cache_on_pipe - apply_cache_on_pipe(pipe, residual_diff_threshold=0.12) - - # release memory - transformer = None - text_encoder = None - vae = None - for k in kwargs.keys(): - kwargs[k] = None - sd_hijack_te.init_hijack(pipe) - devices.torch_gc(force=True, reason='load') - return pipe, allow_post_quant + return pipe diff --git a/pipelines/model_cogview.py b/pipelines/model_cogview.py index 400038dc3..d3ac6f274 100644 --- a/pipelines/model_cogview.py +++ b/pipelines/model_cogview.py @@ -1,34 +1,19 @@ import transformers import diffusers -from modules import shared, devices, sd_models, model_quant, modelloader +from modules import shared, devices, sd_models, model_quant, sd_hijack_te +from pipelines import generic def load_cogview3(checkpoint_info, diffusers_load_config={}): - modelloader.hf_login() repo_id = sd_models.path_to_repo(checkpoint_info) + sd_models.hf_auth_check(checkpoint_info) - load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='Model') - shared.log.debug(f'Load model: type=CogView3 transformer="{repo_id}" quant="{model_quant.get_quant_type(quant_args)}" args={load_args}') - transformer = diffusers.CogView3PlusTransformer2DModel.from_pretrained( - repo_id, - subfolder="transformer", - cache_dir=shared.opts.diffusers_dir, - **load_args, - **quant_args, - ) + load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config) + shared.log.debug(f'Load model: type=CogView3 repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') - load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='TE', device_map=True) - shared.log.debug(f'Load model: type=CogView3 te="{repo_id}" quant="{model_quant.get_quant_type(quant_args)}" args={load_args}') - text_encoder = transformers.T5EncoderModel.from_pretrained( - repo_id, - subfolder="text_encoder", - cache_dir=shared.opts.diffusers_dir, - **diffusers_load_config, - **quant_args, - ) + transformer = generic.load_transformer(repo_id, cls_name=diffusers.CogView3PlusTransformer2DModel, load_config=diffusers_load_config, subfolder="transformer") + text_encoder = generic.load_text_encoder(repo_id, cls_name=transformers.T5EncoderModel, load_config=diffusers_load_config, subfolder="text_encoder") - load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config, allow_quant=False) - shared.log.debug(f'Load model: type=CogView3 model="{checkpoint_info.name}" repo="{repo_id}" offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') pipe = diffusers.CogView3PlusPipeline.from_pretrained( repo_id, text_encoder=text_encoder, @@ -36,36 +21,23 @@ def load_cogview3(checkpoint_info, diffusers_load_config={}): cache_dir=shared.opts.diffusers_dir, **load_args, ) + sd_hijack_te.init_hijack(pipe) + del transformer + del text_encoder devices.torch_gc() return pipe def load_cogview4(checkpoint_info, diffusers_load_config={}): - modelloader.hf_login() repo_id = sd_models.path_to_repo(checkpoint_info) + sd_models.hf_auth_check(checkpoint_info) - load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='Model') - shared.log.debug(f'Load model: type=CogView4 transformer="{repo_id}" quant="{model_quant.get_quant_type(quant_args)}" args={load_args}') - transformer = diffusers.CogView4Transformer2DModel.from_pretrained( - repo_id, - subfolder="transformer", - cache_dir=shared.opts.diffusers_dir, - **diffusers_load_config, - **quant_args, - ) + load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config) + shared.log.debug(f'Load model: type=CogView4 repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') - load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='TE', device_map=True) - shared.log.debug(f'Load model: type=CogView4 te="{repo_id}" quant="{model_quant.get_quant_type(quant_args)}" args={load_args}') - text_encoder = transformers.AutoModelForCausalLM.from_pretrained( # TODO model load: cogview4 balanced offload does not work for GlmModel - repo_id, - subfolder="text_encoder", - cache_dir=shared.opts.diffusers_dir, - **load_args, - # **quant_args, - ) + transformer = generic.load_transformer(repo_id, cls_name=diffusers.CogView4Transformer2DModel, load_config=diffusers_load_config, subfolder="transformer") + text_encoder = generic.load_text_encoder(repo_id, cls_name=transformers.GlmModel, load_config=diffusers_load_config, subfolder="text_encoder", allow_quant=True) - load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config, allow_quant=False) - shared.log.debug(f'Load model: type=CogView4 model="{checkpoint_info.name}" repo="{repo_id}" offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') pipe = diffusers.CogView4Pipeline.from_pretrained( repo_id, text_encoder=text_encoder, @@ -73,6 +45,8 @@ def load_cogview4(checkpoint_info, diffusers_load_config={}): cache_dir=shared.opts.diffusers_dir, **load_args, ) - pipe.enable_model_cpu_offload() + sd_hijack_te.init_hijack(pipe) + del transformer + del text_encoder devices.torch_gc() return pipe diff --git a/pipelines/model_cosmos.py b/pipelines/model_cosmos.py index 419dc3f65..839c2d9c5 100644 --- a/pipelines/model_cosmos.py +++ b/pipelines/model_cosmos.py @@ -1,68 +1,22 @@ -import os import transformers import diffusers from modules import shared, devices, sd_models, model_quant, sd_hijack_te - - -def load_transformer(repo_id, diffusers_load_config={}): - load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='Model', device_map=True) - fn = None - - if shared.opts.sd_unet is not None and shared.opts.sd_unet != 'Default': - from modules import sd_unet - if shared.opts.sd_unet not in list(sd_unet.unet_dict): - shared.log.error(f'Load module: type=Transformer not found: {shared.opts.sd_unet}') - return None - fn = sd_unet.unet_dict[shared.opts.sd_unet] if os.path.exists(sd_unet.unet_dict[shared.opts.sd_unet]) else None - - if fn is not None and 'gguf' in fn.lower(): - shared.log.error('Load model: type=Cosmos format="gguf" unsupported') - transformer = None - elif fn is not None and 'safetensors' in fn.lower(): - shared.log.debug(f'Load model: type=Cosmos transformer="{fn}" quant="{model_quant.get_quant(repo_id)}" args={load_args}') - transformer = diffusers.CosmosTransformer3DModel.from_single_file(fn, cache_dir=shared.opts.hfcache_dir, **load_args) - else: - shared.log.debug(f'Load model: type=Cosmos transformer="{repo_id}" quant="{model_quant.get_quant_type(quant_args)}" args={load_args}') - transformer = diffusers.CosmosTransformer3DModel.from_pretrained( - repo_id, - subfolder="transformer", - cache_dir=shared.opts.hfcache_dir, - **load_args, - **quant_args, - ) - if shared.opts.diffusers_offload_mode != 'none' and transformer is not None: - sd_models.move_model(transformer, devices.cpu) - return transformer - - -def load_text_encoder(repo_id, diffusers_load_config={}): - load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='TE', device_map=True) - shared.log.debug(f'Load model: type=Cosmos te="{repo_id}" quant="{model_quant.get_quant_type(quant_args)}" args={load_args}') - text_encoder = transformers.T5EncoderModel.from_pretrained( - repo_id, - subfolder="text_encoder", - cache_dir=shared.opts.hfcache_dir, - **load_args, - **quant_args, - ) - if shared.opts.diffusers_offload_mode != 'none' and text_encoder is not None: - sd_models.move_model(text_encoder, devices.cpu) - return text_encoder +from pipelines import generic def load_cosmos_t2i(checkpoint_info, diffusers_load_config={}): repo_id = sd_models.path_to_repo(checkpoint_info) sd_models.hf_auth_check(checkpoint_info) - transformer = load_transformer(repo_id, diffusers_load_config) - text_encoder = load_text_encoder(repo_id, diffusers_load_config) + load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config, allow_quant=False) + shared.log.debug(f'Load model: type=Cosmos repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') + + transformer = generic.load_transformer(repo_id, cls_name=diffusers.CosmosTransformer3DModel, load_config=diffusers_load_config, subfolder="transformer") + repo_te = 'nvidia/Cosmos-Predict2-2B-Text2Image' if 'Cosmos-Predict2-14B-Text2Image' in repo_id else repo_id + text_encoder = generic.load_text_encoder(repo_te, cls_name=transformers.T5EncoderModel, load_config=diffusers_load_config, subfolder="text_encoder", allow_shared=False) # cosmos does use standard t5 safety_checker = Fake_safety_checker() - load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config, module='Model') - shared.log.debug(f'Load model: type=Cosmos model="{checkpoint_info.name}" repo="{repo_id}" offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') - - cls = diffusers.Cosmos2TextToImagePipeline - pipe = cls.from_pretrained( + pipe = diffusers.Cosmos2TextToImagePipeline.from_pretrained( repo_id, transformer=transformer, text_encoder=text_encoder, diff --git a/pipelines/model_flex.py b/pipelines/model_flex.py index 4a11152ff..f7a285348 100644 --- a/pipelines/model_flex.py +++ b/pipelines/model_flex.py @@ -1,81 +1,32 @@ -import os import transformers import diffusers from modules import shared, devices, sd_models, model_quant, sd_hijack_te - - -def load_transformer(repo_id, diffusers_load_config={}): - load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='Model', device_map=True) - fn = None - - if shared.opts.sd_unet is not None and shared.opts.sd_unet != 'Default': - from modules import sd_unet - if shared.opts.sd_unet not in list(sd_unet.unet_dict): - shared.log.error(f'Load module: type=Transformer not found: {shared.opts.sd_unet}') - return None - fn = sd_unet.unet_dict[shared.opts.sd_unet] if os.path.exists(sd_unet.unet_dict[shared.opts.sd_unet]) else None - - if fn is not None and 'gguf' in fn.lower(): - shared.log.error('Load model: type=HiDream format="gguf" unsupported') - transformer = None - from modules import ggml - transformer = ggml.load_gguf(fn, cls=diffusers.HiDreamImageTransformer2DModel, compute_dtype=devices.dtype) - elif fn is not None and 'safetensors' in fn.lower(): - shared.log.debug(f'Load model: type=FLEX transformer="{repo_id}" quant="{model_quant.get_quant(repo_id)}" args={load_args}') - transformer = diffusers.FluxTransformer2DModel.from_single_file(fn, cache_dir=shared.opts.hfcache_dir, **load_args) - else: - shared.log.debug(f'Load model: type=FLEX transformer="{repo_id}" quant="{model_quant.get_quant_type(quant_args)}" args={load_args}') - transformer = diffusers.FluxTransformer2DModel.from_pretrained( - repo_id, - subfolder="transformer", - cache_dir=shared.opts.hfcache_dir, - **load_args, - **quant_args, - ) - if shared.opts.diffusers_offload_mode != 'none' and transformer is not None: - sd_models.move_model(transformer, devices.cpu) - return transformer - - -def load_text_encoders(repo_id, diffusers_load_config={}): - load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='TE', device_map=True) - shared.log.debug(f'Load model: type=FLEX t5="{repo_id}" quant="{model_quant.get_quant_type(quant_args)}" args={load_args}') - text_encoder_2 = transformers.T5EncoderModel.from_pretrained( - repo_id, - subfolder="text_encoder_2", - cache_dir=shared.opts.hfcache_dir, - **load_args, - **quant_args, - ) - if shared.opts.diffusers_offload_mode != 'none' and text_encoder_2 is not None: - sd_models.move_model(text_encoder_2, devices.cpu) - return text_encoder_2 +from pipelines import generic def load_flex(checkpoint_info, diffusers_load_config={}): repo_id = sd_models.path_to_repo(checkpoint_info) sd_models.hf_auth_check(checkpoint_info) - transformer = load_transformer(repo_id, diffusers_load_config) - text_encoder_2 = load_text_encoders(repo_id, diffusers_load_config) + load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config, allow_quant=False) + shared.log.debug(f'Load model: type=Flex repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') - load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config, module='Model') - shared.log.debug(f'Load model: type=FLEX model="{checkpoint_info.name}" repo="{repo_id}" offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') + transformer = generic.load_transformer(repo_id, cls_name=diffusers.FluxTransformer2DModel, load_config=diffusers_load_config) + text_encoder_2 = generic.load_text_encoder(repo_id, cls_name=transformers.T5EncoderModel, load_config=diffusers_load_config, subfolder="text_encoder_2") from pipelines.flex2 import Flex2Pipeline pipe = Flex2Pipeline.from_pretrained( repo_id, - # custom_pipeline=repo_id, transformer=transformer, text_encoder_2=text_encoder_2, cache_dir=shared.opts.diffusers_dir, **load_args, ) - sd_hijack_te.init_hijack(pipe) diffusers.pipelines.auto_pipeline.AUTO_TEXT2IMAGE_PIPELINES_MAPPING["flex2"] = Flex2Pipeline diffusers.pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING["flex2"] = Flex2Pipeline diffusers.pipelines.auto_pipeline.AUTO_INPAINT_PIPELINES_MAPPING["flex2"] = Flex2Pipeline + sd_hijack_te.init_hijack(pipe) del text_encoder_2 del transformer diff --git a/pipelines/model_flite.py b/pipelines/model_flite.py index 9c1426fb4..ad883a564 100644 --- a/pipelines/model_flite.py +++ b/pipelines/model_flite.py @@ -1,52 +1,26 @@ import sys +import diffusers import transformers from modules import shared, devices, sd_models, model_quant, sd_hijack_te - - -def load_dit(repo_id, diffusers_load_config={}): - load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='Model', device_map=True) - shared.log.debug(f'Load model: type=FLite dit="{repo_id}" quant="{model_quant.get_quant_type(quant_args)}" args={load_args}') - import pipelines.f_lite - sys.modules['f_lite'] = pipelines.f_lite - transformer = pipelines.f_lite.DiT.from_pretrained( - repo_id, - subfolder="dit_model", - cache_dir=shared.opts.hfcache_dir, - **load_args, - **quant_args, - ) - if shared.opts.diffusers_offload_mode != 'none' and transformer is not None: - sd_models.move_model(transformer, devices.cpu) - return transformer - - -def load_text_encoder(repo_id, diffusers_load_config={}): - load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='TE', device_map=True) - shared.log.debug(f'Load model: type=FLite te="{repo_id}" quant="{model_quant.get_quant_type(quant_args)}" args={load_args}') - text_encoder = transformers.T5EncoderModel.from_pretrained( - repo_id, - subfolder="text_encoder", - cache_dir=shared.opts.hfcache_dir, - **load_args, - **quant_args, - ) - if shared.opts.diffusers_offload_mode != 'none' and text_encoder is not None: - sd_models.move_model(text_encoder, devices.cpu) - return text_encoder +from pipelines import generic def load_flite(checkpoint_info, diffusers_load_config={}): repo_id = sd_models.path_to_repo(checkpoint_info) sd_models.hf_auth_check(checkpoint_info) - from pipelines.f_lite import FLitePipeline - dit_model = load_dit(repo_id, diffusers_load_config) - text_encoder = load_text_encoder(repo_id, diffusers_load_config) + load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config, allow_quant=False) + shared.log.debug(f'Load model: type=FLite repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') - load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config, module='Model') - shared.log.debug(f'Load model: type=FLite model="{checkpoint_info.name}" repo="{repo_id}" offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') - pipe = FLitePipeline.from_pretrained( - repo_id, + from pipelines import f_lite + diffusers.FLitePipeline = f_lite.FLitePipeline + sys.modules['f_lite'] = f_lite + + dit_model = generic.load_transformer(repo_id, cls_name=f_lite.DiT, load_config=diffusers_load_config, subfolder="dit_model") + text_encoder = generic.load_text_encoder(repo_id, cls_name=transformers.T5EncoderModel, load_config=diffusers_load_config, subfolder="text_encoder") + + pipe = f_lite.FLitePipeline.from_pretrained( + "Freepik/F-Lite", # pr only exists on main repo revision="refs/pr/8", dit_model=dit_model, text_encoder=text_encoder, diff --git a/pipelines/model_flux.py b/pipelines/model_flux.py index 5c1ba745b..34de4da83 100644 --- a/pipelines/model_flux.py +++ b/pipelines/model_flux.py @@ -1,360 +1,76 @@ -import os -import json -import torch import diffusers import transformers -from safetensors.torch import load_file -from huggingface_hub import hf_hub_download -from modules import shared, errors, devices, sd_models, sd_unet, model_te, model_quant, sd_hijack_te +from modules import shared, devices, sd_models, model_quant +from pipelines import generic -debug = shared.log.trace if os.environ.get('SD_LOAD_DEBUG', None) is not None else lambda *args, **kwargs: None - - -def load_flux_quanto(checkpoint_info): - transformer, text_encoder_2 = None, None - quanto = model_quant.load_quanto('Load model: type=FLUX') - - if isinstance(checkpoint_info, str): - repo_path = checkpoint_info - else: - repo_path = checkpoint_info.path - - try: - quantization_map = os.path.join(repo_path, "transformer", "quantization_map.json") - debug(f'Load model: type=FLUX quantization map="{quantization_map}" repo="{checkpoint_info.name}" component="transformer"') - if not os.path.exists(quantization_map): - repo_id = sd_models.path_to_repo(checkpoint_info) - quantization_map = hf_hub_download(repo_id, subfolder='transformer', filename='quantization_map.json', cache_dir=shared.opts.diffusers_dir) - with open(quantization_map, "r", encoding='utf8') as f: - quantization_map = json.load(f) - state_dict = load_file(os.path.join(repo_path, "transformer", "diffusion_pytorch_model.safetensors")) - dtype = state_dict['context_embedder.bias'].dtype - with torch.device("meta"): - transformer = diffusers.FluxTransformer2DModel.from_config(os.path.join(repo_path, "transformer", "config.json")).to(dtype=dtype) - quanto.requantize(transformer, state_dict, quantization_map, device=torch.device("cpu")) - transformer_dtype = transformer.dtype - if transformer_dtype != devices.dtype: - try: - transformer = transformer.to(dtype=devices.dtype) - except Exception: - shared.log.error(f"Load model: type=FLUX Failed to cast transformer to {devices.dtype}, set dtype to {transformer_dtype}") - except Exception as e: - shared.log.error(f"Load model: type=FLUX failed to load Quanto transformer: {e}") - if debug: - errors.display(e, 'FLUX Quanto:') - - try: - quantization_map = os.path.join(repo_path, "text_encoder_2", "quantization_map.json") - debug(f'Load model: type=FLUX quantization map="{quantization_map}" repo="{checkpoint_info.name}" component="text_encoder_2"') - if not os.path.exists(quantization_map): - repo_id = sd_models.path_to_repo(checkpoint_info) - quantization_map = hf_hub_download(repo_id, subfolder='text_encoder_2', filename='quantization_map.json', cache_dir=shared.opts.diffusers_dir) - with open(quantization_map, "r", encoding='utf8') as f: - quantization_map = json.load(f) - with open(os.path.join(repo_path, "text_encoder_2", "config.json"), encoding='utf8') as f: - t5_config = transformers.T5Config(**json.load(f)) - state_dict = load_file(os.path.join(repo_path, "text_encoder_2", "model.safetensors")) - dtype = state_dict['encoder.block.0.layer.0.SelfAttention.relative_attention_bias.weight'].dtype - with torch.device("meta"): - text_encoder_2 = transformers.T5EncoderModel(t5_config).to(dtype=dtype) - quanto.requantize(text_encoder_2, state_dict, quantization_map, device=torch.device("cpu")) - text_encoder_2_dtype = text_encoder_2.dtype - if text_encoder_2_dtype != devices.dtype: - try: - text_encoder_2 = text_encoder_2.to(dtype=devices.dtype) - except Exception: - shared.log.error(f"Load model: type=FLUX Failed to cast text encoder to {devices.dtype}, set dtype to {text_encoder_2_dtype}") - except Exception as e: - shared.log.error(f"Load model: type=FLUX failed to load Quanto text encoder: {e}") - if debug: - errors.display(e, 'FLUX Quanto:') - - return transformer, text_encoder_2 - - -def load_flux_bnb(checkpoint_info, diffusers_load_config): # pylint: disable=unused-argument - transformer, text_encoder_2 = None, None - if isinstance(checkpoint_info, str): - repo_path = checkpoint_info - else: - repo_path = checkpoint_info.path - model_quant.load_bnb('Load model: type=FLUX') - quant = model_quant.get_quant(repo_path) - try: - if quant == 'fp8': - quantization_config = transformers.BitsAndBytesConfig(load_in_8bit=True, bnb_4bit_compute_dtype=devices.dtype) - debug(f'Quantization: {quantization_config}') - transformer = diffusers.FluxTransformer2DModel.from_single_file(repo_path, **diffusers_load_config, quantization_config=quantization_config) - elif quant == 'fp4': - quantization_config = transformers.BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=devices.dtype, bnb_4bit_quant_type= 'fp4') - debug(f'Quantization: {quantization_config}') - transformer = diffusers.FluxTransformer2DModel.from_single_file(repo_path, **diffusers_load_config, quantization_config=quantization_config) - elif quant == 'nf4': - quantization_config = transformers.BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=devices.dtype, bnb_4bit_quant_type= 'nf4') - debug(f'Quantization: {quantization_config}') - transformer = diffusers.FluxTransformer2DModel.from_single_file(repo_path, **diffusers_load_config, quantization_config=quantization_config) - else: - transformer = diffusers.FluxTransformer2DModel.from_single_file(repo_path, **diffusers_load_config) - except Exception as e: - shared.log.error(f"Load model: type=FLUX failed to load BnB transformer: {e}") - transformer, text_encoder_2 = None, None - if debug: - errors.display(e, 'FLUX:') - return transformer, text_encoder_2 - - -def load_quants(kwargs, repo_id, cache_dir, allow_quant): # pylint: disable=unused-argument - try: - diffusers_load_config = { - "torch_dtype": devices.dtype, - "cache_dir": cache_dir, - } - if 'transformer' not in kwargs and model_quant.check_nunchaku('Model'): - import nunchaku - nunchaku_precision = nunchaku.utils.get_precision() - nunchaku_repo = None - if 'flux.1-kontext' in repo_id.lower(): - nunchaku_repo = f"mit-han-lab/nunchaku-flux.1-kontext-dev/svdq-{nunchaku_precision}_r32-flux.1-kontext-dev.safetensors" - elif 'flux.1-dev' in repo_id.lower(): - nunchaku_repo = f"mit-han-lab/nunchaku-flux.1-dev/svdq-{nunchaku_precision}_r32-flux.1-dev.safetensors" - elif 'flux.1-schnell' in repo_id.lower(): - nunchaku_repo = f"mit-han-lab/nunchaku-flux.1-schnell/svdq-{nunchaku_precision}_r32-flux.1-schnell.safetensors" - elif 'flux.1-fill' in repo_id.lower(): - nunchaku_repo = f"mit-han-lab/svdq-fp4-flux.1-fill-dev/svdq-{nunchaku_precision}_r32-flux.1-schnell.safetensors" - elif 'flux.1-depth' in repo_id.lower(): - nunchaku_repo = f"mit-han-lab/svdq-int4-flux.1-depth-dev/svdq-{nunchaku_precision}_r32-flux.1-schnell.safetensors" - elif 'shuttle' in repo_id.lower(): - nunchaku_repo = f"mit-han-lab/nunchaku-shuttle-jaguar/svdq-{nunchaku_precision}_r32-shuttle-jaguar.safetensors" - else: - shared.log.error(f'Load module: quant=Nunchaku module=transformer repo="{repo_id}" unsupported') - if nunchaku_repo is not None: - shared.log.debug(f'Load module: quant=Nunchaku module=transformer repo="{nunchaku_repo}" precision={nunchaku_precision} offload={shared.opts.nunchaku_offload} attention={shared.opts.nunchaku_attention}') - kwargs['transformer'] = nunchaku.NunchakuFluxTransformer2dModel.from_pretrained(nunchaku_repo, offload=shared.opts.nunchaku_offload, torch_dtype=devices.dtype) - kwargs['transformer'].quantization_method = 'SVDQuant' - if shared.opts.nunchaku_attention: - kwargs['transformer'].set_attention_impl("nunchaku-fp16") - if 'transformer' not in kwargs and model_quant.check_quant('Model'): - load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='Model', device_map=True) - kwargs['transformer'] = diffusers.FluxTransformer2DModel.from_pretrained(repo_id, subfolder="transformer", **load_args, **quant_args) - if 'text_encoder_2' not in kwargs and model_quant.check_nunchaku('TE'): - import nunchaku - nunchaku_precision = nunchaku.utils.get_precision() - nunchaku_repo = 'mit-han-lab/nunchaku-t5/awq-int4-flux.1-t5xxl.safetensors' - shared.log.debug(f'Load module: quant=Nunchaku module=t5 repo="{nunchaku_repo}" precision={nunchaku_precision}') - kwargs['text_encoder_2'] = nunchaku.NunchakuT5EncoderModel.from_pretrained(nunchaku_repo, torch_dtype=devices.dtype) - kwargs['text_encoder_2'].quantization_method = 'SVDQuant' - if 'text_encoder_2' not in kwargs and model_quant.check_quant('TE'): - load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='TE', device_map=True) - kwargs['text_encoder_2'] = transformers.T5EncoderModel.from_pretrained(repo_id, subfolder="text_encoder_2", **load_args, **quant_args) - except Exception as e: - shared.log.error(f'Quantization: {e}') - errors.display(e, 'Quantization:') - return kwargs - - -def load_transformer(file_path): # triggered by opts.sd_unet change - if file_path is None or not os.path.exists(file_path): - return None - transformer = None - quant = model_quant.get_quant(file_path) - diffusers_load_config = { - "torch_dtype": devices.dtype, - "cache_dir": shared.opts.hfcache_dir, - } - if quant is not None and quant != 'none': - shared.log.info(f'Load module: type=UNet/Transformer file="{file_path}" offload={shared.opts.diffusers_offload_mode} prequant={quant} dtype={devices.dtype}') - if 'gguf' in file_path.lower(): - from modules import ggml - _transformer = ggml.load_gguf(file_path, cls=diffusers.FluxTransformer2DModel, compute_dtype=devices.dtype) - if _transformer is not None: - transformer = _transformer - elif quant == "fp8": - _transformer = model_quant.load_fp8_model_layerwise(file_path, diffusers.FluxTransformer2DModel.from_single_file, diffusers_load_config) - if _transformer is not None: - transformer = _transformer - elif quant in {'qint8', 'qint4'}: - _transformer, _text_encoder_2 = load_flux_quanto(file_path) - if _transformer is not None: - transformer = _transformer - elif quant in {'fp8', 'fp4', 'nf4'}: - _transformer, _text_encoder_2 = load_flux_bnb(file_path, diffusers_load_config) - if _transformer is not None: - transformer = _transformer - elif 'nf4' in quant: # TODO flux: loader for civitai nf4 models - from pipelines.model_flux_nf4 import load_flux_nf4 - _transformer, _text_encoder_2 = load_flux_nf4(file_path, prequantized=True) - if _transformer is not None: - transformer = _transformer - else: - quant_args = model_quant.create_bnb_config({}) - if quant_args: - shared.log.info(f'Load module: type=Flux transformer file="{file_path}" offload={shared.opts.diffusers_offload_mode} quant=bnb dtype={devices.dtype}') - from pipelines.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 - load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='Model', device_map=True) - shared.log.debug(f'Load model: type=Flux transformer file="{file_path}" offload={shared.opts.diffusers_offload_mode} args={load_args}') - transformer = diffusers.FluxTransformer2DModel.from_single_file(file_path, **load_args, **quant_args) - if transformer is None: - shared.log.error('Failed to load UNet model') - shared.opts.sd_unet = 'Default' - return transformer - - -def load_flux(checkpoint_info, diffusers_load_config): # triggered by opts.sd_checkpoint change +def load_flux(checkpoint_info, diffusers_load_config={}): repo_id = sd_models.path_to_repo(checkpoint_info) sd_models.hf_auth_check(checkpoint_info) - allow_post_quant = False - prequantized = model_quant.get_quant(checkpoint_info.path) - 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}') + if 'Fill' in repo_id: + cls_name = diffusers.FluxFillPipeline + elif 'Canny' in repo_id: + cls_name = diffusers.FluxControlPipeline + elif 'Depth' in repo_id: + cls_name = diffusers.FluxControlPipeline + elif 'Kontext' in repo_id: + cls_name = diffusers.FluxKontextPipeline + diffusers.pipelines.auto_pipeline.AUTO_TEXT2IMAGE_PIPELINES_MAPPING["flux1kontext"] = diffusers.FluxKontextPipeline + diffusers.pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING["flux1kontext"] = diffusers.FluxKontextPipeline + diffusers.pipelines.auto_pipeline.AUTO_INPAINT_PIPELINES_MAPPING["flux1kontext"] = diffusers.FluxKontextInpaintPipeline + else: + cls_name = diffusers.FluxPipeline - transformer = None - text_encoder_1 = None - text_encoder_2 = None - vae = None + load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config, allow_quant=False) + shared.log.debug(f'Load model: type=Flux repo="{repo_id}" cls={cls_name.__name__} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') - # unload current model - sd_models.unload_model_weights() - shared.sd_model = None - devices.torch_gc(force=True, reason='load') - - if shared.opts.teacache_enabled: + # optional teacache patch + if shared.opts.teacache_enabled and not model_quant.check_nunchaku('Model'): from modules import teacache shared.log.debug(f'Transformers cache: type=teacache patch=forward cls={diffusers.FluxTransformer2DModel.__name__}') diffusers.FluxTransformer2DModel.forward = teacache.teacache_flux_forward # patch must be done before transformer is loaded - # load overrides if any - 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 = '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 = 'Default' - if debug: - errors.display(e, 'FLUX UNet:') - 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 - if 'vit-l' in shared.opts.sd_text_encoder.lower(): - text_encoder_1 = load_vit_l() - else: - 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 = 'Default' - if debug: - errors.display(e, 'FLUX T5:') - 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 - # vae = sd_vae.load_vae_diffusers(None, sd_vae.vae_dict[shared.opts.sd_vae], 'override') - vae_file = sd_vae.vae_dict[shared.opts.sd_vae] - if os.path.exists(vae_file): - vae_config = os.path.join('configs', 'flux', 'vae', 'config.json') - 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 = 'Default' - if debug: - errors.display(e, 'FLUX VAE:') + transformer = None + text_encoder_2 = None - # load quantized components if any - if prequantized == 'nf4': - try: - from pipelines.model_flux_nf4 import load_flux_nf4 - _transformer, _text_encoder = load_flux_nf4(checkpoint_info) - if _transformer is not None: - transformer = _transformer - if _text_encoder is not None: - text_encoder_2 = _text_encoder - except Exception as e: - shared.log.error(f"Load model: type=FLUX failed to load NF4 components: {e}") - if debug: - errors.display(e, 'FLUX NF4:') - if prequantized == 'qint8' or prequantized == 'qint4': - try: - _transformer, _text_encoder = load_flux_quanto(checkpoint_info) - if _transformer is not None: - transformer = _transformer - if _text_encoder is not None: - text_encoder_2 = _text_encoder - except Exception as e: - shared.log.error(f"Load model: type=FLUX failed to load Quanto components: {e}") - if debug: - errors.display(e, 'FLUX Quanto:') + # handle transformer svdquant if available, t5 is handled inside load_text_encoder + prequantized = model_quant.get_quant(checkpoint_info.path) + if model_quant.check_nunchaku('Model'): + from pipelines.flux.flux_nunchaku import load_flux_nunchaku + transformer = load_flux_nunchaku(repo_id) + # handle prequantized models + elif prequantized == 'nf4': + from pipelines.flux.flux_nf4 import load_flux_nf4 + transformer, text_encoder_2 = load_flux_nf4(checkpoint_info) + elif prequantized == 'qint8' or prequantized == 'qint4': + from pipelines.flux.flux_quanto import load_flux_quanto + transformer, text_encoder_2 = load_flux_quanto(checkpoint_info) + elif prequantized == 'fp4' or prequantized == 'fp8': + from pipelines.flux.flux_bnb import load_flux_bnb + transformer = load_flux_bnb(checkpoint_info, diffusers_load_config) - # initialize pipeline with pre-loaded components - kwargs = {} - if transformer is not None: - kwargs['transformer'] = transformer - sd_unet.loaded_unet = shared.opts.sd_unet - if text_encoder_1 is not None: - kwargs['text_encoder'] = text_encoder_1 - model_te.loaded_te = shared.opts.sd_text_encoder - if text_encoder_2 is not None: - kwargs['text_encoder_2'] = text_encoder_2 - model_te.loaded_te = shared.opts.sd_text_encoder - if vae is not None: - kwargs['vae'] = vae - 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 - if 'Fill' in repo_id: - cls = diffusers.FluxFillPipeline - elif 'Canny' in repo_id: - cls = diffusers.FluxControlPipeline - elif 'Depth' in repo_id: - cls = diffusers.FluxControlPipeline - elif 'Kontext' in repo_id: - cls = diffusers.FluxKontextPipeline - from diffusers import pipelines - pipelines.auto_pipeline.AUTO_TEXT2IMAGE_PIPELINES_MAPPING["flux1kontext"] = diffusers.FluxKontextPipeline - pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING["flux1kontext"] = diffusers.FluxKontextPipeline - pipelines.auto_pipeline.AUTO_INPAINT_PIPELINES_MAPPING["flux1kontext"] = diffusers.FluxKontextInpaintPipeline + # finally load transformer and text encoder if not already loaded + if transformer is None: + transformer = generic.load_transformer(repo_id, cls_name=diffusers.FluxTransformer2DModel, load_config=diffusers_load_config) + if text_encoder_2 is None: + text_encoder_2 = generic.load_text_encoder(repo_id, cls_name=transformers.T5EncoderModel, load_config=diffusers_load_config) - else: - cls = diffusers.FluxPipeline - shared.log.debug(f'Load model: type=FLUX cls={cls.__name__} preloaded={list(kwargs)} revision={diffusers_load_config.get("revision", None)}') - for c in kwargs: - if getattr(kwargs[c], 'quantization_method', None) is not None or getattr(kwargs[c], 'gguf', None) is not None: - shared.log.debug(f'Load model: type=FLUX component={c} dtype={kwargs[c].dtype} quant={getattr(kwargs[c], "quantization_method", None) or getattr(kwargs[c], "gguf", None)}') - if kwargs[c].dtype == torch.float32 and devices.dtype != torch.float32: - try: - kwargs[c] = kwargs[c].to(dtype=devices.dtype) - shared.log.warning(f'Load model: type=FLUX component={c} dtype={kwargs[c].dtype} cast dtype={devices.dtype} recast') - except Exception: - pass + pipe = cls_name.from_pretrained( + repo_id, + transformer=transformer, + text_encoder_2=text_encoder_2, + cache_dir=shared.opts.diffusers_dir, + **load_args, + ) - 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) - if fn.endswith('.safetensors') and os.path.isfile(fn): - pipe = cls.from_single_file(fn, cache_dir=shared.opts.diffusers_dir, **kwargs, **diffusers_load_config) - allow_post_quant = True - else: - pipe = cls.from_pretrained(repo_id, cache_dir=shared.opts.diffusers_dir, **kwargs, **diffusers_load_config) + del text_encoder_2 + del transformer + # optional first-block patch if shared.opts.teacache_enabled and model_quant.check_nunchaku('Model'): from nunchaku.caching.diffusers_adapters import apply_cache_on_pipe apply_cache_on_pipe(pipe, residual_diff_threshold=0.12) - # release memory - transformer = None - text_encoder_1 = None - text_encoder_2 = None - vae = None - for k in kwargs.keys(): - kwargs[k] = None - sd_hijack_te.init_hijack(pipe) devices.torch_gc(force=True, reason='load') - return pipe, allow_post_quant + return pipe diff --git a/pipelines/model_hidream.py b/pipelines/model_hidream.py index 6df9092e2..0a0880103 100644 --- a/pipelines/model_hidream.py +++ b/pipelines/model_hidream.py @@ -1,65 +1,15 @@ -import os import transformers import diffusers from modules import shared, devices, sd_models, model_quant, sd_hijack_te +from pipelines import generic -def load_transformer(repo_id, diffusers_load_config={}): - load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='Model', device_map=True) - fn = None - - if shared.opts.sd_unet is not None and shared.opts.sd_unet != 'Default': - from modules import sd_unet - if shared.opts.sd_unet not in list(sd_unet.unet_dict): - shared.log.error(f'Load module: type=Transformer not found: {shared.opts.sd_unet}') - return None - fn = sd_unet.unet_dict[shared.opts.sd_unet] if os.path.exists(sd_unet.unet_dict[shared.opts.sd_unet]) else None - - if fn is not None and 'gguf' in fn.lower(): - shared.log.error('Load model: type=HiDream format="gguf" unsupported') - transformer = None - # from modules import ggml - # transformer = ggml.load_gguf(fn, cls=diffusers.HiDreamImageTransformer2DModel, compute_dtype=devices.dtype) - elif fn is not None and 'safetensors' in fn.lower(): - shared.log.debug(f'Load model: type=HiDream transformer="{repo_id}" offload={shared.opts.diffusers_offload_mode} quant="{model_quant.get_quant(repo_id)}" args={load_args}') - transformer = diffusers.HiDreamImageTransformer2DModel.from_single_file(fn, cache_dir=shared.opts.hfcache_dir, **load_args, **quant_args) - # elif model_quant.check_nunchaku('Model'): - # shared.log.error(f'Load model: type=HiDream transformer="{repo_id}" quant="Nunchaku" unsupported') - # transformer = None - else: - shared.log.debug(f'Load model: type=HiDream transformer="{repo_id}" quant="{model_quant.get_quant_type(quant_args)}" args={load_args}') - transformer = diffusers.HiDreamImageTransformer2DModel.from_pretrained( - repo_id, - subfolder="transformer", - cache_dir=shared.opts.hfcache_dir, - **load_args, - **quant_args, - ) - if shared.opts.diffusers_offload_mode != 'none' and transformer is not None: - sd_models.move_model(transformer, devices.cpu) - return transformer - - -def load_text_encoders(repo_id, diffusers_load_config={}): - if repo_id == 'HiDream-ai/HiDream-E1-Full': - repo_id = 'HiDream-ai/HiDream-I1-Full' # use I1 for t5 and llm - load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='TE', device_map=True) - shared.log.debug(f'Load model: type=HiDream te3="{repo_id}" quant="{model_quant.get_quant_type(quant_args)}" args={load_args}') - text_encoder_3 = transformers.T5EncoderModel.from_pretrained( - repo_id, - subfolder="text_encoder_3", - cache_dir=shared.opts.hfcache_dir, - **load_args, - **quant_args, - ) - if shared.opts.diffusers_offload_mode != 'none' and text_encoder_3 is not None: - sd_models.move_model(text_encoder_3, devices.cpu) - +def load_llama(diffusers_load_config={}): load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='TE', device_map=True) llama_repo = shared.opts.model_h1_llama_repo if shared.opts.model_h1_llama_repo != 'Default' else 'meta-llama/Meta-Llama-3.1-8B-Instruct' shared.log.debug(f'Load model: type=HiDream te4="{llama_repo}" quant="{model_quant.get_quant_type(quant_args)}" args={load_args}') - sd_models.hf_auth_check(llama_repo) + text_encoder_4 = transformers.LlamaForCausalLM.from_pretrained( llama_repo, output_hidden_states=True, @@ -75,18 +25,19 @@ def load_text_encoders(repo_id, diffusers_load_config={}): ) if shared.opts.diffusers_offload_mode != 'none' and text_encoder_4 is not None: sd_models.move_model(text_encoder_4, devices.cpu) - return text_encoder_3, text_encoder_4, tokenizer_4 + return text_encoder_4, tokenizer_4 def load_hidream(checkpoint_info, diffusers_load_config={}): repo_id = sd_models.path_to_repo(checkpoint_info) sd_models.hf_auth_check(checkpoint_info) - transformer = load_transformer(repo_id, diffusers_load_config) - text_encoder_3, text_encoder_4, tokenizer_4 = load_text_encoders(repo_id, diffusers_load_config) + load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config, allow_quant=False) + shared.log.debug(f'Load model: type=HiDream repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') - load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config, module='Model') - shared.log.debug(f'Load model: type=HiDream model="{checkpoint_info.name}" repo="{repo_id}" offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') + transformer = generic.load_transformer(repo_id, cls_name=diffusers.HiDreamImageTransformer2DModel, load_config=diffusers_load_config, subfolder="transformer") + text_encoder_3 = generic.load_text_encoder(repo_id, cls_name=transformers.T5EncoderModel, load_config=diffusers_load_config, subfolder="text_encoder_3") + text_encoder_4, tokenizer_4 = load_llama(diffusers_load_config) if shared.opts.teacache_enabled: from modules import teacache diff --git a/pipelines/model_hunyuandit.py b/pipelines/model_hunyuandit.py new file mode 100644 index 000000000..335bc7d7e --- /dev/null +++ b/pipelines/model_hunyuandit.py @@ -0,0 +1,29 @@ +import transformers +import diffusers +from modules import shared, sd_models, devices, model_quant +from pipelines import generic + + +def load_hunyuandit(checkpoint_info, diffusers_load_config={}): + repo_id = sd_models.path_to_repo(checkpoint_info) + sd_models.hf_auth_check(checkpoint_info) + + load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config) + shared.log.debug(f'Load model: type=HunyuanDiT repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') + + transformer = generic.load_transformer(repo_id, cls_name=diffusers.HunyuanDiT2DModel, load_config=diffusers_load_config) + repo_te = 'Tencent-Hunyuan/HunyuanDiT-v1.2-Diffusers' if 'HunyuanDiT-v1' in repo_id else repo_id + text_encoder_2 = generic.load_text_encoder(repo_te, cls_name=transformers.T5EncoderModel, load_config=diffusers_load_config, subfolder="text_encoder_2", allow_shared=False) # this is not normal t5 + + pipe = diffusers.HunyuanDiTPipeline.from_pretrained( + repo_id, + transformer=transformer, + text_encoder_2=text_encoder_2, + cache_dir=shared.opts.diffusers_dir, + **load_args, + ) + + del text_encoder_2 + del transformer + devices.torch_gc(force=True, reason='load') + return pipe diff --git a/pipelines/model_kandinsky.py b/pipelines/model_kandinsky.py new file mode 100644 index 000000000..0d5ad0013 --- /dev/null +++ b/pipelines/model_kandinsky.py @@ -0,0 +1,65 @@ +import transformers +import diffusers +from modules import shared, sd_models, devices, model_quant, sd_hijack_te +from pipelines import generic + + +def load_kandinsky21(checkpoint_info, diffusers_load_config={}): + repo_id = sd_models.path_to_repo(checkpoint_info) + sd_models.hf_auth_check(checkpoint_info) + + load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config) + shared.log.debug(f'Load model: type=Kandinsky21 repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') + pipe = diffusers.KandinskyCombinedPipeline.from_pretrained( + repo_id, + cache_dir=shared.opts.diffusers_dir, + **load_args, + ) + sd_hijack_te.init_hijack(pipe) + devices.torch_gc(force=True, reason='load') + return pipe + + +def load_kandinsky22(checkpoint_info, diffusers_load_config={}): + repo_id = sd_models.path_to_repo(checkpoint_info) + sd_models.hf_auth_check(checkpoint_info) + + load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config) + shared.log.debug(f'Load model: type=Kandinsky22 repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') + pipe = diffusers.KandinskyV22CombinedPipeline.from_pretrained( + repo_id, + cache_dir=shared.opts.diffusers_dir, + **load_args, + ) + sd_hijack_te.init_hijack(pipe) + devices.torch_gc(force=True, reason='load') + return pipe + + +def load_kandinsky3(checkpoint_info, diffusers_load_config={}): + repo_id = sd_models.path_to_repo(checkpoint_info) + sd_models.hf_auth_check(checkpoint_info) + + load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config) + shared.log.debug(f'Load model: type=Kandinsky30 repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') + + unet = generic.load_transformer(repo_id, cls_name=diffusers.Kandinsky3UNet, load_config=diffusers_load_config, subfolder="unet", variant="fp16") + text_encoder = generic.load_text_encoder(repo_id, cls_name=transformers.T5EncoderModel, load_config=diffusers_load_config, subfolder="text_encoder", variant="fp16") + + pipe = diffusers.Kandinsky3Pipeline.from_pretrained( + repo_id, + unet=unet, + text_encoder=text_encoder, + variant="fp16", + cache_dir=shared.opts.diffusers_dir, + **load_args, + ) + pipe.task_args = { + 'output_type': 'np', + } + + del text_encoder + del unet + sd_hijack_te.init_hijack(pipe) + devices.torch_gc(force=True, reason='load') + return pipe diff --git a/pipelines/model_kolors.py b/pipelines/model_kolors.py index 8add20664..26fcc8497 100644 --- a/pipelines/model_kolors.py +++ b/pipelines/model_kolors.py @@ -2,21 +2,14 @@ import torch import diffusers -repo_id = 'Kwai-Kolors/Kolors-diffusers' - - def load_kolors(_checkpoint_info, diffusers_load_config={}): from modules import shared, devices diffusers_load_config['variant'] = "fp16" if 'torch_dtype' not in diffusers_load_config: diffusers_load_config['torch_dtype'] = torch.float16 - # import torch - # import transformers - # encoder_id = 'THUDM/chatglm3-6b' - # text_encoder = transformers.AutoModel.from_pretrained(encoder_id, torch_dtype=torch.float16, trust_remote_code=True, cache_dir=shared.opts.diffusers_dir) - # text_encoder = transformers.AutoModel.from_pretrained("THUDM/chatglm3-6b", torch_dtype=torch.float16, trust_remote_code=True).quantize(4).cuda() - # tokenizer = transformers.AutoTokenizer.from_pretrained(encoder_id, trust_remote_code=True, cache_dir=shared.opts.diffusers_dir) + repo_id = 'Kwai-Kolors/Kolors-diffusers' + shared.log.debug(f'Load model: type=Kolors repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={diffusers_load_config}') pipe = diffusers.KolorsPipeline.from_pretrained( repo_id, cache_dir = shared.opts.diffusers_dir, diff --git a/pipelines/model_lumina.py b/pipelines/model_lumina.py index 60b681881..217bbcfa3 100644 --- a/pipelines/model_lumina.py +++ b/pipelines/model_lumina.py @@ -1,9 +1,7 @@ -import os import transformers import diffusers -from modules import errors, shared, sd_models, sd_unet, sd_hijack_te, devices, modelloader, model_quant - -debug = shared.log.trace if os.environ.get('SD_LOAD_DEBUG', None) is not None else lambda *args, **kwargs: None +from modules import shared, sd_models, sd_hijack_te, devices, modelloader, model_quant +from pipelines import generic def load_lumina(_checkpoint_info, diffusers_load_config={}): @@ -19,7 +17,6 @@ def load_lumina(_checkpoint_info, diffusers_load_config={}): def load_lumina2(checkpoint_info, diffusers_load_config={}): - transformer, text_encoder, vae = None, None, None repo_id = sd_models.path_to_repo(checkpoint_info) if shared.opts.teacache_enabled: @@ -27,61 +24,10 @@ def load_lumina2(checkpoint_info, diffusers_load_config={}): shared.log.debug(f'Transformers cache: type=teacache patch=forward cls={diffusers.Lumina2Transformer2DModel.__name__}') diffusers.Lumina2Transformer2DModel.forward = teacache.teacache_lumina2_forward # patch must be done before transformer is loaded - load_config, quant_config = model_quant.get_dit_args(diffusers_load_config, module='Model') - if shared.opts.sd_unet != 'Default': - try: - debug(f'Load model: type=Lumina2 unet="{shared.opts.sd_unet}"') - transformer = diffusers.Lumina2Transformer2DModel.from_single_file( - sd_unet.unet_dict[shared.opts.sd_unet], - cache_dir=shared.opts.diffusers_dir, - **load_config, - **quant_config - ) - if transformer is 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=Lumina2 failed to load UNet: {e}") - shared.opts.sd_unet = 'Default' - if debug: - errors.display(e, 'Lumina2 UNet:') + transformer = generic.load_transformer(repo_id, cls_name=diffusers.Lumina2Transformer2DModel, load_config=diffusers_load_config) + text_encoder = generic.load_text_encoder(repo_id, cls_name=transformers.Gemma2Model, load_config=diffusers_load_config) - if shared.opts.sd_vae != 'Default' and shared.opts.sd_vae != 'Automatic': - try: - debug(f'Load model: type=Lumina2 vae="{shared.opts.sd_vae}"') - from modules import sd_vae - # vae = sd_vae.load_vae_diffusers(None, sd_vae.vae_dict[shared.opts.sd_vae], 'override') - vae_file = sd_vae.vae_dict[shared.opts.sd_vae] - if os.path.exists(vae_file): - vae_config = os.path.join('configs', 'flux', 'vae', 'config.json') - 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=Lumina2 failed to load VAE: {e}") - shared.opts.sd_vae = 'Default' - if debug: - errors.display(e, 'Lumina2 VAE:') - - if transformer is None: - transformer = diffusers.Lumina2Transformer2DModel.from_pretrained( - repo_id, - subfolder="transformer", - cache_dir=shared.opts.diffusers_dir, - **load_config, - **quant_config, - ) - - load_config, quant_config = model_quant.get_dit_args(diffusers_load_config, module='TE', device_map=True) - text_encoder = transformers.AutoModel.from_pretrained( - repo_id, - subfolder="text_encoder", - cache_dir=shared.opts.diffusers_dir, - **load_config, - **quant_config, - ) - - load_config, quant_config = model_quant.get_dit_args(diffusers_load_config, allow_quant=False) - if vae is not None: - load_config['vae'] = vae + load_config, _quant_args = model_quant.get_dit_args(diffusers_load_config, allow_quant=False) pipe = diffusers.Lumina2Pipeline.from_pretrained( repo_id, cache_dir=shared.opts.diffusers_dir, @@ -90,6 +36,8 @@ def load_lumina2(checkpoint_info, diffusers_load_config={}): **load_config, ) + del transformer + del text_encoder sd_hijack_te.init_hijack(pipe) devices.torch_gc(force=True, reason='load') return pipe diff --git a/pipelines/model_meissonic.py b/pipelines/model_meissonic.py index 30671e350..b045f006c 100644 --- a/pipelines/model_meissonic.py +++ b/pipelines/model_meissonic.py @@ -6,10 +6,10 @@ def load_meissonic(checkpoint_info, diffusers_load_config={}): from modules import shared, devices, modelloader, sd_models, shared_items from pipelines.meissonic.transformer import Transformer2DModel as TransformerMeissonic from pipelines.meissonic.scheduler import Scheduler as MeissonicScheduler - from pipelines.meissonic.pipeline import Pipeline as PipelineMeissonic - from pipelines.meissonic.pipeline_img2img import Img2ImgPipeline as PipelineMeissonicImg2Img - from pipelines.meissonic.pipeline_inpaint import InpaintPipeline as PipelineMeissonicInpaint - shared_items.pipelines['Meissonic'] = PipelineMeissonic + from pipelines.meissonic.pipeline import MeissonicPipeline + from pipelines.meissonic.pipeline_img2img import MeissonicImg2ImgPipeline + from pipelines.meissonic.pipeline_inpaint import MeissonicInpaintPipeline + shared_items.pipelines['Meissonic'] = MeissonicPipeline modelloader.hf_login() fn = sd_models.path_to_repo(checkpoint_info) @@ -41,7 +41,7 @@ def load_meissonic(checkpoint_info, diffusers_load_config={}): cache_dir=cache_dir, ) scheduler = MeissonicScheduler.from_pretrained(fn, subfolder="scheduler", cache_dir=cache_dir) - pipe = PipelineMeissonic( + pipe = MeissonicPipeline( vqvae=vqvae.to(devices.dtype), text_encoder=text_encoder.to(devices.dtype), transformer=model.to(devices.dtype), @@ -49,8 +49,8 @@ def load_meissonic(checkpoint_info, diffusers_load_config={}): scheduler=scheduler, ) - 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 + diffusers.pipelines.auto_pipeline.AUTO_TEXT2IMAGE_PIPELINES_MAPPING["meissonic"] = MeissonicPipeline + diffusers.pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING["meissonic"] = MeissonicImg2ImgPipeline + diffusers.pipelines.auto_pipeline.AUTO_INPAINT_PIPELINES_MAPPING["meissonic"] = MeissonicInpaintPipeline devices.torch_gc(force=True, reason='load') return pipe diff --git a/pipelines/model_omnigen.py b/pipelines/model_omnigen.py index 596fe4dbb..b8e8d7fd0 100644 --- a/pipelines/model_omnigen.py +++ b/pipelines/model_omnigen.py @@ -1,9 +1,6 @@ -import os import diffusers from modules import shared, devices, sd_models, model_quant -debug = shared.log.trace if os.environ.get('SD_LOAD_DEBUG', None) is not None else lambda *args, **kwargs: None - def load_omnigen(checkpoint_info, diffusers_load_config={}): # pylint: disable=unused-argument repo_id = sd_models.path_to_repo(checkpoint_info) diff --git a/pipelines/model_omnigen2.py b/pipelines/model_omnigen2.py index 94488ae0a..6f2f48e7b 100644 --- a/pipelines/model_omnigen2.py +++ b/pipelines/model_omnigen2.py @@ -1,8 +1,5 @@ -import os from modules import shared, devices, sd_models, model_quant -debug = shared.log.trace if os.environ.get('SD_LOAD_DEBUG', None) is not None else lambda *args, **kwargs: None - def load_omnigen2(checkpoint_info, diffusers_load_config={}): # pylint: disable=unused-argument repo_id = sd_models.path_to_repo(checkpoint_info) diff --git a/pipelines/model_pixart.py b/pipelines/model_pixart.py index 254326abc..2de1be95d 100644 --- a/pipelines/model_pixart.py +++ b/pipelines/model_pixart.py @@ -1,12 +1,14 @@ import transformers import diffusers from huggingface_hub import file_exists +from modules import shared, devices, sd_models, model_quant +from pipelines import generic def load_pixart(checkpoint_info, diffusers_load_config={}): - from modules import shared, devices, modelloader, sd_models, model_quant - modelloader.hf_login() repo_id = sd_models.path_to_repo(checkpoint_info) + sd_models.hf_auth_check(checkpoint_info) + repo_id_tenc = repo_id repo_id_pipe = repo_id @@ -15,30 +17,21 @@ def load_pixart(checkpoint_info, diffusers_load_config={}): if not file_exists(repo_id_pipe, "model_index.json"): repo_id_pipe = "PixArt-alpha/PixArt-Sigma-XL-2-1024-MS" - load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='Model') - transformer = diffusers.PixArtTransformer2DModel.from_pretrained( - repo_id, - subfolder='transformer', - cache_dir=shared.opts.hfcache_dir, - **load_args, - **quant_args, - ) - load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='TE', device_map=True) - text_encoder = transformers.T5EncoderModel.from_pretrained( - repo_id_tenc, - subfolder="text_encoder", - cache_dir=shared.opts.hfcache_dir, - **load_args, - **quant_args, - ) - load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config, allow_quant=False) + shared.log.debug(f'Load model: type=AuraFlow repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') + + transformer = generic.load_transformer(repo_id, cls_name=diffusers.PixArtTransformer2DModel, load_config=diffusers_load_config) + text_encoder = generic.load_text_encoder(repo_id_tenc, cls_name=transformers.T5EncoderModel, load_config=diffusers_load_config) + pipe = diffusers.PixArtSigmaPipeline.from_pretrained( repo_id_pipe, - cache_dir=shared.opts.diffusers_dir, transformer=transformer, text_encoder=text_encoder, + cache_dir=shared.opts.diffusers_dir, **load_args, ) + + del text_encoder + del transformer devices.torch_gc(force=True, reason='load') return pipe diff --git a/pipelines/model_qwen.py b/pipelines/model_qwen.py new file mode 100644 index 000000000..63ee989e3 --- /dev/null +++ b/pipelines/model_qwen.py @@ -0,0 +1,42 @@ +import transformers +import diffusers +from modules import shared, devices, sd_models, model_quant, sd_hijack_te +from pipelines import generic + + +def load_qwen(checkpoint_info, diffusers_load_config={}): + repo_id = sd_models.path_to_repo(checkpoint_info) + sd_models.hf_auth_check(checkpoint_info) + + load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config, module='Model') + shared.log.debug(f'Load model: type=Qwen model="{checkpoint_info.name}" repo="{repo_id}" offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') + + transformer = generic.load_transformer(repo_id, cls_name=diffusers.QwenImageTransformer2DModel, load_config=diffusers_load_config, modules_dtype_dict={"minimum_6bit": ["img_mod", "pos_embed", "time_text_embed", "img_in", "txt_in", "norm_out"]}) + repo_te = 'Qwen/Qwen-Image' if 'Qwen-Lightning' in repo_id else repo_id + text_encoder = generic.load_text_encoder(repo_te, cls_name=transformers.Qwen2_5_VLForConditionalGeneration, load_config=diffusers_load_config) + + pipe = diffusers.QwenImagePipeline.from_pretrained( + repo_id, + transformer=transformer, + text_encoder=text_encoder, + cache_dir=shared.opts.diffusers_dir, + **load_args, + ) + pipe.task_args = { + 'output_type': 'np', + } + + diffusers.pipelines.auto_pipeline.AUTO_TEXT2IMAGE_PIPELINES_MAPPING["qwen-image"] = diffusers.QwenImagePipeline + diffusers.pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING["qwen-image"] = diffusers.QwenImageImg2ImgPipeline + diffusers.pipelines.auto_pipeline.AUTO_INPAINT_PIPELINES_MAPPING["qwen-image"] = diffusers.QwenImageInpaintPipeline + + del text_encoder + del transformer + + sd_hijack_te.init_hijack(pipe) + from modules.video_models import video_vae + pipe.vae.orig_decode = pipe.vae.decode + pipe.vae.decode = video_vae.hijack_vae_decode + + devices.torch_gc() + return pipe diff --git a/pipelines/model_sd3.py b/pipelines/model_sd3.py index 6130ad81c..e488e89f1 100644 --- a/pipelines/model_sd3.py +++ b/pipelines/model_sd3.py @@ -1,128 +1,35 @@ -import os import diffusers import transformers -from modules import shared, devices, errors, sd_models, sd_unet, model_quant, model_tools +from modules import shared, devices, sd_models, model_quant +from pipelines import generic -def load_overrides(kwargs, cache_dir): - if shared.opts.sd_unet != 'Default': - try: - fn = sd_unet.unet_dict[shared.opts.sd_unet] - if fn.endswith('.safetensors'): - kwargs['transformer'] = diffusers.SD3Transformer2DModel.from_single_file(fn, cache_dir=cache_dir, torch_dtype=devices.dtype) - sd_unet.loaded_unet = shared.opts.sd_unet - shared.log.debug(f'Load model: type=SD3 unet="{shared.opts.sd_unet}" fmt=safetensors') - elif fn.endswith('.gguf'): - from modules import ggml - kwargs['transformer'] = ggml.load_gguf(fn, cls=diffusers.SD3Transformer2DModel, compute_dtype=devices.dtype) - sd_unet.loaded_unet = shared.opts.sd_unet - shared.log.debug(f'Load model: type=SD3 unet="{shared.opts.sd_unet}" fmt=gguf') - except Exception as e: - shared.log.error(f"Load model: type=SD3 failed to load UNet: {e}") - errors.display(e, 'UNet') - shared.opts.sd_unet = 'Default' - sd_unet.failed_unet.append(shared.opts.sd_unet) - - if shared.opts.sd_text_encoder != 'Default': - try: - from modules.model_te import load_t5, load_vit_l, load_vit_g - if 'vit-l' in shared.opts.sd_text_encoder.lower(): - kwargs['text_encoder'] = load_vit_l() - shared.log.debug(f'Load model: type=SD3 variant="vit-l" te="{shared.opts.sd_text_encoder}"') - elif 'vit-g' in shared.opts.sd_text_encoder.lower(): - kwargs['text_encoder_2'] = load_vit_g() - shared.log.debug(f'Load model: type=SD3 variant="vit-g" te="{shared.opts.sd_text_encoder}"') - else: - kwargs['text_encoder_3'] = load_t5(name=shared.opts.sd_text_encoder, cache_dir=shared.opts.diffusers_dir) - shared.log.debug(f'Load model: type=SD3 variant="t5" te="{shared.opts.sd_text_encoder}"') - except Exception as e: - shared.log.error(f"Load model: type=SD3 failed to load T5: {e}") - errors.display(e, 'TE') - shared.opts.sd_text_encoder = 'Default' - - if shared.opts.sd_vae != 'Default' and shared.opts.sd_vae != 'Automatic': - try: - from modules import sd_vae - vae_file = sd_vae.vae_dict[shared.opts.sd_vae] - if os.path.exists(vae_file): - vae_config = os.path.join('configs', 'sd3', 'vae', 'config.json') - kwargs['vae'] = diffusers.AutoencoderKL.from_single_file(vae_file, config=vae_config, cache_dir=cache_dir, torch_dtype=devices.dtype) - shared.log.debug(f'Load model: type=SD3 vae="{shared.opts.sd_vae}"') - except Exception as e: - shared.log.error(f"Load model: type=SD3 failed to load VAE: {e}") - errors.display(e, 'VAE') - shared.opts.sd_vae = 'Default' - return kwargs - - -def load_quants(kwargs, repo_id, cache_dir): - quant_args = model_quant.create_config(module='Model') - if quant_args and 'quantization_config' in quant_args: - kwargs['transformer'] = diffusers.SD3Transformer2DModel.from_pretrained(repo_id, subfolder="transformer", cache_dir=cache_dir, torch_dtype=devices.dtype, **quant_args) - quant_args = model_quant.create_config(module='TE') - if quant_args and 'quantization_config' in quant_args: - kwargs['text_encoder_3'] = transformers.T5EncoderModel.from_pretrained(repo_id, subfolder="text_encoder_3", variant='fp16', cache_dir=cache_dir, torch_dtype=devices.dtype, **quant_args) - return kwargs - - -def load_missing(kwargs, fn, cache_dir): - keys = model_tools.get_safetensor_keys(fn) - size = os.stat(fn).st_size // 1024 // 1024 - if size > 15000: - repo_id = 'stabilityai/stable-diffusion-3.5-large' - else: - repo_id = 'stabilityai/stable-diffusion-3-medium-diffusers' - if 'text_encoder' not in kwargs and 'text_encoder' not in keys: - kwargs['text_encoder'] = transformers.CLIPTextModelWithProjection.from_pretrained(repo_id, subfolder='text_encoder', cache_dir=cache_dir, torch_dtype=devices.dtype) - shared.log.debug(f'Load model: type=SD3 missing=te1 repo="{repo_id}"') - if 'text_encoder_2' not in kwargs and 'text_encoder_2' not in keys: - kwargs['text_encoder_2'] = transformers.CLIPTextModelWithProjection.from_pretrained(repo_id, subfolder='text_encoder_2', cache_dir=cache_dir, torch_dtype=devices.dtype) - shared.log.debug(f'Load model: type=SD3 missing=te2 repo="{repo_id}"') - if 'text_encoder_3' not in kwargs and 'text_encoder_3' not in keys: - load_args, quant_args = model_quant.get_dit_args({}, module='TE', device_map=True) - kwargs['text_encoder_3'] = transformers.T5EncoderModel.from_pretrained(repo_id, subfolder="text_encoder_3", variant='fp16', cache_dir=cache_dir, **load_args, **quant_args) - shared.log.debug(f'Load model: type=SD3 missing=te3 repo="{repo_id}"') - if 'vae' not in kwargs and 'vae' not in keys: - kwargs['vae'] = diffusers.AutoencoderKL.from_pretrained(repo_id, subfolder='vae', cache_dir=cache_dir, torch_dtype=devices.dtype) - shared.log.debug(f'Load model: type=SD3 missing=vae repo="{repo_id}"') - return kwargs - - -def load_sd3(checkpoint_info, cache_dir=None, config=None): +def load_sd3(checkpoint_info, diffusers_load_config={}): repo_id = sd_models.path_to_repo(checkpoint_info) sd_models.hf_auth_check(checkpoint_info) - fn = checkpoint_info.path - kwargs = {} - kwargs = load_overrides(kwargs, cache_dir) - if (fn is None) or (not os.path.exists(fn) or os.path.isdir(fn)): - kwargs = load_quants(kwargs, repo_id, cache_dir) - - loader = diffusers.StableDiffusion3Pipeline.from_pretrained - if fn is not None and os.path.exists(fn) and os.path.isfile(fn): - if fn.endswith('.safetensors'): - loader = diffusers.StableDiffusion3Pipeline.from_single_file - repo_id = fn - elif fn.endswith('.gguf'): - from modules import ggml - kwargs['transformer'] = ggml.load_gguf(fn, cls=diffusers.SD3Transformer2DModel, compute_dtype=devices.dtype) - kwargs = load_missing(kwargs, fn, cache_dir) - kwargs['variant'] = 'fp16' - else: - kwargs['variant'] = 'fp16' - - shared.log.debug(f'Load model: type=SD3 kwargs={list(kwargs)} repo="{repo_id}"') + load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config, allow_quant=False) + shared.log.debug(f'Load model: type=SD3 repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}') + transformer = generic.load_transformer(repo_id, cls_name=diffusers.SD3Transformer2DModel, load_config=diffusers_load_config) + # text_encoder = generic.load_text_encoder(repo_id, cls_name=transformers.CLIPTextModelWithProjection, load_config=diffusers_load_config, subfolder="text_encoder") + # text_encoder_2 = generic.load_text_encoder(repo_id, cls_name=transformers.CLIPTextModelWithProjection, load_config=diffusers_load_config, subfolder="text_encoder_2") if shared.opts.model_sd3_disable_te5: - shared.log.debug('Load model: type=SD3 option="disable-te5"') - kwargs['text_encoder_3'] = None + text_encoder_3 = None + else: + text_encoder_3 = generic.load_text_encoder(repo_id, cls_name=transformers.T5EncoderModel, load_config=diffusers_load_config, subfolder="text_encoder_3") - pipe = loader( + pipe = diffusers.StableDiffusion3Pipeline.from_pretrained( repo_id, - torch_dtype=devices.dtype, - cache_dir=cache_dir, - config=config, - **kwargs, + transformer=transformer, + # text_encoder=text_encoder, + # text_encoder_2=text_encoder_2, + text_encoder_3=text_encoder_3, + cache_dir=shared.opts.diffusers_dir, + **load_args, ) + + del text_encoder_3 + del transformer devices.torch_gc(force=True, reason='load') return pipe diff --git a/pipelines/model_stablecascade.py b/pipelines/model_stablecascade.py index fb780c0f2..60d8ca87c 100644 --- a/pipelines/model_stablecascade.py +++ b/pipelines/model_stablecascade.py @@ -190,6 +190,7 @@ class StableCascadeDecoderPipelineFixed(diffusers.StableCascadeDecoderPipeline): ): shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model) # 0. Define commonly used variables + guidance_scale = guidance_scale or 0.0 self.guidance_scale = guidance_scale self.do_classifier_free_guidance = self.guidance_scale > 1 device = self._execution_device diff --git a/pipelines/model_wanai.py b/pipelines/model_wanai.py index 49503c670..3edbfa67d 100644 --- a/pipelines/model_wanai.py +++ b/pipelines/model_wanai.py @@ -41,6 +41,7 @@ def load_transformer(repo_id, diffusers_load_config={}, subfolder='transformer') def load_text_encoder(repo_id, diffusers_load_config={}): load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='TE', device_map=True) + repo_id = 'Wan-AI/Wan2.1-T2V-1.3B-Diffusers' if 'Wan2.' in repo_id else repo_id # always use shared umt5 shared.log.debug(f'Load model: type=WanAI te="{repo_id}" quant="{model_quant.get_quant_type(quant_args)}" args={load_args}') text_encoder = transformers.UMT5EncoderModel.from_pretrained( repo_id, diff --git a/requirements.txt b/requirements.txt index 863edcccc..20fe134c3 100644 --- a/requirements.txt +++ b/requirements.txt @@ -32,7 +32,7 @@ pi-heif # versioned rich==14.1.0 -safetensors==0.5.3 +safetensors==0.6.2 tensordict==0.8.3 peft==0.17.0 httpx==0.24.1 @@ -41,16 +41,16 @@ torchsde==0.2.6 antlr4-python3-runtime==4.9.3 requests==2.32.4 tqdm==4.67.1 -accelerate==1.9.0 +accelerate==1.10.0 opencv-contrib-python-headless==4.11.0.86 einops==0.8.1 -huggingface_hub==0.34.3 +huggingface_hub==0.34.4 numexpr==2.11.0 numpy==2.1.2 -pandas==2.3.0 +pandas==2.3.1 numba==0.61.2 protobuf==4.25.3 -pytorch_lightning==2.5.2 +pytorch_lightning==2.5.3 tokenizers==0.21.4 urllib3==1.26.19 Pillow==10.4.0 diff --git a/scripts/nudenet/imageguard.py b/scripts/nudenet/imageguard.py index 55bcfaf86..edb2b7e8a 100644 --- a/scripts/nudenet/imageguard.py +++ b/scripts/nudenet/imageguard.py @@ -106,7 +106,7 @@ def image_guard(image, policy:str=None) -> str: attn_implementation='flash_attention_2', torch_dtype=devices.dtype, device_map="auto", - cache_dir='/mnt/models/huggingface', + cache_dir=shared.opts.hfcache_dir, ) processor = transformers.AutoProcessor.from_pretrained(repo_id, cache_dir=shared.opts.hfcache_dir) shared.log.info(f'NudeNet load: model="{repo_id}"') diff --git a/scripts/nudenet/nudenet.py b/scripts/nudenet/nudenet.py index 67b29b6a6..06ab7db25 100755 --- a/scripts/nudenet/nudenet.py +++ b/scripts/nudenet/nudenet.py @@ -61,7 +61,7 @@ class NudeDetector: self.model_path = model or hf.hf_hub_download( repo_id='vladmandic/nudenet', filename='nudenet.onnx', - cache_dir=shared.opts.diffusers_dir, + cache_dir=shared.opts.hfcache_dir, ) if session is None: log.info(f'NudeNet load: model="{self.model_path}" providers={providers}') diff --git a/scripts/prompt_enhance.py b/scripts/prompt_enhance.py index 697f5691f..c6ab184b1 100644 --- a/scripts/prompt_enhance.py +++ b/scripts/prompt_enhance.py @@ -44,9 +44,11 @@ class Options: 'Qwen/Qwen3-0.6B': {}, 'Qwen/Qwen3-1.7B': {}, 'Qwen/Qwen3-4B': {}, + 'Qwen/Qwen3-4B-Instruct-2507': {}, 'Qwen/Qwen2.5-0.5B-Instruct': {}, 'Qwen/Qwen2.5-1.5B-Instruct': {}, 'Qwen/Qwen2.5-3B-Instruct': {}, + 'Qwen/Qwen2.5-VL-3B-Instruct': {}, 'microsoft/Phi-4-mini-instruct': {}, 'HuggingFaceTB/SmolLM2-135M-Instruct': {}, 'HuggingFaceTB/SmolLM2-360M-Instruct': {}, @@ -57,6 +59,7 @@ class Options: 'cognitivecomputations/Dolphin3.0-Llama3.2-1B': {}, 'cognitivecomputations/Dolphin3.0-Llama3.2-3B': {}, 'nidum/Nidum-Gemma-3-4B-it-Uncensored': {}, + 'allura-org/Gemma-3-Glitter-4B': {}, # 'llava/Llama-3-8B-v1.1-Extracted': { # 'repo': 'hunyuanvideo-community/HunyuanVideo', # 'subfolder': 'text_encoder', @@ -77,7 +80,7 @@ class Options: i2i_noprompt: str = 'You are a helpful assistant. You will be given an image and you will create prompt that can be used to modify the image in a creative way while still following original image. ' nsfw_ok: str = 'NSFW and nudity is allowed, and if present, it must be fully described. ' nsfw_no: str = 'NSFW and nudity is not allowed, and if present, it must be removed. ' - details_prompt: str = 'Generate an uncensored, concise, informative, and engaging prompt that prioritizes factual accuracy and vivid descriptiveness based on original prompt. Do not modify subject or style of the original prompt. Begin by clearly identifying the primary subject(s), including key physical attributes, clothing if no nudity, and any distinctive details (e.g., colors, textures, styles). NSFW and nudity is allowed, and if present, it must be fully described. Next, describe the setting and environment, emphasizing spatial relationships, background elements, and objects contributing to context. Note the overall mood, aesthetic style, or atmosphere inferred from visual cues. Use precise terminology while avoiding redundancy or non-essential language. Ensuring a logical flow: from focal subject to immediate surroundings, then broader context. Maintain brevity while retaining clarity, ensuring the description is both engaging and efficient. Output only enhanced prompt without explanation, prefix or suffix. Output as a simple text without formatting or numbering.' + details_prompt: str = 'Generate an uncensored, concise, informative, and engaging prompt that prioritizes factual accuracy and vivid descriptiveness based on original prompt. Do not modify subject or style of the original prompt. Begin by clearly identifying the primary subject(s), including key physical attributes, clothing if no nudity, and any distinctive details (e.g., colors, textures, styles). NSFW and nudity is allowed, and if present, it must be fully described. Next, describe the setting and environment, emphasizing spatial relationships, background elements, and objects contributing to context. Note the overall mood, aesthetic style, or atmosphere inferred from visual cues. Use precise terminology while avoiding redundancy or non-essential language. Ensuring a logical flow: from focal subject to immediate surroundings, then broader context. Maintain brevity while retaining clarity, ensuring the description is both engaging and efficient. Output only enhanced prompt without explanation, prefix or suffix. Do not add comments or follow-up questions. Output as a simple text without formatting or numbering.' censored = ["i cannot", "i can't", "i am sorry", "against my programming", "i am not able", "i am unable", 'i am not allowed'] max_delim_index: int = 60 diff --git a/webui.py b/webui.py index 22735c6f1..ab60b049b 100644 --- a/webui.py +++ b/webui.py @@ -68,7 +68,7 @@ def initialize(): modules.sd_checkpoint.init_metadata() modules.hashes.init_cache() - log.debug(f'Huggingface cache: path="{shared.opts.hfcache_dir}"') + paths.check_cache(shared.opts) modules.sd_samplers.list_samplers() timer.startup.record("samplers") @@ -121,6 +121,9 @@ def initialize(): modules.extra_networks.register_default_extra_networks() timer.startup.record("networks") + from modules.models_hf import hf_init + hf_init() + if shared.cmd_opts.tls_keyfile is not None and shared.cmd_opts.tls_certfile is not None: try: if not os.path.exists(shared.cmd_opts.tls_keyfile): @@ -413,23 +416,5 @@ def webui(restart=False): return shared.demo.server -def api_only(): - start_common() - from fastapi import FastAPI - app = FastAPI(**fastapi_args) - modules.api.middleware.setup_middleware(app, shared.cmd_opts) - shared.api = create_api(app) - shared.api.register() - shared.api.wants_restart = False - modules.script_callbacks.app_started_callback(None, app) - modules.sd_models.write_metadata() - log.info(f"Startup time: {timer.startup.summary()}") - server = shared.api.launch() - return server - - if __name__ == "__main__": - if shared.cmd_opts.api_only: - api_only() - else: - webui() + webui() diff --git a/wiki b/wiki index 79b18f2c5..f91e819d2 160000 --- a/wiki +++ b/wiki @@ -1 +1 @@ -Subproject commit 79b18f2c5e3438f3f564fd264fdb27bed76b0f72 +Subproject commit f91e819d22603f34be0c3e8fb674d4ab89421622