nodes/lodestone-rock-ComfyUI_FluxMod
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

add v2 approximator

Browse Source
This commit is contained in:
lodestone-rock
2024-11-24 11:11:09 +07:00
parent bb2dc04229
commit 2da24a8f7e
1 changed files with 6 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
flux_mod/loader.py
+6 -2
View File
@@ -85,7 +85,11 @@ def load_flux_mod(model_path, timestep_guidance_path, linear_dtypes=torch.bfloat
state_dict = load_selected_keys(model_path, ["mod", "time_in", "guidance_in", "vector_in"])
timestep_state_dict = comfy.utils.load_torch_file(timestep_guidance_path)
if "v2" in timestep_guidance_path:
n_layers = 5
else:
n_layers = 4
param_count = sum([x.numel() for x in state_dict.values()])
unet_dtype = torch.bfloat16
@@ -120,7 +124,7 @@ def load_flux_mod(model_path, timestep_guidance_path, linear_dtypes=torch.bfloat
model.diffusion_model = FluxMod(params=params)
model.diffusion_model.load_state_dict(state_dict)
model.diffusion_model.distilled_guidance_layer = Approximator(64, 3072, 5120, 4)
model.diffusion_model.distilled_guidance_layer = Approximator(64, 3072, 5120, n_layers)
model.diffusion_model.distilled_guidance_layer.load_state_dict(timestep_state_dict)
model.diffusion_model.dtype = unet_dtype
model.diffusion_model.eval()