Modifying the load for the models

main.py

@@ -8,8 +8,10 @@ import re
 import io
 import base64
 import matplotlib.pyplot as plt 
-from transformers import pipeline
-
+import torch
+from transformers import pipeline,VisionEncoderDecoderModel,ViTImageProcessor,AutoTokenizer
+from transformers import BartForConditionalGeneration, BartTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer
 import fitz
 from docx import Document
 from pptx import Presentation
@@ -29,23 +31,45 @@ app.add_middleware(
     allow_headers=["*"],
 )
 try:
-    interpreter = pipeline("image-to-text", model="nlpconnect/vit-gpt2-image-captioning")
+    #interpreter = pipeline("image-to-text", model="nlpconnect/vit-gpt2-image-captioning")
+    interpreter_model = VisionEncoderDecoderModel.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
+    interpreter_processor = ViTImageProcessor.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
+    interpreter_tokenizer = AutoTokenizer.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
 except Exception as exp:
-    print("[ERROR] Can't load nlpconnect/vit-gpt2-image-captioning ")
+    print("[ERROR] Can't load nlpconnect/vit-gpt2-image-captioning")
     print(str(exp))
+    
+#try:
+#    summarizer = pipeline("summarization", model="facebook/bart-large-cnn",device=0)
+#except Exception as exp:
+#    print("[ERROR] Can't load facebook/bart-large-cnn ")
+#    print(str(exp))
+
 try:
-    summarizer = pipeline("summarization", model="facebook/bart-large-cnn",device=0)
-except Exception as exp:
-    print("[ERROR] Can't load facebook/bart-large-cnn ")
-    print(str(exp))
+    summarizer_model = BartForConditionalGeneration.from_pretrained("facebook/bart-large-cnn")
+except OSError as e:
+    print(f"[INFO] PyTorch weights not found. Falling back to TensorFlow weights.\n{e}")
+    summarizer_model = BartForConditionalGeneration.from_pretrained("facebook/bart-large-cnn", from_tf=True)
+
+summarizer_tokenizer = BartTokenizer.from_pretrained("facebook/bart-large-cnn")
+
+
+
+#try:
+#    generator = pipeline("text-generation", model="deepseek-ai/deepseek-coder-1.3b-instruct", device_map="auto")
+#except Exception as exp:
+#    print("[ERROR] Can't load deepseek-ai/deepseek-coder-1.3b-instruct ")
+#    print(str(exp))
+
+
 try:
-    generator = pipeline("text-generation", model="deepseek-ai/deepseek-coder-1.3b-instruct", device_map="auto")
-except Exception as exp:
+    generator_model = AutoModelForCausalLM.from_pretrained("deepseek-ai/deepseek-coder-1.3b-instruct", trust_remote_code=True)
+    tokengenerator_modelizer = AutoTokenizer.from_pretrained("deepseek-ai/deepseek-coder-1.3b-instruct", trust_remote_code=True)
+except Exception as exp :
     print("[ERROR] Can't load deepseek-ai/deepseek-coder-1.3b-instruct ")
     print(str(exp))
 
 
-
 app.mount("/static",StaticFiles(directory='static'),'static')
 templates = Jinja2Templates(directory='templates')
 
@@ -71,9 +95,12 @@ def caption(file:UploadFile=File(...)):
     if extension not in Supported_extensions:
         return {"error": "Unsupported file type"}
     image = Image.open(file.file)
-    
-    caption = interpreter(image)
-    return {"caption": caption[0]['generated_text']}   
+    #caption = interpreter(image)
+    pixel_values = interpreter_processor(images=image, return_tensors="pt").pixel_values
+    output_ids = interpreter_model.generate(pixel_values, max_length=16, num_beams=4)
+    caption = interpreter_tokenizer.decode(output_ids[0], skip_special_tokens=True)
+    return {"caption":caption}
+    #return {"caption": caption[0]['generated_text']}   
 
 @app.post("/summerize")
 def summerzation(file:UploadFile=File(...)):
@@ -93,8 +120,8 @@ def summerzation(file:UploadFile=File(...)):
       return {"error": "File is empty"}
     
     result=""
-    for i in range(0,len(text),1024):
-        result+=summarizer(text, max_length=150, min_length=30, do_sample=False)[0]['summary_text']
+    #for i in range(0,len(text),1024):
+    #    result+=summarizer(text, max_length=150, min_length=30, do_sample=False)[0]['summary_text']
     return {"summary": result}
 
 
@@ -133,7 +160,7 @@ error.
 ##Prompt: {prompt}.
         """
         
-        output = generator(message, max_length=1000)
+        output = [{"generated_text":""}]#generator(message, max_length=1000)
         match = re.search(r'```python(.*?)```', output[0]["generated_text"], re.DOTALL)
         code =''
         if not match: