sayakpaul HF staff commited on
Commit
370d4c1
1 Parent(s): c733976

add: files.

Browse files
Files changed (5) hide show
  1. __init__.py +0 -0
  2. app.py +4 -0
  3. requirements.txt +4 -0
  4. text_to_image_inpainting.py +58 -0
  5. tool_config.json +5 -0
__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_inpainting import InpatingTool
3
+
4
+ launch_gradio_demo(InpatingTool)
requirements.txt ADDED
@@ -0,0 +1,4 @@
 
 
 
 
 
1
+ transformers
2
+ diffusers
3
+ accelerate
4
+ torch
text_to_image_inpainting.py ADDED
@@ -0,0 +1,58 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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 StableDiffusionInpaintPipeline
6
+
7
+
8
+ INPAINTING_DESCRIPTION = (
9
+ "This is a tool that inpaints some parts of an image StableDiffusionInpaintPipeline according to a prompt."
10
+ " It takes three inputs: `image`, which should be the original image which will be inpainted,"
11
+ " `mask_image`, which should be used to determine which parts of the original image"
12
+ " (stored in the `image` variable) should be inpainted,"
13
+ " and `prompt`, which should be the prompt to use to guide the inpainting process. It returns the"
14
+ " inpainted image."
15
+ )
16
+
17
+
18
+ class InpatingTool(Tool):
19
+ default_checkpoint = "stabilityai/stable-diffusion-2-inpainting"
20
+ description = INPAINTING_DESCRIPTION
21
+ inputs = ['text']
22
+ outputs = ['image']
23
+
24
+ def __init__(self, device=None, **hub_kwargs) -> None:
25
+ if not is_accelerate_available():
26
+ raise ImportError("Accelerate should be installed in order to use tools.")
27
+
28
+ super().__init__()
29
+
30
+ self.device = device
31
+ self.pipeline = None
32
+ self.hub_kwargs = hub_kwargs
33
+
34
+ def setup(self):
35
+ if self.device is None:
36
+ self.device = get_default_device()
37
+
38
+ self.pipeline = StableDiffusionInpaintPipeline.from_pretrained(self.default_checkpoint)
39
+ self.pipeline.to(self.device)
40
+
41
+ if self.device.type == "cuda":
42
+ self.pipeline.to(torch_dtype=torch.float16)
43
+
44
+ self.is_initialized = True
45
+
46
+ def __call__(self, image, mask_image, prompt):
47
+ if not self.is_initialized:
48
+ self.setup()
49
+
50
+ negative_prompt = "low quality, bad quality, deformed, low resolution"
51
+ added_prompt = " , highest quality, highly realistic, very high resolution"
52
+
53
+ return self.pipeline(
54
+ prompt=prompt + added_prompt,
55
+ negative_prompt=negative_prompt,
56
+ image=image,
57
+ mask_image=mask_image
58
+ ).images[0]
tool_config.json ADDED
@@ -0,0 +1,5 @@
 
 
 
 
 
 
1
+ {
2
+ "tool_class": "text_to_image_inpainting.InpatingTool",
3
+ "description": "This is a tool that inpaints some parts of an image StableDiffusionInpaintPipeline according to a prompt. It takes three inputs: `image`, which should be the original image which will be inpainted, `mask_image`, which should be used to determine which parts of the original image (stored in the `image` variable) should be inpainted, and `prompt`, which should be the prompt to use to guide the inpainting process. It returns the inpainted image.",
4
+ "name": "image_inpainter"
5
+ }