Update app.py

app.py

@@ -6,12 +6,12 @@ import random
 from diffusers import DiffusionPipeline
 import torch
 import transformers
+from transformers import AutoModelForVision2Seq, AutoProcessor, BitsAndBytesConfig
 transformers.utils.move_cache()
 device = "cuda" if torch.cuda.is_available() else "cpu"
 import os
 password1=os.environ["password"]
 model_id = "peterpeter8585/ai2"
-client2=InferenceClient("peterpeter8585/ai2")
 
 def respond1(multimodal_input, history: list[tuple[str, str]],system_message,max_tokens,temperature,top_p,password):
     if password==password1:
@@ -26,16 +26,24 @@ def respond1(multimodal_input, history: list[tuple[str, str]],system_message,max
                 messages.append({"role": "assistant", "content": [{"type":"text", "text":val[1]}]})
         response = ""
 
-        for message in client2.chat_completion(
-            messages,
-            max_tokens=max_tokens,
-            stream=True,
-            temperature=temperature,
-            top_p=top_p,
-        ):
-            token = message
-            response += token
-            yield response
+        model_id = "HuggingFaceM4/idefics2-8b"
+
+        quantization_config = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_use_double_quant=True, bnb_4bit_compute_dtype=torch.float16)
+        processor = AutoProcessor.from_pretrained(model_id)
+        model = AutoModelForVision2Seq.from_pretrained(model_id, torch_dtype=torch.float16, quantization_config=quantization_config)
+        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()}
+        num_tokens = len(inputs["input_ids"][0])
+        with torch.inference_mode():
+            generated_ids = model.generate(**inputs, max_new_tokens=max_tokens,top_p=top_p, temperature=1.0,)
+
+        new_tokens = generated_ids[:, num_tokens:]
+        generated_text = processor.batch_decode(new_tokens, skip_special_tokens=True)[0]
+
+        token = generated_text
+        response+=token
+        yield response