have doubts about flux model load

[mobilejammer]

I'm sorry for always posting to ask for help or seek solutions. This time, I have a code problem that I need help with.

`elif (
args.model_family.lower() == "flux" and not args.flux_attention_masked_training
):
#from helpers.models.flux.transformer import (
# FluxTransformer2DModelWithMasking as FluxTransformer2DModel,
#)
from diffusers import FluxTransformer2DModel
import torch

    if torch.cuda.is_available():
        rank = (
            torch.distributed.get_rank()
            if torch.distributed.is_initialized()
            else 0
        )
        primary_device = torch.cuda.get_device_properties(0)
        if primary_device.major >= 9:
            try:
                from flash_attn_interface import flash_attn_func
                import diffusers

                from helpers.models.flux.attention import (
                    FluxAttnProcessor3_0,
                    FluxSingleAttnProcessor3_0,
                )

                diffusers.models.attention_processor.FluxSingleAttnProcessor2_0 = (
                    FluxSingleAttnProcessor3_0
                )
                diffusers.models.attention_processor.FluxAttnProcessor2_0 = (
                    FluxAttnProcessor3_0
                )
                if rank == 0:
                    print("Using FlashAttention3_0 for H100 GPU (Single block)")
            except:
                if rank == 0:
                    logger.warning(
                        "No flash_attn is available, using slower FlashAttention_2_0. Install flash_attn to make use of FA3 for Hopper or newer arch."
                    )

    transformer = FluxTransformer2DModel.from_pretrained(
        args.pretrained_transformer_model_name_or_path
        or args.pretrained_model_name_or_path,
        subfolder=determine_subfolder(args.pretrained_transformer_subfolder),
        **pretrained_load_args,
    )
elif args.model_family.lower() == "flux" and args.flux_attention_masked_training:
    from helpers.models.flux.transformer import (
        FluxTransformer2DModelWithMasking,
    )

    transformer = FluxTransformer2DModelWithMasking.from_pretrained(
        args.pretrained_transformer_model_name_or_path
        or args.pretrained_model_name_or_path,
        subfolder=determine_subfolder(args.pretrained_transformer_subfolder),
        **pretrained_load_args,
    )`

This is the code starting from line 53 of the diffusion_model. py. I've looked at the two branches related to flux, and ultimately they are both about mask. Then I modified your code to the following:

#from helpers.models.flux.transformer import (

FluxTransformer2DModelWithMasking as FluxTransformer2DModel,

#)
from diffusers import FluxTransformer2DModel

I'm not sure if my modification has misunderstood the meaning of your code. I'm afraid that I haven't fully grasped the intention of your code, which may cause more problems for the subsequent program. Thank you!

[bghira]

i can't even tell what the modification is. but you can't switch back to the Diffusers modeling code, it lacks a bunch of custom features that SimpleTuner relies on.

[bghira]

historically, the masking transformer modeling class was only used for masked training, to avoid introducing any errors to users that did not enable this option.

eventually, the masking modeling class became more fully featured than the Diffusers implementation.

now we rely on it pretty much everywhere that the training weights are instantiated. we don't HAVE to use it for other things where we're not training the model, eg. saving the weights.

[mobilejammer]

I'm terribly sorry that I might have misled you in helping me solve the problem. The FLUX model I downloaded is from https://huggingface.co/black-forest-labs/FLUX.1-dev. However, the "config.json" inside it is as follows:
{
"_class_name": "FluxTransformer2DModel",
"_diffusers_version": "0.30.0.dev0",
"_name_or_path": "../checkpoints/flux-dev/transformer",
"attention_head_dim": 128,
"guidance_embeds": true,
"in_channels": 64,
"joint_attention_dim": 4096,
"num_attention_heads": 24,
"num_layers": 19,
"num_single_layers": 38,
"patch_size": 1,
"pooled_projection_dim": 768
}
It seems that there is no "mask" in it. It appears that regardless of whether args.flux_attention_masked_training is set or not, Simple Tuner always loads FluxTransformer2DModelWithMasking. That's why I have this question.

[bghira]

this is fine. it's just a local copy of Flux transformer with upgrades.