Spaces:
Running
on
Zero
Running
on
Zero
Update app.py
Browse files
app.py
CHANGED
@@ -6,13 +6,15 @@ from PIL import Image
|
|
6 |
|
7 |
import spaces
|
8 |
|
|
|
9 |
from transformers import pipeline
|
10 |
|
11 |
# Initialize Caption Generation Model
|
12 |
-
get_caption = pipeline("image-to-text",model="Salesforce/blip-image-captioning-base")
|
13 |
|
14 |
# Initialize Image Generation Model
|
15 |
-
generate_pipeline = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5")
|
|
|
16 |
|
17 |
def captioner(input: Image.Image) -> str:
|
18 |
"""
|
@@ -37,7 +39,7 @@ def generate(prompt: str) -> Image.Image:
|
|
37 |
Returns:
|
38 |
Image.Image: The generated image corresponding to the given prompt.
|
39 |
"""
|
40 |
-
return
|
41 |
|
42 |
@spaces.GPU(duration=300)
|
43 |
def caption_and_generate(image: Image.Image) -> list:
|
|
|
6 |
|
7 |
import spaces
|
8 |
|
9 |
+
import torch
|
10 |
from transformers import pipeline
|
11 |
|
12 |
# Initialize Caption Generation Model
|
13 |
+
get_caption = pipeline("image-to-text",model="Salesforce/blip-image-captioning-base", device=0)
|
14 |
|
15 |
# Initialize Image Generation Model
|
16 |
+
generate_pipeline = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", )
|
17 |
+
pipe = generate_pipeline.to("cuda")
|
18 |
|
19 |
def captioner(input: Image.Image) -> str:
|
20 |
"""
|
|
|
39 |
Returns:
|
40 |
Image.Image: The generated image corresponding to the given prompt.
|
41 |
"""
|
42 |
+
return pipe(prompt).images[0]
|
43 |
|
44 |
@spaces.GPU(duration=300)
|
45 |
def caption_and_generate(image: Image.Image) -> list:
|