[Python] AttributeError: Can't get attribute 'RivaGAN' on <module '__main__' from...

I run the code in pytest and error at import a pt file as name "rivagan_32bit_model.pt".It's also normally running at main like

if __name__ == "__main__":

filePath = 'test_visible.mp4'
output_path = 'output.mp4'

orgCode = '0000111011111110'

addVideoImLabel(filePath=filePath, savePath="dist",orgCode=orgCode, label='Ai')


And the project tree like this


/data/home/username/ai_gcmark_alg_video (19.89MB)
├── visible_video_watermark.py (4.75KB)
├── rivagan_32bit_model.pt (1.26MB)
├── .git (7.23MB)
├── test_invisible_video_watermark.py (288b)
├── invisible_video_watermark.py (30.27KB)
├── __pycache__ (83.17KB)
├── .gitignore (47b)
├── test_visible_video_watermark.py (4b)
├── .pytest_cache (798b)
├── __init__.py (0b)
├── test_visible.mp4 (1.25MB)
├── requirements.txt (4.63KB)
├── SimSun.ttf (10.03MB)
├── README.md (2.64KB)
└── .vscode (473b)


And the test_invisible_video_watermark.py content :


from .invisible_video_watermark import *
import pytest

SAVE_DIR = "dist"
VIDEO_DIR = "test_visible.mp4"

def test_load_model():
model = RivaGAN(data_dim=32)
model = RivaGAN.load(f'rivagan_32bit_model.pt')
assert model != None




RivaGAN is a class define at invisible_video_watermark.py

At Last My Pytorch is

pytorch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2 pytorch-cuda=11.7


first, I tried


def test_load_model():
from .invisible_video_watermark import RivaGAN
model = RivaGAN(data_dim=32)
model = RivaGAN.load(f'rivagan_32bit_model.pt')
assert model != None



I run this code at termianl with

pytest


and get error like this


self = <torch.serialization._load.<locals>.UnpicklerWrapper object at 0x7f812eb97810>, mod_name = '__main__', name = 'RivaGAN'

def find_class(self, mod_name, name):
if type(name) is str and 'Storage' in name:
try:
return StorageType(name)
except KeyError:
pass
mod_name = load_module_mapping.get(mod_name, mod_name)
> return super().find_class(mod_name, name)
E AttributeError: Can't get attribute 'RivaGAN' on <module '__main__' from '/data/home/Xinsheng/miniconda3/envs/video/bin/pytest'>

../miniconda3/envs/video/lib/python3.10/site-packages/torch/serialization.py:1165: AttributeError
======================================================================================= short test summary info =======================================================================================
FAILED test_invisible_video_watermark.py::test_load_model - AttributeError: Can't get attribute 'RivaGAN' on <module '__main__' from '/data/home/Xinsheng/miniconda3/envs/video/bin/pytest'>
========================================================================================== 1 failed in 6.30s ==========================================================================================


How can I load this model at another file or other application like FastAPI or Flask ?

Continue reading...