在低配显卡下利用diffusers库驱动Flux模型

望峰
1,241 阅读4分钟

配置信息

  • 机器资源配置

    • 显卡 >= 15G
    • 内存 >= 显存

  • diffusers版本>=0.31.0

概述

本文旨在探讨在中低端显卡环境下，利用加速技术驱动Flux dev fp8模型，通过diffusers实现高效推理的方法、遇到的问题及相应策略。

通过本地模型驱动需要哪些文件

模型文件说明：由于需要单独使用fp8量化模型，下载时可以去除transformer文件夹

flux fp8地址：huggingface.co/Kijai/flux-…

flux地址：huggingface.co/black-fores…

image.png

加载优化

方式一：组装式加载

为了节省显存，需要将transformer单独加载和量化，并开启enable_model_cpu_offload以减少显存占用

from optimum.quanto import freeze, qfloat8, quantize
dtype = torch.bfloat16

print('loading flux transformer……')
transformer = FluxTransformer2DModel.from_single_file(
    pretrained_model_name_or_path, # 下载的fp8模型地址
    torch_dtype=dtype,
    config=config_path,
    local_files_only=True
)
print('loaded flux transformer')

print('optimized flux transformer')
quantize(transformer, weights=qfloat8) # 对模型进行量化
freeze(transformer)

print('loading flux text_encoder_2……')
text_encoder_2 = T5EncoderModel.from_pretrained(bfl_repo, subfolder="text_encoder_2", torch_dtype=dtype)
quantize(text_encoder_2, weights=qfloat8)
freeze(text_encoder_2)

print('loaded text encoder')

pipe = FluxPipeline.from_pretrained(flux_dir, # flux 模型对应的文件夹地址
                                    transformer=None, #单独加载transformer和text_encoder_2单元
                                    text_encoder_2=None, 
                                    torch_dtype=dtype)
pipe.transformer = transformer
pipe.text_encoder_2 = text_encoder_2
print('loaded flux pipeline')

pipeline.enable_model_cpu_offload()
return pipe

持续优化手段：

  • 如果显存还是不足以支撑模型加载，可以尝试将enable_model_cpu_offload切换为enable_sequential_cpu_offload，将不同模型按执行顺序加载，但是会拖慢生图的整体速度（在A30测试，速度降低80%: 20s -> 100s）
  • 开启vae优化，可继续将显存降低（会降低生图质量）
pipeline.vae.enable_slicing() # 多批次生图优化
pipeline.vae.enable_tiling() # 大尺寸图像优化

方式二：分段式加载

基于方式一，可以将处理prompt_embeds和扩散过程拆分进行，两个阶段组装不同的单元

def flush():
    gc.collect()
    torch.cuda.empty_cache()
    torch.cuda.reset_max_memory_allocated()
    torch.cuda.reset_peak_memory_stats()

text_encoder = CLIPTextModel.from_pretrained(
    flux_dir, subfolder="text_encoder", torch_dtype=torch.bfloat16
)
text_encoder_2 = T5EncoderModel.from_pretrained(
    flux_dir, subfolder="text_encoder_2", torch_dtype=torch.bfloat16
)
tokenizer = CLIPTokenizer.from_pretrained(flux_dir, subfolder="tokenizer")
tokenizer_2 = T5TokenizerFast.from_pretrained(flux_dir, subfolder="tokenizer_2")

# 第一阶段：加载text_encoder 和 tokenizer处理prompt
pipeline = FluxPipeline.from_pretrained(
    flux_dir,
    text_encoder=text_encoder,
    text_encoder_2=text_encoder_2,
    tokenizer=tokenizer,
    tokenizer_2=tokenizer_2,
    transformer=None,
    vae=None,
).to("cuda")

with torch.no_grad():
    print("Encoding prompts.")
    prompt_embeds, pooled_prompt_embeds, text_ids = pipeline.encode_prompt(
        prompt=prompt, prompt_2=None, max_sequence_length=256
    )

del text_encoder
del text_encoder_2
del tokenizer
del tokenizer_2
del pipeline

flush()

# 第二阶段：加载vae、transformer，执行扩散处理
pipeline = FluxPipeline.from_pretrained(
    flux_dir,
    text_encoder=None,
    text_encoder_2=None,
    tokenizer=None,
    tokenizer_2=None,
    torch_dtype=torch.bfloat16,
)

Lora支持

Lora不生效（也不会报错）

问题表现：

存在大量的key缺失的Warning，信息如下：

for single_transformer_blocks.30.attn.to_v.lora_A.digit_art.weight: copying from a non-meta parameter in the checkpoint to a meta parameter in the current model, which is a no-op. (Did you mean to pass assign=True to assign items in the state dictionary to their corresponding key in the module instead of copying them in place?)

原因：

在diffusers的0.30.0版本中，未支持xlabs和kohya的lora，导致lora中缺失大量的key未加入到transformer和text_encoder中，具体代码如下：

image.png 代码位置可参考：github.com/huggingface…

解决方式：

方式一（推荐）：diffusers库更新到0.31.0

方式二：根据最新的diffusers源码适配0.30.0版本

koyha适配代码：github.com/huggingface…

xlabs适配代码：github.com/huggingface…

quanto和lora不适配

问题表现：

抛出键值异常：KeyError: 'time_text_embed.timestep_embedder.linear_1.weight_qtype'

原因：

在使用optimum-quanto优化网络时，如果版本<=0.2.2, 会在加载lora时尝试从权重矩阵中获取weight_qtype属性，而这个属性lora中大部分是缺失的

0.2.2版本代码如下：

image.png 在0.2.3版本之后，weight_qtype会直接从初始化的权重矩阵中自动读取，代码如下：

image.png

image.png

解决方式：

更新optimum-quanto版本>=0.2.3

Prompt limitation

在flux中仍然有77 tokens的限制，之前的SD模型中，diffusers提供了基于Compel库的解决方案，SDXL示例如下：

from compel import Compel, ReturnedEmbeddingsType

compel = Compel(
            tokenizer=[pipeline.tokenizer, pipeline.tokenizer_2] ,
            text_encoder=[pipeline.text_encoder, pipeline.text_encoder_2],
            returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED,
            requires_pooled=[False, True]
         )
prompt_embeds, pooled_prompt_embeds = compel(prompt)

但是Compel现阶段还不支持flux（来自作者的回复：github.com/damian0815/…）

可改用sd_embed支持long prompt，它支持AUTOMATIC1111 prompt规范，可无缝通过sd webui迁移，具体使用方式如下：

from sd_embed.embedding_funcs import get_weighted_text_embeddings_flux1

prompt_embeds, pooled_prompt_embeds = get_weighted_text_embeddings_flux1(
    pipe=pipeline,
    prompt=prompt
)

## 支持批量生图num_images_per_prompt，需要张开张量
bs_embed, seq_len, _ = prompt_embeds.shape
# duplicate text embeddings for each generation per prompt, using mps friendly method
prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1)

pooled_prompt_embeds = pooled_prompt_embeds.repeat(1, num_images_per_prompt).view(
    bs_embed * num_images_per_prompt, -1
)

image = pipe(
    prompt_embeds               = prompt_embeds
    , pooled_prompt_embeds      = pooled_prompt_embeds
    , width                     = 896
    , height                    = 1280
    , num_inference_steps       = 20
    , guidance_scale            = 4.0
    , generator                 = torch.Generator().manual_seed(1234)
).images[0]
display(image)

参数支持

flux不支持参数列表

名称说明
clip_skip
eta用于DDIMScheduler
sigmas用于部分scheduler
negative_promptflux无需负向词支持，可直接在prompt中描述

flux扩充支持的参数列表

名称说明
prompt_2支持扩展更长的prompt
avatar
文章
阅读
粉丝