twodgirl / flux-magcache-pulid-diffusers

Flux MagCache and PuLID

The merge of two projects ( 1 , 2 ).

The transformer model supports block swapping during inference. This saves more resources in addition to the enable_model_cpu_offload function.

Using MagCache with float8 precision provides an additional speed benefit.

Inference

from diffusers import FluxPipeline
from flux_magcache_model import FluxTransformer2DModel
from flux_pulid_loader import PuLID
from PIL import Image
import torch

class FluxGenerator:
    def __init__(self, pulid=False, compute_fp8=False, **kwargs):
        pipe_args = {}
        if compute_fp8:
            transformer = FluxTransformer2DModel.from_pretrained('repo-name/flux.1-dev-fp8-diffusers',
                                                                 subfolder='transformer',
                                                                 torch_dtype=torch.float8_e4m3fn)
            replace_regular_linears(transformer)
            pipe_args['transformer'] = transformer
        self.pipe = MagPipeline.from_pretrained('black-forest-labs/FLUX.1-dev', torch_dtype=torch.bfloat16, **pipe_args)
        self.pipe.enable_model_cpu_offload()
        if pulid:
            self.pulid_model = PuLID()
            self.pulid_model.setup(self.pipe.transformer)

    def clear_id(self):
        self.pipe.transformer.pul_id = None
        self.pipe.transformer.pul_weight = 1.0

    def set_id(self,
               id_image: Image.Image,
               id_weight=1.0,
               true_cfg=1.0):
        use_true_cfg = abs(true_cfg - 1.0) > 1e-2
        id_embeddings, uncond_id_embeddings = self.pulid_model.get_id_embedding_by_pil(id_image, true_cfg=true_cfg)
        self.pipe.transformer.pul_id = uncond_id_embeddings if use_true_cfg else id_embeddings
        self.pipe.transformer.pul_id_weight = id_weight

class MagPipeline(FluxPipeline):
    def __call__(self, **kwargs):
        num_inference_steps = kwargs.get('num_inference_steps', 28)
        self.transformer.retention_steps = int(num_inference_steps * self.transformer.retention_ratio)

        return super().__call__(**kwargs)

if __name__ == '__main__':
    gen = FluxGenerator()
    gen.set_id(Image.open('reference_image.png'))
    prompt = """A stylish woman walks down a street filled with warm glowing neon and animated city signs."""
    image = gen.pipe(prompt=prompt).images[0]
    image.save('preview.png')

Disclaimer

Use of this code and the model requires citation and attribution to the author via a link to their Hugging Face profile in all resulting work.