Update app.py

app.py

@@ -7,13 +7,14 @@ CAPTION_MODELS = {
     'blip-large': 'Salesforce/blip-image-captioning-large',
     'vit-gpt2-coco-en': 'ydshieh/vit-gpt2-coco-en',
     'blip2-2.7b-fp16': 'Mediocreatmybest/blip2-opt-2.7b-fp16-sharded',
+    'blip2-2.7b': 'Salesforce/blip2-opt-2.7b',
 }
 
 # Create a dictionary to store loaded models
 loaded_models = {}
 
 # Simple caption creation
-def caption_image(model_choice, image_input, url_input):
+def caption_image(model_choice, image_input, url_input, load_in_8bit):
     if image_input is not None:
         input_data = image_input
     else:
@@ -23,10 +24,11 @@ def caption_image(model_choice, image_input, url_input):
     if model_choice in loaded_models:
         captioner = loaded_models[model_choice]
     else:
+        model_kwargs = {"load_in_8bit": load_in_8bit} if load_in_8bit else {}
         captioner = pipeline(task="image-to-text",
                             model=CAPTION_MODELS[model_choice],
                             max_new_tokens=30,
-                            device_map="cpu", use_fast=True
+                            device_map="cpu", model_kwargs=model_kwargs, use_fast=True
                             )
         # Store the loaded model
         loaded_models[model_choice] = captioner
@@ -34,12 +36,13 @@ def caption_image(model_choice, image_input, url_input):
     caption = captioner(input_data)[0]['generated_text']
     return str(caption).strip()
 
-def launch(model_choice, image_input, url_input):
-    return caption_image(model_choice, image_input, url_input)
+def launch(model_choice, image_input, url_input, load_in_8bit):
+    return caption_image(model_choice, image_input, url_input, load_in_8bit)
 
 model_dropdown = gr.Dropdown(choices=list(CAPTION_MODELS.keys()), label='Select Caption Model')
 image_input = gr.Image(type="pil", label="Input Image")
 url_input = gr.Text(label="Input URL")
+load_in_8bit = gr.Checkbox(label="Load model in 8bit")
 
-iface = gr.Interface(launch, inputs=[model_dropdown, image_input, url_input], outputs="text")
-iface.launch()
+iface = gr.Interface(launch, inputs=[model_dropdown, image_input, url_input, load_in_8bit], outputs="text")
+iface.launch()