Update main.py

main.py

@@ -22,6 +22,7 @@ import asyncio
 from typing import Optional
 from dotenv import load_dotenv
 import boto3
+from groq import Groq
 
 app = FastAPI()
 
@@ -33,11 +34,13 @@ app.add_middleware(
     allow_headers=["*"],  
 )
 
+groqClient = Groq (api_key=os.environ.get("GROQ_API_KEY"))
+
 load_dotenv()
 token = os.environ.get("HF_TOKEN")
 login(token)
 
-prompt_model = "meta-llama/Meta-Llama-3.1-8B-Instruct"
+prompt_model = "llama-3.1-8b-instant"
 magic_prompt_model = "Gustavosta/MagicPrompt-Stable-Diffusion"
 options = {"use_cache": False, "wait_for_model": True}
 parameters = {"return_full_text":False, "max_new_tokens":300}
@@ -88,33 +91,37 @@ async def core():
     
 
 def getPrompt(prompt, modelID, attempts=1):
-    input = prompt
     if modelID != magic_prompt_model:
-        tokenizer = AutoTokenizer.from_pretrained(modelID)
         chat = [
             {"role": "user", "content": prompt_base},
             {"role": "assistant", "content": prompt_assistant},
             {"role": "user", "content": prompt},
             ]
-        input = tokenizer.apply_chat_template(chat, tokenize=False, add_generation_prompt=True)
-    try:
-        print(modelID)
-        apiData={"inputs":input, "parameters": parameters, "options": options, "timeout": 45}
-        response = requests.post(API_URL + modelID, headers=headers, data=json.dumps(apiData))
-        if response.status_code == 200:
-            try:
-                responseData = response.json()
-                return responseData
-            except ValueError as e:
-                print(f"Error parsing JSON: {e}")
-        else:
-            print(f"Error from API: {response.status_code} - {response.text}")
+        try:
+            response = client.chat.completions.create(messages=chat, temperature=1, max_tokens=2048, top_p=1, stream=False, stop=None, model=modelID) 
+        except Exception as e:
+            print(f"An error occurred: {e}")
+            if attempts < 3:
+                getPrompt(prompt, modelID, attempts + 1) 
+    else:
+        try:
+            print(modelID)
+            apiData={"inputs":input, "parameters": parameters, "options": options, "timeout": 45}
+            response = requests.post(API_URL + modelID, headers=headers, data=json.dumps(apiData))
+            if response.status_code == 200:
+                try:
+                    responseData = response.json()
+                    return responseData
+                except ValueError as e:
+                    print(f"Error parsing JSON: {e}")
+            else:
+                print(f"Error from API: {response.status_code} - {response.text}")
+                if attempts < 3:
+                    getPrompt(prompt, modelID, attempts + 1)
+        except Exception as e:
+            print(f"An error occurred: {e}")
             if attempts < 3:
                 getPrompt(prompt, modelID, attempts + 1)
-    except Exception as e:
-        print(f"An error occurred: {e}")
-        if attempts < 3:
-            getPrompt(prompt, modelID, attempts + 1)
     return response.json()
 
 @app.post("/inferencePrompt")
@@ -229,7 +236,7 @@ def lambda_image(prompt, modelID):
     return response_data['body']
 
 def inferenceAPI(model, item, attempts = 1):
-    print(model)
+    print(f'Inference model {model}')
     if attempts > 5:
         return 'An error occured when Processing', model
     prompt = item.prompt
@@ -285,7 +292,8 @@ def get_random_model(models):
         print("Choosing randomly")
         model = random.choice(models)
     last_two_models.append(model)
-    last_two_models = last_two_models[-5:]        
+    last_two_models = last_two_models[-5:]
+    
     return model
    
 def nsfw_check(item, attempts=1):
@@ -324,6 +332,7 @@ async def inference(item: Item):
     print(activeModels['text-to-image'])
     base64_img = ""
     model = item.modelID
+    print(f'Start Model {model}')
     NSFW = False
     try:
         if item.image: