bgfs

app.py

@@ -13,7 +13,8 @@ from pdf2image import convert_from_path
 from PIL import Image
 from torch.utils.data import DataLoader
 from tqdm import tqdm
-from transformers import AutoProcessor, Idefics3ForConditionalGeneration
+from transformers import Qwen2VLForConditionalGeneration, AutoTokenizer, AutoProcessor
+from qwen_vl_utils import process_vision_info
 import re
 import time
 from PIL import Image
@@ -28,76 +29,70 @@ subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENT
 
 @spaces.GPU
 def model_inference(
-    images, text, assistant_prefix= "Réfléchis step by step. Répond uniquement avec les informations du document fourni.", decoding_strategy = "Greedy", temperature= 0.4, max_new_tokens=512,
-    repetition_penalty=1.2, top_p=0.8
+    images, text,
 ):
-    ## Load idefics
-    id_processor = AutoProcessor.from_pretrained("HuggingFaceM4/Idefics3-8B-Llama3")
-
-    id_model = Idefics3ForConditionalGeneration.from_pretrained("HuggingFaceM4/Idefics3-8B-Llama3", 
-            torch_dtype=torch.bfloat16,
-            #_attn_implementation="flash_attention_2"
-                                                            ).to("cuda")
-
-    BAD_WORDS_IDS = id_processor.tokenizer(["<image>", "<fake_token_around_image>"], add_special_tokens=False).input_ids
-    EOS_WORDS_IDS = [id_processor.tokenizer.eos_token_id]
+    
     print(type(images))
     print(images[0])
     images = Image.open(images[0][0])
     print(images)
     print(type(images))
-    if text == "" and not images:
-        gr.Error("Please input a query and optionally image(s).")
-
-    if text == "" and images:
-        gr.Error("Please input a text query along the image(s).")
+    # model = Qwen2VLForConditionalGeneration.from_pretrained(
+    # "Qwen/Qwen2-VL-7B-Instruct", torch_dtype="auto", device_map="auto"
+    # )
+
+    #We recommend enabling flash_attention_2 for better acceleration and memory saving, especially in multi-image and video scenarios.
+    model = Qwen2VLForConditionalGeneration.from_pretrained(
+        "Qwen/Qwen2-VL-7B-Instruct",
+        torch_dtype=torch.bfloat16,
+        attn_implementation="flash_attention_2",
+        device_map="auto",
+    )
 
-    if isinstance(images, Image.Image):
-        images = [images]
+    # default processer
+    processor = AutoProcessor.from_pretrained("Qwen/Qwen2-VL-7B-Instruct")
 
+    # The default range for the number of visual tokens per image in the model is 4-16384. You can set min_pixels and max_pixels according to your needs, such as a token count range of 256-1280, to balance speed and memory usage.
+    # min_pixels = 256*28*28
+    # max_pixels = 1280*28*28
+    # processor = AutoProcessor.from_pretrained("Qwen/Qwen2-VL-7B-Instruct", min_pixels=min_pixels, max_pixels=max_pixels)
 
-    resulting_messages = [
+    messages = [
+        {
+            "role": "user",
+            "content": [
                 {
-                    "role": "user",
-                    "content": [{"type": "image"}] + [
-                        {"type": "text", "text": text}
-                    ]
-                }
-            ]
-
-    if assistant_prefix:
-      text = f"{assistant_prefix} {text}"
-
-
-    prompt = id_processor.apply_chat_template(resulting_messages, add_generation_prompt=True)
-    inputs = id_processor(text=prompt, images=[images], return_tensors="pt")
-    inputs = {k: v.to("cuda") for k, v in inputs.items()}
-
-    generation_args = {
-        "max_new_tokens": max_new_tokens,
-        "repetition_penalty": repetition_penalty,
-
-    }
-
-    assert decoding_strategy in [
-        "Greedy",
-        "Top P Sampling",
+                    "type": "image",
+                    "image": images,
+                },
+                {"type": "text", "text": text},
+            ],
+        }
     ]
-    if decoding_strategy == "Greedy":
-        generation_args["do_sample"] = False
-    elif decoding_strategy == "Top P Sampling":
-        generation_args["temperature"] = temperature
-        generation_args["do_sample"] = True
-        generation_args["top_p"] = top_p
-
-
-    generation_args.update(inputs)
 
-    # Generate
-    generated_ids = id_model.generate(**generation_args)
+    # Preparation for inference
+    text = processor.apply_chat_template(
+        messages, tokenize=False, add_generation_prompt=True
+    )
+    image_inputs, video_inputs = process_vision_info(messages)
+    inputs = processor(
+        text=[text],
+        images=image_inputs,
+        videos=video_inputs,
+        padding=True,
+        return_tensors="pt",
+    )
+    inputs = inputs.to("cuda")
 
-    generated_texts = id_processor.batch_decode(generated_ids[:, generation_args["input_ids"].size(1):], skip_special_tokens=True)
-    return generated_texts[0]
+    # Inference: Generation of the output
+    generated_ids = model.generate(**inputs, max_new_tokens=128)
+    generated_ids_trimmed = [
+        out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
+    ]
+    output_text = processor.batch_decode(
+        generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
+    )
+    return output_text[0]