Upload handler.py
Browse files- handler.py +1 -1
handler.py
CHANGED
|
@@ -179,7 +179,7 @@ def load_pipeline_opt(repo_id: str, dtype: torch.dtype) -> Any:
|
|
| 179 |
|
| 180 |
def load_pipeline_para(repo_id: str, dtype: torch.dtype) -> Any:
|
| 181 |
weight = int8_dynamic_activation_int4_weight() if IS_4BIT else int8_dynamic_activation_int8_weight()
|
| 182 |
-
pipe = FluxPipeline.from_pretrained(repo_id, torch_dtype=dtype)
|
| 183 |
if IS_PARA:
|
| 184 |
if IS_MGPU:
|
| 185 |
from para_attn.context_parallel import init_context_parallel_mesh
|
|
|
|
| 179 |
|
| 180 |
def load_pipeline_para(repo_id: str, dtype: torch.dtype) -> Any:
|
| 181 |
weight = int8_dynamic_activation_int4_weight() if IS_4BIT else int8_dynamic_activation_int8_weight()
|
| 182 |
+
pipe = FluxPipeline.from_pretrained(repo_id, torch_dtype=dtype).to("cuda")
|
| 183 |
if IS_PARA:
|
| 184 |
if IS_MGPU:
|
| 185 |
from para_attn.context_parallel import init_context_parallel_mesh
|