Upload joycaption.py

joycaption.py

@@ -11,12 +11,15 @@ import os
 import gc
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
+HF_TOKEN = os.environ.get("HF_TOKEN", None)
+use_inference_client = False
 
 llm_models = {
     "Sao10K/Llama-3.1-8B-Stheno-v3.4": None,
     "unsloth/Meta-Llama-3.1-8B-bnb-4bit": None,
+    "DevQuasar/HermesNova-Llama-3.1-8B": None,
     "mergekit-community/L3.1-Boshima-b-FIX": None,
-    "meta-llama/Meta-Llama-3.1-8B": None,
+    "meta-llama/Meta-Llama-3.1-8B": None, # gated
 }
 
 CLIP_PATH = "google/siglip-so400m-patch14-384"
@@ -25,9 +28,6 @@ MODEL_PATH = list(llm_models.keys())[0]
 CHECKPOINT_PATH = Path("wpkklhc6")
 TITLE = "<h1><center>JoyCaption Pre-Alpha (2024-07-30a)</center></h1>"
 
-HF_TOKEN = os.environ.get("HF_TOKEN", None)
-use_inference_client = False
-
 class ImageAdapter(nn.Module):
     def __init__(self, input_features: int, output_features: int):
         super().__init__()
@@ -200,8 +200,6 @@ def stream_chat_mod(input_image: Image.Image, max_new_tokens: int=300, top_k: in
     #generate_ids = text_model.generate(input_ids, inputs_embeds=inputs_embeds, attention_mask=attention_mask, max_new_tokens=300, do_sample=False, suppress_tokens=None)
     generate_ids = text_model.generate(input_ids, inputs_embeds=inputs_embeds, attention_mask=attention_mask,
                                        max_new_tokens=max_new_tokens, do_sample=True, top_k=top_k, temperature=temperature, suppress_tokens=None)
-    
-    print(prompt)
 
     # Trim off the prompt
     generate_ids = generate_ids[:, input_ids.shape[1]:]