Update src/pipeline.py

src/pipeline.py

@@ -20,9 +20,9 @@ def clear():
 @torch.inference_mode()
 def load_pipeline() -> Pipeline:    
     clear()
-    dtype, device = torch.bfloat16, "cuda"
+    dtype, device = torch.float8_e5m2, "cuda"
     
-    vae = AutoencoderTiny.from_pretrained("manbeast3b/quantized6", torch_dtype=torch.float8_e5m2)
+    vae = AutoencoderTiny.from_pretrained("manbeast3b/quantized6", torch_dtype=dtype)
     pipeline = DiffusionPipeline.from_pretrained(
         MODEL_ID,
         vae=vae,