import torch
from safetensors import safe_open
from safetensors.torch import save_file
import numpy as np
input_path = "pornmasterZImage_turboV01.safetensors"
target_dtype = 'fp8'
output_path = input_path.split('.safetensors')[0] + f"_{target_dtype}.safetensors"
exclude_patterns = []
dtype_map = {'fp32': torch.float32, 'fp16': torch.float16, 'bf16': torch.bfloat16, 'fp8': torch.float8_e4m3fn}
target_torch_dtype = dtype_map[target_dtype.lower()]
tensors = {}
with safe_open(input_path, framework="pt", device="cuda") as f:
  for key in f.keys():
    tensor = f.get_tensor(key)
    original_dtype = tensor.dtype
    should_exclude = any((p.endswith('*') and key.startswith(p[:-1])) or key == p for p in exclude_patterns)
    if should_exclude:
      print(f"Skipped {key}: {original_dtype} (excluded from conversion)")
      tensors[key] = tensor
    else:
      if target_dtype.lower() == 'fp8':
        if tensor.dtype not in [torch.float16, torch.float32, torch.bfloat16]: tensor = tensor.to(torch.float32)
        tensor = torch.clamp(tensor, -448.0, 448.0)
        tensor = tensor.to(target_torch_dtype)
      else: tensor = tensor.to(target_torch_dtype)
      #print(f"Converted {key}: {original_dtype} > {tensor.dtype}")
      tensors[key] = tensor
save_file(tensors, output_path)
print(f"Saved in {output_path}")