github/ai-toolkit
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Phase 2

Browse Source
This commit is contained in:
Jaret Burkett
2026-08-27 18:18:31 -06:00
parent 45886f01b2
commit 520d96aac3
27 changed files with 333 additions and 253 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
extensions_built_in/diffusion_models/omnigen2/__init__.py
View File

@@ -13,7 +13,7 @@ from toolkit.samplers.custom_flowmatch_sampler import (
)
from toolkit.accelerator import unwrap_model
from optimum.quanto import freeze
from toolkit.util.quantize import quantize, get_qtype
from toolkit.util.quantize import quantize, get_qtype, quantize_model
from .src.pipelines.omnigen2.pipeline_omnigen2 import OmniGen2Pipeline
from .src.models.transformers import OmniGen2Transformer2DModel
from .src.models.transformers.repo import OmniGen2RotaryPosEmbed
@@ -105,9 +105,9 @@ class OmniGen2Model(BaseModel):
if self.model_config.quantize:
self.print_and_status_update("Quantizing transformer")
quantization_type = get_qtype(self.model_config.qtype)
quantize(transformer, weights=quantization_type)
freeze(transformer)
quantize_model(self, transformer)
if not self.low_vram:
transformer.to(self.device_torch)
if self.low_vram:
# unload it for now