Upload app.py

app.py

@@ -38,7 +38,7 @@ from transformers import BlipProcessor, BlipForConditionalGeneration
 def generate_caption(processor, blip_model, raw_image):
     # unconditional image captioning
     inputs = processor(raw_image, return_tensors="pt").to(
-        "cuda", torch.float16)
+        device) #fp 16
     out = blip_model.generate(**inputs)
     caption = processor.decode(out[0], skip_special_tokens=True)
     return caption
@@ -153,8 +153,6 @@ sam_checkpoint = 'sam_vit_b_01ec64.pth'
 asam_checkpoint = 'asam_vit_b.pth'
 output_dir = "outputs"
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
-
-
 blip_processor = None
 blip_model = None
 groundingdino_model = None
@@ -185,7 +183,7 @@ def run_grounded_sam(input_image, text_prompt, task_type, box_threshold, text_th
         blip_processor = blip_processor or BlipProcessor.from_pretrained(
             "Salesforce/blip-image-captioning-large")
         blip_model = blip_model or BlipForConditionalGeneration.from_pretrained(
-            "Salesforce/blip-image-captioning-large", torch_dtype=torch.float16).to("cuda")
+            "Salesforce/blip-image-captioning-large").to(device) #torch_dtype=torch.float16
         text_prompt = generate_caption(blip_processor, blip_model, image_pil)
         print(f"Caption: {text_prompt}")
 
@@ -281,7 +279,7 @@ def run_grounded_sam(input_image, text_prompt, task_type, box_threshold, text_th
             boxes=transformed_boxes,
             multimask_output=False,
         )
-        print(torch.sum(masks))
+        print(torch.sum(masks), masks.device)
         # masks: [1, 1, 512, 512]
         mask_image = Image.new('RGBA', size, color=(0, 0, 0, 0))
         mask_draw = ImageDraw.Draw(mask_image)
@@ -449,7 +447,7 @@ if __name__ == "__main__":
                 input_image = gr.Image(
                     source='upload', type="pil", value="example9.jpg", tool="sketch",brush_radius=20)
                 task_type = gr.Dropdown(
-                    ["default_box","automatic", "scribble_point", "scribble_box", "text"], value="default_box", label="task_type")
+                    ["default_box","automatic", "scribble_point", "scribble_box"], value="default_box", label="task_type")
                 text_prompt = gr.Textbox(label="Text Prompt", placeholder="bench .", visible=False)
                 run_button = gr.Button(label="Run")
                 with gr.Accordion("Advanced options", open=False):