Update src/pipeline.py

src/pipeline.py

@@ -37,6 +37,7 @@ class SDXLQuantization:
 
 Pipeline = None
 MODEL_ID = "black-forest-labs/FLUX.1-schnell"
+DTYPE = torch.bfloat16
 def clear():
     gc.collect()
     torch.cuda.empty_cache()
@@ -47,12 +48,12 @@ def load_pipeline() -> Pipeline:
     clear()
     dtype, device = torch.bfloat16, "cuda"
     vae = AutoencoderKL.from_pretrained(
-        MODEL_ID, subfolder="vae", torch_dtype=torch.bfloat16
+        MODEL_ID, subfolder="vae", torch_dtype=DTYPE
     )
     instance = SDXLQuantization(vae, 8)
     instance.quantize_model()
 
-    pipeline = FluxPipeline.from_pretrained(FLUX_CHECKPOINT,vae=vae,
+    pipeline = FluxPipeline.from_pretrained(MODEL_ID,vae=vae,
                                         torch_dtype=DTYPE)
     pipeline.text_encoder.to(memory_format=torch.channels_last)
     pipeline.text_encoder_2.to(memory_format=torch.channels_last)