test_quantize_custom_model

[ccssu]

doc: OneDiff Quant 🚀 NEW

[ccssu]

python faster_quantize.py

"""
CUDA_VISIBLE_DEVICES=7 python faster_quant.py --model_id /share_nfs/hf_models/stable-diffusion-xl-base-1.0 \
                                        --seed 1 \
                                        --backend torch \
                                        --output_file sdxl_pt.png \
                                        --prompt "street style, detailed, raw photo, woman, face, shot on CineStill 800T" 

CUDA_VISIBLE_DEVICES=7 python faster_quant.py --model_id /share_nfs/hf_models/stable-diffusion-xl-base-1.0 \
                                        --seed 1 \
                                        --backend onediff \
                                        --quantize \
                                        --prompt "street style, detailed, raw photo, woman, face, shot on CineStill 800T" \
                                        --cache_dir ./run_sdxl_quant 

CUDA_VISIBLE_DEVICES=7 python faster_quant.py --model_id /share_nfs/hf_models/stable-diffusion-v1-5 \
                                        --seed 1 \
                                        --backend torch \
                                        --output_file sd-v1-5_pt.png \
                                        --height 512 \
                                        --width 512

CUDA_VISIBLE_DEVICES=7 python faster_quant.py --model_id  /share_nfs/hf_models/stable-diffusion-v1-5  \
                                        --seed 1 \
                                        --backend onediff \
                                        --quantize \
                                        --cache_dir ./run_sd-v1-5_quant \
                                        --output_file sd-v1-5_of.png   \
                                        --height 512 \
                                        --width 512
"""
import os 
from diffusers import AutoPipelineForText2Image
import torch
import argparse
from onediff.infer_compiler import oneflow_compile
import time
from onediff_quant.quantization import QuantizationConfig
class Time:
    def __enter__(self):
        self.start_time = time.time()
        return self 
    
    def __exit__(self, *args, **kwargs):
        self.end_time = time.time()
        self.duration = self.end_time - self.start_time
        print(f"Duration: {self.duration:.2f} seconds")
        return self
    
def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model_id", type=str, default="/share_nfs/hf_models/stable-diffusion-xl-base-1.0 ")
    parser.add_argument("--prompt", type=str, default="a photo of an astronaut riding a horse on mars")
    parser.add_argument("--output_file", type=str, default="astronaut_rides_horse_onediff_quant.png")
    parser.add_argument("--seed", type=int, default=1)
    parser.add_argument("--quality_level", type=int, default=2, choices=list(range(8)))
    parser.add_argument("--compute_density_threshold", type=int, default=100)
    parser.add_argument("--backend", type=str, default="onediff", choices=["onediff", "torch"])
    parser.add_argument("--quantize", action="store_true")
    parser.add_argument("--cache_dir", type=str, default="./run_sdxl")
    parser.add_argument("--height", type=int, default=1024)
    parser.add_argument("--width", type=int, default=1024)
    args = parser.parse_args()
    return args
args = parse_args()
pipe = AutoPipelineForText2Image.from_pretrained(args.model_id, torch_dtype=torch.float16, use_safetensors=True, variant="fp16")
pipe.to("cuda")

def auto_pipe(pipe, cache_dir: str):
    compiled_options = {
            "graph_file": os.path.join(cache_dir, "tmp.graph"),
        }
    
    pipe.unet = oneflow_compile(pipe.unet, options=compiled_options)
    config = QuantizationConfig.from_settings(
        conv_mae_threshold=0.0060,
        linear_mae_threshold=0.007,
        conv_compute_density_threshold=900,
        linear_compute_density_threshold=300,
        cache_dir=args.cache_dir, plot_calibrate_info=False)
    pipe.unet.apply_online_quant(quant_config=config)
    return pipe

graph_dir = os.path.join(args.cache_dir, "graph")

torch.manual_seed(args.seed)
if args.backend == "onediff":
    if args.quantize:
        pipe = auto_pipe(pipe, graph_dir)

# warmup
for _ in range(1):
    images = pipe(prompt=args.prompt, height = 1024, width = 1024, num_inference_steps=30)

torch.manual_seed(args.seed)
images = pipe(prompt=args.prompt, height = args.height, width = args.width).images[0]

print(f'save to {args.output_file}')
images.save(args.output_file)

stable-diffusion-xl-base-1.0:
    pipe.unet = oneflow_compile(pipe.unet, options=compiled_options)
    config = QuantizationConfig.from_settings(
        conv_mae_threshold=0.007,
        linear_mae_threshold=0.007,
        conv_compute_density_threshold=300,
        linear_compute_density_threshold=300,
        cache_dir="./run_sdxl_quant", plot_calibrate_info=True)
    pipe.unet.apply_online_quant(quant_config=config)

stable-diffusion-v1-5
    pipe.unet = oneflow_compile(pipe.unet, options=compiled_options)
    config = QuantizationConfig.from_settings(
        conv_mae_threshold=0.0060,
        linear_mae_threshold=0.007,
        conv_compute_density_threshold=900,
        linear_compute_density_threshold=300,
        cache_dir=args.cache_dir, plot_calibrate_info=False)
    pipe.unet.apply_online_quant(quant_config=config)
    return pipe

下图量化 sd1.5 结果：{"quantized_conv_count": 33, "quantized_linear_count": 113, "conv_count": 98, "linear_count": 184, }

$ python -m oneflow --doctor
path: ['/data/home/miniconda3/envs/py310/lib/python3.10/site-packages/oneflow']
version: 0.9.1.dev20240325+cu121
git_commit: 0523481
cmake_build_type: Release
rdma: True
mlir: True
enterprise: True