caption using an auxiliary space if on spaces

app.py

@@ -4,18 +4,10 @@ from typing import Union
 from huggingface_hub import whoami
 is_spaces = True if os.environ.get("SPACE_ID") else False
 is_canonical = True if os.environ.get("SPACE_ID") == "autotrain-projects/train-flux-lora-ease" else False
-
-if is_spaces:
-    subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
-    import spaces
     
 os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
 import sys
 
-from dotenv import load_dotenv
-
-load_dotenv()
-
 # Add the current working directory to the Python path
 sys.path.insert(0, os.getcwd())
 
@@ -28,9 +20,13 @@ import shutil
 import json
 import yaml
 from slugify import slugify
-from transformers import AutoProcessor, AutoModelForCausalLM
 
+if is_spaces:
+    from gradio_client import Client, handle_file
+    client = Client("multimodalart/Florence-2-l4")
+    
 if not is_spaces:
+    from transformers import AutoProcessor, AutoModelForCausalLM
     sys.path.insert(0, "ai-toolkit")
     from toolkit.job import get_job
     gr.OAuthProfile = None
@@ -38,7 +34,6 @@ if not is_spaces:
     
 MAX_IMAGES = 150
 
-
 def load_captioning(uploaded_files, concept_sentence):
     uploaded_images = [file for file in uploaded_files if not file.endswith('.txt')]
     txt_files = [file for file in uploaded_files if file.endswith('.txt')]
@@ -71,7 +66,6 @@ def load_captioning(uploaded_files, concept_sentence):
             print(base_name)
             print(image_value)
             if base_name in txt_files_dict:
-                print("entrou")
                 with open(txt_files_dict[base_name], 'r') as file:
                     corresponding_caption = file.read()
                     
@@ -112,13 +106,13 @@ def create_dataset(*inputs):
     return destination_folder
 
 
-def run_captioning(images, concept_sentence, *captions):
+def run_captioning_local(images, concept_sentence, *captions):
     device = "cuda" if torch.cuda.is_available() else "cpu"
     torch_dtype = torch.float16
     model = AutoModelForCausalLM.from_pretrained(
-        "microsoft/Florence-2-large", torch_dtype=torch_dtype, trust_remote_code=True
+        "multimodalart/Florence-2-large-no-flash-attn", torch_dtype=torch_dtype, trust_remote_code=True
     ).to(device)
-    processor = AutoProcessor.from_pretrained("microsoft/Florence-2-large", trust_remote_code=True)
+    processor = AutoProcessor.from_pretrained("multimodalart/Florence-2-large-no-flash-attn", trust_remote_code=True)
 
     captions = list(captions)
     for i, image_path in enumerate(images):
@@ -147,8 +141,26 @@ def run_captioning(images, concept_sentence, *captions):
     del model
     del processor
 
-if is_spaces:
-    run_captioning = spaces.GPU()(run_captioning)
+def run_captioning_spaces(images, concept_sentence, *captions):
+    captions = list(captions)
+    for i, image_path in enumerate(images):
+        print(captions[i])
+        if isinstance(image_path, str):  # If image is a file path
+            image = Image.open(image_path).convert("RGB")
+        
+        answer = client.predict(
+    		image=handle_file(image_path),
+    		task_prompt="Detailed Caption",
+    		text_input=None,
+    		api_name="/process_image"
+        )[0].replace("'", '"')
+        parsed_answer = json.loads(answer)
+        caption_text = parsed_answer["<DETAILED_CAPTION>"].replace("The image shows ", "")
+        if concept_sentence:
+            caption_text = f"{caption_text} [trigger]"
+        captions[i] = caption_text
+
+        yield captions
 
 def recursive_update(d, u):
     for k, v in u.items():
@@ -548,7 +560,7 @@ with gr.Blocks(theme=theme, css=css) as demo:
         outputs=progress_area,
     )
 
-    do_captioning.click(fn=run_captioning, inputs=[images, concept_sentence] + caption_list, outputs=caption_list)
+    do_captioning.click(fn=run_captioning_spaces if is_spaces else run_captioning_local, inputs=[images, concept_sentence] + caption_list, outputs=caption_list)
     demo.load(fn=swap_visibilty, outputs=main_ui)
 
 if __name__ == "__main__":