Update app.py

app.py

@@ -1,18 +1,17 @@
 import gradio as gr
-import spaces
 import torch
 import os
 import numpy as np
 from groq import Groq
-from transformers import AutoModel, AutoTokenizer
+from transformers import AutoModel, AutoTokenizer, BitsAndBytesConfig
 from diffusers import StableDiffusionXLPipeline, UNet2DConditionModel, EulerDiscreteScheduler
 from parler_tts import ParlerTTSForConditionalGeneration
 import soundfile as sf
-from langchain.embeddings.openai import OpenAIEmbeddings
-from langchain.vectorstores import Chroma
+from langchain.embeddings import OpenAIEmbeddings
+from langchain.vectorstores import Chroma 
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain.chains import RetrievalQA
-from langchain.llms import OpenAI
+from langchain.llms import OpenAI 
 from PIL import Image
 from decord import VideoReader, cpu
 import requests
@@ -20,9 +19,17 @@ import requests
 client = Groq(api_key=os.environ.get("GROQ_API_KEY"))
 MODEL = 'llama3-groq-70b-8192-tool-use-preview'
 
+# Configure transformers to load the model with 4-bit quantization
+bnb_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_use_double_quant=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_compute_dtype=torch.bfloat16
+)
+
 # Load models for text, speech, and image processing
 text_model = AutoModel.from_pretrained('openbmb/MiniCPM-V-2_6', trust_remote_code=True,
-    attn_implementation='sdpa', torch_dtype=torch.bfloat16).eval().cuda()
+    quantization_config=bnb_config, device_map="auto") 
 tokenizer = AutoTokenizer.from_pretrained('openbmb/MiniCPM-V-2_6', trust_remote_code=True)
 
 tts_model = ParlerTTSForConditionalGeneration.from_pretrained("parler-tts/parler-tts-large-v1").to('cuda')