Update src/pipeline.py
Browse files- src/pipeline.py +1 -1
src/pipeline.py
CHANGED
@@ -32,7 +32,7 @@ def load_pipeline() -> Pipeline:
|
|
32 |
quantize_(vae, int8_weight_only())
|
33 |
text_encoder_2 = T5EncoderModel.from_pretrained("city96/t5-v1_1-xxl-encoder-bf16", torch_dtype=torch.bfloat16)
|
34 |
quantize_(text_encoder_2, int8_weight_only())
|
35 |
-
model = FluxTransformer2DModel.from_pretrained(
|
36 |
pipeline = DiffusionPipeline.from_pretrained(ckpt_id, text_encoder=text_encoder, transformer=model, text_encoder_2=text_encoder_2, torch_dtype=torch.bfloat16)
|
37 |
pipeline.to("cuda")
|
38 |
|
|
|
32 |
quantize_(vae, int8_weight_only())
|
33 |
text_encoder_2 = T5EncoderModel.from_pretrained("city96/t5-v1_1-xxl-encoder-bf16", torch_dtype=torch.bfloat16)
|
34 |
quantize_(text_encoder_2, int8_weight_only())
|
35 |
+
model = FluxTransformer2DModel.from_pretrained(ckpt_id, subfolder="transformer", torch_dtype=torch.bfloat16, use_safetensors=False)
|
36 |
pipeline = DiffusionPipeline.from_pretrained(ckpt_id, text_encoder=text_encoder, transformer=model, text_encoder_2=text_encoder_2, torch_dtype=torch.bfloat16)
|
37 |
pipeline.to("cuda")
|
38 |
|