Upload 2 files

prompt.py

@@ -14,6 +14,8 @@ import random
 
 subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
 
+huggingface_token = os.getenv("HF_TOKEN")
+
 # Initialize Florence model
 device = "cuda" if torch.cuda.is_available() else "cpu"
 florence_model = AutoModelForCausalLM.from_pretrained('microsoft/Florence-2-base', trust_remote_code=True).to("cpu").eval()
@@ -318,11 +320,13 @@ class PromptGenerator:
         return prompt
 
 import os
+from openai import OpenAI
 
 class HuggingFaceInferenceNode:
 
     def __init__(self):
-        self.client = InferenceClient("meta-llama/Meta-Llama-3.1-70B-Instruct") #
+        #self.client = InferenceClient("meta-llama/Meta-Llama-3.1-70B-Instruct") #
+        self.client = OpenAI(base_url="https://api-inference.huggingface.co/v1/", api_key=huggingface_token)
         self.prompts_dir = "./prompts"
         os.makedirs(self.prompts_dir, exist_ok=True)
 
@@ -379,7 +383,8 @@ You are allowed to make up film and branding names, and do them like 80's, 90's
                 {"role": "user", "content": user_message}
             ]
 
-            response = self.client.chat_completion( #
+            #response = self.client.chat_completion( #
+            response = self.client.chat.completions.create(
                 model="meta-llama/Meta-Llama-3.1-70B-Instruct",
                 max_tokens=1024,
                 temperature=0.7,

requirements.txt

@@ -13,4 +13,5 @@ controlnet-aux
 kornia
 numpy
 opencv-python
-deepspeed
+deepspeed
+openai==1.37.0