ostris/ai-toolkit
1
0
Fork 0
mirror of https://github.com/ostris/ai-toolkit.git synced 2026-04-27 09:44:02 +00:00
Code Issues Packages Projects Releases Wiki Activity

Allow DFE to not have a VAE

Browse Source
This commit is contained in:
Jaret Burkett
2025-03-30 09:23:01 -06:00
parent 860d892214
commit c083a0e5ea
2 changed files with 10 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
toolkit/models/diffusion_feature_extraction.py
View File

@@ -154,11 +154,12 @@ class DiffusionFeatureExtractor(nn.Module):
class DiffusionFeatureExtractor3(nn.Module):
def __init__(self, device=torch.device("cuda"), dtype=torch.bfloat16):
def __init__(self, device=torch.device("cuda"), dtype=torch.bfloat16, vae=None):
super().__init__()
self.version = 3
vae = AutoencoderTiny.from_pretrained(
"madebyollin/taef1", torch_dtype=torch.bfloat16)
if vae is None:
vae = AutoencoderTiny.from_pretrained(
"madebyollin/taef1", torch_dtype=torch.bfloat16)
self.vae = vae
image_encoder_path = "google/siglip-so400m-patch14-384"
try:
@@ -342,9 +343,9 @@ class DiffusionFeatureExtractor3(nn.Module):
return total_loss
def load_dfe(model_path) -> DiffusionFeatureExtractor:
def load_dfe(model_path, vae=None) -> DiffusionFeatureExtractor:
if model_path == "v3":
dfe = DiffusionFeatureExtractor3()
dfe = DiffusionFeatureExtractor3(vae=vae)
dfe.eval()
return dfe
if not os.path.exists(model_path):