from huggingface_hub import hf_hub_download
import torch, importlib.util, json
model_file = hf_hub_download(repo_id="olaverse/prism-steganography", filename="model.py")
ckpt_file = hf_hub_download(repo_id="olaverse/prism-steganography", filename="pytorch_model.pt")
config_file = hf_hub_download(repo_id="olaverse/prism-steganography", filename="config.json")
spec = importlib.util.spec_from_file_location("model", model_file)
model_module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(model_module)
config = json.load(open(config_file))
checkpoint = torch.load(ckpt_file, map_location="cpu")
encoder = model_module.StegEncoder(**config)
encoder.load_state_dict(checkpoint["encoder"])
encoder.eval()
decoder = model_module.StegDecoder(**config)
decoder.load_state_dict(checkpoint["decoder"])
decoder.eval()