Update app.py

app.py

@@ -6,6 +6,18 @@ import random
 from diffusers import DiffusionPipeline
 import torch
 import transformers
+
+from transformers import BitsAndBytesConfig
+
+quantization_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_use_double_quant=True,
+    bnb_4bit_compute_dtype=torch.float16
+)
+
+
+
 from transformers import AutoModelForVision2Seq, AutoProcessor
 transformers.utils.move_cache()
 device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -59,7 +71,11 @@ def respond0(multimodal_input,password):
             model_id = "HuggingFaceM4/idefics2-8b"
 
             processor = AutoProcessor.from_pretrained(model_id)
-            model = AutoModelForVision2Seq.from_pretrained(model_id, torch_dtype=torch.float16)
+            model = AutoModelForVision2Seq.from_pretrained(
+                "HuggingFaceM4/idefics2-8b",
+                torch_dtype=torch.float16,    
+                quantization_config=quantization_config
+            ).to("cpu")           
             prompt = processor.apply_chat_template(messages, add_generation_prompt=True)
             inputs = processor(text=prompt, images=[images], return_tensors="pt")
             inputs = {k: v.to(model.device) for k, v in inputs.items()}
@@ -83,7 +99,11 @@ def respond0(multimodal_input,password):
             model_id = "HuggingFaceM4/idefics2-8b"
 
             processor = AutoProcessor.from_pretrained(model_id)
-            model = AutoModelForVision2Seq.from_pretrained(model_id, torch_dtype=torch.float16)
+            model = AutoModelForVision2Seq.from_pretrained(
+                "HuggingFaceM4/idefics2-8b",
+                torch_dtype=torch.float16,    
+                quantization_config=quantization_config
+            ).to("cpu")            
             prompt = processor.apply_chat_template(messages, add_generation_prompt=True)
             inputs = processor(text=prompt, images=[images], return_tensors="pt")
             inputs = {k: v.to(model.device) for k, v in inputs.items()}