sayakpaul/inference_with_torchao_serialized.py

from diffusers import DiffusionPipeline, FluxTransformer2DModel, AutoencoderKL
from transformers import T5EncoderModel, CLIPTextModel
from torchao.quantization import quantize_, int8_weight_only
import torch
import os

ckpt_id = "sayakpaul/FLUX.1-merged"

# Function to load and quantize a component
def load_and_quantize(model_class, subfolder):
    model = model_class.from_pretrained(
        ckpt_id, subfolder=subfolder, torch_dtype=torch.float16
    )
    quantize_(model, int8_weight_only())
    return model

# Load and quantize components
transformer = load_and_quantize(FluxTransformer2DModel, "transformer")
text_encoder = load_and_quantize(CLIPTextModel, "text_encoder")
text_encoder_2 = load_and_quantize(T5EncoderModel, "text_encoder_2")
vae = load_and_quantize(AutoencoderKL, "vae")

# Initialize the pipeline
pipeline = DiffusionPipeline.from_pretrained(
    ckpt_id,
    transformer=transformer,
    vae=vae,
    text_encoder=text_encoder,
    text_encoder_2=text_encoder_2,
    torch_dtype=torch.float16
)

# Enable CPU offloading to save GPU memory
# pipeline.enable_sequential_cpu_offload()

# Save the quantized pipeline
save_directory = "quantized_flux_pipeline"
os.makedirs(save_directory, exist_ok=True)
pipeline.save_pretrained(save_directory)
print(f"Quantized pipeline saved to {save_directory}")

from diffusers import DiffusionPipeline, FluxTransformer2DModel, AutoencoderKL
from transformers import T5EncoderModel, CLIPTextModel
from torchao.quantization import quantize_, int8_weight_only
import torch
import os

ckpt_id = "sayakpaul/FLUX.1-merged"

def load_and_quantize(model_class, subfolder):
    model = model_class.from_pretrained(
        ckpt_id, subfolder=subfolder, torch_dtype=torch.float16
    )
    quantize_(model, int8_weight_only())
    return model

# Load and quantize components
transformer = load_and_quantize(FluxTransformer2DModel, "transformer")
text_encoder = load_and_quantize(CLIPTextModel, "text_encoder")
text_encoder_2 = load_and_quantize(T5EncoderModel, "text_encoder_2")
vae = load_and_quantize(AutoencoderKL, "vae")

# Initialize the pipeline
pipeline = DiffusionPipeline.from_pretrained(
    ckpt_id,
    transformer=transformer,
    vae=vae,
    text_encoder=text_encoder,
    text_encoder_2=text_encoder_2,
    torch_dtype=torch.float16
)

# Save the quantized pipeline components separately
save_directory = "quantized_flux_pipeline"
os.makedirs(save_directory, exist_ok=True)

# Move all components to CPU before saving
pipeline.to("cpu")

# Save each component separately
torch.save(pipeline.transformer.state_dict(), os.path.join(save_directory, "transformer.pt"))
torch.save(pipeline.vae.state_dict(), os.path.join(save_directory, "vae.pt"))
torch.save(pipeline.text_encoder.state_dict(), os.path.join(save_directory, "text_encoder.pt"))
torch.save(pipeline.text_encoder_2.state_dict(), os.path.join(save_directory, "text_encoder_2.pt"))

#forget those 
# Save the config files
# pipeline.transformer.save_config(save_directory)
# pipeline.vae.save_config(save_directory)
#pipeline.text_encoder.save_config(save_directory)
# pipeline.text_encoder_2.save_config(save_directory)

print(f"Quantized pipeline components saved to {save_directory}")

# Optional: Clear CUDA cache and print memory usage
torch.cuda.empty_cache()
memory = (torch.cuda.memory_allocated() / 1024 / 1024 / 1024)
print(f"GPU memory usage: {memory:.3f} GB")

import torch
from diffusers import FluxTransformer2DModel, DiffusionPipeline, AutoencoderKL
from transformers import T5EncoderModel, CLIPTextModel, CLIPTextConfig ,T5Config 
import os

dtype, device = torch.bfloat16, "cuda"
ckpt_id = "sayakpaul/FLUX.1-merged"
save_directory = "quantized_flux_pipeline"

# Load transformer
with torch.device("meta"):
    config = FluxTransformer2DModel.load_config(ckpt_id, subfolder="transformer")
    transformer = FluxTransformer2DModel.from_config(config).to(dtype)

transformer_state_dict = torch.load(os.path.join(save_directory, "transformer.pt"), map_location="cpu")
transformer.load_state_dict(transformer_state_dict, assign=True)

# Load VAE
with torch.device("meta"):
    vae_config = AutoencoderKL.load_config(ckpt_id, subfolder="vae")
    vae = AutoencoderKL.from_config(vae_config).to(dtype)

vae_state_dict = torch.load(os.path.join(save_directory, "vae.pt"), map_location="cpu")
vae.load_state_dict(vae_state_dict, assign=True)

text_encoder_config = CLIPTextConfig(ckpt_id, subfolder="text_encoder")
text_encoder = CLIPTextModel(text_encoder_config)

text_encoder_state_dict = torch.load(os.path.join(save_directory, "text_encoder.pt"), map_location="cpu")
text_encoder.load_state_dict(text_encoder_state_dict, assign=True)

# Load text encoder 2
# with torch.device("meta"):
text_encoder_2_config = T5Config(ckpt_id, subfolder="text_encoder_2")
text_encoder_2 = T5EncoderModel(text_encoder_2_config)

text_encoder_2_state_dict = torch.load(os.path.join(save_directory, "text_encoder_2.pt"), map_location="cpu",)
text_encoder_2.load_state_dict(text_encoder_2_state_dict, assign=True)

# Initialize the pipeline
pipeline = DiffusionPipeline.from_pretrained(
    ckpt_id,
    transformer=transformer,
    vae=vae,
    text_encoder=text_encoder,
    text_encoder_2=text_encoder_2,
    torch_dtype=dtype
).to(device)

# Generate image
image = pipeline(
    "cat",
    guidance_scale=0.0,
    num_inference_steps=4,
    max_sequence_length=256
).images[0]

image.save("flux_local_int8.png")
print("Image generated and saved as 'flux_local_int8.png'")