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()
@@ -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