update

app.py

@@ -1,23 +1,26 @@
 import gradio as gr
-from transformers import CLIPSegProcessor, CLIPSegForImageSegmentation
-import torch
+from lavis.models import load_model_and_preprocess
 from PIL import Image
-import numpy as np
 
 def process(input_image, prompt):
-    inputs = processor(text=prompt, images=input_image, padding="max_length", return_tensors="pt")
-    # predict
-    with torch.no_grad():
-        outputs = model(**inputs)
-    preds = torch.sigmoid(outputs.logits).squeeze().detach().cpu().numpy()
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+    model, vis_processors, txt_processors = load_model_and_preprocess(name="blip_feature_extractor", model_type="base", is_eval=True, device=device)
+
+    input_image = input_image.resize((256, 256), Image.ANTIALIAS)
+    image = vis_processors["eval"](input_image).unsqueeze(0).to(device)
+    text_input = txt_processors["eval"](prompt)
+    sample = {"image": image, "text_input": [text_input]}
+
+    features_multimodal = model.extract_features(sample, mode="multimodal")
+    preds = features_multimodal.multimodal_embeds.squeeze().detach().cpu().numpy()
     preds = np.where(preds > 0.3, 255, 0).astype(np.uint8)
     preds = Image.fromarray(preds.astype(np.uint8))
     preds = np.array(preds.resize((input_image.width, input_image.height)))
-    return preds
     
+    return preds
+
 if __name__ == '__main__':
-    processor = CLIPSegProcessor.from_pretrained("CIDAS/clipseg-rd64-refined")
-    model = CLIPSegForImageSegmentation.from_pretrained("CIDAS/clipseg-rd64-refined")
     input_image = gr.inputs.Image(label='image', type='pil')
     prompt = gr.Textbox(label='Prompt')
     ips = [
@@ -31,4 +34,4 @@ if __name__ == '__main__':
                          outputs=outputs,
                          input_size=input_size, 
                          output_size=output_size)
-    iface.launch()
+    iface.launch()

requirements.txt

@@ -1,4 +1,5 @@
 gradio
 transformers
 torch
-pillow
+pillow
+salesforce-lavis==1.0.2