Upload 2 files

README.md

@@ -5,7 +5,7 @@ language:
 ---
 # Image Captioning App
 
-This is a mod of [Wi-zz/joy-caption-pre-alpha](https://huggingface.co/Wi-zz/joy-caption-pre-alpha) and [fancyfeast/joy-caption-alpha-two](https://huggingface.co/spaces/fancyfeast/joy-caption-alpha-two). Thanks to [dominic1021](https://huggingface.co/dominic1021).
+This is a mod of [Wi-zz/joy-caption-pre-alpha](https://huggingface.co/Wi-zz/joy-caption-pre-alpha) and [fancyfeast/joy-caption-alpha-two](https://huggingface.co/spaces/fancyfeast/joy-caption-alpha-two). Thanks to [dominic1021](https://huggingface.co/dominic1021), [IceHibiki](https://huggingface.co/IceHibiki).
 
 # Notice: I will contribute to Wi-zz after shaping the code.
 

app.py

@@ -12,7 +12,7 @@ from torch import nn
 from transformers import AutoModel, AutoProcessor, AutoTokenizer, PreTrainedTokenizer, PreTrainedTokenizerFast, AutoModelForCausalLM
 from typing import List, Union
 import torchvision.transforms.functional as TVF
-from peft import PeftConfig
+from peft import PeftModel
 import gc
 
 # Constants
@@ -127,6 +127,7 @@ class ImageAdapter(nn.Module):
 
 # Global Variables
 IS_NF4 = True
+IS_LORA = True
 MODEL_PATH = DEFAULT_MODEL_PATH
 device = "cuda" if torch.cuda.is_available() else "cpu"
 print(f"Running on {device}")
@@ -196,7 +197,7 @@ class ImageAdapter(nn.Module):
         return self.other_tokens(torch.tensor([2], device=self.other_tokens.weight.device)).squeeze(0)
 
 def load_models():
-    global MODEL_PATH, IS_NF4
+    global MODEL_PATH, IS_NF4, IS_LORA
     try:
         if IS_NF4:
             from transformers import BitsAndBytesConfig
@@ -222,11 +223,11 @@ def load_models():
             print(f"Loading LLM: {MODEL_PATH} 🤖")
             text_model = AutoModelForCausalLM.from_pretrained(MODEL_PATH, quantization_config=nf4_config, device_map=device, torch_dtype=torch.bfloat16).eval()
 
-            if LORA_PATH.exists():
+            if False and IS_LORA and LORA_PATH.exists(): # omitted
                 print("Loading VLM's custom text model 🤖")
-                peft_config = PeftConfig.from_pretrained(LORA_PATH, device_map=device, quantization_config=nf4_config)
-                text_model.add_adapter(peft_config)
-                text_model.enable_adapters()
+                text_model = PeftModel.from_pretrained(model=text_model, model_id=LORA_PATH, device_map=device, quantization_config=nf4_config)
+                text_model = text_model.merge_and_unload(safe_merge=True) # to avoid PEFT bug https://github.com/huggingface/transformers/issues/28515
+            else: print("VLM's custom text model isn't loaded 🤖")
 
             print("Loading image adapter 🖼️")
             image_adapter = ImageAdapter(clip_model.config.hidden_size, text_model.config.hidden_size, False, False, 38, False).eval().to("cpu")
@@ -250,13 +251,13 @@ def load_models():
             assert isinstance(tokenizer, (PreTrainedTokenizer, PreTrainedTokenizerFast)), f"Tokenizer is of type {type(tokenizer)}"
 
             print(f"Loading LLM: {MODEL_PATH} 🤖")
-            text_model = AutoModelForCausalLM.from_pretrained(MODEL_PATH, device_map="auto", torch_dtype=torch.bfloat16).eval()
+            text_model = AutoModelForCausalLM.from_pretrained(MODEL_PATH, device_map="auto", torch_dtype=torch.bfloat16).eval() # device_map="auto" may cause LoRA issue
 
-            if LORA_PATH.exists():
+            if IS_LORA and LORA_PATH.exists():
                 print("Loading VLM's custom text model 🤖")
-                peft_config = PeftConfig.from_pretrained(LORA_PATH, device_map=device)
-                text_model.add_adapter(peft_config)
-                text_model.enable_adapters()
+                text_model = PeftModel.from_pretrained(model=text_model, model_id=LORA_PATH, device_map=device)
+                text_model = text_model.merge_and_unload(safe_merge=True) # to avoid PEFT bug https://github.com/huggingface/transformers/issues/28515
+            else: print("VLM's custom text model isn't loaded 🤖")
 
             print("Loading image adapter 🖼️")
             image_adapter = ImageAdapter(clip_model.config.hidden_size, text_model.config.hidden_size, False, False, 38, False).eval().to("cpu")
@@ -461,7 +462,8 @@ def parse_arguments():
     parser.add_argument("--prompt", type=str, default="", help='Custom Prompt (optional, will override all other settings)')
     parser.add_argument("--model", type=str, default=DEFAULT_MODEL_PATH,
                         help='Huggingface LLM repo (default: "unsloth/Meta-Llama-3.1-8B-bnb-4bit")')
-    parser.add_argument("--bf16", action="store_true", help="Use bfloat16 (default: NF4)")
+    parser.add_argument("--bf16", action="store_true", default=False, help="Use bfloat16 (default: NF4)")
+    parser.add_argument("--nolora", action="store_true", default=False, help="Disable VLM's custom text model (default: Enable)")
     parser.add_argument("--tokens", type=int, default=300, help="Max tokens (default: 300)")
     parser.add_argument("--topp", type=float, default=0.9, help="Top-P (default: 0.9)")
     parser.add_argument("--temp", type=float, default=0.6, help="Temperature (default: 0.6)")
@@ -480,7 +482,7 @@ def is_valid_repo(repo_id):
         return False
 
 def main():
-    global MODEL_PATH, IS_NF4
+    global MODEL_PATH, IS_NF4, IS_LORA
     args = parse_arguments()
     input_paths = [Path(input_path) for input_path in args.input]
     batch_size = args.bs
@@ -492,8 +494,8 @@ def main():
     max_new_tokens = args.tokens
     top_p = args.topp
     temperature = args.temp
-    if args.bf16: IS_NF4 = False
-    else: IS_NF4 = True
+    IS_NF4 = False if args.bf16 else True
+    IS_LORA = False if args.nolora else True
     if is_valid_repo(args.model): MODEL_PATH = args.model
     else: sys.exit(1)
     models = load_models()