Update app.py

app.py

@@ -9,6 +9,9 @@ CAPTION_MODELS = {
     'blip2-2.7b-fp16': 'Mediocreatmybest/blip2-opt-2.7b-fp16-sharded',
 }
 
+# Create a dictionary to store loaded models
+loaded_models = {}
+
 # Simple caption creation
 def caption_image(model_choice, image_input, url_input):
     if image_input is not None:
@@ -16,11 +19,18 @@ def caption_image(model_choice, image_input, url_input):
     else:
         input_data = url_input
 
-    captioner = pipeline(task="image-to-text",
-                         model=CAPTION_MODELS[model_choice],
-                         max_new_tokens=30,
-                         device_map="cpu", use_fast=True
-                         )
+    # Check if the model is already loaded
+    if model_choice in loaded_models:
+        captioner = loaded_models[model_choice]
+    else:
+        captioner = pipeline(task="image-to-text",
+                            model=CAPTION_MODELS[model_choice],
+                            max_new_tokens=30,
+                            device_map="cpu", use_fast=True
+                            )
+        # Store the loaded model
+        loaded_models[model_choice] = captioner
+
     caption = captioner(input_data)[0]['generated_text']
     return str(caption).strip()
 
@@ -32,4 +42,4 @@ image_input = gr.Image(type="pil", label="Input Image")
 url_input = gr.Text(label="Input URL")
 
 iface = gr.Interface(launch, inputs=[model_dropdown, image_input, url_input], outputs="text")
-iface.launch()
+iface.launch()