v1

app.py

@@ -1,5 +1,5 @@
 # A100 Zero GPU
-import spaces
+# import spaces
 
 # TroL Package
 import torch
@@ -18,8 +18,8 @@ from transformers import TextIteratorStreamer
 from torchvision.transforms.functional import pil_to_tensor
 
 # flash attention
-import subprocess
-subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
+# import subprocess
+# subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
 
 # accel
 accel = Accelerator()
@@ -33,10 +33,10 @@ question="What is the troll doing? Provide the detail in the image and imagine w
 model_1_8, tokenizer_1_8 = load_trol(link='TroL-1.8B')
 
 # loading model
-model_3_8, tokenizer_3_8 = load_trol(link='TroL-3.8B')
+# model_3_8, tokenizer_3_8 = load_trol(link='TroL-3.8B')
 
 # loading model
-model_7, tokenizer_7 = load_trol(link='TroL-7B')
+# model_7, tokenizer_7 = load_trol(link='TroL-7B')
 
 def threading_function(inputs, image_token_number, streamer, device, model, tokenizer, temperature, new_max_token, top_p):
 
@@ -55,7 +55,7 @@ def threading_function(inputs, image_token_number, streamer, device, model, toke
     generation_kwargs.update({'use_cache': True})
     return model.generate(**generation_kwargs)
 
-@spaces.GPU
+# @spaces.GPU
 def bot_streaming(message, history, link, temperature, new_max_token, top_p):
 
     # model selection
@@ -135,7 +135,7 @@ def bot_streaming(message, history, link, temperature, new_max_token, top_p):
         yield buffer
 
 demo = gr.ChatInterface(fn=bot_streaming,
-                        additional_inputs = [gr.Radio(["1.8B", "3.8B"], label="Size", info="Select one model size", value="3.8B"), gr.Slider(0, 1, 0.9, label="temperature"), gr.Slider(1, 1024, 128, label="new_max_token"), gr.Slider(0, 1, 0.95, label="top_p")],
+                        additional_inputs = [gr.Radio(["1.8B"], label="Size", info="Select one model size", value="1.8B"), gr.Slider(0, 1, 0.9, label="temperature"), gr.Slider(1, 1024, 128, label="new_max_token"), gr.Slider(0, 1, 0.95, label="top_p")],
                         additional_inputs_accordion="Generation Hyperparameters",
                         theme=gr.themes.Soft(),
                         title="TroL",

trol/load_trol.py

@@ -14,14 +14,14 @@ def load_trol(link):
     if link == 'TroL-1.8B':
         from .arch_internlm2.modeling_trol import TroLForCausalLM
         from .arch_internlm2.tokenization_internlm2 import InternLM2Tokenizer as TroLTokenizer
-        bits = 16
+        bits = 4
         path = TROL_1_8B
         bit_quant_skip = ["vit", "vision_proj", "ffn", "output"]
 
     elif link == 'TroL-3.8B':
         from trol.arch_phi3.modeling_trol import TroLForCausalLM 
         from transformers import LlamaTokenizerFast as TroLTokenizer
-        bits = 16
+        bits = 8
         path = TROL_3_8B
         bit_quant_skip = ["vision_model", "vision_proj", "lm_head"]