Add support to Flex.1- alpha

[elen07zz]

https://huggingface.co/ostris/Flex.1-alpha

Flex.1 alpha is a pre-trained base 8 billion parameter rectified flow transformer capable of generating images from text descriptions. It has a similar architecture to FLUX.1-dev, but with fewer double transformer blocks (8 vs 19). It began as a finetune of FLUX.1-schnell which allows the model to retain the Apache 2.0 license. A guidance embedder has been trained for it so that it no longer requires CFG to generate images.

[stepfunction83]

I think this should be fairly straightforward to do by training the model as if it were Schnell, similar to what is done with the de-distill models. I'm going to experiment with that approach and see if it works.

[stepfunction83]

I can confirm that forcing the model to be trained as schnell by changing:

is_schnell = not ("guidance_in.in_layer.bias" in keys or "time_text_embed.guidance_embedder.linear_1.bias" in keys)

to

is_schnell = True

In library/flux_utils.py, line 62 seems to allow it to finetune directly.

[stepfunction83]

I can also confirm that the new guidance module that Ostris created gets destroyed in the process. It will need to be explicitly excluded.

[stepfunction83]

This seems to be the key snippet of code to bypass the guidance block in ai-toolkit:

def guidance_embed_bypass_forward(self, timestep, guidance, pooled_projection):
    timesteps_proj = self.time_proj(timestep)
    timesteps_emb = self.timestep_embedder(
        timesteps_proj.to(dtype=pooled_projection.dtype))  # (N, D)
    pooled_projections = self.text_embedder(pooled_projection)
    conditioning = timesteps_emb + pooled_projections
    return conditioning

# bypass the forward function

def bypass_flux_guidance(transformer):
    if hasattr(transformer.time_text_embed, '_bfg_orig_forward'):
        return
    # dont bypass if it doesnt have the guidance embedding
    if not hasattr(transformer.time_text_embed, 'guidance_embedder'):
        return
    transformer.time_text_embed._bfg_orig_forward = transformer.time_text_embed.forward
    transformer.time_text_embed.forward = partial(
        guidance_embed_bypass_forward, transformer.time_text_embed
    )

# restore the forward function

def restore_flux_guidance(transformer):
    if not hasattr(transformer.time_text_embed, '_bfg_orig_forward'):
        return
    transformer.time_text_embed.forward = transformer.time_text_embed._bfg_orig_forward
    del transformer.time_text_embed._bfg_orig_forward

Where transformer is the unet