Update app.py

app.py

@@ -16,7 +16,7 @@ CAPTION_MODELS = {
 loaded_models = {}
 
 # Simple caption creation
-def caption_image(model_choice, image_input, url_inputs, load_in_8bit):
+def caption_image(model_choice, image_input, url_inputs, load_in_8bit, device):
     if image_input is not None:
         input_data = [image_input]
     else:
@@ -34,7 +34,7 @@ def caption_image(model_choice, image_input, url_inputs, load_in_8bit):
         captioner = pipeline(task="image-to-text",
                             model=CAPTION_MODELS[model_choice],
                             max_new_tokens=30,
-                            device='cpu', # Set the device as CPU
+                            device=device,  # Use selected device
                             model_kwargs=model_kwargs, 
                             torch_dtype=dtype,  # Set the floating point
                             use_fast=True