Spaces:
Running
Running
Update testlcm/backend/pipelines/lcm.py
Browse files
testlcm/backend/pipelines/lcm.py
CHANGED
@@ -33,10 +33,8 @@ def _get_lcm_pipeline_from_base_model(
|
|
33 |
base_model_id,
|
34 |
unet=unet,
|
35 |
torch_dtype=torch.float32,
|
36 |
-
low_cpu_mem_usage=True,
|
37 |
safety_checker=None,
|
38 |
-
|
39 |
-
local_files_only=use_local_model,
|
40 |
)
|
41 |
pipeline.scheduler = LCMScheduler.from_config(pipeline.scheduler.config)
|
42 |
return pipeline
|
@@ -46,18 +44,17 @@ def load_taesd(
|
|
46 |
pipeline: Any,
|
47 |
use_local_model: bool = False,
|
48 |
torch_data_type: torch.dtype = torch.float32,
|
49 |
-
|
50 |
safety_checker=None,
|
51 |
-
|
52 |
):
|
53 |
vae_model = get_tiny_decoder_vae_model(pipeline.__class__.__name__)
|
54 |
pipeline.vae = AutoencoderTiny.from_pretrained(
|
55 |
vae_model,
|
56 |
torch_dtype=torch_data_type,
|
57 |
local_files_only=use_local_model,
|
58 |
-
low_cpu_mem_usage=True,
|
59 |
safety_checker=None,
|
60 |
-
|
61 |
)
|
62 |
|
63 |
|
|
|
33 |
base_model_id,
|
34 |
unet=unet,
|
35 |
torch_dtype=torch.float32,
|
|
|
36 |
safety_checker=None,
|
37 |
+
|
|
|
38 |
)
|
39 |
pipeline.scheduler = LCMScheduler.from_config(pipeline.scheduler.config)
|
40 |
return pipeline
|
|
|
44 |
pipeline: Any,
|
45 |
use_local_model: bool = False,
|
46 |
torch_data_type: torch.dtype = torch.float32,
|
47 |
+
|
48 |
safety_checker=None,
|
49 |
+
|
50 |
):
|
51 |
vae_model = get_tiny_decoder_vae_model(pipeline.__class__.__name__)
|
52 |
pipeline.vae = AutoencoderTiny.from_pretrained(
|
53 |
vae_model,
|
54 |
torch_dtype=torch_data_type,
|
55 |
local_files_only=use_local_model,
|
|
|
56 |
safety_checker=None,
|
57 |
+
|
58 |
)
|
59 |
|
60 |
|