Unverified Commit d603ccb6 authored by Dhruv Nair's avatar Dhruv Nair Committed by GitHub
Browse files

Small change to download in dance diffusion convert script (#7070)

* update

* make style
parent fd0f4695
...@@ -4,6 +4,7 @@ import math ...@@ -4,6 +4,7 @@ import math
import os import os
from copy import deepcopy from copy import deepcopy
import requests
import torch import torch
from audio_diffusion.models import DiffusionAttnUnet1D from audio_diffusion.models import DiffusionAttnUnet1D
from diffusion import sampling from diffusion import sampling
...@@ -73,9 +74,14 @@ class DiffusionUncond(nn.Module): ...@@ -73,9 +74,14 @@ class DiffusionUncond(nn.Module):
def download(model_name): def download(model_name):
url = MODELS_MAP[model_name]["url"] url = MODELS_MAP[model_name]["url"]
os.system(f"wget {url} ./") r = requests.get(url, stream=True)
return f"./{model_name}.ckpt" local_filename = f"./{model_name}.ckpt"
with open(local_filename, "wb") as fp:
for chunk in r.iter_content(chunk_size=8192):
fp.write(chunk)
return local_filename
DOWN_NUM_TO_LAYER = { DOWN_NUM_TO_LAYER = {
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment