Spaces:
Running
on
Zero
Running
on
Zero
Update app.py
Browse files
app.py
CHANGED
@@ -195,15 +195,15 @@ def main(args):
|
|
195 |
config = OmegaConf.load("configs/config_customnet.yaml")
|
196 |
model = instantiate_from_config(config.model)
|
197 |
|
198 |
-
|
199 |
-
|
200 |
-
|
201 |
|
202 |
-
|
203 |
-
|
204 |
-
|
205 |
model = model.to(device)
|
206 |
-
sampler =
|
207 |
|
208 |
# load demo
|
209 |
demo = gr.Blocks()
|
|
|
195 |
config = OmegaConf.load("configs/config_customnet.yaml")
|
196 |
model = instantiate_from_config(config.model)
|
197 |
|
198 |
+
model_path='./customnet_v1.pt?download=true'
|
199 |
+
if not os.path.exists(model_path):
|
200 |
+
os.system(f'wget https://huggingface.co/TencentARC/CustomNet/resolve/main/customnet_v1.pt?download=true -P .')
|
201 |
|
202 |
+
ckpt = torch.load(model_path, map_location="cpu")
|
203 |
+
model.load_state_dict(ckpt)
|
204 |
+
del ckpt
|
205 |
model = model.to(device)
|
206 |
+
sampler = None
|
207 |
|
208 |
# load demo
|
209 |
demo = gr.Blocks()
|