ComfyUI

#6
by nhathoangfoto - opened

How to use this model in ComfyUI?

you need to merge models

from diffusers import FluxTransformer2DModel
from huggingface_hub import snapshot_download
from accelerate import init_empty_weights
from diffusers.models.model_loading_utils import load_model_dict_into_meta
import safetensors.torch
import glob
import torch

with init_empty_weights():
config = FluxTransformer2DModel.load_config("black-forest-labs/FLUX.1-dev", subfolder="transformer")
model = FluxTransformer2DModel.from_config(config)

dev_ckpt = snapshot_download(repo_id="black-forest-labs/FLUX.1-dev", allow_patterns="transformer/")
schnell_ckpt = snapshot_download(repo_id="black-forest-labs/FLUX.1-schnell", allow_patterns="transformer/
")

dev_shards = sorted(glob.glob(f"{dev_ckpt}/transformer/.safetensors"))
schnell_shards = sorted(glob.glob(f"{schnell_ckpt}/transformer/
.safetensors"))

merged_state_dict = {}
guidance_state_dict = {}

for i in range(len(dev_shards)):
state_dict_dev_temp = safetensors.torch.load_file(dev_shards[i])
state_dict_schnell_temp = safetensors.torch.load_file(schnell_shards[i])

keys = list(state_dict_dev_temp.keys())
for k in keys:
    if "guidance" not in k:
        merged_state_dict[k] = (state_dict_dev_temp.pop(k) + state_dict_schnell_temp.pop(k)) / 2
    else:
        guidance_state_dict[k] = state_dict_dev_temp.pop(k)

if len(state_dict_dev_temp) > 0:
    raise ValueError(f"There should not be any residue but got: {list(state_dict_dev_temp.keys())}.")
if len(state_dict_schnell_temp) > 0:
    raise ValueError(f"There should not be any residue but got: {list(state_dict_dev_temp.keys())}.")

merged_state_dict.update(guidance_state_dict)
load_model_dict_into_meta(model, merged_state_dict)

model.to(torch.bfloat16).save_pretrained("./merged-flux")

@jeje3869
u can download the merged one from here
but than its in blobls as 3 huge chunks of safetensors
how to load it in comfyui
what node/s ?

When I run the merge script it dies here:
load_model_dict_into_meta(model, merged_state_dict)

Doesn't give an error just says Killed.
(Guessing that's the process that was killed)

I don't know how to run the ComfyUI conversion and won't have time to do that either.

sayakpaul changed discussion status to closed

When I run the merge script it dies here:
load_model_dict_into_meta(model, merged_state_dict)

Doesn't give an error just says Killed.
(Guessing that's the process that was killed)

https://huggingface.co/drbaph/FLUX.1-schnell-dev-merged

@jeje3869
u can download the merged one from here
but than its in blobls as 3 huge chunks of safetensors
how to load it in comfyui
what node/s ?

https://huggingface.co/drbaph/FLUX.1-schnell-dev-merged

No description provided.

I can confim on 16GB VRAM WORKS like CHARM:

1 file safetensors with comfy workflow
https://huggingface.co/drbaph/FLUX.1-schnell-dev-merged
thx to @sayakpaul @drbaph

Sign up or log in to comment