root
commited on
Commit
·
22af2ed
1
Parent(s):
2a94fc4
trying nanother
Browse files- handler.py +4 -3
handler.py
CHANGED
|
@@ -50,9 +50,10 @@ class EndpointHandler():
|
|
| 50 |
pretrained_base_model_path_unet = os.path.join(base_dir, 'pretrained_weights', 'stable-diffusion-v1-5', 'unet')
|
| 51 |
|
| 52 |
print("model path is " + pretrained_base_model_path_unet)
|
| 53 |
-
|
| 54 |
-
|
| 55 |
-
|
|
|
|
| 56 |
|
| 57 |
inference_config_path = os.path.join(base_dir, 'configs', 'inference', 'inference_v2.yaml')
|
| 58 |
motion_module_path = os.path.join(base_dir, 'pretrained_weights', 'motion_module.pth')
|
|
|
|
| 50 |
pretrained_base_model_path_unet = os.path.join(base_dir, 'pretrained_weights', 'stable-diffusion-v1-5', 'unet')
|
| 51 |
|
| 52 |
print("model path is " + pretrained_base_model_path_unet)
|
| 53 |
+
reference_unet = UNet2DConditionModel.from_pretrained(
|
| 54 |
+
self.config.pretrained_base_model_path,
|
| 55 |
+
subfolder="unet",
|
| 56 |
+
).to(dtype=self.weight_dtype, device="cuda")
|
| 57 |
|
| 58 |
inference_config_path = os.path.join(base_dir, 'configs', 'inference', 'inference_v2.yaml')
|
| 59 |
motion_module_path = os.path.join(base_dir, 'pretrained_weights', 'motion_module.pth')
|