Update app.py

app.py

@@ -1,4 +1,4 @@
-from transformers import MllamaForConditionalGeneration, AutoProcessor, BitsAndBytesConfig
+from transformers import MllamaForConditionalGeneration, AutoProcessor
 from PIL import Image
 import torch
 import gradio as gr
@@ -6,18 +6,10 @@ import spaces
 
 # Initialize model and processor
 ckpt = "unsloth/Llama-3.2-11B-Vision-Instruct"
-bnb_config = BitsAndBytesConfig(
-    load_in_4bit=True,                   # Quantification 4-bit pour économiser la VRAM
-    bnb_4bit_use_double_quant=True,
-    bnb_4bit_quant_type="nf4",
-    bnb_4bit_compute_dtype=torch.float16
-)
-
 model = MllamaForConditionalGeneration.from_pretrained(
     ckpt,
-    quantization_config=bnb_config,
-    device_map="auto"
-)
+    torch_dtype=torch.bfloat16
+).to("cuda")
 processor = AutoProcessor.from_pretrained(ckpt)
 
 @spaces.GPU
@@ -37,9 +29,6 @@ def extract_text(image):
     "- Do NOT separate handwritten and printed text.\n"
     "- Do NOT confuse '.' (a period) with '|' (a border).\n\n"
     "Only extract the text that is actually visible in the image, and nothing else.")
-                                                                                      
-
-    
     # Create message structure
     messages = [
         {
@@ -67,9 +56,9 @@ def extract_text(image):
         result = result[result.lower().find("assistant") + len("assistant"):].strip()
     
     # Remove any remaining conversation markers
-    cleaned_result = result.replace("user", "").replace(prompt, "").strip()
+    result = result.replace("user", "").replace(prompt, "").strip()
 
-    print(cleaned_result)
+    print(result)
     
     return result