how to run flux on your 16gb potato

flux_on_potato.py

# First, in your terminal.
#
# $ python3 -m virtualenv env
# $ source env/bin/activate
# $ pip install torch torchvision transformers sentencepiece protobuf accelerate
# $ pip install git+https://github.com/huggingface/diffusers.git
# $ pip install optimum-quanto

import torch

from optimum.quanto import freeze, qfloat8, quantize

from diffusers import FlowMatchEulerDiscreteScheduler, AutoencoderKL
from diffusers.models.transformers.transformer_flux import FluxTransformer2DModel
from diffusers.pipelines.flux.pipeline_flux import FluxPipeline
from transformers import CLIPTextModel, CLIPTokenizer,T5EncoderModel, T5TokenizerFast

dtype = torch.bfloat16

# schnell is the distilled turbo model. For the CFG distilled model, use:
# bfl_repo = "black-forest-labs/FLUX.1-dev"
# revision = "refs/pr/3"
#
# The undistilled model that uses CFG ("pro") which can use negative prompts
# was not released.
bfl_repo = "black-forest-labs/FLUX.1-schnell"
revision = "refs/pr/1"

scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(bfl_repo, subfolder="scheduler", revision=revision)
text_encoder = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14", torch_dtype=dtype)
tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14", torch_dtype=dtype)
text_encoder_2 = T5EncoderModel.from_pretrained(bfl_repo, subfolder="text_encoder_2", torch_dtype=dtype, revision=revision)
tokenizer_2 = T5TokenizerFast.from_pretrained(bfl_repo, subfolder="tokenizer_2", torch_dtype=dtype, revision=revision)
vae = AutoencoderKL.from_pretrained(bfl_repo, subfolder="vae", torch_dtype=dtype, revision=revision)
transformer = FluxTransformer2DModel.from_pretrained(bfl_repo, subfolder="transformer", torch_dtype=dtype, revision=revision)

# Experimental: Try this to load in 4-bit for <16GB cards.
#
# from optimum.quanto import qint4
# quantize(transformer, weights=qint4, exclude=["proj_out", "x_embedder", "norm_out", "context_embedder"])
# freeze(transformer)
quantize(transformer, weights=qfloat8)
freeze(transformer)

quantize(text_encoder_2, weights=qfloat8)
freeze(text_encoder_2)

pipe = FluxPipeline(
    scheduler=scheduler,
    text_encoder=text_encoder,
    tokenizer=tokenizer,
    text_encoder_2=None,
    tokenizer_2=tokenizer_2,
    vae=vae,
    transformer=None,
)
pipe.text_encoder_2 = text_encoder_2
pipe.transformer = transformer
pipe.enable_model_cpu_offload()

generator = torch.Generator().manual_seed(12345)
image = pipe(
    prompt='nekomusume cat girl, digital painting', 
    width=1024,
    height=1024,
    num_inference_steps=4, 
    generator=generator,
    guidance_scale=3.5,
).images[0]
image.save('test_flux_distilled.png')

transformer = torch.load('D:/models/transformer.pt')
transformer.eval()

Thanks. That gets the run time down from 2m45s to 1m01s.
Using the same logic to save text_encoders_2 and load it gets the total runtime down to 25 seconds.

Can you modify the script to support the various Flux LoRAs out there at the moment?

I tried adding in
pipeline.load_lora_weights(adapter_id)
from here
https://huggingface.co/zouzoumaki/flux-loras
but that has issues with optimum.

LoRA support would be good. Especially now that with the loading of the models the whole script takes only 25 seconds.

You need to

pipe.load_lora_weights("./pytorch_lora_weights.safetensors")
pipe.fuse_lora(lora_scale=0.125)
pipe.unload_lora_weights()

But it takes a long time. I am still waiting for a better way to do this with quanto.

You need to

pipe.load_lora_weights("./pytorch_lora_weights.safetensors")
pipe.fuse_lora(lora_scale=0.125)
pipe.unload_lora_weights()

But it takes a long time. I am still waiting for a better way to do this with quanto.

Like this?

pipe = FluxPipeline(
    scheduler=scheduler,
    text_encoder=text_encoder,
    tokenizer=tokenizer,
    text_encoder_2=text_encoder_2,
    tokenizer_2=tokenizer_2,
    vae=vae,
    transformer=transformer,
)

pipe.load_lora_weights("./pytorch_lora_weights.safetensors")
pipe.fuse_lora(lora_scale=0.125)
pipe.unload_lora_weights()

because that gives these errors

Traceback (most recent call last):
  File "flux_on_potato.py", line 146, in <module>
    pipe.load_lora_weights("./pytorch_lora_weights.safetensors")
  File "venv\lib\site-packages\diffusers\loaders\lora_pipeline.py", line 1620, in load_lora_weights
    self.load_lora_into_transformer(
  File "venv\lib\site-packages\diffusers\loaders\lora_pipeline.py", line 1700, in load_lora_into_transformer
    incompatible_keys = set_peft_model_state_dict(transformer, state_dict, adapter_name)
  File "venv\lib\site-packages\peft\utils\save_and_load.py", line 395, in set_peft_model_state_dict
    load_result = model.load_state_dict(peft_model_state_dict, strict=False)
  File "venv\lib\site-packages\torch\nn\modules\module.py", line 2201, in load_state_dict
    load(self, state_dict)
  File "venv\lib\site-packages\torch\nn\modules\module.py", line 2189, in load
    load(child, child_state_dict, child_prefix)  # noqa: F821
  File "venv\lib\site-packages\torch\nn\modules\module.py", line 2189, in load
    load(child, child_state_dict, child_prefix)  # noqa: F821
  File "venv\lib\site-packages\torch\nn\modules\module.py", line 2189, in load
    load(child, child_state_dict, child_prefix)  # noqa: F821
  File "venv\lib\site-packages\torch\nn\modules\module.py", line 2183, in load
    module._load_from_state_dict(
  File "venv\lib\site-packages\optimum\quanto\nn\qmodule.py", line 159, in _load_from_state_dict
    deserialized_weight = QBytesTensor.load_from_state_dict(
  File "venv\lib\site-packages\optimum\quanto\tensor\qbytes.py", line 90, in load_from_state_dict
    inner_tensors_dict[name] = state_dict.pop(prefix + name)
KeyError: 'time_text_embed.timestep_embedder.linear_1.weight._data'

Yeah, something about that lora isn't working :( You can open an issue at https://github.com/huggingface/diffusers/issues

Yeah, something about that lora isn't working :( You can open an issue at https://github.com/huggingface/diffusers/issues

So you can use other LoRAs with that syntax? Can you share one that works?

I tried making the edit in the script for quantize, ie

#quantize(transformer, weights=qfloat8)
quantize(transformer, weights=qint4, exclude=["proj_out", "x_embedder", "norm_out", "context_embedder"])

You also need to edit the import for optimum.quanto to from optimum.quanto import freeze, qfloat8, qint4, quantize so it know what qint4 is. But the script was even slower as it seems to only use CPU now? I gave up waiting after iterations 0/4 just sat there.

I noticed the same thing, too, 0% GPU usage and only CPU usage. Were you ever able to resolve this?

Edit: I should note that this seems to happen only on Windows. On my Ubuntu machine, the GPU does get utilized.

I never got an answer so I gave up on LoRA support.

you need to add lora before quantizing

In my case, qfloat8 works well on V100. But 4090 needs qfloat8_e5m2, don't really know the reson tho :)

My generated image looks like this. what is the issue with the script?

My generated image looks like this. what is the issue with the script?

I got a similar result on my rtx4060ti (ubuntu), using qfloat8_e5m2 instead of qfloat8 solved it for me (not sure why).