Update app.py

app.py

@@ -7,9 +7,10 @@
 import gradio as gr
 from transformers import AutoProcessor, Idefics3ForConditionalGeneration, image_utils
 import torch
+import spaces
 
 # Determine the device (GPU or CPU) to run the model on
-device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
 print(f"Using device: {device}")  # Log the device being used
 
 # Define the model ID and base model path
@@ -28,6 +29,7 @@ model = Idefics3ForConditionalGeneration.from_pretrained(
 model.load_adapter(model_id, device_map="auto")
 
 # Define a function to infer a description from an image
+@spaces.GPU
 def infere(image):
     """
     Generate a description of a medical image.
@@ -77,7 +79,7 @@ def infere(image):
 title = f"<a href='https://huggingface.co/eltorio/IDEFICS3_ROCO'>IDEFICS3_ROCO</a>: Medical Image to Text <b>running on {device}</b>"
 desc = "This model generates a description of a medical image."
 
-device_desc = f"This model is running on {device} 🚀." if device == torch.device('cuda') else f"🐢 This model is running on {device} it will be very (very) slow. If you can donate some GPU time it will be usable 🐢. <a href='https://huggingface.co/eltorio/IDEFICS3_ROCO/discussions'>Please contact us.</a>"
+device_desc = f"This model is running on {device} 🚀." if device == 'cuda' else f"🐢 This model is running on {device} it will be very (very) slow. If you can donate some GPU time it will be usable 🐢. <a href='https://huggingface.co/eltorio/IDEFICS3_ROCO/discussions'>Please contact us.</a>"
 
 # Define the long description for the Gradio interface
 long_desc = f"This demo is based on the <a href='https://huggingface.co/eltorio/IDEFICS3_ROCO'>IDEFICS3_ROCO model</a>, which is a multimodal model that can generate text from images. It has been fine-tuned on <a href='https://huggingface.co/datasets/eltorio/ROCO-radiology'>eltorio/ROCO-radiology</a>&nbsp;a dataset of medical images and can generate descriptions of medical images. Try uploading an image of a medical image and see what the model generates!<br><b>{device_desc}</b><br> 2024 - Ronan Le Meillat"