Add again bnb

app.py

@@ -4,7 +4,7 @@ from transformers import (
 )
 from transformers import (
     AutoProcessor,
-    # BitsAndBytesConfig,
+    BitsAndBytesConfig,
     LlavaForConditionalGeneration,
 )
 from PIL import Image
@@ -38,12 +38,12 @@ USER_ID = None
 model_id = "blanchon/PixDiet-pixtral-nutrition-v2"
 
 # BitsAndBytesConfig int-4 config
-# bnb_config = BitsAndBytesConfig(
-#     load_in_4bit=True,
-#     bnb_4bit_use_double_quant=True,
-#     bnb_4bit_quant_type="nf4",
-#     bnb_4bit_compute_dtype=torch.bfloat16,
-# )
+bnb_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_use_double_quant=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_compute_dtype=torch.bfloat16,
+)
 
 # Modify the model and processor initialization
 if TESTING:
@@ -55,7 +55,7 @@ else:
         model_id,
         device_map="auto",
         torch_dtype=torch.bfloat16,
-        # quantization_config=bnb_config,
+        quantization_config=bnb_config,
     )
     processor = AutoProcessor.from_pretrained(model_id)