ported to chatgpt api

app.py

@@ -5,39 +5,45 @@ from transformers import AutoModelForCausalLM, AutoProcessor
 import torch
 import gradio as gr
 import re
+import openai
+import base64
+from io import BytesIO
 
 # model_name = "arjunanand13/Florence-enphase2"
-model_name = "Stardragon2099/florence-adlp-40e"
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-print(f"Using device: {device}")
+# model_name = "Stardragon2099/florence-adlp-40e"
+# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+# print(f"Using device: {device}")
 
-device = "cuda:0" if torch.cuda.is_available() else "cpu"
-torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
-model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch_dtype, trust_remote_code=True).to(device)
-processor = AutoProcessor.from_pretrained(model_name, trust_remote_code=True)
+# device = "cuda:0" if torch.cuda.is_available() else "cpu"
+# torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+# model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch_dtype, trust_remote_code=True).to(device)
+# processor = AutoProcessor.from_pretrained(model_name, trust_remote_code=True)
 
-torch.cuda.empty_cache()
+# Retrieve the API key from Hugging Face Secrets
+openai.api_key = os.environ.get("Chatgpt_api")
+
+# torch.cuda.empty_cache()
 
 DEFAULT_PROMPT = ("You are a Leg Lift Classifier. There is an image of a throughput component "
                   "and we need to identify if the leg is inserted in the hole or not. Return 'True' "
                   "if any leg is not completely seated in the hole; return 'False' if the leg is inserted "
                   "in the hole. Return only the required JSON in this format: {Leg_lift: , Reason: }.")
 
-def predict(img, prompt):
-    inputs = processor(text=prompt, images=img, return_tensors="pt").to(device, torch_dtype)
+# def predict(img, prompt):
+#     inputs = processor(text=prompt, images=img, return_tensors="pt").to(device, torch_dtype)
 
-    generated_ids = model.generate(
-        input_ids=inputs["input_ids"],
-        pixel_values=inputs["pixel_values"],
-        max_new_tokens=1024,
-        do_sample=False,
-        num_beams=3
-    )
-    generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+#     generated_ids = model.generate(
+#         input_ids=inputs["input_ids"],
+#         pixel_values=inputs["pixel_values"],
+#         max_new_tokens=1024,
+#         do_sample=False,
+#         num_beams=3
+#     )
+#     generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
 
-    parsed_answer = processor.post_process_generation(generated_text, task= prompt, image_size=(img.width, img.height))
+#     parsed_answer = processor.post_process_generation(generated_text, task= prompt, image_size=(img.width, img.height))
 
-    return parsed_answer
+#     return parsed_answer
 
 
     # if not isinstance(image, Image.Image):
@@ -52,6 +58,34 @@ def predict(img, prompt):
     # answer = processor.batch_decode(outputs, skip_special_tokens=True)[0]
     # return answer
 
+def predict(img, prompt):
+    # Encode the Pillow image in base64
+    buffer = BytesIO()
+    img.save(buffer, format="PNG")  # Save the image in PNG format (can change to JPEG if needed)
+    buffer.seek(0)
+    encoded_image = base64.b64encode(buffer.read()).decode('utf-8')
+
+    # Combine prompt and image
+    # system_prompt = (
+    #     "You are an AI that can process text and images. The user has uploaded an image encoded in base64 "
+    #     "format along with a text prompt. You need to consider both the image and the text while responding."
+    # )
+
+    messages = [
+        {"role": "system", "content": prompt},
+        {"role": "user", "content": f"Image (base64): {encoded_image}"}
+    ]
+
+    # Make API call
+    try:
+        response = openai.ChatCompletion.create(
+            model="gpt-4",
+            messages=messages
+        )
+        return response['choices'][0]['message']['content']
+    except Exception as e:
+        return f"Error: {e}"
+
 def get_prompt(method, component_label=None):
     global prompt
     if method == "leg_lift":

requirements.txt

@@ -2,4 +2,5 @@ torch
 pillow
 transformers
 timm
-einops
+einops
+openai