Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

FLUX.1-dev FP8 Example Code Cleanup #9690

Open
hashnimo opened this issue Oct 16, 2024 · 1 comment
Open

FLUX.1-dev FP8 Example Code Cleanup #9690

hashnimo opened this issue Oct 16, 2024 · 1 comment

Comments

@hashnimo
Copy link

hashnimo commented Oct 16, 2024

I was looking at the FLUX.1-dev FP8 example code in the documentation and noticed some unnecessary imports and variable declarations, which make it appear longer and cluttered. Here is the shorter version of the code I'm currently using. Please note that I'm not an expert, just trying to make it easier for anyone else looking to run the code.

import torch
from diffusers import FluxTransformer2DModel, FluxPipeline
from transformers import T5EncoderModel
from optimum.quanto import quantize, qfloat8, freeze

bfl_repo = "black-forest-labs/FLUX.1-dev"

transformer = FluxTransformer2DModel.from_single_file("https://huggingface.co/Kijai/flux-fp8/blob/main/flux1-dev-fp8.safetensors", torch_dtype=torch.bfloat16)
quantize(transformer, weights=qfloat8)
freeze(transformer)

text_encoder_2 = T5EncoderModel.from_pretrained(bfl_repo, subfolder="text_encoder_2", torch_dtype=torch.bfloat16)
quantize(text_encoder_2, weights=qfloat8)
freeze(text_encoder_2)

pipe = FluxPipeline.from_pretrained(bfl_repo, transformer=transformer, text_encoder_2=text_encoder_2, torch_dtype=torch.bfloat16)
pipe.enable_model_cpu_offload()

prompt = "A cat holding a sign that says hello world"
image = pipe(
    prompt,
    generator=torch.Generator("cpu").manual_seed(0)
).images[0]

image.save("flux-dev-fp8.png")

Optimizations made:
Removed unnecessary code and reduced the line count from 32 to 26.

@a-r-r-o-w
Copy link
Member

Thanks for making it leaner! Would you like to open a PR? You can tag @sayakpaul for reviews if you do

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

No branches or pull requests

2 participants