Update app.py

app.py

@@ -1,60 +1,81 @@
 """
-This module provides an interface for image captioning using the BLIP model.
+This module provides an interface for image captioning using the BLIP-2 model.
 The interface allows users to upload an image and receive a caption.
 """
 
 import gradio as gr
 import spaces
-from transformers import BlipProcessor, BlipForConditionalGeneration
+from transformers import BlipProcessor, BlipForConditionalGeneration, BitsAndBytesConfig
 from PIL import Image
 
+# Define device
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# Define quantization configuration
+quantization_config = BitsAndBytesConfig(load_in_8bit=True)
+
 # Initialize the processor and model
-processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
-model = (
-    BlipForConditionalGeneration
-    .from_pretrained("Salesforce/blip-image-captioning-base")
-    .to("cuda")
-)
+try:
+    processor = Blip2Processor.from_pretrained("Salesforce/blip2-opt-6.7b-coco")
+    model = Blip2ForConditionalGeneration.from_pretrained(
+        "Salesforce/blip2-opt-6.7b-coco",
+        quantization_config=quantization_config,  # Quantize model to 8-bit
+        device_map="auto",  # Efficient GPU utilization
+        torch_dtype=torch.float16  # Load weights in float16 to save memory
+    ).to(device)
+except Exception as error:
+    print(f"Error initializing model: {error}")
 
-def generate_caption(image: Image) -> str:
+def generate_caption(image: Image.Image) -> str:
     """
-    Generates a caption for a given image using the BLIP model.
+    Generates a caption for the given image using the BLIP-2 model.
 
     Args:
-        image (Image): The input image as a PIL Image object.
+        image (PIL.Image): The input image to generate a caption for.
 
     Returns:
-        str: The generated caption.
+        str: The generated caption as a string.
     """
-    inputs = processor(images=image, return_tensors="pt").to("cuda")
-    outputs = model.generate(**inputs)
-    caption = processor.decode(outputs[0], skip_special_tokens=True)
-    return caption
+    if not isinstance(image, Image.Image):
+        raise ValueError("Input must be a PIL Image.")
+    
+    try:
+        inputs = processor(images=image, return_tensors="pt").to(device)
+        outputs = model.generate(**inputs)
+        caption = processor.decode(outputs[0], skip_special_tokens=True)
+        return caption
+    except Exception as error:
+        return f"Error generating caption: {str(error)}"
 
 @spaces.GPU
-def caption_image(image: Image) -> str:
+def caption_image(image: Image.Image) -> str:
     """
     Takes a PIL Image input and returns a caption.
 
     Args:
-        image (Image): The input image as a PIL Image object.
+        image (PIL.Image): The input image to generate a caption for.
 
     Returns:
-        str: The generated caption or an error message.
+        str: The generated caption, or an error message if something goes wrong.
     """
     try:
-        return generate_caption(image)
-    except Exception as e:
-        return f"An error occurred: {str(e)}"
+        caption = generate_caption(image)
+        return caption
+    except Exception as error:
+        return f"An error occurred: {str(error)}"
+
+# Constants for Gradio interface configuration
+IMAGE_TYPE = "pil"
+OUTPUT_TYPE = "text"
 
-# Define the Gradio interface
+# Define the Gradio interface for image captioning
 demo = gr.Interface(
     fn=caption_image,
-    inputs=gr.Image(type="pil"),
-    outputs="text",
+    inputs=gr.Image(type=IMAGE_TYPE),
+    outputs=OUTPUT_TYPE,
     title="Image Captioning with BLIP",
     description="Upload an image to generate a caption."
 )
 
-# Launch the interface
+# Launch the Gradio interface
 demo.launch()