compress_diffusion_aggressive.py

#!/usr/bin/env python3
"""
AGGRESSIVE compression for diffusion model that's already in FP16.
This will use INT8 quantization to achieve ~50% more reduction.
"""

import os
import sys
import json
import torch
import numpy as np
from safetensors import safe_open
from safetensors.torch import save_file

print("=" * 80)
print("[*] AGGRESSIVE DIFFUSION MODEL COMPRESSION")
print("   FP16 -> INT8 Quantization")
print("=" * 80)

MODEL_PATH = "diffusion_pytorch_model.safetensors"
OUTPUT_DIR = "."

# Check model exists
if not os.path.exists(MODEL_PATH):
    print(f"\n[ERROR] Model not found at {MODEL_PATH}")
    sys.exit(1)

# Get original size
original_size_mb = os.path.getsize(MODEL_PATH) / (1024 ** 2)
original_size_gb = original_size_mb / 1024

print(f"\n[FILES] Original Model:")
print(f"   Size: {original_size_gb:.2f} GB ({original_size_mb:.0f} MB)")
print(f"   Format: SafeTensors (FP16)")

print("\n" + "=" * 80)
print("STEP 1: LOADING MODEL")
print("=" * 80)

print("\n[LOAD] Loading SafeTensors weights...")
weights = {}
metadata = {}

with safe_open(MODEL_PATH, framework="pt", device="cpu") as f:
    keys = f.keys()
    print(f"   Found {len(keys)} tensors")
    
    # Check data types
    fp16_count = 0
    fp32_count = 0
    
    for key in keys:
        tensor = f.get_tensor(key)
        weights[key] = tensor
        
        if tensor.dtype == torch.float16:
            fp16_count += 1
        elif tensor.dtype == torch.float32:
            fp32_count += 1

print(f"\n[OK] Loaded {len(weights)} tensors")
print(f"   FP16 tensors: {fp16_count}")
print(f"   FP32 tensors: {fp32_count}")

print("\n" + "=" * 80)
print("STEP 2: INT8 QUANTIZATION")
print("=" * 80)

print("\n[COMPRESS] Converting FP16 -> INT8 with scale factors...")
print("   (This will reduce size by ~50% more)")

compressed_weights = {}
scale_factors = {}
quantized_count = 0

for key, tensor in weights.items():
    if tensor.dtype in [torch.float16, torch.float32]:
        # Quantize to INT8
        # Calculate scale factor (min-max quantization)
        tensor_min = tensor.min().item()
        tensor_max = tensor.max().item()
        
        # Avoid division by zero
        if abs(tensor_max - tensor_min) < 1e-8:
            scale = 1.0
            zero_point = 0
        else:
            # Scale to [-127, 127] range
            scale = (tensor_max - tensor_min) / 254.0
            zero_point = int(-tensor_min / scale - 127)
        
        # Quantize
        quantized = ((tensor.float() / scale) + zero_point).clamp(-128, 127).to(torch.int8)
        
        # Store quantized weights and scale factors
        compressed_weights[key] = quantized
        scale_factors[key + ".scale"] = torch.tensor([scale], dtype=torch.float32)
        scale_factors[key + ".zero_point"] = torch.tensor([zero_point], dtype=torch.int32)
        
        quantized_count += 1
    else:
        # Keep non-float tensors as-is
        compressed_weights[key] = tensor

print(f"\n[OK] Quantization complete:")
print(f"   Quantized {quantized_count} tensors to INT8")

print("\n" + "=" * 80)
print("STEP 3: SAVING COMPRESSED MODEL")
print("=" * 80)

# Combine weights and scale factors
all_tensors = {**compressed_weights, **scale_factors}

# Save as safetensors
output_int8_safetensors = os.path.join(OUTPUT_DIR, "compressed_diffusion_int8.safetensors")
print(f"\n[SAVE] Saving to: {output_int8_safetensors}")

try:
    save_file(all_tensors, output_int8_safetensors)
    print("[OK] Saved in SafeTensors format")
except Exception as e:
    print(f"[ERROR] Failed to save: {e}")
    sys.exit(1)

# Also save as PyTorch for easy loading
output_int8_pt = os.path.join(OUTPUT_DIR, "compressed_diffusion_int8.pt")
torch.save({
    'weights': compressed_weights,
    'scale_factors': scale_factors,
    'quantization': 'int8'
}, output_int8_pt)
print(f"[OK] Also saved as PyTorch: {output_int8_pt}")

print("\n" + "=" * 80)
print("STEP 4: COMPRESSION METRICS")
print("=" * 80)

compressed_size_mb = os.path.getsize(output_int8_safetensors) / (1024 ** 2)
compressed_size_gb = compressed_size_mb / 1024

size_reduction = ((original_size_mb - compressed_size_mb) / original_size_mb * 100)
compression_ratio = original_size_mb / compressed_size_mb if compressed_size_mb > 0 else 0
space_saved_mb = original_size_mb - compressed_size_mb
space_saved_gb = space_saved_mb / 1024

print(f"\n[METRICS] Compression Results:")
print(f"   Original Size:     {original_size_gb:.2f} GB ({original_size_mb:.0f} MB)")
print(f"   Compressed Size:   {compressed_size_gb:.2f} GB ({compressed_size_mb:.0f} MB)")
print(f"   Size Reduction:    {size_reduction:.2f}%")
print(f"   Compression Ratio: {compression_ratio:.2f}x")
print(f"   Space Saved:       {space_saved_gb:.2f} GB ({space_saved_mb:.0f} MB)")

# Save metrics
metrics = {
    "model_info": {
        "original_path": MODEL_PATH,
        "total_tensors": len(weights),
        "original_format": "FP16"
    },
    "compression_config": {
        "method": "INT8 Quantization",
        "tensors_quantized": quantized_count,
        "target_precision": "int8"
    },
    "size_metrics": {
        "original_size_mb": round(original_size_mb, 2),
        "original_size_gb": round(original_size_gb, 2),
        "compressed_size_mb": round(compressed_size_mb, 2),
        "compressed_size_gb": round(compressed_size_gb, 2),
        "size_reduction_percent": round(size_reduction, 2),
        "compression_ratio": round(compression_ratio, 2),
        "space_saved_mb": round(space_saved_mb, 2),
        "space_saved_gb": round(space_saved_gb, 2)
    },
    "output_files": {
        "int8_safetensors": output_int8_safetensors,
        "int8_pytorch": output_int8_pt
    },
    "notes": [
        "INT8 quantization applied to all FP16 tensors",
        "Scale factors stored for dequantization",
        "Model needs to be dequantized before use",
        "May have slight quality loss compared to FP16"
    ]
}

metrics_file = os.path.join(OUTPUT_DIR, "int8_compression_metrics.json")
with open(metrics_file, 'w') as f:
    json.dump(metrics, f, indent=2)

print(f"\n[OK] Metrics saved to: {metrics_file}")

# Create dequantization helper script
dequant_script = """# How to load and dequantize the INT8 model

import torch
from safetensors.torch import load_file

# Load INT8 model
checkpoint = torch.load('compressed_diffusion_int8.pt')
weights_int8 = checkpoint['weights']
scale_factors = checkpoint['scale_factors']

# Dequantize back to FP16
weights_fp16 = {}
for key, tensor in weights_int8.items():
    if tensor.dtype == torch.int8:
        scale_key = key + ".scale"
        zero_point_key = key + ".zero_point"
        
        if scale_key in scale_factors:
            scale = scale_factors[scale_key].item()
            zero_point = scale_factors[zero_point_key].item()
            
            # Dequantize: (quantized - zero_point) * scale
            dequantized = (tensor.float() - zero_point) * scale
            weights_fp16[key] = dequantized.half()  # Convert to FP16
        else:
            weights_fp16[key] = tensor
    else:
        weights_fp16[key] = tensor

print(f"Dequantized {len(weights_fp16)} tensors back to FP16")

# Now you can use weights_fp16 with your diffusion model
"""

helper_file = os.path.join(OUTPUT_DIR, "dequantize_int8_model.py")
with open(helper_file, 'w') as f:
    f.write(dequant_script)

print(f"[OK] Dequantization helper saved to: {helper_file}")

print("\n" + "=" * 80)
print("[SUCCESS] AGGRESSIVE COMPRESSION COMPLETE!")
print("=" * 80)

print(f"\n[FILES] Generated Files:")
print(f"   1. {output_int8_safetensors}")
print(f"   2. {output_int8_pt}")
print(f"   3. {metrics_file}")
print(f"   4. {helper_file}")

print(f"\n[METRICS] Summary:")
print(f"   Saved: {space_saved_gb:.2f} GB ({space_saved_mb:.0f} MB)")
print(f"   Compression: {compression_ratio:.2f}x smaller")
print(f"   Reduction: {size_reduction:.1f}%")

print(f"\n[NEXT] Next Steps:")
print(f"   1. Test the compressed model using the dequantization script")
print(f"   2. Compare image generation quality (INT8 vs FP16)")
print(f"   3. Measure inference speed")
print(f"   4. If quality is acceptable, deploy INT8 model")

print("\n" + "=" * 80)
print("[DONE] INT8 COMPRESSION SUCCESSFUL!")
print("=" * 80)