Modelpublic

Silan10/flux-torchao-int8

sourceHugging Faceotherupdated 8mo agoView on Hugging Face

0likes4downloads

Model Card

Model Overview

Silan10/flux-torchao-int8 is an 8-bit quantized version of the `black-forest-labs/FLUX.1-dev` text-to-image model. In this version, the `transformer` component has been quantized to 8-bit precision using TorchAO.

TorchAO quantization uses 8-bit integer weight-only representation with optimized kernels. This provides substantial memory savings while maintaining high image quality. TorchAO is compatible with torch.compile for additional inference speedups.

Usage

python

import torch
from diffusers import FluxPipeline, AutoModel

REPO_ID = "Silan10/flux-torchao-int8"
FLUX_MODEL_PATH = "black-forest-labs/FLUX.1-dev"  # Or local path

# Load quantized transformer
print("Loading quantized transformer...")
transformer = AutoModel.from_pretrained(
    REPO_ID,
    torch_dtype=torch.bfloat16,
    use_safetensors=False,
)
transformer.to("cuda")

# Load rest of pipeline
print("Loading pipeline...")
pipe = FluxPipeline.from_pretrained(
    FLUX_MODEL_PATH,
    transformer=None,
    torch_dtype=torch.bfloat16
)
pipe.to("cuda")

pipe.transformer = transformer
print("✓ Pipeline ready.")

prompt = "Ultra-detailed nighttime cyberpunk city street, several pedestrians in modern clothes, one person in the foreground looking toward the camera, sharp facial features and detailed hair, wet pavement reflecting colorful neon signs, shop windows with small readable text on signs, a gradient sky fading from deep blue to purple, a mix of strong highlights and deep shadows, highly detailed, 4K, cinematic lighting."
print("Generating image...")

image = pipe(
    prompt,
    num_inference_steps=20,
    guidance_scale=3.5,
    max_sequence_length=512,
    width=1024,
    height=1024,
    generator=torch.Generator("cpu").manual_seed(42)
).images[0]

image.save("output_torchao_int8.png")
print("✓ Image generated successfully.")
print("DONE!")

Sample Output

Credits

—Quantization: TorchAO by PyTorch Team