Update app.py

app.py

@@ -1,4 +1,5 @@
 import spaces
+from install_flsh_attn import attn_implementation, dtype
 from transformers import (
     AutoImageProcessor,
     AutoModelForCausalLM,
@@ -12,16 +13,18 @@ from PIL import Image
 import zipfile
 import io
 import json
-import os
 
 DEVICE = Accelerator().device
 MODEL_NAME = "qihoo360/fg-clip2-so400m"
-BATCH_SIZE = 64
+BATCH_SIZE = 128
 
 
-model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, trust_remote_code=True).to(
-    DEVICE
-)
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_NAME,
+    trust_remote_code=True,
+    dtype=dtype,
+    attn_implementation=attn_implementation,
+).to(DEVICE)
 image_processor = AutoImageProcessor.from_pretrained(MODEL_NAME)