Update app.py

app.py

@@ -8,15 +8,39 @@ import spaces
 import subprocess
 subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
 
+DEVICE = "cuda:0"
+USE_LORA = False
+USE_QLORA = True
 
-processor = AutoProcessor.from_pretrained("HuggingFaceM4/idefics2-8b")
+processor = AutoProcessor.from_pretrained("HuggingFaceM4/idefics2-8b", do_image_splitting=False)
 
-model = Idefics2ForConditionalGeneration.from_pretrained(
+if USE_QLORA or USE_LORA:
+    lora_config = LoraConfig(
+        r=8,
+        lora_alpha=8,
+        lora_dropout=0.1,
+        target_modules='.*(text_model|modality_projection|perceiver_resampler).*(down_proj|gate_proj|up_proj|k_proj|q_proj|v_proj|o_proj).*',
+        use_dora=False if USE_QLORA else True,
+        init_lora_weights="gaussian"
+    )
+    bnb_config = BitsAndBytesConfig(
+        load_in_4bit=True,
+        bnb_4bit_quant_type="nf4",
+        bnb_4bit_compute_dtype=torch.float16
+    ) if USE_QLORA else None
+    model = Idefics2ForConditionalGeneration.from_pretrained(
         "HuggingFaceM4/idefics2-8b",
-        torch_dtype=torch.bfloat16,
-        # _attn_implementation="flash_attention_2",
-        # trust_remote_code=True
-).to("cuda")
+        torch_dtype=torch.float16,
+        quantization_config=bnb_config,
+    )
+    model.add_adapter(lora_config)
+    model.enable_adapters()
+else:
+    model = Idefics2ForConditionalGeneration.from_pretrained(
+        "HuggingFaceM4/idefics2-8b",
+        torch_dtype=torch.float16,
+        _attn_implementation="flash_attention_2"
+    ).to("cuda")
 
 import gradio as gr
 from huggingface_hub import InferenceApi