Small fixes & clean up

.gitignore

@@ -9,4 +9,8 @@ __pycache__
 *.mp3
 *.mp3
 *.txt
-.copilotignore
+.copilotignore
+.misc
+BFL-flux-diffusers
+.env
+.env.*

image_encoder.py

@@ -54,7 +54,7 @@ def test_real_img():
     img_hwc = torch.from_numpy(im).cuda().type(torch.float32)
     img_chw = img_hwc.permute(2, 0, 1).contiguous()
     img_gray = img_hwc.mean(dim=2, keepdim=False).contiguous().clamp(0, 255)
-    tj = TurboImage()
+    tj = ImageEncoder()
     o = tj.encode_torch(img_chw)
     o2 = tj.encode_torch(img_hwc)
     o3 = tj.encode_torch(img_gray)
@@ -64,7 +64,6 @@ def test_real_img():
         f.write(o)
     with open("out_gray.jpg", "wb") as f:
         f.write(o3)
-    # print(o)
 
 
 if __name__ == "__main__":

main_gr.py

@@ -1,7 +1,7 @@
 import torch
 
 from flux_pipeline import FluxPipeline
-import gradio as gr
+import gradio as gr  # type: ignore
 from PIL import Image
 
 

util.py

@@ -236,18 +236,13 @@ def load_autoencoder(config: ModelSpec) -> AutoEncoder:
         missing, unexpected = ae.load_state_dict(sd, strict=False, assign=True)
         print_load_warning(missing, unexpected)
     if config.ae_quantization_dtype is not None:
-        from quantize_swap_and_dispatch import _full_quant, into_qtype
-
         ae.to(into_device(config.ae_device))
-        _full_quant(
-            ae,
-            max_quants=8000,
-            current_quants=0,
-            quantization_dtype=into_qtype(config.ae_quantization_dtype),
-        )
-        if config.offload_vae:
-            ae.to("cpu")
-            torch.cuda.empty_cache()
+        from float8_quantize import recursive_swap_linears
+
+        recursive_swap_linears(ae)
+    if config.offload_vae:
+        ae.to("cpu")
+        torch.cuda.empty_cache()
     return ae