ostris/ai-toolkit
1
0
Fork 0
mirror of https://github.com/ostris/ai-toolkit.git synced 2026-04-24 16:29:26 +00:00
Code Issues Packages Projects Releases Wiki Activity

Added ability to quantize with torchao

Browse Source
This commit is contained in:
Jaret Burkett
2025-03-20 16:28:54 -06:00
parent 3a6b24f4c8
commit f5aa4232fa
7 changed files with 57 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
toolkit/models/base_model.py
View File

@@ -23,27 +23,22 @@ from toolkit.models.decorator import Decorator
from toolkit.paths import KEYMAPS_ROOT
from toolkit.prompt_utils import inject_trigger_into_prompt, PromptEmbeds, concat_prompt_embeds
from toolkit.reference_adapter import ReferenceAdapter
from toolkit.saving import save_ldm_model_from_diffusers
from toolkit.sd_device_states_presets import empty_preset
from toolkit.train_tools import get_torch_dtype, apply_noise_offset
import torch
from toolkit.pipelines import CustomStableDiffusionXLPipeline
from diffusers import StableDiffusionPipeline, StableDiffusionXLPipeline, T2IAdapter, DDPMScheduler, \
LCMScheduler, Transformer2DModel, AutoencoderTiny, ControlNetModel, \
FluxTransformer2DModel
from toolkit.models.lumina2 import Lumina2Transformer2DModel
LCMScheduler, Transformer2DModel, AutoencoderTiny, ControlNetModel
import diffusers
from diffusers import \
AutoencoderKL, \
UNet2DConditionModel
from diffusers import PixArtAlphaPipeline
from transformers import T5EncoderModel, UMT5EncoderModel
from transformers import CLIPTextModel, CLIPTokenizer, CLIPTextModelWithProjection
from toolkit.accelerator import get_accelerator, unwrap_model
from typing import TYPE_CHECKING
from toolkit.print import print_acc
from transformers import Gemma2Model, Qwen2Model, LlamaModel
if TYPE_CHECKING:
from toolkit.lora_special import LoRASpecialNetwork