keplersj commited on
Commit
4dba69f
·
1 Parent(s): 7d3d31c

use transfromer pipeline

Browse files
Files changed (1) hide show
  1. app.py +12 -8
app.py CHANGED
@@ -1,10 +1,8 @@
1
  import streamlit as st
2
  from PIL import Image
3
- from transformers import BlipProcessor, BlipForConditionalGeneration
4
  from diffusers import StableDiffusionPipeline
5
 
6
- processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-large")
7
- model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-large")
8
  pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4")
9
 
10
  captions = []
@@ -12,6 +10,13 @@ captions = []
12
  with st.sidebar:
13
  files = st.file_uploader("Upload images to blend", accept_multiple_files=True)
14
  st.divider()
 
 
 
 
 
 
 
15
  image_gen_guidance = st.slider("Stable Diffusion: Guidance Scale", value=7.5)
16
  image_gen_steps = st.slider("stable Diffusion: Inference Steps", value=50)
17
 
@@ -22,12 +27,11 @@ with col1:
22
  image = Image.open(file_name)
23
 
24
  with st.spinner('Captioning Provided Image'):
25
- inputs = processor(image, return_tensors="pt")
26
- out = model.generate(**inputs)
27
- description = processor.decode(out[0], skip_special_tokens=True)
28
- captions.append(description)
29
 
30
- st.image(image, caption=description)
 
31
 
32
  with col2:
33
  if len(captions) > 0:
 
1
  import streamlit as st
2
  from PIL import Image
3
+ from transformers import pipeline as transformer
4
  from diffusers import StableDiffusionPipeline
5
 
 
 
6
  pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4")
7
 
8
  captions = []
 
10
  with st.sidebar:
11
  files = st.file_uploader("Upload images to blend", accept_multiple_files=True)
12
  st.divider()
13
+ caption_model = st.selectbox("Caption Model", [
14
+ "ydshieh/vit-gpt2-coco-en",
15
+ "Salesforce/blip-image-captioning-large",
16
+ "nlpconnect/vit-gpt2-image-captioning",
17
+ "microsoft/git-base"
18
+ ])
19
+ st.divider()
20
  image_gen_guidance = st.slider("Stable Diffusion: Guidance Scale", value=7.5)
21
  image_gen_steps = st.slider("stable Diffusion: Inference Steps", value=50)
22
 
 
27
  image = Image.open(file_name)
28
 
29
  with st.spinner('Captioning Provided Image'):
30
+ captioner = transformer(model=caption_model)
31
+ caption = captioner(image)[0].generated_text
 
 
32
 
33
+ captions.append(caption)
34
+ st.image(image, caption=caption)
35
 
36
  with col2:
37
  if len(captions) > 0: