Update app.py

app.py

@@ -28,12 +28,13 @@ def caption_image(model_choice, image_input, url_input, load_in_8bit):
         captioner = loaded_models[model_key]
     else:
         model_kwargs = {"load_in_8bit": load_in_8bit} if load_in_8bit else {}
+        dtype = torch.float16 if load_in_8bit else torch.float32  # Set dtype based on the value of load_in_8bit
         captioner = pipeline(task="image-to-text",
                             model=CAPTION_MODELS[model_choice],
                             max_new_tokens=30,
-                            device=0, # Set the device as CPU
+                            device='cpu', # Set the device as CPU
                             model_kwargs=model_kwargs, 
-                            torch_dtype=torch.float16,  # Set the floating point to Float16
+                            torch_dtype=dtype,  # Set the floating point
                             use_fast=True
                             )
         # Store the loaded model