fix bug with importing generator

OlgaChaganova · Jun 28, 2021 · dda2c05 · dda2c05
1 parent b031189
commit dda2c05
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/pix2art_streamlit.py b/pix2art_streamlit.py
@@ -6,13 +6,13 @@
 import requests
 import numpy as np
 
-import models
+from models.generator import Generator
 from utils import config
 
 
 @st.cache
 def load_model(style_type):
-    model = models.Generator().to(config.DEVICE)
+    model = Generator().to(config.DEVICE)
     if style_type == 'Ayvazovsky':
         PATH = config.PATH_AYVAZOVSKY
     else: