Spaces:
Running
on
A10G
Running
on
A10G
Zhouyan248
commited on
Commit
•
6e5c15c
1
Parent(s):
08c8308
Update base/text_to_video/__init__.py
Browse files
base/text_to_video/__init__.py
CHANGED
@@ -24,7 +24,7 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
|
|
24 |
def model_t2v_fun(args):
|
25 |
# sd_path = args.pretrained_path + "/stable-diffusion-v1-4"
|
26 |
sd_path = args.pretrained_path
|
27 |
-
|
28 |
# state_dict = find_model(args.pretrained_path + "/lavie_base.pt")
|
29 |
state_dict = find_model("./pretrained_models/lavie_base.pt")
|
30 |
unet.load_state_dict(state_dict)
|
|
|
24 |
def model_t2v_fun(args):
|
25 |
# sd_path = args.pretrained_path + "/stable-diffusion-v1-4"
|
26 |
sd_path = args.pretrained_path
|
27 |
+
unet = get_models(args, sd_path).to(device, dtype=torch.float16)
|
28 |
# state_dict = find_model(args.pretrained_path + "/lavie_base.pt")
|
29 |
state_dict = find_model("./pretrained_models/lavie_base.pt")
|
30 |
unet.load_state_dict(state_dict)
|