update

app.py

@@ -4,10 +4,16 @@ import re
 import time
 from PIL import Image
 import torch
-import spaces
+#import spaces
 import subprocess
 subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
 
+if torch.cuda.is_available():
+    device = "cuda"
+elif torch.backends.mps.is_available():
+    device = "mps"
+else:
+    device = "cpu"
 
 processor = AutoProcessor.from_pretrained("HuggingFaceM4/idefics2-8b")
 
@@ -15,9 +21,9 @@ model = Idefics2ForConditionalGeneration.from_pretrained(
         "HuggingFaceM4/idefics2-8b",
         torch_dtype=torch.bfloat16,
         #_attn_implementation="flash_attention_2",
-        trust_remote_code=True).to("cuda")
+        trust_remote_code=True).to(device)
 
-@spaces.GPU(duration=180)
+#@spaces.GPU(duration=180)
 def model_inference(
     image, text, decoding_strategy, temperature,
     max_new_tokens, repetition_penalty, top_p
@@ -40,7 +46,7 @@ def model_inference(
 
     prompt = processor.apply_chat_template(resulting_messages, add_generation_prompt=True)
     inputs = processor(text=prompt, images=[image], return_tensors="pt")
-    inputs = {k: v.to("cuda") for k, v in inputs.items()}
+    inputs = {k: v.to(device) for k, v in inputs.items()}
 
     generation_args = {
         "max_new_tokens": max_new_tokens,
@@ -173,4 +179,4 @@ with gr.Blocks(fill_height=True) as demo:
                                                       max_new_tokens, repetition_penalty, top_p], outputs=output)
 
 
-demo.launch(debug=True)
+demo.launch(debug=True)

requirements.txt

@@ -1,2 +1,2 @@
-spaces
+#spaces
 git+https://github.com/huggingface/transformers.git