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

Test model loading

Browse Source
This commit is contained in:
Jaret Burkett
2026-08-27 14:31:53 -06:00
parent 8db198ec0a
commit 9ed2e0b8e7
5 changed files with 326 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
extensions_built_in/diffusion_models/chroma/chroma_model.py
View File

@@ -50,10 +50,12 @@ class FakeConfig:
self.patch_size = 1
class FakeCLIP(torch.nn.Module):
def __init__(self):
def __init__(self, device='cuda'):
super().__init__()
self.dtype = torch.bfloat16
self.device = 'cuda'
# the pipeline derives its execution device from this attribute;
# nn.Module.to() does not update it
self.device = device
self.text_model = None
self.tokenizer = None
self.model_max_length = 77
@@ -180,8 +182,8 @@ class ChromaModel(BaseModel):
self.prepare_text_encoder(text_encoder_2, dtype=dtype)
# self.print_and_status_update("Loading CLIP")
text_encoder = FakeCLIP()
tokenizer = FakeCLIP()
text_encoder = FakeCLIP(device=self.device_torch)
tokenizer = FakeCLIP(device=self.device_torch)
text_encoder.to(self.device_torch, dtype=dtype)
self.noise_scheduler = ChromaModel.get_train_scheduler()

10
extensions_built_in/diffusion_models/chroma/chroma_radiance_model.py
View File

@@ -50,10 +50,12 @@ class FakeConfig:
self.patch_size = 1
class FakeCLIP(torch.nn.Module):
def __init__(self):
def __init__(self, device='cuda'):
super().__init__()
self.dtype = torch.bfloat16
self.device = 'cuda'
# the pipeline derives its execution device from this attribute;
# nn.Module.to() does not update it
self.device = device
self.text_model = None
self.tokenizer = None
self.model_max_length = 77
@@ -182,8 +184,8 @@ class ChromaRadianceModel(BaseModel):
self.prepare_text_encoder(text_encoder_2, dtype=dtype)
# self.print_and_status_update("Loading CLIP")
text_encoder = FakeCLIP()
tokenizer = FakeCLIP()
text_encoder = FakeCLIP(device=self.device_torch)
tokenizer = FakeCLIP(device=self.device_torch)
text_encoder.to(self.device_torch, dtype=dtype)
self.noise_scheduler = ChromaRadianceModel.get_train_scheduler()