freddyaboulton HF staff lysandre HF staff commited on
Commit
f640be4
0 Parent(s):

Duplicate from huggingface-tools/text-to-image

Browse files

Co-authored-by: Lysandre <lysandre@users.noreply.huggingface.co>

Files changed (6) hide show
  1. README.md +13 -0
  2. __init__.py +0 -0
  3. app.py +4 -0
  4. requirements.txt +4 -0
  5. text_to_image.py +51 -0
  6. tool_config.json +5 -0
README.md ADDED
@@ -0,0 +1,13 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ title: Text to Image
3
+ emoji: ⚡
4
+ colorFrom: blue
5
+ colorTo: green
6
+ sdk: gradio
7
+ sdk_version: 3.27.0
8
+ app_file: app.py
9
+ pinned: false
10
+ tags:
11
+ - tool
12
+ duplicated_from: huggingface-tools/text-to-image
13
+ ---
__init__.py ADDED
File without changes
app.py ADDED
@@ -0,0 +1,4 @@
 
 
 
 
 
1
+ from transformers.tools.base import launch_gradio_demo
2
+ from text_to_image import TextToImageTool
3
+
4
+ launch_gradio_demo(TextToImageTool)
requirements.txt ADDED
@@ -0,0 +1,4 @@
 
 
 
 
 
1
+ transformers>=4.29.0
2
+ diffusers
3
+ accelerate
4
+ torch
text_to_image.py ADDED
@@ -0,0 +1,51 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from transformers.tools.base import Tool, get_default_device
2
+ from transformers.utils import is_accelerate_available
3
+ import torch
4
+
5
+ from diffusers import DiffusionPipeline, DPMSolverMultistepScheduler
6
+
7
+
8
+ TEXT_TO_IMAGE_DESCRIPTION = (
9
+ "This is a tool that creates an image according to a prompt, which is a text description. It takes an input named `prompt` which "
10
+ "contains the image description and outputs an image."
11
+ )
12
+
13
+
14
+ class TextToImageTool(Tool):
15
+ default_checkpoint = "runwayml/stable-diffusion-v1-5"
16
+ description = TEXT_TO_IMAGE_DESCRIPTION
17
+ inputs = ['text']
18
+ outputs = ['image']
19
+
20
+ def __init__(self, device=None, **hub_kwargs) -> None:
21
+ if not is_accelerate_available():
22
+ raise ImportError("Accelerate should be installed in order to use tools.")
23
+
24
+ super().__init__()
25
+
26
+ self.device = device
27
+ self.pipeline = None
28
+ self.hub_kwargs = hub_kwargs
29
+
30
+ def setup(self):
31
+ if self.device is None:
32
+ self.device = get_default_device()
33
+
34
+ self.pipeline = DiffusionPipeline.from_pretrained(self.default_checkpoint)
35
+ self.pipeline.scheduler = DPMSolverMultistepScheduler.from_config(self.pipeline.scheduler.config)
36
+ self.pipeline.to(self.device)
37
+
38
+ if self.device.type == "cuda":
39
+ self.pipeline.to(torch_dtype=torch.float16)
40
+
41
+ self.is_initialized = True
42
+
43
+ def __call__(self, prompt):
44
+ if not self.is_initialized:
45
+ self.setup()
46
+
47
+ negative_prompt = "low quality, bad quality, deformed, low resolution"
48
+ added_prompt = " , highest quality, highly realistic, very high resolution"
49
+
50
+ return self.pipeline(prompt + added_prompt, negative_prompt=negative_prompt, num_inference_steps=25).images[0]
51
+
tool_config.json ADDED
@@ -0,0 +1,5 @@
 
 
 
 
 
 
1
+ {
2
+ "description": "This is a tool that creates an image according to a prompt, which is a text description. It takes an input named `prompt` which contains the image description and outputs an image.",
3
+ "name": "image_generator",
4
+ "tool_class": "text_to_image.TextToImageTool"
5
+ }