Update app.py

app.py

@@ -3,11 +3,11 @@ from transformers import AutoProcessor, AutoModelForCausalLM
 import spaces
 import re
 from PIL import Image 
-
+import torch
 import subprocess
-subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
-
-model = AutoModelForCausalLM.from_pretrained('gokaygokay/Florence-2-SD3-Captioner', trust_remote_code=True).eval()
+#subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model = AutoModelForCausalLM.from_pretrained('gokaygokay/Florence-2-SD3-Captioner', trust_remote_code=True).to(device).eval()
 
 processor = AutoProcessor.from_pretrained('gokaygokay/Florence-2-SD3-Captioner', trust_remote_code=True)
 
@@ -84,14 +84,6 @@ with gr.Blocks(css=css) as demo:
             with gr.Column():
                 output_text = gr.Textbox(label="Output Text")
 
-        gr.Examples(
-            [["image1.jpg"], ["image2.jpg"], ["image3.png"], ["image4.jpg"], ["image5.jpg"], ["image6.PNG"]],
-            inputs = [input_img],
-            outputs = [output_text],
-            fn=run_example,
-            label='Try captioning on below examples'
-            )
-
         submit_btn.click(run_example, [input_img], [output_text])
 
 demo.launch(debug=True)