Only Blip2 model active for testing on cpu

app.py

@@ -1,13 +1,18 @@
 import gradio as gr
-from transformers import AutoProcessor, AutoTokenizer, AutoImageProcessor, AutoModelForCausalLM, BlipForConditionalGeneration, Blip2ForConditionalGeneration, VisionEncoderDecoderModel
+from transformers import AutoProcessor, AutoTokenizer, AutoImageProcessor, AutoModelForCausalLM, BlipForConditionalGeneration, Blip2ForConditionalGeneration, VisionEncoderDecoderModel, BitsAndBytesConfig
 import torch
 import open_clip
 
 from huggingface_hub import hf_hub_download
 
+quantizer_config = BitsAndBytesConfig(llm_int8_enable_fp32_cpu_offload=True)
 # Use when running on a CPU
 device_map = {
-    "embedding": "cpu",
+    "transformer.word_embeddings": 0,
+    "transformer.word_embeddings_layernorm": 0,
+    "lm_head": "cpu",
+    "transformer.h": 0,
+    "transformer.ln_f": 0,
 }
 
 # Load the Blip2 model
@@ -15,30 +20,30 @@ preprocessor_blip2_8_bit = AutoProcessor.from_pretrained("Salesforce/blip2-opt-6
 model_blip2_8_bit = Blip2ForConditionalGeneration.from_pretrained("Salesforce/blip2-opt-6.7b", device_map=device_map, load_in_8bit_fp32_cpu_offload=True)
 
 # Load the Blip base model
-preprocessor_blip_base = AutoProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
-model_blip_base = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")
+# preprocessor_blip_base = AutoProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
+# model_blip_base = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")
 
-# Load the Blip large model
-preprocessor_blip_large = AutoProcessor.from_pretrained("Salesforce/blip-image-captioning-large")
-model_blip_large = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-large")
+# # Load the Blip large model
+# preprocessor_blip_large = AutoProcessor.from_pretrained("Salesforce/blip-image-captioning-large")
+# model_blip_large = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-large")
 
-# Load the GIT coco model
-preprocessor_git_large_coco = AutoProcessor.from_pretrained("microsoft/git-large-coco")
-model_git_large_coco = AutoModelForCausalLM.from_pretrained("microsoft/git-large-coco")
+# # Load the GIT coco model
+# preprocessor_git_large_coco = AutoProcessor.from_pretrained("microsoft/git-large-coco")
+# model_git_large_coco = AutoModelForCausalLM.from_pretrained("microsoft/git-large-coco")
 
-# Load the CLIP model
-model_oc_coca, _, transform_oc_coca = open_clip.create_model_and_transforms(
-	model_name="coca_ViT-L-14",
-	pretrained="mscoco_finetuned_laion2B-s13B-b90k"
-)
+# # Load the CLIP model
+# model_oc_coca, _, transform_oc_coca = open_clip.create_model_and_transforms(
+# 	model_name="coca_ViT-L-14",
+# 	pretrained="mscoco_finetuned_laion2B-s13B-b90k"
+# )
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 # Transfer the models to the device
 model_blip2_8_bit.to(device)
-model_blip_base.to(device)
-model_blip_large.to(device)
-model_git_large_coco.to(device)
-model_oc_coca.to(device)
+# model_blip_base.to(device)
+# model_blip_large.to(device)
+# model_git_large_coco.to(device)
+# model_oc_coca.to(device)
 
 
 def generate_caption(
@@ -139,18 +144,18 @@ def generate_captions(
 	caption_blip2_8_bit = generate_caption(preprocessor_blip2_8_bit, model_blip2_8_bit, image, use_float_16=True).strip()
 
 	# Generate captions for the image using the Blip base model
-	caption_blip_base = generate_caption(preprocessor_blip_base, model_blip_base, image).strip()
+	# caption_blip_base = generate_caption(preprocessor_blip_base, model_blip_base, image).strip()
 
-	# Generate captions for the image using the Blip large model
-	caption_blip_large = generate_caption(preprocessor_blip_large, model_blip_large, image).strip()
+	# # Generate captions for the image using the Blip large model
+	# caption_blip_large = generate_caption(preprocessor_blip_large, model_blip_large, image).strip()
 	
-	# Generate captions for the image using the GIT coco model
-	caption_git_large_coco = generate_caption(preprocessor_git_large_coco, model_git_large_coco, image).strip()
+	# # Generate captions for the image using the GIT coco model
+	# caption_git_large_coco = generate_caption(preprocessor_git_large_coco, model_git_large_coco, image).strip()
 	
-	# Generate captions for the image using the CLIP model
-	caption_oc_coca = generate_captions_clip(model_oc_coca, transform_oc_coca, image).strip()
+	# # Generate captions for the image using the CLIP model
+	# caption_oc_coca = generate_captions_clip(model_oc_coca, transform_oc_coca, image).strip()
 
-	return caption_blip2_8_bit, caption_blip_base, caption_blip_large, caption_git_large_coco, caption_oc_coca
+	return caption_blip2_8_bit
 
 
 # Create the interface
@@ -165,10 +170,10 @@ iface = gr.Interface(
 	# Define the outputs
 	outputs=[
 		gr.outputs.Textbox(label="Blip2 8-bit"),
-		gr.outputs.Textbox(label="Blip base"),
-		gr.outputs.Textbox(label="Blip large"),
-		gr.outputs.Textbox(label="GIT large coco"),
-		gr.outputs.Textbox(label="CLIP"),
+		# gr.outputs.Textbox(label="Blip base"),
+		# gr.outputs.Textbox(label="Blip large"),
+		# gr.outputs.Textbox(label="GIT large coco"),
+		# gr.outputs.Textbox(label="CLIP"),
 	],
 	title="Image Captioning",
 	description="Generate captions for images using the Blip2 model, the Blip base model, the Blip large model, the GIT large coco model, and the CLIP model.",