LanHarmony commited on
Commit
78df1b1
1 Parent(s): ae5d7c3
Files changed (4) hide show
  1. README.md +1 -1
  2. app.py +28 -18
  3. visual_chat_diffuser.py +0 -952
  4. visual_foundation_models.py +107 -872
README.md CHANGED
@@ -1,6 +1,6 @@
1
  ---
2
  title: Visual Chatgpt
3
- emoji: 🐢
4
  colorFrom: yellow
5
  colorTo: yellow
6
  sdk: gradio
 
1
  ---
2
  title: Visual Chatgpt
3
+ emoji: 🎨
4
  colorFrom: yellow
5
  colorTo: yellow
6
  sdk: gradio
app.py CHANGED
@@ -47,12 +47,11 @@ from langchain.agents.initialize import initialize_agent
47
  from langchain.agents.tools import Tool
48
  from langchain.chains.conversation.memory import ConversationBufferMemory
49
  from langchain.llms.openai import OpenAI
50
- from langchain.vectorstores import Weaviate
51
  import re
52
  import gradio as gr
53
 
54
 
55
- def cut_dialogue_history(history_memory, keep_last_n_words=500):
56
  tokens = history_memory.split()
57
  n_tokens = len(tokens)
58
  print(f"hitory_memory:{history_memory}, n_tokens: {n_tokens}")
@@ -73,24 +72,24 @@ class ConversationBot:
73
  self.edit = ImageEditing(device="cuda:0")
74
  self.i2t = ImageCaptioning(device="cuda:0")
75
  self.t2i = T2I(device="cuda:0")
76
- self.image2canny = image2canny_new()
77
- self.canny2image = canny2image_new(device="cuda:0")
78
- # self.image2line = image2line_new()
79
- # self.line2image = line2image_new(device="cuda:0")
80
- # self.image2hed = image2hed_new()
81
- # self.hed2image = hed2image_new(device="cuda:0")
82
- # self.image2scribble = image2scribble_new()
83
- # self.scribble2image = scribble2image_new(device="cuda:0")
84
- # self.image2pose = image2pose_new()
85
- # self.pose2image = pose2image_new(device="cuda:0")
86
  self.BLIPVQA = BLIPVQA(device="cuda:0")
 
 
 
 
 
 
 
 
 
 
 
87
  # self.image2seg = image2seg_new()
88
  # self.seg2image = seg2image_new(device="cuda:0")
89
  # self.image2depth = image2depth_new()
90
  # self.depth2image = depth2image_new(device="cuda:0")
91
  # self.image2normal = image2normal_new()
92
  # self.normal2image = normal2image_new(device="cuda:0")
93
- self.pix2pix = Pix2Pix(device="cuda:0")
94
  self.memory = ConversationBufferMemory(memory_key="chat_history", output_key='output')
95
  self.tools = [
96
  Tool(name="Get Photo Description", func=self.i2t.inference,
@@ -178,7 +177,7 @@ class ConversationBot:
178
  print("===============Running run_text =============")
179
  print("Inputs:", text, state)
180
  print("======>Previous memory:\n %s" % self.agent.memory)
181
- self.agent.memory.buffer = cut_dialogue_history(self.agent.memory.buffer, keep_last_n_words=500)
182
  res = self.agent({"input": text})
183
  print("======>Current memory:\n %s" % self.agent.memory)
184
  response = re.sub('(image/\S*png)', lambda m: f'![](/file={m.group(0)})*{m.group(0)}*', res['output'])
@@ -215,8 +214,9 @@ with gr.Blocks(css="#chatbot .overflow-y-auto{height:500px}") as demo:
215
  with gr.Row():
216
  gr.Markdown("<h3><center>Visual ChatGPT</center></h3>")
217
 
 
218
  openai_api_key_textbox = gr.Textbox(
219
- placeholder="Paste your OpenAI API key (sk-...)",
220
  show_label=False,
221
  lines=1,
222
  type="password",
@@ -233,15 +233,25 @@ with gr.Blocks(css="#chatbot .overflow-y-auto{height:500px}") as demo:
233
  with gr.Column(scale=0.15, min_width=0):
234
  btn = gr.UploadButton("Upload", file_types=["image"])
235
 
 
 
 
 
 
 
 
 
 
 
 
 
 
236
  openai_api_key_textbox.submit(bot.init_agent, [openai_api_key_textbox], [input_raws])
237
  txt.submit(bot.run_text, [txt, state], [chatbot, state])
238
  txt.submit(lambda: "", None, txt)
239
-
240
  btn.upload(bot.run_image, [btn, state, txt], [chatbot, state, txt])
241
-
242
  clear.click(bot.memory.clear)
243
  clear.click(lambda: [], None, chatbot)
244
  clear.click(lambda: [], None, state)
245
 
246
-
247
  demo.launch(server_name="0.0.0.0", server_port=7860)
 
47
  from langchain.agents.tools import Tool
48
  from langchain.chains.conversation.memory import ConversationBufferMemory
49
  from langchain.llms.openai import OpenAI
 
50
  import re
51
  import gradio as gr
52
 
53
 
54
+ def cut_dialogue_history(history_memory, keep_last_n_words=400):
55
  tokens = history_memory.split()
56
  n_tokens = len(tokens)
57
  print(f"hitory_memory:{history_memory}, n_tokens: {n_tokens}")
 
72
  self.edit = ImageEditing(device="cuda:0")
73
  self.i2t = ImageCaptioning(device="cuda:0")
74
  self.t2i = T2I(device="cuda:0")
 
 
 
 
 
 
 
 
 
 
75
  self.BLIPVQA = BLIPVQA(device="cuda:0")
76
+ self.pix2pix = Pix2Pix(device="cuda:0")
77
+ self.image2canny = image2canny()
78
+ self.canny2image = canny2image(device="cuda:0")
79
+ # self.image2line = image2line()
80
+ # self.line2image = line2image(device="cuda:0")
81
+ # self.image2hed = image2hed()
82
+ # self.hed2image = hed2image(device="cuda:0")
83
+ # self.image2scribble = image2scribble()
84
+ # self.scribble2image = scribble2image(device="cuda:0")
85
+ # self.image2pose = image2pose()
86
+ # self.pose2image = pose2image(device="cuda:0")
87
  # self.image2seg = image2seg_new()
88
  # self.seg2image = seg2image_new(device="cuda:0")
89
  # self.image2depth = image2depth_new()
90
  # self.depth2image = depth2image_new(device="cuda:0")
91
  # self.image2normal = image2normal_new()
92
  # self.normal2image = normal2image_new(device="cuda:0")
 
93
  self.memory = ConversationBufferMemory(memory_key="chat_history", output_key='output')
94
  self.tools = [
95
  Tool(name="Get Photo Description", func=self.i2t.inference,
 
177
  print("===============Running run_text =============")
178
  print("Inputs:", text, state)
179
  print("======>Previous memory:\n %s" % self.agent.memory)
180
+ self.agent.memory.buffer = cut_dialogue_history(self.agent.memory.buffer, keep_last_n_words=400)
181
  res = self.agent({"input": text})
182
  print("======>Current memory:\n %s" % self.agent.memory)
183
  response = re.sub('(image/\S*png)', lambda m: f'![](/file={m.group(0)})*{m.group(0)}*', res['output'])
 
214
  with gr.Row():
215
  gr.Markdown("<h3><center>Visual ChatGPT</center></h3>")
216
 
217
+ with gr.Row():
218
  openai_api_key_textbox = gr.Textbox(
219
+ placeholder="Paste your OpenAI API key here to start Visual ChatGPT(sk-...)",
220
  show_label=False,
221
  lines=1,
222
  type="password",
 
233
  with gr.Column(scale=0.15, min_width=0):
234
  btn = gr.UploadButton("Upload", file_types=["image"])
235
 
236
+ gr.Examples(
237
+ examples=["Generate a figure of a lovely cat",
238
+ "Replace the cat with a lovely dog",
239
+ "Remove the cat in this image",
240
+ "Can you detect the canny edge of this image?",
241
+ "Can you use this canny image to generate a oil painting of a lovely dog",
242
+ "Make it like water-color painting",
243
+ "What is the background color"
244
+ "Describe this image"],
245
+ inputs=txt
246
+ )
247
+
248
+
249
  openai_api_key_textbox.submit(bot.init_agent, [openai_api_key_textbox], [input_raws])
250
  txt.submit(bot.run_text, [txt, state], [chatbot, state])
251
  txt.submit(lambda: "", None, txt)
 
252
  btn.upload(bot.run_image, [btn, state, txt], [chatbot, state, txt])
 
253
  clear.click(bot.memory.clear)
254
  clear.click(lambda: [], None, chatbot)
255
  clear.click(lambda: [], None, state)
256
 
 
257
  demo.launch(server_name="0.0.0.0", server_port=7860)
visual_chat_diffuser.py DELETED
@@ -1,952 +0,0 @@
1
- VISUAL_CHATGPT_PREFIX = """Visual ChatGPT is designed to be able to assist with a wide range of text and visual related tasks, from answering simple questions to providing in-depth explanations and discussions on a wide range of topics. Visual ChatGPT is able to generate human-like text based on the input it receives, allowing it to engage in natural-sounding conversations and provide responses that are coherent and relevant to the topic at hand.
2
-
3
- Visual ChatGPT is able to process and understand large amounts of text and image. As a language model, Visual ChatGPT can not directly read images, but it has a list of tools to finish different visual tasks. Each image will have a file name formed as "image/xxx.png", and Visual ChatGPT can invoke different tools to indirectly understand pictures. When talking about images, Visual ChatGPT is very strict to the file name and will never fabricate nonexistent files. When using tools to generate new image files, Visual ChatGPT is also known that the image may not be the same as user's demand, and will use other visual question answering tools or description tools to observe the real image. Visual ChatGPT is able to use tools in a sequence, and is loyal to the tool observation outputs rather than faking the image content and image file name. It will remember to provide the file name from the last tool observation, if a new image is generated.
4
-
5
- Human may provide new figures to Visual ChatGPT with a description. The description helps Visual ChatGPT to understand this image, but Visual ChatGPT should use tools to finish following tasks, rather than directly imagine from the description.
6
-
7
- Overall, Visual ChatGPT is a powerful visual dialogue assistant tool that can help with a wide range of tasks and provide valuable insights and information on a wide range of topics.
8
-
9
-
10
- TOOLS:
11
- ------
12
-
13
- Visual ChatGPT has access to the following tools:"""
14
-
15
- VISUAL_CHATGPT_FORMAT_INSTRUCTIONS = """To use a tool, please use the following format:
16
-
17
- ```
18
- Thought: Do I need to use a tool? Yes
19
- Action: the action to take, should be one of [{tool_names}]
20
- Action Input: the input to the action
21
- Observation: the result of the action
22
- ```
23
-
24
- When you have a response to say to the Human, or if you do not need to use a tool, you MUST use the format:
25
-
26
- ```
27
- Thought: Do I need to use a tool? No
28
- {ai_prefix}: [your response here]
29
- ```
30
- """
31
-
32
- VISUAL_CHATGPT_SUFFIX = """You are very strict to the filename correctness and will never fake a file name if not exists.
33
- You will remember to provide the image file name loyally if it's provided in the last tool observation.
34
-
35
- Begin!
36
-
37
- Previous conversation history:
38
- {chat_history}
39
-
40
- New input: {input}
41
- Since Visual ChatGPT is a text language model, Visual ChatGPT must use tools to observe images rather than imagination.
42
- The thoughts and observations are only visible for Visual ChatGPT, Visual ChatGPT should remember to repeat important information in the final response for Human.
43
- Thought: Do I need to use a tool? {agent_scratchpad}"""
44
-
45
- import uuid
46
- import os
47
- import cv2
48
- import random
49
- from PIL import Image
50
- import torch
51
- import numpy as np
52
- from pytorch_lightning import seed_everything
53
- import re
54
- import gradio as gr
55
-
56
- from diffusers import StableDiffusionPipeline, StableDiffusionInpaintPipeline, StableDiffusionInstructPix2PixPipeline
57
- from diffusers import EulerAncestralDiscreteScheduler
58
- from diffusers import StableDiffusionControlNetPipeline, ControlNetModel, UniPCMultistepScheduler
59
- from controlnet_aux import OpenposeDetector, MLSDdetector, HEDdetector
60
-
61
- from transformers import AutoModelForCausalLM, AutoTokenizer, CLIPSegProcessor, CLIPSegForImageSegmentation
62
- from transformers import pipeline, BlipProcessor, BlipForConditionalGeneration, BlipForQuestionAnswering
63
- from transformers import AutoImageProcessor, UperNetForSemanticSegmentation
64
-
65
- from langchain.agents.initialize import initialize_agent
66
- from langchain.agents.tools import Tool
67
- from langchain.chains.conversation.memory import ConversationBufferMemory
68
- from langchain.llms.openai import OpenAI
69
- from langchain.vectorstores import Weaviate
70
-
71
-
72
- def cut_dialogue_history(history_memory, keep_last_n_words=500):
73
- tokens = history_memory.split()
74
- n_tokens = len(tokens)
75
- print(f"hitory_memory:{history_memory}, n_tokens: {n_tokens}")
76
- if n_tokens < keep_last_n_words:
77
- return history_memory
78
- else:
79
- paragraphs = history_memory.split('\n')
80
- last_n_tokens = n_tokens
81
- while last_n_tokens >= keep_last_n_words:
82
- last_n_tokens = last_n_tokens - len(paragraphs[0].split(' '))
83
- paragraphs = paragraphs[1:]
84
- return '\n' + '\n'.join(paragraphs)
85
-
86
- def get_new_image_name(org_img_name, func_name="update"):
87
- head_tail = os.path.split(org_img_name)
88
- head = head_tail[0]
89
- tail = head_tail[1]
90
- name_split = tail.split('.')[0].split('_')
91
- this_new_uuid = str(uuid.uuid4())[0:4]
92
- if len(name_split) == 1:
93
- most_org_file_name = name_split[0]
94
- recent_prev_file_name = name_split[0]
95
- new_file_name = '{}_{}_{}_{}.png'.format(this_new_uuid, func_name, recent_prev_file_name, most_org_file_name)
96
- else:
97
- assert len(name_split) == 4
98
- most_org_file_name = name_split[3]
99
- recent_prev_file_name = name_split[0]
100
- new_file_name = '{}_{}_{}_{}.png'.format(this_new_uuid, func_name, recent_prev_file_name, most_org_file_name)
101
- return os.path.join(head, new_file_name)
102
-
103
- def ade_palette():
104
- return [[120, 120, 120], [180, 120, 120], [6, 230, 230], [80, 50, 50],
105
- [4, 200, 3], [120, 120, 80], [140, 140, 140], [204, 5, 255],
106
- [230, 230, 230], [4, 250, 7], [224, 5, 255], [235, 255, 7],
107
- [150, 5, 61], [120, 120, 70], [8, 255, 51], [255, 6, 82],
108
- [143, 255, 140], [204, 255, 4], [255, 51, 7], [204, 70, 3],
109
- [0, 102, 200], [61, 230, 250], [255, 6, 51], [11, 102, 255],
110
- [255, 7, 71], [255, 9, 224], [9, 7, 230], [220, 220, 220],
111
- [255, 9, 92], [112, 9, 255], [8, 255, 214], [7, 255, 224],
112
- [255, 184, 6], [10, 255, 71], [255, 41, 10], [7, 255, 255],
113
- [224, 255, 8], [102, 8, 255], [255, 61, 6], [255, 194, 7],
114
- [255, 122, 8], [0, 255, 20], [255, 8, 41], [255, 5, 153],
115
- [6, 51, 255], [235, 12, 255], [160, 150, 20], [0, 163, 255],
116
- [140, 140, 140], [250, 10, 15], [20, 255, 0], [31, 255, 0],
117
- [255, 31, 0], [255, 224, 0], [153, 255, 0], [0, 0, 255],
118
- [255, 71, 0], [0, 235, 255], [0, 173, 255], [31, 0, 255],
119
- [11, 200, 200], [255, 82, 0], [0, 255, 245], [0, 61, 255],
120
- [0, 255, 112], [0, 255, 133], [255, 0, 0], [255, 163, 0],
121
- [255, 102, 0], [194, 255, 0], [0, 143, 255], [51, 255, 0],
122
- [0, 82, 255], [0, 255, 41], [0, 255, 173], [10, 0, 255],
123
- [173, 255, 0], [0, 255, 153], [255, 92, 0], [255, 0, 255],
124
- [255, 0, 245], [255, 0, 102], [255, 173, 0], [255, 0, 20],
125
- [255, 184, 184], [0, 31, 255], [0, 255, 61], [0, 71, 255],
126
- [255, 0, 204], [0, 255, 194], [0, 255, 82], [0, 10, 255],
127
- [0, 112, 255], [51, 0, 255], [0, 194, 255], [0, 122, 255],
128
- [0, 255, 163], [255, 153, 0], [0, 255, 10], [255, 112, 0],
129
- [143, 255, 0], [82, 0, 255], [163, 255, 0], [255, 235, 0],
130
- [8, 184, 170], [133, 0, 255], [0, 255, 92], [184, 0, 255],
131
- [255, 0, 31], [0, 184, 255], [0, 214, 255], [255, 0, 112],
132
- [92, 255, 0], [0, 224, 255], [112, 224, 255], [70, 184, 160],
133
- [163, 0, 255], [153, 0, 255], [71, 255, 0], [255, 0, 163],
134
- [255, 204, 0], [255, 0, 143], [0, 255, 235], [133, 255, 0],
135
- [255, 0, 235], [245, 0, 255], [255, 0, 122], [255, 245, 0],
136
- [10, 190, 212], [214, 255, 0], [0, 204, 255], [20, 0, 255],
137
- [255, 255, 0], [0, 153, 255], [0, 41, 255], [0, 255, 204],
138
- [41, 0, 255], [41, 255, 0], [173, 0, 255], [0, 245, 255],
139
- [71, 0, 255], [122, 0, 255], [0, 255, 184], [0, 92, 255],
140
- [184, 255, 0], [0, 133, 255], [255, 214, 0], [25, 194, 194],
141
- [102, 255, 0], [92, 0, 255]]
142
-
143
- def HWC3(x):
144
- assert x.dtype == np.uint8
145
- if x.ndim == 2:
146
- x = x[:, :, None]
147
- assert x.ndim == 3
148
- H, W, C = x.shape
149
- assert C == 1 or C == 3 or C == 4
150
- if C == 3:
151
- return x
152
- if C == 1:
153
- return np.concatenate([x, x, x], axis=2)
154
- if C == 4:
155
- color = x[:, :, 0:3].astype(np.float32)
156
- alpha = x[:, :, 3:4].astype(np.float32) / 255.0
157
- y = color * alpha + 255.0 * (1.0 - alpha)
158
- y = y.clip(0, 255).astype(np.uint8)
159
- return y
160
-
161
- def resize_image(input_image, resolution):
162
- H, W, C = input_image.shape
163
- H = float(H)
164
- W = float(W)
165
- k = float(resolution) / min(H, W)
166
- H *= k
167
- W *= k
168
- H = int(np.round(H / 64.0)) * 64
169
- W = int(np.round(W / 64.0)) * 64
170
- img = cv2.resize(input_image, (W, H), interpolation=cv2.INTER_LANCZOS4 if k > 1 else cv2.INTER_AREA)
171
- return img
172
-
173
- class MaskFormer:
174
- def __init__(self, device):
175
- self.device = device
176
- self.processor = CLIPSegProcessor.from_pretrained("CIDAS/clipseg-rd64-refined")
177
- self.model = CLIPSegForImageSegmentation.from_pretrained("CIDAS/clipseg-rd64-refined").to(device)
178
-
179
- def inference(self, image_path, text):
180
- threshold = 0.5
181
- min_area = 0.02
182
- padding = 20
183
- original_image = Image.open(image_path)
184
- image = original_image.resize((512, 512))
185
- inputs = self.processor(text=text, images=image, padding="max_length", return_tensors="pt",).to(self.device)
186
- with torch.no_grad():
187
- outputs = self.model(**inputs)
188
- mask = torch.sigmoid(outputs[0]).squeeze().cpu().numpy() > threshold
189
- area_ratio = len(np.argwhere(mask)) / (mask.shape[0] * mask.shape[1])
190
- if area_ratio < min_area:
191
- return None
192
- true_indices = np.argwhere(mask)
193
- mask_array = np.zeros_like(mask, dtype=bool)
194
- for idx in true_indices:
195
- padded_slice = tuple(slice(max(0, i - padding), i + padding + 1) for i in idx)
196
- mask_array[padded_slice] = True
197
- visual_mask = (mask_array * 255).astype(np.uint8)
198
- image_mask = Image.fromarray(visual_mask)
199
- return image_mask.resize(image.size)
200
-
201
- class ImageEditing:
202
- def __init__(self, device):
203
- print("Initializing StableDiffusionInpaint to %s" % device)
204
- self.device = device
205
- self.mask_former = MaskFormer(device=self.device)
206
- self.inpainting = StableDiffusionInpaintPipeline.from_pretrained( "runwayml/stable-diffusion-inpainting",).to(device)
207
-
208
- def remove_part_of_image(self, input):
209
- image_path, to_be_removed_txt = input.split(",")
210
- print(f'remove_part_of_image: to_be_removed {to_be_removed_txt}')
211
- return self.replace_part_of_image(f"{image_path},{to_be_removed_txt},background")
212
-
213
- def replace_part_of_image(self, input):
214
- image_path, to_be_replaced_txt, replace_with_txt = input.split(",")
215
- print(f'replace_part_of_image: replace_with_txt {replace_with_txt}')
216
- original_image = Image.open(image_path)
217
- mask_image = self.mask_former.inference(image_path, to_be_replaced_txt)
218
- updated_image = self.inpainting(prompt=replace_with_txt, image=original_image, mask_image=mask_image).images[0]
219
- updated_image_path = get_new_image_name(image_path, func_name="replace-something")
220
- updated_image.save(updated_image_path)
221
- return updated_image_path
222
-
223
- class Pix2Pix:
224
- def __init__(self, device):
225
- print("Initializing Pix2Pix to %s" % device)
226
- self.device = device
227
- self.pipe = StableDiffusionInstructPix2PixPipeline.from_pretrained("timbrooks/instruct-pix2pix", torch_dtype=torch.float16, safety_checker=None).to(device)
228
- self.pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(self.pipe.scheduler.config)
229
-
230
- def inference(self, inputs):
231
- """Change style of image."""
232
- print("===>Starting Pix2Pix Inference")
233
- image_path, instruct_text = inputs.split(",")[0], ','.join(inputs.split(',')[1:])
234
- original_image = Image.open(image_path)
235
- image = self.pipe(instruct_text,image=original_image,num_inference_steps=40,image_guidance_scale=1.2,).images[0]
236
- updated_image_path = get_new_image_name(image_path, func_name="pix2pix")
237
- image.save(updated_image_path)
238
- return updated_image_path
239
-
240
- class T2I:
241
- def __init__(self, device):
242
- print("Initializing T2I to %s" % device)
243
- self.device = device
244
- self.pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16)
245
- self.text_refine_tokenizer = AutoTokenizer.from_pretrained("Gustavosta/MagicPrompt-Stable-Diffusion")
246
- self.text_refine_model = AutoModelForCausalLM.from_pretrained("Gustavosta/MagicPrompt-Stable-Diffusion")
247
- self.text_refine_gpt2_pipe = pipeline("text-generation", model=self.text_refine_model, tokenizer=self.text_refine_tokenizer, device=self.device)
248
- self.pipe.to(device)
249
-
250
- def inference(self, text):
251
- image_filename = os.path.join('image', str(uuid.uuid4())[0:8] + ".png")
252
- refined_text = self.text_refine_gpt2_pipe(text)[0]["generated_text"]
253
- print(f'{text} refined to {refined_text}')
254
- image = self.pipe(refined_text).images[0]
255
- image.save(image_filename)
256
- print(f"Processed T2I.run, text: {text}, image_filename: {image_filename}")
257
- return image_filename
258
-
259
- class ImageCaptioning:
260
- def __init__(self, device):
261
- print("Initializing ImageCaptioning to %s" % device)
262
- self.device = device
263
- self.processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
264
- self.model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base").to(self.device)
265
-
266
- def inference(self, image_path):
267
- inputs = self.processor(Image.open(image_path), return_tensors="pt").to(self.device)
268
- out = self.model.generate(**inputs)
269
- captions = self.processor.decode(out[0], skip_special_tokens=True)
270
- return captions
271
-
272
- class image2canny:
273
- def __init__(self):
274
- print("Direct detect canny.")
275
- self.low_threshold = 100
276
- self.high_threshold = 200
277
-
278
- def inference(self, inputs):
279
- print("===>Starting image2canny Inference")
280
- image = Image.open(inputs)
281
- image = np.array(image)
282
- canny = cv2.Canny(image, self.low_threshold, self.high_threshold)
283
- canny = canny[:, :, None]
284
- canny = np.concatenate([canny, canny, canny], axis=2)
285
- canny = 255 - canny
286
- canny = Image.fromarray(canny)
287
- updated_image_path = get_new_image_name(inputs, func_name="edge")
288
- canny.save(updated_image_path)
289
- return updated_image_path
290
-
291
- class canny2image:
292
- def __init__(self, device):
293
- self.controlnet = ControlNetModel.from_pretrained(
294
- "fusing/stable-diffusion-v1-5-controlnet-canny"
295
- )
296
-
297
- self.pipe = StableDiffusionControlNetPipeline.from_pretrained(
298
- "runwayml/stable-diffusion-v1-5", controlnet=self.controlnet, safety_checker=None
299
- )
300
-
301
- self.pipe.scheduler = UniPCMultistepScheduler.from_config(self.pipe.scheduler.config)
302
- self.pipe.to(device)
303
- self.image_resolution = 512
304
- self.num_inference_steps = 20
305
- self.seed = -1
306
- self.unconditional_guidance_scale = 9.0
307
- self.a_prompt = 'best quality, extremely detailed'
308
- self.n_prompt = 'longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality'
309
-
310
- def inference(self, inputs):
311
- print("===>Starting canny2image Inference")
312
- image_path, instruct_text = inputs.split(",")[0], ','.join(inputs.split(',')[1:])
313
- image = Image.open(image_path)
314
- image = np.array(image)
315
- image = 255 - image
316
- prompt = instruct_text
317
- img = resize_image(HWC3(image), self.image_resolution)
318
- img = Image.fromarray(img)
319
-
320
- self.seed = random.randint(0, 65535)
321
- seed_everything(self.seed)
322
- prompt = prompt + ', ' + self.a_prompt
323
- image = self.pipe(prompt, img, num_inference_steps=self.num_inference_steps, eta=0.0, negative_prompt=self.n_prompt, guidance_scale=self.unconditional_guidance_scale).images[0]
324
- updated_image_path = get_new_image_name(image_path, func_name="canny2image")
325
- image.save(updated_image_path)
326
- return updated_image_path
327
-
328
- class image2line:
329
- def __init__(self):
330
- self.detector = MLSDdetector.from_pretrained('lllyasviel/ControlNet')
331
- self.value_thresh = 0.1
332
- self.dis_thresh = 0.1
333
- self.resolution = 512
334
-
335
- def inference(self, inputs):
336
- print("===>Starting image2line Inference")
337
- image = Image.open(inputs)
338
- image = np.array(image)
339
- image = HWC3(image)
340
- mlsd = self.detector(resize_image(image, self.resolution), thr_v=self.value_thresh, thr_d=self.dis_thresh)
341
- mlsd = np.array(mlsd)
342
- mlsd = 255 - mlsd
343
- mlsd = Image.fromarray(mlsd)
344
- updated_image_path = get_new_image_name(inputs, func_name="line-of")
345
- mlsd.save(updated_image_path)
346
- return updated_image_path
347
-
348
- class line2image:
349
- def __init__(self, device):
350
- print("Initialize the line2image model...")
351
- self.controlnet = ControlNetModel.from_pretrained(
352
- "fusing/stable-diffusion-v1-5-controlnet-mlsd"
353
- )
354
-
355
- self.pipe = StableDiffusionControlNetPipeline.from_pretrained(
356
- "runwayml/stable-diffusion-v1-5", controlnet=self.controlnet, safety_checker=None
357
- )
358
-
359
- self.pipe.scheduler = UniPCMultistepScheduler.from_config(self.pipe.scheduler.config)
360
- self.pipe.to(device)
361
- self.image_resolution = 512
362
- self.num_inference_steps = 20
363
- self.seed = -1
364
- self.unconditional_guidance_scale = 9.0
365
- self.a_prompt = 'best quality, extremely detailed'
366
- self.n_prompt = 'longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality'
367
-
368
- def inference(self, inputs):
369
- print("===>Starting line2image Inference")
370
- image_path, instruct_text = inputs.split(",")[0], ','.join(inputs.split(',')[1:])
371
- image = Image.open(image_path)
372
- image = np.array(image)
373
- image = 255 - image
374
- prompt = instruct_text
375
- img = resize_image(HWC3(image), self.image_resolution)
376
- img = Image.fromarray(img)
377
-
378
- self.seed = random.randint(0, 65535)
379
- seed_everything(self.seed)
380
-
381
- prompt = prompt + ', ' + self.a_prompt
382
- image = self.pipe(prompt, img, num_inference_steps=self.num_inference_steps, eta=0.0, negative_prompt=self.n_prompt, guidance_scale=self.unconditional_guidance_scale).images[0]
383
- updated_image_path = get_new_image_name(image_path, func_name="line2image")
384
- image.save(updated_image_path)
385
- return updated_image_path
386
-
387
- class image2hed:
388
- def __init__(self):
389
- print("Direct detect soft HED boundary...")
390
- self.detector = HEDdetector.from_pretrained('lllyasviel/ControlNet')
391
- self.resolution = 512
392
-
393
- def inference(self, inputs):
394
- print("===>Starting image2hed Inference")
395
- image = Image.open(inputs)
396
- image = np.array(image)
397
- image = HWC3(image)
398
- image = Image.fromarray(resize_image(image, self.resolution))
399
- hed = self.detector(image)
400
-
401
- updated_image_path = get_new_image_name(inputs, func_name="hed-boundary")
402
- hed.save(updated_image_path)
403
- return updated_image_path
404
-
405
- class hed2image:
406
- def __init__(self, device):
407
- print("Initialize the hed2image model...")
408
- self.controlnet = ControlNetModel.from_pretrained(
409
- "fusing/stable-diffusion-v1-5-controlnet-hed"
410
- )
411
-
412
- self.pipe = StableDiffusionControlNetPipeline.from_pretrained(
413
- "runwayml/stable-diffusion-v1-5", controlnet=self.controlnet, safety_checker=None
414
- )
415
-
416
- self.pipe.scheduler = UniPCMultistepScheduler.from_config(self.pipe.scheduler.config)
417
- self.pipe.to(device)
418
- self.image_resolution = 512
419
- self.num_inference_steps = 20
420
- self.seed = -1
421
- self.unconditional_guidance_scale = 9.0
422
- self.a_prompt = 'best quality, extremely detailed'
423
- self.n_prompt = 'longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality'
424
-
425
- def inference(self, inputs):
426
- print("===>Starting hed2image Inference")
427
- image_path, instruct_text = inputs.split(",")[0], ','.join(inputs.split(',')[1:])
428
- image = Image.open(image_path)
429
- image = np.array(image)
430
- img = resize_image(HWC3(image), self.image_resolution)
431
- img = Image.fromarray(img)
432
-
433
- self.seed = random.randint(0, 65535)
434
- seed_everything(self.seed)
435
-
436
- prompt = instruct_text
437
- prompt = prompt + ', ' + self.a_prompt
438
- image = \
439
- self.pipe(prompt, img, num_inference_steps=self.num_inference_steps, eta=0.0, negative_prompt=self.n_prompt,
440
- guidance_scale=self.unconditional_guidance_scale).images[0]
441
- updated_image_path = get_new_image_name(image_path, func_name="hed2image")
442
- image.save(updated_image_path)
443
- return updated_image_path
444
-
445
- class image2scribble:
446
- def __init__(self):
447
- print("Direct detect scribble.")
448
- self.detector = HEDdetector.from_pretrained('lllyasviel/ControlNet')
449
- self.resolution = 512
450
-
451
- def inference(self, inputs):
452
- print("===>Starting image2scribble Inference")
453
- image = Image.open(inputs)
454
- image = np.array(image)
455
- image = HWC3(image)
456
- image = resize_image(image, self.resolution)
457
- image = Image.fromarray(image)
458
- scribble = self.detector(image, scribble=True)
459
- scribble = np.array(scribble)
460
- scribble = 255 - scribble
461
- scribble = Image.fromarray(scribble)
462
- updated_image_path = get_new_image_name(inputs, func_name="scribble")
463
- scribble.save(updated_image_path)
464
- return updated_image_path
465
-
466
- class scribble2image:
467
- def __init__(self, device):
468
- self.controlnet = ControlNetModel.from_pretrained(
469
- "fusing/stable-diffusion-v1-5-controlnet-scribble"
470
- )
471
-
472
- self.pipe = StableDiffusionControlNetPipeline.from_pretrained(
473
- "runwayml/stable-diffusion-v1-5", controlnet=self.controlnet, safety_checker=None
474
- )
475
-
476
- self.pipe.scheduler = UniPCMultistepScheduler.from_config(self.pipe.scheduler.config)
477
- self.pipe.to(device)
478
- self.image_resolution = 512
479
- self.num_inference_steps = 20
480
- self.seed = -1
481
- self.unconditional_guidance_scale = 9.0
482
- self.a_prompt = 'best quality, extremely detailed'
483
- self.n_prompt = 'longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality'
484
-
485
- def inference(self, inputs):
486
- print("===>Starting scribble2image Inference")
487
- image_path, instruct_text = inputs.split(",")[0], ','.join(inputs.split(',')[1:])
488
- image = Image.open(image_path)
489
- image = np.array(image)
490
- image = 255 - image
491
- img = resize_image(HWC3(image), self.image_resolution)
492
- img = Image.fromarray(img)
493
-
494
- self.seed = random.randint(0, 65535)
495
- seed_everything(self.seed)
496
-
497
- prompt = instruct_text
498
- prompt = prompt + ', ' + self.a_prompt
499
- image = \
500
- self.pipe(prompt, img, num_inference_steps=self.num_inference_steps, eta=0.0, negative_prompt=self.n_prompt,
501
- guidance_scale=self.unconditional_guidance_scale).images[0]
502
- updated_image_path = get_new_image_name(image_path, func_name="scribble2image")
503
- image.save(updated_image_path)
504
- return updated_image_path
505
-
506
- class image2pose:
507
- def __init__(self):
508
- self.detector = OpenposeDetector.from_pretrained('lllyasviel/ControlNet')
509
- self.resolution = 512
510
-
511
- def inference(self, inputs):
512
- print("===>Starting image2pose Inference")
513
- image = Image.open(inputs)
514
- image = np.array(image)
515
- image = HWC3(image)
516
- image = resize_image(image, self.resolution)
517
- image = Image.fromarray(image)
518
- pose = self.detector(image)
519
-
520
- updated_image_path = get_new_image_name(inputs, func_name="human-pose")
521
- pose.save(updated_image_path)
522
- return updated_image_path
523
-
524
- class pose2image:
525
- def __init__(self, device):
526
- self.controlnet = ControlNetModel.from_pretrained(
527
- "fusing/stable-diffusion-v1-5-controlnet-openpose"
528
- )
529
-
530
- self.pipe = StableDiffusionControlNetPipeline.from_pretrained(
531
- "runwayml/stable-diffusion-v1-5", controlnet=self.controlnet, safety_checker=None
532
- )
533
-
534
- self.pipe.scheduler = UniPCMultistepScheduler.from_config(self.pipe.scheduler.config)
535
- self.pipe.to(device)
536
- self.image_resolution = 512
537
- self.num_inference_steps = 20
538
- self.seed = -1
539
- self.unconditional_guidance_scale = 9.0
540
- self.a_prompt = 'best quality, extremely detailed'
541
- self.n_prompt = 'longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality'
542
-
543
- def inference(self, inputs):
544
- print("===>Starting pose2image Inference")
545
- image_path, instruct_text = inputs.split(",")[0], ','.join(inputs.split(',')[1:])
546
- image = Image.open(image_path)
547
- image = np.array(image)
548
- img = resize_image(HWC3(image), self.image_resolution)
549
- img = Image.fromarray(img)
550
-
551
- self.seed = random.randint(0, 65535)
552
- seed_everything(self.seed)
553
-
554
- prompt = instruct_text
555
- prompt = prompt + ', ' + self.a_prompt
556
- image = \
557
- self.pipe(prompt, img, num_inference_steps=self.num_inference_steps, eta=0.0, negative_prompt=self.n_prompt,
558
- guidance_scale=self.unconditional_guidance_scale).images[0]
559
- updated_image_path = get_new_image_name(image_path, func_name="pose2image")
560
- image.save(updated_image_path)
561
- return updated_image_path
562
-
563
- class image2seg:
564
- def __init__(self):
565
- print("Initialize image2segmentation Inference")
566
- self.image_processor = AutoImageProcessor.from_pretrained("openmmlab/upernet-convnext-small")
567
- self.image_segmentor = UperNetForSemanticSegmentation.from_pretrained("openmmlab/upernet-convnext-small")
568
- self.resolution = 512
569
-
570
- def inference(self, inputs):
571
- image = Image.open(inputs)
572
- image = np.array(image)
573
- image = HWC3(image)
574
- image = resize_image(image, self.resolution)
575
- image = Image.fromarray(image)
576
- pixel_values = self.image_processor(image, return_tensors="pt").pixel_values
577
-
578
- with torch.no_grad():
579
- outputs = self.image_segmentor(pixel_values)
580
-
581
- seg = self.image_processor.post_process_semantic_segmentation(outputs, target_sizes=[image.size[::-1]])[0]
582
-
583
- color_seg = np.zeros((seg.shape[0], seg.shape[1], 3), dtype=np.uint8) # height, width, 3
584
-
585
- palette = np.array(ade_palette())
586
-
587
- for label, color in enumerate(palette):
588
- color_seg[seg == label, :] = color
589
-
590
- color_seg = color_seg.astype(np.uint8)
591
-
592
- segmentation = Image.fromarray(color_seg)
593
- updated_image_path = get_new_image_name(inputs, func_name="segmentation")
594
- segmentation.save(updated_image_path)
595
- return updated_image_path
596
-
597
- class seg2image:
598
- def __init__(self, device):
599
- self.controlnet = ControlNetModel.from_pretrained(
600
- "fusing/stable-diffusion-v1-5-controlnet-seg"
601
- )
602
-
603
- self.pipe = StableDiffusionControlNetPipeline.from_pretrained(
604
- "runwayml/stable-diffusion-v1-5", controlnet=self.controlnet, safety_checker=None
605
- )
606
-
607
- self.pipe.scheduler = UniPCMultistepScheduler.from_config(self.pipe.scheduler.config)
608
- self.pipe.to(device)
609
- self.image_resolution = 512
610
- self.num_inference_steps = 20
611
- self.seed = -1
612
- self.unconditional_guidance_scale = 9.0
613
- self.a_prompt = 'best quality, extremely detailed'
614
- self.n_prompt = 'longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality'
615
-
616
- def inference(self, inputs):
617
- print("===>Starting seg2image Inference")
618
- image_path, instruct_text = inputs.split(",")[0], ','.join(inputs.split(',')[1:])
619
- image = Image.open(image_path)
620
- image = np.array(image)
621
- img = resize_image(HWC3(image), self.image_resolution)
622
- img = Image.fromarray(img)
623
-
624
- self.seed = random.randint(0, 65535)
625
- seed_everything(self.seed)
626
-
627
- prompt = instruct_text
628
- prompt = prompt + ', ' + self.a_prompt
629
- image = \
630
- self.pipe(prompt, img, num_inference_steps=self.num_inference_steps, eta=0.0, negative_prompt=self.n_prompt,
631
- guidance_scale=self.unconditional_guidance_scale).images[0]
632
- updated_image_path = get_new_image_name(image_path, func_name="segment2image")
633
- image.save(updated_image_path)
634
- return updated_image_path
635
-
636
- class image2depth:
637
- def __init__(self):
638
- print("initialize depth estimation")
639
- self.depth_estimator = pipeline('depth-estimation')
640
- self.resolution = 512
641
-
642
- def inference(self, inputs):
643
- image = Image.open(inputs)
644
- image = np.array(image)
645
- image = HWC3(image)
646
- image = resize_image(image, self.resolution)
647
- image = Image.fromarray(image)
648
- depth = self.depth_estimator(image)['depth']
649
- depth = np.array(depth)
650
- depth = depth[:, :, None]
651
- depth = np.concatenate([depth, depth, depth], axis=2)
652
- depth = Image.fromarray(depth)
653
- updated_image_path = get_new_image_name(inputs, func_name="depth")
654
- depth.save(updated_image_path)
655
- return updated_image_path
656
-
657
- class depth2image:
658
- def __init__(self, device):
659
- self.controlnet = ControlNetModel.from_pretrained(
660
- "fusing/stable-diffusion-v1-5-controlnet-depth"
661
- )
662
-
663
- self.pipe = StableDiffusionControlNetPipeline.from_pretrained(
664
- "runwayml/stable-diffusion-v1-5", controlnet=self.controlnet, safety_checker=None
665
- )
666
-
667
- self.pipe.scheduler = UniPCMultistepScheduler.from_config(self.pipe.scheduler.config)
668
- self.pipe.to(device)
669
- self.image_resolution = 512
670
- self.num_inference_steps = 20
671
- self.seed = -1
672
- self.unconditional_guidance_scale = 9.0
673
- self.a_prompt = 'best quality, extremely detailed'
674
- self.n_prompt = 'longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality'
675
-
676
- def inference(self, inputs):
677
- print("===>Starting depth2image Inference")
678
- image_path, instruct_text = inputs.split(",")[0], ','.join(inputs.split(',')[1:])
679
- image = Image.open(image_path)
680
- image = np.array(image)
681
- img = resize_image(HWC3(image), self.image_resolution)
682
- img = Image.fromarray(img)
683
-
684
- self.seed = random.randint(0, 65535)
685
- seed_everything(self.seed)
686
-
687
- prompt = instruct_text
688
- prompt = prompt + ', ' + self.a_prompt
689
- image = \
690
- self.pipe(prompt, img, num_inference_steps=self.num_inference_steps, eta=0.0, negative_prompt=self.n_prompt,
691
- guidance_scale=self.unconditional_guidance_scale).images[0]
692
- updated_image_path = get_new_image_name(image_path, func_name="depth2image")
693
- image.save(updated_image_path)
694
- return updated_image_path
695
-
696
- class image2normal:
697
- def __init__(self):
698
- print("normal estimation")
699
- self.depth_estimator = pipeline("depth-estimation", model="Intel/dpt-hybrid-midas")
700
- self.resolution = 512
701
- self.bg_threhold = 0.4
702
-
703
- def inference(self, inputs):
704
- image = Image.open(inputs)
705
- image = np.array(image)
706
- image = HWC3(image)
707
- image = resize_image(image, self.resolution)
708
- image = Image.fromarray(image)
709
- image = self.depth_estimator(image)['predicted_depth'][0]
710
-
711
- image = image.numpy()
712
-
713
- image_depth = image.copy()
714
- image_depth -= np.min(image_depth)
715
- image_depth /= np.max(image_depth)
716
-
717
- bg_threhold = 0.4
718
-
719
- x = cv2.Sobel(image, cv2.CV_32F, 1, 0, ksize=3)
720
- x[image_depth < bg_threhold] = 0
721
-
722
- y = cv2.Sobel(image, cv2.CV_32F, 0, 1, ksize=3)
723
- y[image_depth < bg_threhold] = 0
724
-
725
- z = np.ones_like(x) * np.pi * 2.0
726
-
727
- image = np.stack([x, y, z], axis=2)
728
- image /= np.sum(image ** 2.0, axis=2, keepdims=True) ** 0.5
729
- image = (image * 127.5 + 127.5).clip(0, 255).astype(np.uint8)
730
- image = Image.fromarray(image)
731
- updated_image_path = get_new_image_name(inputs, func_name="normal-map")
732
- image.save(updated_image_path)
733
- return updated_image_path
734
-
735
- class normal2image:
736
- def __init__(self, device):
737
- self.controlnet = ControlNetModel.from_pretrained(
738
- "fusing/stable-diffusion-v1-5-controlnet-normal"
739
- )
740
-
741
- self.pipe = StableDiffusionControlNetPipeline.from_pretrained(
742
- "runwayml/stable-diffusion-v1-5", controlnet=self.controlnet, safety_checker=None
743
- )
744
-
745
- self.pipe.scheduler = UniPCMultistepScheduler.from_config(self.pipe.scheduler.config)
746
- self.pipe.to(device)
747
- self.image_resolution = 512
748
- self.num_inference_steps = 20
749
- self.seed = -1
750
- self.unconditional_guidance_scale = 9.0
751
- self.a_prompt = 'best quality, extremely detailed'
752
- self.n_prompt = 'longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality'
753
-
754
- def inference(self, inputs):
755
- print("===>Starting normal2image Inference")
756
- image_path, instruct_text = inputs.split(",")[0], ','.join(inputs.split(',')[1:])
757
- image = Image.open(image_path)
758
- image = np.array(image)
759
- img = resize_image(HWC3(image), self.image_resolution)
760
- img = Image.fromarray(img)
761
-
762
- self.seed = random.randint(0, 65535)
763
- seed_everything(self.seed)
764
-
765
- prompt = instruct_text
766
- prompt = prompt + ', ' + self.a_prompt
767
- image = \
768
- self.pipe(prompt, img, num_inference_steps=self.num_inference_steps, eta=0.0, negative_prompt=self.n_prompt,
769
- guidance_scale=self.unconditional_guidance_scale).images[0]
770
- updated_image_path = get_new_image_name(image_path, func_name="normal2image")
771
- image.save(updated_image_path)
772
- return updated_image_path
773
-
774
- class BLIPVQA:
775
- def __init__(self, device):
776
- print("Initializing BLIP VQA to %s" % device)
777
- self.device = device
778
- self.processor = BlipProcessor.from_pretrained("Salesforce/blip-vqa-base")
779
- self.model = BlipForQuestionAnswering.from_pretrained("Salesforce/blip-vqa-base").to(self.device)
780
-
781
- def get_answer_from_question_and_image(self, inputs):
782
- image_path, question = inputs.split(",")
783
- raw_image = Image.open(image_path).convert('RGB')
784
- print(F'BLIPVQA :question :{question}')
785
- inputs = self.processor(raw_image, question, return_tensors="pt").to(self.device)
786
- out = self.model.generate(**inputs)
787
- answer = self.processor.decode(out[0], skip_special_tokens=True)
788
- return answer
789
-
790
- class ConversationBot:
791
- def __init__(self):
792
- print("Initializing VisualChatGPT")
793
- self.llm = OpenAI(temperature=0)
794
- self.edit = ImageEditing(device="cuda:0")
795
- self.i2t = ImageCaptioning(device="cuda:0")
796
- self.t2i = T2I(device="cuda:0")
797
- self.image2canny = image2canny()
798
- self.canny2image = canny2image(device="cuda:1")
799
- self.image2line = image2line()
800
- self.line2image = line2image(device="cuda:1")
801
- self.image2hed = image2hed()
802
- self.hed2image = hed2image(device="cuda:1")
803
- self.image2scribble = image2scribble()
804
- self.scribble2image = scribble2image(device="cuda:2")
805
- self.image2pose = image2pose()
806
- self.pose2image = pose2image(device="cuda:2")
807
- self.BLIPVQA = BLIPVQA(device="cuda:2")
808
- self.image2seg = image2seg()
809
- self.seg2image = seg2image(device="cuda:3")
810
- self.image2depth = image2depth()
811
- self.depth2image = depth2image(device="cuda:3")
812
- self.image2normal = image2normal()
813
- self.normal2image = normal2image(device="cuda:3")
814
- self.pix2pix = Pix2Pix(device="cuda:3")
815
- self.memory = ConversationBufferMemory(memory_key="chat_history", output_key='output')
816
- self.tools = [
817
- Tool(name="Get Photo Description", func=self.i2t.inference,
818
- description="useful for when you want to know what is inside the photo. receives image_path as input. "
819
- "The input to this tool should be a string, representing the image_path. "),
820
- Tool(name="Generate Image From User Input Text", func=self.t2i.inference,
821
- description="useful for when you want to generate an image from a user input text and it saved it to a file. like: generate an image of an object or something, or generate an image that includes some objects. "
822
- "The input to this tool should be a string, representing the text used to generate image. "),
823
- Tool(name="Remove Something From The Photo", func=self.edit.remove_part_of_image,
824
- description="useful for when you want to remove and object or something from the photo from its description or location. "
825
- "The input to this tool should be a comma seperated string of two, representing the image_path and the object need to be removed. "),
826
- Tool(name="Replace Something From The Photo", func=self.edit.replace_part_of_image,
827
- description="useful for when you want to replace an object from the object description or location with another object from its description. "
828
- "The input to this tool should be a comma seperated string of three, representing the image_path, the object to be replaced, the object to be replaced with "),
829
- Tool(name="Instruct Image Using Text", func=self.pix2pix.inference,
830
- description="useful for when you want to the style of the image to be like the text. like: make it look like a painting. or make it like a robot. "
831
- "The input to this tool should be a comma seperated string of two, representing the image_path and the text. "),
832
- Tool(name="Answer Question About The Image", func=self.BLIPVQA.get_answer_from_question_and_image,
833
- description="useful for when you need an answer for a question based on an image. like: what is the background color of the last image, how many cats in this figure, what is in this figure. "
834
- "The input to this tool should be a comma seperated string of two, representing the image_path and the question"),
835
- Tool(name="Edge Detection On Image", func=self.image2canny.inference,
836
- description="useful for when you want to detect the edge of the image. like: detect the edges of this image, or canny detection on image, or peform edge detection on this image, or detect the canny image of this image. "
837
- "The input to this tool should be a string, representing the image_path"),
838
- Tool(name="Generate Image Condition On Canny Image", func=self.canny2image.inference,
839
- description="useful for when you want to generate a new real image from both the user desciption and a canny image. like: generate a real image of a object or something from this canny image, or generate a new real image of a object or something from this edge image. "
840
- "The input to this tool should be a comma seperated string of two, representing the image_path and the user description. "),
841
- Tool(name="Line Detection On Image", func=self.image2line.inference,
842
- description="useful for when you want to detect the straight line of the image. like: detect the straight lines of this image, or straight line detection on image, or peform straight line detection on this image, or detect the straight line image of this image. "
843
- "The input to this tool should be a string, representing the image_path"),
844
- Tool(name="Generate Image Condition On Line Image", func=self.line2image.inference,
845
- description="useful for when you want to generate a new real image from both the user desciption and a straight line image. like: generate a real image of a object or something from this straight line image, or generate a new real image of a object or something from this straight lines. "
846
- "The input to this tool should be a comma seperated string of two, representing the image_path and the user description. "),
847
- Tool(name="Hed Detection On Image", func=self.image2hed.inference,
848
- description="useful for when you want to detect the soft hed boundary of the image. like: detect the soft hed boundary of this image, or hed boundary detection on image, or peform hed boundary detection on this image, or detect soft hed boundary image of this image. "
849
- "The input to this tool should be a string, representing the image_path"),
850
- Tool(name="Generate Image Condition On Soft Hed Boundary Image", func=self.hed2image.inference,
851
- description="useful for when you want to generate a new real image from both the user desciption and a soft hed boundary image. like: generate a real image of a object or something from this soft hed boundary image, or generate a new real image of a object or something from this hed boundary. "
852
- "The input to this tool should be a comma seperated string of two, representing the image_path and the user description"),
853
- Tool(name="Segmentation On Image", func=self.image2seg.inference,
854
- description="useful for when you want to detect segmentations of the image. like: segment this image, or generate segmentations on this image, or peform segmentation on this image. "
855
- "The input to this tool should be a string, representing the image_path"),
856
- Tool(name="Generate Image Condition On Segmentations", func=self.seg2image.inference,
857
- description="useful for when you want to generate a new real image from both the user desciption and segmentations. like: generate a real image of a object or something from this segmentation image, or generate a new real image of a object or something from these segmentations. "
858
- "The input to this tool should be a comma seperated string of two, representing the image_path and the user description"),
859
- Tool(name="Predict Depth On Image", func=self.image2depth.inference,
860
- description="useful for when you want to detect depth of the image. like: generate the depth from this image, or detect the depth map on this image, or predict the depth for this image. "
861
- "The input to this tool should be a string, representing the image_path"),
862
- Tool(name="Generate Image Condition On Depth", func=self.depth2image.inference,
863
- description="useful for when you want to generate a new real image from both the user desciption and depth image. like: generate a real image of a object or something from this depth image, or generate a new real image of a object or something from the depth map. "
864
- "The input to this tool should be a comma seperated string of two, representing the image_path and the user description"),
865
- Tool(name="Predict Normal Map On Image", func=self.image2normal.inference,
866
- description="useful for when you want to detect norm map of the image. like: generate normal map from this image, or predict normal map of this image. "
867
- "The input to this tool should be a string, representing the image_path"),
868
- Tool(name="Generate Image Condition On Normal Map", func=self.normal2image.inference,
869
- description="useful for when you want to generate a new real image from both the user desciption and normal map. like: generate a real image of a object or something from this normal map, or generate a new real image of a object or something from the normal map. "
870
- "The input to this tool should be a comma seperated string of two, representing the image_path and the user description"),
871
- Tool(name="Sketch Detection On Image", func=self.image2scribble.inference,
872
- description="useful for when you want to generate a scribble of the image. like: generate a scribble of this image, or generate a sketch from this image, detect the sketch from this image. "
873
- "The input to this tool should be a string, representing the image_path"),
874
- Tool(name="Generate Image Condition On Sketch Image", func=self.scribble2image.inference,
875
- description="useful for when you want to generate a new real image from both the user desciption and a scribble image or a sketch image. "
876
- "The input to this tool should be a comma seperated string of two, representing the image_path and the user description"),
877
- Tool(name="Pose Detection On Image", func=self.image2pose.inference,
878
- description="useful for when you want to detect the human pose of the image. like: generate human poses of this image, or generate a pose image from this image. "
879
- "The input to this tool should be a string, representing the image_path"),
880
- Tool(name="Generate Image Condition On Pose Image", func=self.pose2image.inference,
881
- description="useful for when you want to generate a new real image from both the user desciption and a human pose image. like: generate a real image of a human from this human pose image, or generate a new real image of a human from this pose. "
882
- "The input to this tool should be a comma seperated string of two, representing the image_path and the user description")
883
- ]
884
- self.agent = initialize_agent(
885
- self.tools,
886
- self.llm,
887
- agent="conversational-react-description",
888
- verbose=True,
889
- memory=self.memory,
890
- return_intermediate_steps=True,
891
- agent_kwargs={'prefix': VISUAL_CHATGPT_PREFIX, 'format_instructions': VISUAL_CHATGPT_FORMAT_INSTRUCTIONS, 'suffix': VISUAL_CHATGPT_SUFFIX}, )
892
-
893
- def run_text(self, text, state):
894
- print("===============Running run_text =============")
895
- print("Inputs:", text, state)
896
- print("======>Previous memory:\n %s" % self.agent.memory)
897
- self.agent.memory.buffer = cut_dialogue_history(self.agent.memory.buffer, keep_last_n_words=500)
898
- res = self.agent({"input": text})
899
- print("======>Current memory:\n %s" % self.agent.memory)
900
- response = re.sub('(image/\S*png)', lambda m: f'![](/file={m.group(0)})*{m.group(0)}*', res['output'])
901
- state = state + [(text, response)]
902
- print("Outputs:", state)
903
- return state, state
904
-
905
- def run_image(self, image, state, txt):
906
- print("===============Running run_image =============")
907
- print("Inputs:", image, state)
908
- print("======>Previous memory:\n %s" % self.agent.memory)
909
- image_filename = os.path.join('image', str(uuid.uuid4())[0:8] + ".png")
910
- print("======>Auto Resize Image...")
911
- img = Image.open(image.name)
912
- width, height = img.size
913
- ratio = min(512 / width, 512 / height)
914
- width_new, height_new = (round(width * ratio), round(height * ratio))
915
- img = img.resize((width_new, height_new))
916
- img = img.convert('RGB')
917
- img.save(image_filename, "PNG")
918
- print(f"Resize image form {width}x{height} to {width_new}x{height_new}")
919
- description = self.i2t.inference(image_filename)
920
- Human_prompt = "\nHuman: provide a figure named {}. The description is: {}. This information helps you to understand this image, but you should use tools to finish following tasks, " \
921
- "rather than directly imagine from my description. If you understand, say \"Received\". \n".format(image_filename, description)
922
- AI_prompt = "Received. "
923
- self.agent.memory.buffer = self.agent.memory.buffer + Human_prompt + 'AI: ' + AI_prompt
924
- print("======>Current memory:\n %s" % self.agent.memory)
925
- state = state + [(f"![](/file={image_filename})*{image_filename}*", AI_prompt)]
926
- print("Outputs:", state)
927
- return state, state, txt + ' ' + image_filename + ' '
928
-
929
- bot = ConversationBot()
930
- with gr.Blocks(css="#chatbot .overflow-y-auto{height:500px}") as demo:
931
- chatbot = gr.Chatbot(elem_id="chatbot", label="Visual ChatGPT")
932
- state = gr.State([])
933
-
934
- with gr.Row():
935
- with gr.Column(scale=0.7):
936
- txt = gr.Textbox(show_label=False, placeholder="Enter text and press enter, or upload an image").style(container=False)
937
- with gr.Column(scale=0.15, min_width=0):
938
- clear = gr.Button("Clear️")
939
- with gr.Column(scale=0.15, min_width=0):
940
- btn = gr.UploadButton("Upload", file_types=["image"])
941
-
942
- txt.submit(bot.run_text, [txt, state], [chatbot, state])
943
- txt.submit(lambda: "", None, txt)
944
-
945
- btn.upload(bot.run_image, [btn, state, txt], [chatbot, state, txt])
946
-
947
- clear.click(bot.memory.clear)
948
- clear.click(lambda: [], None, chatbot)
949
- clear.click(lambda: [], None, state)
950
-
951
-
952
- demo.launch(server_name="0.0.0.0", server_port=7860)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
visual_foundation_models.py CHANGED
@@ -1,6 +1,5 @@
1
- from diffusers import StableDiffusionPipeline
2
- from diffusers import StableDiffusionInpaintPipeline
3
- from diffusers import StableDiffusionInstructPix2PixPipeline, EulerAncestralDiscreteScheduler
4
  from diffusers import StableDiffusionControlNetPipeline, ControlNetModel, UniPCMultistepScheduler
5
  from controlnet_aux import OpenposeDetector, MLSDdetector, HEDdetector
6
 
@@ -8,84 +7,14 @@ from transformers import AutoModelForCausalLM, AutoTokenizer, CLIPSegProcessor,
8
  from transformers import pipeline, BlipProcessor, BlipForConditionalGeneration, BlipForQuestionAnswering
9
  from transformers import AutoImageProcessor, UperNetForSemanticSegmentation
10
 
11
- from PIL import Image
 
12
  import torch
13
- import numpy as np
14
  import uuid
 
 
15
  from pytorch_lightning import seed_everything
16
- import cv2
17
- import random
18
- import os
19
-
20
- def ade_palette():
21
- return [[120, 120, 120], [180, 120, 120], [6, 230, 230], [80, 50, 50],
22
- [4, 200, 3], [120, 120, 80], [140, 140, 140], [204, 5, 255],
23
- [230, 230, 230], [4, 250, 7], [224, 5, 255], [235, 255, 7],
24
- [150, 5, 61], [120, 120, 70], [8, 255, 51], [255, 6, 82],
25
- [143, 255, 140], [204, 255, 4], [255, 51, 7], [204, 70, 3],
26
- [0, 102, 200], [61, 230, 250], [255, 6, 51], [11, 102, 255],
27
- [255, 7, 71], [255, 9, 224], [9, 7, 230], [220, 220, 220],
28
- [255, 9, 92], [112, 9, 255], [8, 255, 214], [7, 255, 224],
29
- [255, 184, 6], [10, 255, 71], [255, 41, 10], [7, 255, 255],
30
- [224, 255, 8], [102, 8, 255], [255, 61, 6], [255, 194, 7],
31
- [255, 122, 8], [0, 255, 20], [255, 8, 41], [255, 5, 153],
32
- [6, 51, 255], [235, 12, 255], [160, 150, 20], [0, 163, 255],
33
- [140, 140, 140], [250, 10, 15], [20, 255, 0], [31, 255, 0],
34
- [255, 31, 0], [255, 224, 0], [153, 255, 0], [0, 0, 255],
35
- [255, 71, 0], [0, 235, 255], [0, 173, 255], [31, 0, 255],
36
- [11, 200, 200], [255, 82, 0], [0, 255, 245], [0, 61, 255],
37
- [0, 255, 112], [0, 255, 133], [255, 0, 0], [255, 163, 0],
38
- [255, 102, 0], [194, 255, 0], [0, 143, 255], [51, 255, 0],
39
- [0, 82, 255], [0, 255, 41], [0, 255, 173], [10, 0, 255],
40
- [173, 255, 0], [0, 255, 153], [255, 92, 0], [255, 0, 255],
41
- [255, 0, 245], [255, 0, 102], [255, 173, 0], [255, 0, 20],
42
- [255, 184, 184], [0, 31, 255], [0, 255, 61], [0, 71, 255],
43
- [255, 0, 204], [0, 255, 194], [0, 255, 82], [0, 10, 255],
44
- [0, 112, 255], [51, 0, 255], [0, 194, 255], [0, 122, 255],
45
- [0, 255, 163], [255, 153, 0], [0, 255, 10], [255, 112, 0],
46
- [143, 255, 0], [82, 0, 255], [163, 255, 0], [255, 235, 0],
47
- [8, 184, 170], [133, 0, 255], [0, 255, 92], [184, 0, 255],
48
- [255, 0, 31], [0, 184, 255], [0, 214, 255], [255, 0, 112],
49
- [92, 255, 0], [0, 224, 255], [112, 224, 255], [70, 184, 160],
50
- [163, 0, 255], [153, 0, 255], [71, 255, 0], [255, 0, 163],
51
- [255, 204, 0], [255, 0, 143], [0, 255, 235], [133, 255, 0],
52
- [255, 0, 235], [245, 0, 255], [255, 0, 122], [255, 245, 0],
53
- [10, 190, 212], [214, 255, 0], [0, 204, 255], [20, 0, 255],
54
- [255, 255, 0], [0, 153, 255], [0, 41, 255], [0, 255, 204],
55
- [41, 0, 255], [41, 255, 0], [173, 0, 255], [0, 245, 255],
56
- [71, 0, 255], [122, 0, 255], [0, 255, 184], [0, 92, 255],
57
- [184, 255, 0], [0, 133, 255], [255, 214, 0], [25, 194, 194],
58
- [102, 255, 0], [92, 0, 255]]
59
-
60
- def HWC3(x):
61
- assert x.dtype == np.uint8
62
- if x.ndim == 2:
63
- x = x[:, :, None]
64
- assert x.ndim == 3
65
- H, W, C = x.shape
66
- assert C == 1 or C == 3 or C == 4
67
- if C == 3:
68
- return x
69
- if C == 1:
70
- return np.concatenate([x, x, x], axis=2)
71
- if C == 4:
72
- color = x[:, :, 0:3].astype(np.float32)
73
- alpha = x[:, :, 3:4].astype(np.float32) / 255.0
74
- y = color * alpha + 255.0 * (1.0 - alpha)
75
- y = y.clip(0, 255).astype(np.uint8)
76
- return y
77
-
78
- def resize_image(input_image, resolution):
79
- H, W, C = input_image.shape
80
- H = float(H)
81
- W = float(W)
82
- k = float(resolution) / min(H, W)
83
- H *= k
84
- W *= k
85
- H = int(np.round(H / 64.0)) * 64
86
- W = int(np.round(W / 64.0)) * 64
87
- img = cv2.resize(input_image, (W, H), interpolation=cv2.INTER_LANCZOS4 if k > 1 else cv2.INTER_AREA)
88
- return img
89
 
90
  def get_new_image_name(org_img_name, func_name="update"):
91
  head_tail = os.path.split(org_img_name)
@@ -104,6 +33,7 @@ def get_new_image_name(org_img_name, func_name="update"):
104
  new_file_name = '{}_{}_{}_{}.png'.format(this_new_uuid, func_name, recent_prev_file_name, most_org_file_name)
105
  return os.path.join(head, new_file_name)
106
 
 
107
  class MaskFormer:
108
  def __init__(self, device):
109
  self.device = device
@@ -130,7 +60,7 @@ class MaskFormer:
130
  mask_array[padded_slice] = True
131
  visual_mask = (mask_array * 255).astype(np.uint8)
132
  image_mask = Image.fromarray(visual_mask)
133
- return image_mask.resize(image.size)
134
 
135
  class ImageEditing:
136
  def __init__(self, device):
@@ -148,9 +78,11 @@ class ImageEditing:
148
  image_path, to_be_replaced_txt, replace_with_txt = input.split(",")
149
  print(f'replace_part_of_image: replace_with_txt {replace_with_txt}')
150
  original_image = Image.open(image_path)
 
151
  mask_image = self.mask_former.inference(image_path, to_be_replaced_txt)
152
- updated_image = self.inpainting(prompt=replace_with_txt, image=original_image, mask_image=mask_image).images[0]
153
  updated_image_path = get_new_image_name(image_path, func_name="replace-something")
 
154
  updated_image.save(updated_image_path)
155
  return updated_image_path
156
 
@@ -176,7 +108,7 @@ class T2I:
176
  print("Initializing T2I to %s" % device)
177
  self.device = device
178
  self.pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16)
179
- self.text_refine_tokenizer = AutoTokenizer.from_pretrained("Gustavosta/MagicPrompt-Stable-Diffusion")
180
  self.text_refine_model = AutoModelForCausalLM.from_pretrained("Gustavosta/MagicPrompt-Stable-Diffusion", torch_dtype=torch.float16)
181
  self.text_refine_gpt2_pipe = pipeline("text-generation", model=self.text_refine_model, tokenizer=self.text_refine_tokenizer, device=self.device, torch_dtype=torch.float16)
182
  self.pipe.to(device)
@@ -203,7 +135,7 @@ class ImageCaptioning:
203
  captions = self.processor.decode(out[0], skip_special_tokens=True)
204
  return captions
205
 
206
- class image2canny_new:
207
  def __init__(self):
208
  print("Direct detect canny.")
209
  self.low_threshold = 100
@@ -216,28 +148,20 @@ class image2canny_new:
216
  canny = cv2.Canny(image, self.low_threshold, self.high_threshold)
217
  canny = canny[:, :, None]
218
  canny = np.concatenate([canny, canny, canny], axis=2)
219
- canny = 255 - canny
220
  canny = Image.fromarray(canny)
221
  updated_image_path = get_new_image_name(inputs, func_name="edge")
222
  canny.save(updated_image_path)
223
  return updated_image_path
224
 
225
- class canny2image_new:
226
  def __init__(self, device):
227
- self.controlnet = ControlNetModel.from_pretrained(
228
- "fusing/stable-diffusion-v1-5-controlnet-canny", torch_dtype=torch.float16
229
- )
230
-
231
  self.pipe = StableDiffusionControlNetPipeline.from_pretrained(
232
  "runwayml/stable-diffusion-v1-5", controlnet=self.controlnet, safety_checker=None, torch_dtype=torch.float16
233
  )
234
-
235
  self.pipe.scheduler = UniPCMultistepScheduler.from_config(self.pipe.scheduler.config)
236
  self.pipe.to(device)
237
- self.image_resolution = 512
238
- self.num_inference_steps = 20
239
  self.seed = -1
240
- self.unconditional_guidance_scale = 9.0
241
  self.a_prompt = 'best quality, extremely detailed'
242
  self.n_prompt = 'longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality'
243
 
@@ -245,124 +169,36 @@ class canny2image_new:
245
  print("===>Starting canny2image Inference")
246
  image_path, instruct_text = inputs.split(",")[0], ','.join(inputs.split(',')[1:])
247
  image = Image.open(image_path)
248
- image = np.array(image)
249
- image = 255 - image
250
- prompt = instruct_text
251
- img = resize_image(HWC3(image), self.image_resolution)
252
- img = Image.fromarray(img)
253
-
254
  self.seed = random.randint(0, 65535)
255
  seed_everything(self.seed)
256
- prompt = prompt + ', ' + self.a_prompt
257
- image = self.pipe(prompt, img, num_inference_steps=self.num_inference_steps, eta=0.0, negative_prompt=self.n_prompt, guidance_scale=self.unconditional_guidance_scale).images[0]
258
  updated_image_path = get_new_image_name(image_path, func_name="canny2image")
259
  image.save(updated_image_path)
260
  return updated_image_path
261
 
262
-
263
- # class image2canny:
264
- # def __init__(self):
265
- # print("Direct detect canny.")
266
- # self.detector = CannyDetector()
267
- # self.low_thresh = 100
268
- # self.high_thresh = 200
269
- #
270
- # def inference(self, inputs):
271
- # print("===>Starting image2canny Inference")
272
- # image = Image.open(inputs)
273
- # image = np.array(image)
274
- # canny = self.detector(image, self.low_thresh, self.high_thresh)
275
- # canny = 255 - canny
276
- # image = Image.fromarray(canny)
277
- # updated_image_path = get_new_image_name(inputs, func_name="edge")
278
- # image.save(updated_image_path)
279
- # return updated_image_path
280
- #
281
- # class canny2image:
282
- # def __init__(self, device):
283
- # print("Initialize the canny2image model.")
284
- # model = create_model('ControlNet/models/cldm_v15.yaml', device=device).to(device)
285
- # model.load_state_dict(load_state_dict('ControlNet/models/control_sd15_canny.pth', location='cpu'))
286
- # self.model = model.to(device)
287
- # self.device = device
288
- # self.ddim_sampler = DDIMSampler(self.model)
289
- # self.ddim_steps = 20
290
- # self.image_resolution = 512
291
- # self.num_samples = 1
292
- # self.save_memory = False
293
- # self.strength = 1.0
294
- # self.guess_mode = False
295
- # self.scale = 9.0
296
- # self.seed = -1
297
- # self.a_prompt = 'best quality, extremely detailed'
298
- # self.n_prompt = 'longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality'
299
- #
300
- # def inference(self, inputs):
301
- # print("===>Starting canny2image Inference")
302
- # image_path, instruct_text = inputs.split(",")[0], ','.join(inputs.split(',')[1:])
303
- # image = Image.open(image_path)
304
- # image = np.array(image)
305
- # image = 255 - image
306
- # prompt = instruct_text
307
- # img = resize_image(HWC3(image), self.image_resolution)
308
- # H, W, C = img.shape
309
- # control = torch.from_numpy(img.copy()).float().to(device=self.device) / 255.0
310
- # control = torch.stack([control for _ in range(self.num_samples)], dim=0)
311
- # control = einops.rearrange(control, 'b h w c -> b c h w').clone()
312
- # self.seed = random.randint(0, 65535)
313
- # seed_everything(self.seed)
314
- # if self.save_memory:
315
- # self.model.low_vram_shift(is_diffusing=False)
316
- # cond = {"c_concat": [control], "c_crossattn": [self.model.get_learned_conditioning([prompt + ', ' + self.a_prompt] * self.num_samples)]}
317
- # un_cond = {"c_concat": None if self.guess_mode else [control], "c_crossattn": [self.model.get_learned_conditioning([self.n_prompt] * self.num_samples)]}
318
- # shape = (4, H // 8, W // 8)
319
- # self.model.control_scales = [self.strength * (0.825 ** float(12 - i)) for i in range(13)] if self.guess_mode else ([self.strength] * 13) # Magic number. IDK why. Perhaps because 0.825**12<0.01 but 0.826**12>0.01
320
- # samples, intermediates = self.ddim_sampler.sample(self.ddim_steps, self.num_samples, shape, cond, verbose=False, eta=0., unconditional_guidance_scale=self.scale, unconditional_conditioning=un_cond)
321
- # if self.save_memory:
322
- # self.model.low_vram_shift(is_diffusing=False)
323
- # x_samples = self.model.decode_first_stage(samples)
324
- # x_samples = (einops.rearrange(x_samples, 'b c h w -> b h w c') * 127.5 + 127.5).cpu().numpy().clip(0, 255).astype(np.uint8)
325
- # updated_image_path = get_new_image_name(image_path, func_name="canny2image")
326
- # real_image = Image.fromarray(x_samples[0]) # get default the index0 image
327
- # real_image.save(updated_image_path)
328
- # return updated_image_path
329
- class image2line_new:
330
  def __init__(self):
331
  self.detector = MLSDdetector.from_pretrained('lllyasviel/ControlNet')
332
- self.value_thresh = 0.1
333
- self.dis_thresh = 0.1
334
- self.resolution = 512
335
 
336
  def inference(self, inputs):
337
  print("===>Starting image2line Inference")
338
  image = Image.open(inputs)
339
- image = np.array(image)
340
- image = HWC3(image)
341
- mlsd = self.detector(resize_image(image, self.resolution), thr_v=self.value_thresh, thr_d=self.dis_thresh)
342
- mlsd = np.array(mlsd)
343
- mlsd = 255 - mlsd
344
- mlsd = Image.fromarray(mlsd)
345
  updated_image_path = get_new_image_name(inputs, func_name="line-of")
346
  mlsd.save(updated_image_path)
347
  return updated_image_path
348
 
349
- class line2image_new:
350
  def __init__(self, device):
351
  print("Initialize the line2image model...")
352
- self.controlnet = ControlNetModel.from_pretrained(
353
- "fusing/stable-diffusion-v1-5-controlnet-mlsd"
354
- )
355
-
356
  self.pipe = StableDiffusionControlNetPipeline.from_pretrained(
357
- "runwayml/stable-diffusion-v1-5", controlnet=self.controlnet, safety_checker=None
358
  )
359
-
360
  self.pipe.scheduler = UniPCMultistepScheduler.from_config(self.pipe.scheduler.config)
361
  self.pipe.to(device)
362
- self.image_resolution = 512
363
- self.num_inference_steps = 20
364
  self.seed = -1
365
- self.unconditional_guidance_scale = 9.0
366
  self.a_prompt = 'best quality, extremely detailed'
367
  self.n_prompt = 'longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality'
368
 
@@ -370,129 +206,37 @@ class line2image_new:
370
  print("===>Starting line2image Inference")
371
  image_path, instruct_text = inputs.split(",")[0], ','.join(inputs.split(',')[1:])
372
  image = Image.open(image_path)
373
- image = np.array(image)
374
- image = 255 - image
375
- prompt = instruct_text
376
- img = resize_image(HWC3(image), self.image_resolution)
377
- img = Image.fromarray(img)
378
-
379
  self.seed = random.randint(0, 65535)
380
  seed_everything(self.seed)
381
-
382
- prompt = prompt + ', ' + self.a_prompt
383
- image = self.pipe(prompt, img, num_inference_steps=self.num_inference_steps, eta=0.0, negative_prompt=self.n_prompt, guidance_scale=self.unconditional_guidance_scale).images[0]
384
  updated_image_path = get_new_image_name(image_path, func_name="line2image")
385
  image.save(updated_image_path)
386
  return updated_image_path
387
 
388
-
389
- # class image2line:
390
- # def __init__(self):
391
- # print("Direct detect straight line...")
392
- # self.detector = MLSDdetector()
393
- # self.value_thresh = 0.1
394
- # self.dis_thresh = 0.1
395
- # self.resolution = 512
396
- #
397
- # def inference(self, inputs):
398
- # print("===>Starting image2hough Inference")
399
- # image = Image.open(inputs)
400
- # image = np.array(image)
401
- # image = HWC3(image)
402
- # hough = self.detector(resize_image(image, self.resolution), self.value_thresh, self.dis_thresh)
403
- # updated_image_path = get_new_image_name(inputs, func_name="line-of")
404
- # hough = 255 - cv2.dilate(hough, np.ones(shape=(3, 3), dtype=np.uint8), iterations=1)
405
- # image = Image.fromarray(hough)
406
- # image.save(updated_image_path)
407
- # return updated_image_path
408
- #
409
- #
410
- # class line2image:
411
- # def __init__(self, device):
412
- # print("Initialize the line2image model...")
413
- # model = create_model('ControlNet/models/cldm_v15.yaml', device=device).to(device)
414
- # model.load_state_dict(load_state_dict('ControlNet/models/control_sd15_mlsd.pth', location='cpu'))
415
- # self.model = model.to(device)
416
- # self.device = device
417
- # self.ddim_sampler = DDIMSampler(self.model)
418
- # self.ddim_steps = 20
419
- # self.image_resolution = 512
420
- # self.num_samples = 1
421
- # self.save_memory = False
422
- # self.strength = 1.0
423
- # self.guess_mode = False
424
- # self.scale = 9.0
425
- # self.seed = -1
426
- # self.a_prompt = 'best quality, extremely detailed'
427
- # self.n_prompt = 'longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality'
428
- #
429
- # def inference(self, inputs):
430
- # print("===>Starting line2image Inference")
431
- # image_path, instruct_text = inputs.split(",")[0], ','.join(inputs.split(',')[1:])
432
- # image = Image.open(image_path)
433
- # image = np.array(image)
434
- # image = 255 - image
435
- # prompt = instruct_text
436
- # img = resize_image(HWC3(image), self.image_resolution)
437
- # H, W, C = img.shape
438
- # img = cv2.resize(img, (W, H), interpolation=cv2.INTER_NEAREST)
439
- # control = torch.from_numpy(img.copy()).float().to(device=self.device) / 255.0
440
- # control = torch.stack([control for _ in range(self.num_samples)], dim=0)
441
- # control = einops.rearrange(control, 'b h w c -> b c h w').clone()
442
- # self.seed = random.randint(0, 65535)
443
- # seed_everything(self.seed)
444
- # if self.save_memory:
445
- # self.model.low_vram_shift(is_diffusing=False)
446
- # cond = {"c_concat": [control], "c_crossattn": [self.model.get_learned_conditioning([prompt + ', ' + self.a_prompt] * self.num_samples)]}
447
- # un_cond = {"c_concat": None if self.guess_mode else [control], "c_crossattn": [self.model.get_learned_conditioning([self.n_prompt] * self.num_samples)]}
448
- # shape = (4, H // 8, W // 8)
449
- # self.model.control_scales = [self.strength * (0.825 ** float(12 - i)) for i in range(13)] if self.guess_mode else ([self.strength] * 13) # Magic number. IDK why. Perhaps because 0.825**12<0.01 but 0.826**12>0.01
450
- # samples, intermediates = self.ddim_sampler.sample(self.ddim_steps, self.num_samples, shape, cond, verbose=False, eta=0., unconditional_guidance_scale=self.scale, unconditional_conditioning=un_cond)
451
- # if self.save_memory:
452
- # self.model.low_vram_shift(is_diffusing=False)
453
- # x_samples = self.model.decode_first_stage(samples)
454
- # x_samples = (einops.rearrange(x_samples, 'b c h w -> b h w c') * 127.5 + 127.5).\
455
- # cpu().numpy().clip(0,255).astype(np.uint8)
456
- # updated_image_path = get_new_image_name(image_path, func_name="line2image")
457
- # real_image = Image.fromarray(x_samples[0]) # default the index0 image
458
- # real_image.save(updated_image_path)
459
- # return updated_image_path
460
-
461
- class image2hed_new:
462
  def __init__(self):
463
  print("Direct detect soft HED boundary...")
464
  self.detector = HEDdetector.from_pretrained('lllyasviel/ControlNet')
465
- self.resolution = 512
466
 
467
  def inference(self, inputs):
468
  print("===>Starting image2hed Inference")
469
  image = Image.open(inputs)
470
- image = np.array(image)
471
- image = HWC3(image)
472
- image = Image.fromarray(resize_image(image, self.resolution))
473
  hed = self.detector(image)
474
-
475
  updated_image_path = get_new_image_name(inputs, func_name="hed-boundary")
476
  hed.save(updated_image_path)
477
  return updated_image_path
478
 
479
- class hed2image_new:
480
  def __init__(self, device):
481
  print("Initialize the hed2image model...")
482
- self.controlnet = ControlNetModel.from_pretrained(
483
- "fusing/stable-diffusion-v1-5-controlnet-hed"
484
- )
485
-
486
  self.pipe = StableDiffusionControlNetPipeline.from_pretrained(
487
- "runwayml/stable-diffusion-v1-5", controlnet=self.controlnet, safety_checker=None
488
  )
489
-
490
  self.pipe.scheduler = UniPCMultistepScheduler.from_config(self.pipe.scheduler.config)
491
  self.pipe.to(device)
492
- self.image_resolution = 512
493
- self.num_inference_steps = 20
494
  self.seed = -1
495
- self.unconditional_guidance_scale = 9.0
496
  self.a_prompt = 'best quality, extremely detailed'
497
  self.n_prompt = 'longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality'
498
 
@@ -500,125 +244,36 @@ class hed2image_new:
500
  print("===>Starting hed2image Inference")
501
  image_path, instruct_text = inputs.split(",")[0], ','.join(inputs.split(',')[1:])
502
  image = Image.open(image_path)
503
- image = np.array(image)
504
- img = resize_image(HWC3(image), self.image_resolution)
505
- img = Image.fromarray(img)
506
-
507
  self.seed = random.randint(0, 65535)
508
  seed_everything(self.seed)
509
-
510
- prompt = instruct_text
511
- prompt = prompt + ', ' + self.a_prompt
512
- image = \
513
- self.pipe(prompt, img, num_inference_steps=self.num_inference_steps, eta=0.0, negative_prompt=self.n_prompt,
514
- guidance_scale=self.unconditional_guidance_scale).images[0]
515
  updated_image_path = get_new_image_name(image_path, func_name="hed2image")
516
  image.save(updated_image_path)
517
  return updated_image_path
518
 
519
- # class image2hed:
520
- # def __init__(self):
521
- # print("Direct detect soft HED boundary...")
522
- # self.detector = HEDdetector()
523
- # self.resolution = 512
524
- #
525
- # def inference(self, inputs):
526
- # print("===>Starting image2hed Inference")
527
- # image = Image.open(inputs)
528
- # image = np.array(image)
529
- # image = HWC3(image)
530
- # hed = self.detector(resize_image(image, self.resolution))
531
- # updated_image_path = get_new_image_name(inputs, func_name="hed-boundary")
532
- # image = Image.fromarray(hed)
533
- # image.save(updated_image_path)
534
- # return updated_image_path
535
- #
536
- #
537
- # class hed2image:
538
- # def __init__(self, device):
539
- # print("Initialize the hed2image model...")
540
- # model = create_model('ControlNet/models/cldm_v15.yaml', device=device).to(device)
541
- # model.load_state_dict(load_state_dict('ControlNet/models/control_sd15_hed.pth', location='cpu'))
542
- # self.model = model.to(device)
543
- # self.device = device
544
- # self.ddim_sampler = DDIMSampler(self.model)
545
- # self.ddim_steps = 20
546
- # self.image_resolution = 512
547
- # self.num_samples = 1
548
- # self.save_memory = False
549
- # self.strength = 1.0
550
- # self.guess_mode = False
551
- # self.scale = 9.0
552
- # self.seed = -1
553
- # self.a_prompt = 'best quality, extremely detailed'
554
- # self.n_prompt = 'longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality'
555
- #
556
- # def inference(self, inputs):
557
- # print("===>Starting hed2image Inference")
558
- # image_path, instruct_text = inputs.split(",")[0], ','.join(inputs.split(',')[1:])
559
- # image = Image.open(image_path)
560
- # image = np.array(image)
561
- # prompt = instruct_text
562
- # img = resize_image(HWC3(image), self.image_resolution)
563
- # H, W, C = img.shape
564
- # img = cv2.resize(img, (W, H), interpolation=cv2.INTER_NEAREST)
565
- # control = torch.from_numpy(img.copy()).float().to(device=self.device) / 255.0
566
- # control = torch.stack([control for _ in range(self.num_samples)], dim=0)
567
- # control = einops.rearrange(control, 'b h w c -> b c h w').clone()
568
- # self.seed = random.randint(0, 65535)
569
- # seed_everything(self.seed)
570
- # if self.save_memory:
571
- # self.model.low_vram_shift(is_diffusing=False)
572
- # cond = {"c_concat": [control], "c_crossattn": [self.model.get_learned_conditioning([prompt + ', ' + self.a_prompt] * self.num_samples)]}
573
- # un_cond = {"c_concat": None if self.guess_mode else [control], "c_crossattn": [self.model.get_learned_conditioning([self.n_prompt] * self.num_samples)]}
574
- # shape = (4, H // 8, W // 8)
575
- # self.model.control_scales = [self.strength * (0.825 ** float(12 - i)) for i in range(13)] if self.guess_mode else ([self.strength] * 13)
576
- # samples, intermediates = self.ddim_sampler.sample(self.ddim_steps, self.num_samples, shape, cond, verbose=False, eta=0., unconditional_guidance_scale=self.scale, unconditional_conditioning=un_cond)
577
- # if self.save_memory:
578
- # self.model.low_vram_shift(is_diffusing=False)
579
- # x_samples = self.model.decode_first_stage(samples)
580
- # x_samples = (einops.rearrange(x_samples, 'b c h w -> b h w c') * 127.5 + 127.5).cpu().numpy().clip(0, 255).astype(np.uint8)
581
- # updated_image_path = get_new_image_name(image_path, func_name="hed2image")
582
- # real_image = Image.fromarray(x_samples[0]) # default the index0 image
583
- # real_image.save(updated_image_path)
584
- # return updated_image_path
585
- class image2scribble_new:
586
  def __init__(self):
587
  print("Direct detect scribble.")
588
  self.detector = HEDdetector.from_pretrained('lllyasviel/ControlNet')
589
- self.resolution = 512
590
 
591
  def inference(self, inputs):
592
  print("===>Starting image2scribble Inference")
593
  image = Image.open(inputs)
594
- image = np.array(image)
595
- image = HWC3(image)
596
- image = resize_image(image, self.resolution)
597
- image = Image.fromarray(image)
598
  scribble = self.detector(image, scribble=True)
599
- scribble = np.array(scribble)
600
- scribble = 255 - scribble
601
- scribble = Image.fromarray(scribble)
602
  updated_image_path = get_new_image_name(inputs, func_name="scribble")
603
  scribble.save(updated_image_path)
604
  return updated_image_path
605
 
606
- class scribble2image_new:
607
  def __init__(self, device):
608
- self.controlnet = ControlNetModel.from_pretrained(
609
- "fusing/stable-diffusion-v1-5-controlnet-scribble"
610
- )
611
-
612
  self.pipe = StableDiffusionControlNetPipeline.from_pretrained(
613
- "runwayml/stable-diffusion-v1-5", controlnet=self.controlnet, safety_checker=None
614
  )
615
-
616
  self.pipe.scheduler = UniPCMultistepScheduler.from_config(self.pipe.scheduler.config)
617
  self.pipe.to(device)
618
- self.image_resolution = 512
619
- self.num_inference_steps = 20
620
  self.seed = -1
621
- self.unconditional_guidance_scale = 9.0
622
  self.a_prompt = 'best quality, extremely detailed'
623
  self.n_prompt = 'longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality'
624
 
@@ -626,131 +281,34 @@ class scribble2image_new:
626
  print("===>Starting scribble2image Inference")
627
  image_path, instruct_text = inputs.split(",")[0], ','.join(inputs.split(',')[1:])
628
  image = Image.open(image_path)
629
- image = np.array(image)
630
- image = 255 - image
631
- img = resize_image(HWC3(image), self.image_resolution)
632
- img = Image.fromarray(img)
633
-
634
  self.seed = random.randint(0, 65535)
635
  seed_everything(self.seed)
636
-
637
- prompt = instruct_text
638
- prompt = prompt + ', ' + self.a_prompt
639
- image = \
640
- self.pipe(prompt, img, num_inference_steps=self.num_inference_steps, eta=0.0, negative_prompt=self.n_prompt,
641
- guidance_scale=self.unconditional_guidance_scale).images[0]
642
  updated_image_path = get_new_image_name(image_path, func_name="scribble2image")
643
  image.save(updated_image_path)
644
  return updated_image_path
645
 
646
- # class image2scribble:
647
- # def __init__(self):
648
- # print("Direct detect scribble.")
649
- # self.detector = HEDdetector()
650
- # self.resolution = 512
651
- #
652
- # def inference(self, inputs):
653
- # print("===>Starting image2scribble Inference")
654
- # image = Image.open(inputs)
655
- # image = np.array(image)
656
- # image = HWC3(image)
657
- # detected_map = self.detector(resize_image(image, self.resolution))
658
- # detected_map = HWC3(detected_map)
659
- # image = resize_image(image, self.resolution)
660
- # H, W, C = image.shape
661
- # detected_map = cv2.resize(detected_map, (W, H), interpolation=cv2.INTER_LINEAR)
662
- # detected_map = nms(detected_map, 127, 3.0)
663
- # detected_map = cv2.GaussianBlur(detected_map, (0, 0), 3.0)
664
- # detected_map[detected_map > 4] = 255
665
- # detected_map[detected_map < 255] = 0
666
- # detected_map = 255 - detected_map
667
- # updated_image_path = get_new_image_name(inputs, func_name="scribble")
668
- # image = Image.fromarray(detected_map)
669
- # image.save(updated_image_path)
670
- # return updated_image_path
671
- #
672
- # class scribble2image:
673
- # def __init__(self, device):
674
- # print("Initialize the scribble2image model...")
675
- # model = create_model('ControlNet/models/cldm_v15.yaml', device=device).to(device)
676
- # model.load_state_dict(load_state_dict('ControlNet/models/control_sd15_scribble.pth', location='cpu'))
677
- # self.model = model.to(device)
678
- # self.device = device
679
- # self.ddim_sampler = DDIMSampler(self.model)
680
- # self.ddim_steps = 20
681
- # self.image_resolution = 512
682
- # self.num_samples = 1
683
- # self.save_memory = False
684
- # self.strength = 1.0
685
- # self.guess_mode = False
686
- # self.scale = 9.0
687
- # self.seed = -1
688
- # self.a_prompt = 'best quality, extremely detailed'
689
- # self.n_prompt = 'longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality'
690
- #
691
- # def inference(self, inputs):
692
- # print("===>Starting scribble2image Inference")
693
- # print(f'sketch device {self.device}')
694
- # image_path, instruct_text = inputs.split(",")[0], ','.join(inputs.split(',')[1:])
695
- # image = Image.open(image_path)
696
- # image = np.array(image)
697
- # prompt = instruct_text
698
- # image = 255 - image
699
- # img = resize_image(HWC3(image), self.image_resolution)
700
- # H, W, C = img.shape
701
- # img = cv2.resize(img, (W, H), interpolation=cv2.INTER_NEAREST)
702
- # control = torch.from_numpy(img.copy()).float().to(device=self.device) / 255.0
703
- # control = torch.stack([control for _ in range(self.num_samples)], dim=0)
704
- # control = einops.rearrange(control, 'b h w c -> b c h w').clone()
705
- # self.seed = random.randint(0, 65535)
706
- # seed_everything(self.seed)
707
- # if self.save_memory:
708
- # self.model.low_vram_shift(is_diffusing=False)
709
- # cond = {"c_concat": [control], "c_crossattn": [self.model.get_learned_conditioning([prompt + ', ' + self.a_prompt] * self.num_samples)]}
710
- # un_cond = {"c_concat": None if self.guess_mode else [control], "c_crossattn": [self.model.get_learned_conditioning([self.n_prompt] * self.num_samples)]}
711
- # shape = (4, H // 8, W // 8)
712
- # self.model.control_scales = [self.strength * (0.825 ** float(12 - i)) for i in range(13)] if self.guess_mode else ([self.strength] * 13)
713
- # samples, intermediates = self.ddim_sampler.sample(self.ddim_steps, self.num_samples, shape, cond, verbose=False, eta=0., unconditional_guidance_scale=self.scale, unconditional_conditioning=un_cond)
714
- # if self.save_memory:
715
- # self.model.low_vram_shift(is_diffusing=False)
716
- # x_samples = self.model.decode_first_stage(samples)
717
- # x_samples = (einops.rearrange(x_samples, 'b c h w -> b h w c') * 127.5 + 127.5).cpu().numpy().clip(0, 255).astype(np.uint8)
718
- # updated_image_path = get_new_image_name(image_path, func_name="scribble2image")
719
- # real_image = Image.fromarray(x_samples[0]) # default the index0 image
720
- # real_image.save(updated_image_path)
721
- # return updated_image_path
722
-
723
- class image2pose_new:
724
  def __init__(self):
725
  self.detector = OpenposeDetector.from_pretrained('lllyasviel/ControlNet')
726
- self.resolution = 512
727
 
728
  def inference(self, inputs):
729
  print("===>Starting image2pose Inference")
730
  image = Image.open(inputs)
731
- image = np.array(image)
732
- image = HWC3(image)
733
- image = resize_image(image, self.resolution)
734
- image = Image.fromarray(image)
735
  pose = self.detector(image)
736
-
737
  updated_image_path = get_new_image_name(inputs, func_name="human-pose")
738
  pose.save(updated_image_path)
739
  return updated_image_path
740
 
741
- class pose2image_new:
742
  def __init__(self, device):
743
- self.controlnet = ControlNetModel.from_pretrained(
744
- "fusing/stable-diffusion-v1-5-controlnet-openpose"
745
- )
746
-
747
  self.pipe = StableDiffusionControlNetPipeline.from_pretrained(
748
- "runwayml/stable-diffusion-v1-5", controlnet=self.controlnet, safety_checker=None
749
  )
750
-
751
  self.pipe.scheduler = UniPCMultistepScheduler.from_config(self.pipe.scheduler.config)
752
  self.pipe.to(device)
753
- self.image_resolution = 512
754
  self.num_inference_steps = 20
755
  self.seed = -1
756
  self.unconditional_guidance_scale = 9.0
@@ -761,142 +319,84 @@ class pose2image_new:
761
  print("===>Starting pose2image Inference")
762
  image_path, instruct_text = inputs.split(",")[0], ','.join(inputs.split(',')[1:])
763
  image = Image.open(image_path)
764
- image = np.array(image)
765
- img = resize_image(HWC3(image), self.image_resolution)
766
- img = Image.fromarray(img)
767
-
768
  self.seed = random.randint(0, 65535)
769
  seed_everything(self.seed)
770
-
771
- prompt = instruct_text
772
- prompt = prompt + ', ' + self.a_prompt
773
- image = \
774
- self.pipe(prompt, img, num_inference_steps=self.num_inference_steps, eta=0.0, negative_prompt=self.n_prompt,
775
- guidance_scale=self.unconditional_guidance_scale).images[0]
776
  updated_image_path = get_new_image_name(image_path, func_name="pose2image")
777
  image.save(updated_image_path)
778
  return updated_image_path
779
 
780
-
781
- # class image2pose:
782
- # def __init__(self):
783
- # print("Direct human pose.")
784
- # self.detector = OpenposeDetector()
785
- # self.resolution = 512
786
- #
787
- # def inference(self, inputs):
788
- # print("===>Starting image2pose Inference")
789
- # image = Image.open(inputs)
790
- # image = np.array(image)
791
- # image = HWC3(image)
792
- # detected_map, _ = self.detector(resize_image(image, self.resolution))
793
- # detected_map = HWC3(detected_map)
794
- # image = resize_image(image, self.resolution)
795
- # H, W, C = image.shape
796
- # detected_map = cv2.resize(detected_map, (W, H), interpolation=cv2.INTER_LINEAR)
797
- # updated_image_path = get_new_image_name(inputs, func_name="human-pose")
798
- # image = Image.fromarray(detected_map)
799
- # image.save(updated_image_path)
800
- # return updated_image_path
801
- #
802
- # class pose2image:
803
- # def __init__(self, device):
804
- # print("Initialize the pose2image model...")
805
- # model = create_model('ControlNet/models/cldm_v15.yaml', device=device).to(device)
806
- # model.load_state_dict(load_state_dict('ControlNet/models/control_sd15_openpose.pth', location='cpu'))
807
- # self.model = model.to(device)
808
- # self.device = device
809
- # self.ddim_sampler = DDIMSampler(self.model)
810
- # self.ddim_steps = 20
811
- # self.image_resolution = 512
812
- # self.num_samples = 1
813
- # self.save_memory = False
814
- # self.strength = 1.0
815
- # self.guess_mode = False
816
- # self.scale = 9.0
817
- # self.seed = -1
818
- # self.a_prompt = 'best quality, extremely detailed'
819
- # self.n_prompt = 'longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality'
820
- #
821
- # def inference(self, inputs):
822
- # print("===>Starting pose2image Inference")
823
- # image_path, instruct_text = inputs.split(",")[0], ','.join(inputs.split(',')[1:])
824
- # image = Image.open(image_path)
825
- # image = np.array(image)
826
- # prompt = instruct_text
827
- # img = resize_image(HWC3(image), self.image_resolution)
828
- # H, W, C = img.shape
829
- # img = cv2.resize(img, (W, H), interpolation=cv2.INTER_NEAREST)
830
- # control = torch.from_numpy(img.copy()).float().to(device=self.device) / 255.0
831
- # control = torch.stack([control for _ in range(self.num_samples)], dim=0)
832
- # control = einops.rearrange(control, 'b h w c -> b c h w').clone()
833
- # self.seed = random.randint(0, 65535)
834
- # seed_everything(self.seed)
835
- # if self.save_memory:
836
- # self.model.low_vram_shift(is_diffusing=False)
837
- # cond = {"c_concat": [control], "c_crossattn": [ self.model.get_learned_conditioning([prompt + ', ' + self.a_prompt] * self.num_samples)]}
838
- # un_cond = {"c_concat": None if self.guess_mode else [control], "c_crossattn": [self.model.get_learned_conditioning([self.n_prompt] * self.num_samples)]}
839
- # shape = (4, H // 8, W // 8)
840
- # self.model.control_scales = [self.strength * (0.825 ** float(12 - i)) for i in range(13)] if self.guess_mode else ([self.strength] * 13)
841
- # samples, intermediates = self.ddim_sampler.sample(self.ddim_steps, self.num_samples, shape, cond, verbose=False, eta=0., unconditional_guidance_scale=self.scale, unconditional_conditioning=un_cond)
842
- # if self.save_memory:
843
- # self.model.low_vram_shift(is_diffusing=False)
844
- # x_samples = self.model.decode_first_stage(samples)
845
- # x_samples = (einops.rearrange(x_samples, 'b c h w -> b h w c') * 127.5 + 127.5).cpu().numpy().clip(0, 255).astype(np.uint8)
846
- # updated_image_path = get_new_image_name(image_path, func_name="pose2image")
847
- # real_image = Image.fromarray(x_samples[0]) # default the index0 image
848
- # real_image.save(updated_image_path)
849
- # return updated_image_path
850
- class image2seg_new:
851
  def __init__(self):
852
  print("Initialize image2segmentation Inference")
853
  self.image_processor = AutoImageProcessor.from_pretrained("openmmlab/upernet-convnext-small")
854
  self.image_segmentor = UperNetForSemanticSegmentation.from_pretrained("openmmlab/upernet-convnext-small")
855
- self.resolution = 512
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
856
 
857
  def inference(self, inputs):
858
  image = Image.open(inputs)
859
- image = np.array(image)
860
- image = HWC3(image)
861
- image = resize_image(image, self.resolution)
862
- image = Image.fromarray(image)
863
  pixel_values = self.image_processor(image, return_tensors="pt").pixel_values
864
-
865
  with torch.no_grad():
866
  outputs = self.image_segmentor(pixel_values)
867
-
868
  seg = self.image_processor.post_process_semantic_segmentation(outputs, target_sizes=[image.size[::-1]])[0]
869
-
870
  color_seg = np.zeros((seg.shape[0], seg.shape[1], 3), dtype=np.uint8) # height, width, 3
871
-
872
- palette = np.array(ade_palette())
873
-
874
  for label, color in enumerate(palette):
875
  color_seg[seg == label, :] = color
876
-
877
  color_seg = color_seg.astype(np.uint8)
878
-
879
  segmentation = Image.fromarray(color_seg)
880
  updated_image_path = get_new_image_name(inputs, func_name="segmentation")
881
  segmentation.save(updated_image_path)
882
  return updated_image_path
883
 
884
- class seg2image_new:
885
  def __init__(self, device):
886
- self.controlnet = ControlNetModel.from_pretrained(
887
- "fusing/stable-diffusion-v1-5-controlnet-seg"
888
- )
889
-
890
  self.pipe = StableDiffusionControlNetPipeline.from_pretrained(
891
- "runwayml/stable-diffusion-v1-5", controlnet=self.controlnet, safety_checker=None
892
  )
893
-
894
  self.pipe.scheduler = UniPCMultistepScheduler.from_config(self.pipe.scheduler.config)
895
  self.pipe.to(device)
896
- self.image_resolution = 512
897
- self.num_inference_steps = 20
898
  self.seed = -1
899
- self.unconditional_guidance_scale = 9.0
900
  self.a_prompt = 'best quality, extremely detailed'
901
  self.n_prompt = 'longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality'
902
 
@@ -904,106 +404,21 @@ class seg2image_new:
904
  print("===>Starting seg2image Inference")
905
  image_path, instruct_text = inputs.split(",")[0], ','.join(inputs.split(',')[1:])
906
  image = Image.open(image_path)
907
- image = np.array(image)
908
- img = resize_image(HWC3(image), self.image_resolution)
909
- img = Image.fromarray(img)
910
-
911
  self.seed = random.randint(0, 65535)
912
  seed_everything(self.seed)
913
-
914
- prompt = instruct_text
915
- prompt = prompt + ', ' + self.a_prompt
916
- image = \
917
- self.pipe(prompt, img, num_inference_steps=self.num_inference_steps, eta=0.0, negative_prompt=self.n_prompt,
918
- guidance_scale=self.unconditional_guidance_scale).images[0]
919
  updated_image_path = get_new_image_name(image_path, func_name="segment2image")
920
  image.save(updated_image_path)
921
  return updated_image_path
922
 
923
-
924
-
925
- # class image2seg:
926
- # def __init__(self):
927
- # print("===>Starting image2seg Inference")
928
- # print("Direct segmentations.")
929
- # self.detector = UniformerDetector()
930
- # self.resolution = 512
931
- #
932
- # def inference(self, inputs):
933
- # print("===>Starting image2seg Inference")
934
- # image = Image.open(inputs)
935
- # image = np.array(image)
936
- # image = HWC3(image)
937
- # detected_map = self.detector(resize_image(image, self.resolution))
938
- # detected_map = HWC3(detected_map)
939
- # image = resize_image(image, self.resolution)
940
- # H, W, C = image.shape
941
- # detected_map = cv2.resize(detected_map, (W, H), interpolation=cv2.INTER_LINEAR)
942
- # updated_image_path = get_new_image_name(inputs, func_name="segmentation")
943
- # image = Image.fromarray(detected_map)
944
- # image.save(updated_image_path)
945
- # return updated_image_path
946
- #
947
- # class seg2image:
948
- # def __init__(self, device):
949
- # print("Initialize the seg2image model...")
950
- # model = create_model('ControlNet/models/cldm_v15.yaml', device=device).to(device)
951
- # model.load_state_dict(load_state_dict('ControlNet/models/control_sd15_seg.pth', location='cpu'))
952
- # self.model = model.to(device)
953
- # self.device = device
954
- # self.ddim_sampler = DDIMSampler(self.model)
955
- # self.ddim_steps = 20
956
- # self.image_resolution = 512
957
- # self.num_samples = 1
958
- # self.save_memory = False
959
- # self.strength = 1.0
960
- # self.guess_mode = False
961
- # self.scale = 9.0
962
- # self.seed = -1
963
- # self.a_prompt = 'best quality, extremely detailed'
964
- # self.n_prompt = 'longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality'
965
- #
966
- # def inference(self, inputs):
967
- # print("===>Starting seg2image Inference")
968
- # image_path, instruct_text = inputs.split(",")[0], ','.join(inputs.split(',')[1:])
969
- # image = Image.open(image_path)
970
- # image = np.array(image)
971
- # prompt = instruct_text
972
- # img = resize_image(HWC3(image), self.image_resolution)
973
- # H, W, C = img.shape
974
- # img = cv2.resize(img, (W, H), interpolation=cv2.INTER_NEAREST)
975
- # control = torch.from_numpy(img.copy()).float().to(device=self.device) / 255.0
976
- # control = torch.stack([control for _ in range(self.num_samples)], dim=0)
977
- # control = einops.rearrange(control, 'b h w c -> b c h w').clone()
978
- # self.seed = random.randint(0, 65535)
979
- # seed_everything(self.seed)
980
- # if self.save_memory:
981
- # self.model.low_vram_shift(is_diffusing=False)
982
- # cond = {"c_concat": [control], "c_crossattn": [self.model.get_learned_conditioning([prompt + ', ' + self.a_prompt] * self.num_samples)]}
983
- # un_cond = {"c_concat": None if self.guess_mode else [control], "c_crossattn": [self.model.get_learned_conditioning([self.n_prompt] * self.num_samples)]}
984
- # shape = (4, H // 8, W // 8)
985
- # self.model.control_scales = [self.strength * (0.825 ** float(12 - i)) for i in range(13)] if self.guess_mode else ([self.strength] * 13)
986
- # samples, intermediates = self.ddim_sampler.sample(self.ddim_steps, self.num_samples, shape, cond, verbose=False, eta=0., unconditional_guidance_scale=self.scale, unconditional_conditioning=un_cond)
987
- # if self.save_memory:
988
- # self.model.low_vram_shift(is_diffusing=False)
989
- # x_samples = self.model.decode_first_stage(samples)
990
- # x_samples = (einops.rearrange(x_samples, 'b c h w -> b h w c') * 127.5 + 127.5).cpu().numpy().clip(0, 255).astype(np.uint8)
991
- # updated_image_path = get_new_image_name(image_path, func_name="segment2image")
992
- # real_image = Image.fromarray(x_samples[0]) # default the index0 image
993
- # real_image.save(updated_image_path)
994
- # return updated_image_path
995
- class image2depth_new:
996
  def __init__(self):
997
  print("initialize depth estimation")
998
  self.depth_estimator = pipeline('depth-estimation')
999
- self.resolution = 512
1000
 
1001
  def inference(self, inputs):
1002
  image = Image.open(inputs)
1003
- image = np.array(image)
1004
- image = HWC3(image)
1005
- image = resize_image(image, self.resolution)
1006
- image = Image.fromarray(image)
1007
  depth = self.depth_estimator(image)['depth']
1008
  depth = np.array(depth)
1009
  depth = depth[:, :, None]
@@ -1013,22 +428,15 @@ class image2depth_new:
1013
  depth.save(updated_image_path)
1014
  return updated_image_path
1015
 
1016
- class depth2image_new:
1017
  def __init__(self, device):
1018
- self.controlnet = ControlNetModel.from_pretrained(
1019
- "fusing/stable-diffusion-v1-5-controlnet-depth"
1020
- )
1021
-
1022
  self.pipe = StableDiffusionControlNetPipeline.from_pretrained(
1023
- "runwayml/stable-diffusion-v1-5", controlnet=self.controlnet, safety_checker=None
1024
  )
1025
-
1026
  self.pipe.scheduler = UniPCMultistepScheduler.from_config(self.pipe.scheduler.config)
1027
  self.pipe.to(device)
1028
- self.image_resolution = 512
1029
- self.num_inference_steps = 20
1030
  self.seed = -1
1031
- self.unconditional_guidance_scale = 9.0
1032
  self.a_prompt = 'best quality, extremely detailed'
1033
  self.n_prompt = 'longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality'
1034
 
@@ -1036,147 +444,54 @@ class depth2image_new:
1036
  print("===>Starting depth2image Inference")
1037
  image_path, instruct_text = inputs.split(",")[0], ','.join(inputs.split(',')[1:])
1038
  image = Image.open(image_path)
1039
- image = np.array(image)
1040
- img = resize_image(HWC3(image), self.image_resolution)
1041
- img = Image.fromarray(img)
1042
-
1043
  self.seed = random.randint(0, 65535)
1044
  seed_everything(self.seed)
1045
-
1046
- prompt = instruct_text
1047
- prompt = prompt + ', ' + self.a_prompt
1048
- image = \
1049
- self.pipe(prompt, img, num_inference_steps=self.num_inference_steps, eta=0.0, negative_prompt=self.n_prompt,
1050
- guidance_scale=self.unconditional_guidance_scale).images[0]
1051
  updated_image_path = get_new_image_name(image_path, func_name="depth2image")
1052
  image.save(updated_image_path)
1053
  return updated_image_path
1054
 
1055
- # class image2depth:
1056
- # def __init__(self):
1057
- # print("Direct depth estimation.")
1058
- # self.detector = MidasDetector()
1059
- # self.resolution = 512
1060
- #
1061
- # def inference(self, inputs):
1062
- # print("===>Starting image2depth Inference")
1063
- # image = Image.open(inputs)
1064
- # image = np.array(image)
1065
- # image = HWC3(image)
1066
- # detected_map, _ = self.detector(resize_image(image, self.resolution))
1067
- # detected_map = HWC3(detected_map)
1068
- # image = resize_image(image, self.resolution)
1069
- # H, W, C = image.shape
1070
- # detected_map = cv2.resize(detected_map, (W, H), interpolation=cv2.INTER_LINEAR)
1071
- # updated_image_path = get_new_image_name(inputs, func_name="depth")
1072
- # image = Image.fromarray(detected_map)
1073
- # image.save(updated_image_path)
1074
- # return updated_image_path
1075
- #
1076
- # class depth2image:
1077
- # def __init__(self, device):
1078
- # print("Initialize depth2image model...")
1079
- # model = create_model('ControlNet/models/cldm_v15.yaml', device=device).to(device)
1080
- # model.load_state_dict(load_state_dict('ControlNet/models/control_sd15_depth.pth', location='cpu'))
1081
- # self.model = model.to(device)
1082
- # self.device = device
1083
- # self.ddim_sampler = DDIMSampler(self.model)
1084
- # self.ddim_steps = 20
1085
- # self.image_resolution = 512
1086
- # self.num_samples = 1
1087
- # self.save_memory = False
1088
- # self.strength = 1.0
1089
- # self.guess_mode = False
1090
- # self.scale = 9.0
1091
- # self.seed = -1
1092
- # self.a_prompt = 'best quality, extremely detailed'
1093
- # self.n_prompt = 'longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality'
1094
- #
1095
- # def inference(self, inputs):
1096
- # print("===>Starting depth2image Inference")
1097
- # image_path, instruct_text = inputs.split(",")[0], ','.join(inputs.split(',')[1:])
1098
- # image = Image.open(image_path)
1099
- # image = np.array(image)
1100
- # prompt = instruct_text
1101
- # img = resize_image(HWC3(image), self.image_resolution)
1102
- # H, W, C = img.shape
1103
- # img = cv2.resize(img, (W, H), interpolation=cv2.INTER_NEAREST)
1104
- # control = torch.from_numpy(img.copy()).float().to(device=self.device) / 255.0
1105
- # control = torch.stack([control for _ in range(self.num_samples)], dim=0)
1106
- # control = einops.rearrange(control, 'b h w c -> b c h w').clone()
1107
- # self.seed = random.randint(0, 65535)
1108
- # seed_everything(self.seed)
1109
- # if self.save_memory:
1110
- # self.model.low_vram_shift(is_diffusing=False)
1111
- # cond = {"c_concat": [control], "c_crossattn": [ self.model.get_learned_conditioning([prompt + ', ' + self.a_prompt] * self.num_samples)]}
1112
- # un_cond = {"c_concat": None if self.guess_mode else [control], "c_crossattn": [self.model.get_learned_conditioning([self.n_prompt] * self.num_samples)]}
1113
- # shape = (4, H // 8, W // 8)
1114
- # self.model.control_scales = [self.strength * (0.825 ** float(12 - i)) for i in range(13)] if self.guess_mode else ([self.strength] * 13) # Magic number. IDK why. Perhaps because 0.825**12<0.01 but 0.826**12>0.01
1115
- # samples, intermediates = self.ddim_sampler.sample(self.ddim_steps, self.num_samples, shape, cond, verbose=False, eta=0., unconditional_guidance_scale=self.scale, unconditional_conditioning=un_cond)
1116
- # if self.save_memory:
1117
- # self.model.low_vram_shift(is_diffusing=False)
1118
- # x_samples = self.model.decode_first_stage(samples)
1119
- # x_samples = (einops.rearrange(x_samples, 'b c h w -> b h w c') * 127.5 + 127.5).cpu().numpy().clip(0, 255).astype(np.uint8)
1120
- # updated_image_path = get_new_image_name(image_path, func_name="depth2image")
1121
- # real_image = Image.fromarray(x_samples[0]) # default the index0 image
1122
- # real_image.save(updated_image_path)
1123
- # return updated_image_path
1124
-
1125
- class image2normal_new:
1126
  def __init__(self):
1127
  print("normal estimation")
1128
  self.depth_estimator = pipeline("depth-estimation", model="Intel/dpt-hybrid-midas")
1129
- self.resolution = 512
1130
  self.bg_threhold = 0.4
1131
 
1132
  def inference(self, inputs):
1133
  image = Image.open(inputs)
1134
- image = np.array(image)
1135
- image = HWC3(image)
1136
- image = resize_image(image, self.resolution)
1137
- image = Image.fromarray(image)
1138
  image = self.depth_estimator(image)['predicted_depth'][0]
1139
-
1140
  image = image.numpy()
1141
-
1142
  image_depth = image.copy()
1143
  image_depth -= np.min(image_depth)
1144
  image_depth /= np.max(image_depth)
1145
 
1146
- bg_threhold = 0.4
1147
-
1148
  x = cv2.Sobel(image, cv2.CV_32F, 1, 0, ksize=3)
1149
- x[image_depth < bg_threhold] = 0
1150
 
1151
  y = cv2.Sobel(image, cv2.CV_32F, 0, 1, ksize=3)
1152
- y[image_depth < bg_threhold] = 0
1153
 
1154
  z = np.ones_like(x) * np.pi * 2.0
1155
-
1156
  image = np.stack([x, y, z], axis=2)
1157
  image /= np.sum(image ** 2.0, axis=2, keepdims=True) ** 0.5
1158
  image = (image * 127.5 + 127.5).clip(0, 255).astype(np.uint8)
1159
  image = Image.fromarray(image)
 
1160
  updated_image_path = get_new_image_name(inputs, func_name="normal-map")
1161
  image.save(updated_image_path)
1162
  return updated_image_path
1163
 
1164
- class normal2image_new:
1165
  def __init__(self, device):
1166
- self.controlnet = ControlNetModel.from_pretrained(
1167
- "fusing/stable-diffusion-v1-5-controlnet-normal"
1168
- )
1169
-
1170
  self.pipe = StableDiffusionControlNetPipeline.from_pretrained(
1171
- "runwayml/stable-diffusion-v1-5", controlnet=self.controlnet, safety_checker=None
1172
  )
1173
-
1174
  self.pipe.scheduler = UniPCMultistepScheduler.from_config(self.pipe.scheduler.config)
1175
  self.pipe.to(device)
1176
- self.image_resolution = 512
1177
- self.num_inference_steps = 20
1178
  self.seed = -1
1179
- self.unconditional_guidance_scale = 9.0
1180
  self.a_prompt = 'best quality, extremely detailed'
1181
  self.n_prompt = 'longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality'
1182
 
@@ -1184,100 +499,20 @@ class normal2image_new:
1184
  print("===>Starting normal2image Inference")
1185
  image_path, instruct_text = inputs.split(",")[0], ','.join(inputs.split(',')[1:])
1186
  image = Image.open(image_path)
1187
- image = np.array(image)
1188
- img = resize_image(HWC3(image), self.image_resolution)
1189
- img = Image.fromarray(img)
1190
-
1191
  self.seed = random.randint(0, 65535)
1192
  seed_everything(self.seed)
1193
-
1194
- prompt = instruct_text
1195
- prompt = prompt + ', ' + self.a_prompt
1196
- image = \
1197
- self.pipe(prompt, img, num_inference_steps=self.num_inference_steps, eta=0.0, negative_prompt=self.n_prompt,
1198
- guidance_scale=self.unconditional_guidance_scale).images[0]
1199
  updated_image_path = get_new_image_name(image_path, func_name="normal2image")
1200
  image.save(updated_image_path)
1201
  return updated_image_path
1202
 
1203
- # class image2normal:
1204
- # def __init__(self):
1205
- # print("Direct normal estimation.")
1206
- # self.detector = MidasDetector()
1207
- # self.resolution = 512
1208
- # self.bg_threshold = 0.4
1209
- #
1210
- # def inference(self, inputs):
1211
- # print("===>Starting image2 normal Inference")
1212
- # image = Image.open(inputs)
1213
- # image = np.array(image)
1214
- # image = HWC3(image)
1215
- # _, detected_map = self.detector(resize_image(image, self.resolution), bg_th=self.bg_threshold)
1216
- # detected_map = HWC3(detected_map)
1217
- # image = resize_image(image, self.resolution)
1218
- # H, W, C = image.shape
1219
- # detected_map = cv2.resize(detected_map, (W, H), interpolation=cv2.INTER_LINEAR)
1220
- # updated_image_path = get_new_image_name(inputs, func_name="normal-map")
1221
- # image = Image.fromarray(detected_map)
1222
- # image.save(updated_image_path)
1223
- # return updated_image_path
1224
- #
1225
- # class normal2image:
1226
- # def __init__(self, device):
1227
- # print("Initialize normal2image model...")
1228
- # model = create_model('ControlNet/models/cldm_v15.yaml', device=device).to(device)
1229
- # model.load_state_dict(load_state_dict('ControlNet/models/control_sd15_normal.pth', location='cpu'))
1230
- # self.model = model.to(device)
1231
- # self.device = device
1232
- # self.ddim_sampler = DDIMSampler(self.model)
1233
- # self.ddim_steps = 20
1234
- # self.image_resolution = 512
1235
- # self.num_samples = 1
1236
- # self.save_memory = False
1237
- # self.strength = 1.0
1238
- # self.guess_mode = False
1239
- # self.scale = 9.0
1240
- # self.seed = -1
1241
- # self.a_prompt = 'best quality, extremely detailed'
1242
- # self.n_prompt = 'longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality'
1243
- #
1244
- # def inference(self, inputs):
1245
- # print("===>Starting normal2image Inference")
1246
- # image_path, instruct_text = inputs.split(",")[0], ','.join(inputs.split(',')[1:])
1247
- # image = Image.open(image_path)
1248
- # image = np.array(image)
1249
- # prompt = instruct_text
1250
- # img = image[:, :, ::-1].copy()
1251
- # img = resize_image(HWC3(img), self.image_resolution)
1252
- # H, W, C = img.shape
1253
- # img = cv2.resize(img, (W, H), interpolation=cv2.INTER_NEAREST)
1254
- # control = torch.from_numpy(img.copy()).float().to(device=self.device) / 255.0
1255
- # control = torch.stack([control for _ in range(self.num_samples)], dim=0)
1256
- # control = einops.rearrange(control, 'b h w c -> b c h w').clone()
1257
- # self.seed = random.randint(0, 65535)
1258
- # seed_everything(self.seed)
1259
- # if self.save_memory:
1260
- # self.model.low_vram_shift(is_diffusing=False)
1261
- # cond = {"c_concat": [control], "c_crossattn": [self.model.get_learned_conditioning([prompt + ', ' + self.a_prompt] * self.num_samples)]}
1262
- # un_cond = {"c_concat": None if self.guess_mode else [control], "c_crossattn": [self.model.get_learned_conditioning([self.n_prompt] * self.num_samples)]}
1263
- # shape = (4, H // 8, W // 8)
1264
- # self.model.control_scales = [self.strength * (0.825 ** float(12 - i)) for i in range(13)] if self.guess_mode else ([self.strength] * 13)
1265
- # samples, intermediates = self.ddim_sampler.sample(self.ddim_steps, self.num_samples, shape, cond, verbose=False, eta=0., unconditional_guidance_scale=self.scale, unconditional_conditioning=un_cond)
1266
- # if self.save_memory:
1267
- # self.model.low_vram_shift(is_diffusing=False)
1268
- # x_samples = self.model.decode_first_stage(samples)
1269
- # x_samples = (einops.rearrange(x_samples, 'b c h w -> b h w c') * 127.5 + 127.5).cpu().numpy().clip(0, 255).astype(np.uint8)
1270
- # updated_image_path = get_new_image_name(image_path, func_name="normal2image")
1271
- # real_image = Image.fromarray(x_samples[0]) # default the index0 image
1272
- # real_image.save(updated_image_path)
1273
- # return updated_image_path
1274
-
1275
  class BLIPVQA:
1276
  def __init__(self, device):
1277
  print("Initializing BLIP VQA to %s" % device)
1278
  self.device = device
1279
- self.processor = BlipProcessor.from_pretrained("Salesforce/blip-vqa-base")
1280
- self.model = BlipForQuestionAnswering.from_pretrained("Salesforce/blip-vqa-base").to(self.device)
1281
 
1282
  def get_answer_from_question_and_image(self, inputs):
1283
  image_path, question = inputs.split(",")
 
1
+ from diffusers import StableDiffusionPipeline, StableDiffusionInpaintPipeline, StableDiffusionInstructPix2PixPipeline
2
+ from diffusers import EulerAncestralDiscreteScheduler
 
3
  from diffusers import StableDiffusionControlNetPipeline, ControlNetModel, UniPCMultistepScheduler
4
  from controlnet_aux import OpenposeDetector, MLSDdetector, HEDdetector
5
 
 
7
  from transformers import pipeline, BlipProcessor, BlipForConditionalGeneration, BlipForQuestionAnswering
8
  from transformers import AutoImageProcessor, UperNetForSemanticSegmentation
9
 
10
+ import os
11
+ import random
12
  import torch
13
+ import cv2
14
  import uuid
15
+ from PIL import Image
16
+ import numpy as np
17
  from pytorch_lightning import seed_everything
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
18
 
19
  def get_new_image_name(org_img_name, func_name="update"):
20
  head_tail = os.path.split(org_img_name)
 
33
  new_file_name = '{}_{}_{}_{}.png'.format(this_new_uuid, func_name, recent_prev_file_name, most_org_file_name)
34
  return os.path.join(head, new_file_name)
35
 
36
+
37
  class MaskFormer:
38
  def __init__(self, device):
39
  self.device = device
 
60
  mask_array[padded_slice] = True
61
  visual_mask = (mask_array * 255).astype(np.uint8)
62
  image_mask = Image.fromarray(visual_mask)
63
+ return image_mask.resize(original_image.size)
64
 
65
  class ImageEditing:
66
  def __init__(self, device):
 
78
  image_path, to_be_replaced_txt, replace_with_txt = input.split(",")
79
  print(f'replace_part_of_image: replace_with_txt {replace_with_txt}')
80
  original_image = Image.open(image_path)
81
+ original_size = original_image.size
82
  mask_image = self.mask_former.inference(image_path, to_be_replaced_txt)
83
+ updated_image = self.inpainting(prompt=replace_with_txt, image=original_image.resize((512,512)), mask_image=mask_image.resize((512,512))).images[0]
84
  updated_image_path = get_new_image_name(image_path, func_name="replace-something")
85
+ updated_image = updated_image.resize(original_size)
86
  updated_image.save(updated_image_path)
87
  return updated_image_path
88
 
 
108
  print("Initializing T2I to %s" % device)
109
  self.device = device
110
  self.pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16)
111
+ self.text_refine_tokenizer = AutoTokenizer.from_pretrained("Gustavosta/MagicPrompt-Stable-Diffusion", torch_dtype=torch.float16)
112
  self.text_refine_model = AutoModelForCausalLM.from_pretrained("Gustavosta/MagicPrompt-Stable-Diffusion", torch_dtype=torch.float16)
113
  self.text_refine_gpt2_pipe = pipeline("text-generation", model=self.text_refine_model, tokenizer=self.text_refine_tokenizer, device=self.device, torch_dtype=torch.float16)
114
  self.pipe.to(device)
 
135
  captions = self.processor.decode(out[0], skip_special_tokens=True)
136
  return captions
137
 
138
+ class image2canny:
139
  def __init__(self):
140
  print("Direct detect canny.")
141
  self.low_threshold = 100
 
148
  canny = cv2.Canny(image, self.low_threshold, self.high_threshold)
149
  canny = canny[:, :, None]
150
  canny = np.concatenate([canny, canny, canny], axis=2)
 
151
  canny = Image.fromarray(canny)
152
  updated_image_path = get_new_image_name(inputs, func_name="edge")
153
  canny.save(updated_image_path)
154
  return updated_image_path
155
 
156
+ class canny2image:
157
  def __init__(self, device):
158
+ self.controlnet = ControlNetModel.from_pretrained("fusing/stable-diffusion-v1-5-controlnet-canny", torch_dtype=torch.float16)
 
 
 
159
  self.pipe = StableDiffusionControlNetPipeline.from_pretrained(
160
  "runwayml/stable-diffusion-v1-5", controlnet=self.controlnet, safety_checker=None, torch_dtype=torch.float16
161
  )
 
162
  self.pipe.scheduler = UniPCMultistepScheduler.from_config(self.pipe.scheduler.config)
163
  self.pipe.to(device)
 
 
164
  self.seed = -1
 
165
  self.a_prompt = 'best quality, extremely detailed'
166
  self.n_prompt = 'longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality'
167
 
 
169
  print("===>Starting canny2image Inference")
170
  image_path, instruct_text = inputs.split(",")[0], ','.join(inputs.split(',')[1:])
171
  image = Image.open(image_path)
 
 
 
 
 
 
172
  self.seed = random.randint(0, 65535)
173
  seed_everything(self.seed)
174
+ prompt = instruct_text + ', ' + self.a_prompt
175
+ image = self.pipe(prompt, image, num_inference_steps=20, eta=0.0, negative_prompt=self.n_prompt, guidance_scale=9.0).images[0]
176
  updated_image_path = get_new_image_name(image_path, func_name="canny2image")
177
  image.save(updated_image_path)
178
  return updated_image_path
179
 
180
+ class image2line:
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
181
  def __init__(self):
182
  self.detector = MLSDdetector.from_pretrained('lllyasviel/ControlNet')
 
 
 
183
 
184
  def inference(self, inputs):
185
  print("===>Starting image2line Inference")
186
  image = Image.open(inputs)
187
+ mlsd = self.detector(image)
 
 
 
 
 
188
  updated_image_path = get_new_image_name(inputs, func_name="line-of")
189
  mlsd.save(updated_image_path)
190
  return updated_image_path
191
 
192
+ class line2image:
193
  def __init__(self, device):
194
  print("Initialize the line2image model...")
195
+ self.controlnet = ControlNetModel.from_pretrained("fusing/stable-diffusion-v1-5-controlnet-mlsd", torch_dtype=torch.float16)
 
 
 
196
  self.pipe = StableDiffusionControlNetPipeline.from_pretrained(
197
+ "runwayml/stable-diffusion-v1-5", controlnet=self.controlnet, safety_checker=None, torch_dtype=torch.float16
198
  )
 
199
  self.pipe.scheduler = UniPCMultistepScheduler.from_config(self.pipe.scheduler.config)
200
  self.pipe.to(device)
 
 
201
  self.seed = -1
 
202
  self.a_prompt = 'best quality, extremely detailed'
203
  self.n_prompt = 'longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality'
204
 
 
206
  print("===>Starting line2image Inference")
207
  image_path, instruct_text = inputs.split(",")[0], ','.join(inputs.split(',')[1:])
208
  image = Image.open(image_path)
 
 
 
 
 
 
209
  self.seed = random.randint(0, 65535)
210
  seed_everything(self.seed)
211
+ prompt = instruct_text + ', ' + self.a_prompt
212
+ image = self.pipe(prompt, image, num_inference_steps=20, eta=0.0, negative_prompt=self.n_prompt, guidance_scale=9.0).images[0]
 
213
  updated_image_path = get_new_image_name(image_path, func_name="line2image")
214
  image.save(updated_image_path)
215
  return updated_image_path
216
 
217
+ class image2hed:
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
218
  def __init__(self):
219
  print("Direct detect soft HED boundary...")
220
  self.detector = HEDdetector.from_pretrained('lllyasviel/ControlNet')
 
221
 
222
  def inference(self, inputs):
223
  print("===>Starting image2hed Inference")
224
  image = Image.open(inputs)
 
 
 
225
  hed = self.detector(image)
 
226
  updated_image_path = get_new_image_name(inputs, func_name="hed-boundary")
227
  hed.save(updated_image_path)
228
  return updated_image_path
229
 
230
+ class hed2image:
231
  def __init__(self, device):
232
  print("Initialize the hed2image model...")
233
+ self.controlnet = ControlNetModel.from_pretrained("fusing/stable-diffusion-v1-5-controlnet-hed", torch_dtype=torch.float16)
 
 
 
234
  self.pipe = StableDiffusionControlNetPipeline.from_pretrained(
235
+ "runwayml/stable-diffusion-v1-5", controlnet=self.controlnet, safety_checker=None, torch_dtype=torch.float16
236
  )
 
237
  self.pipe.scheduler = UniPCMultistepScheduler.from_config(self.pipe.scheduler.config)
238
  self.pipe.to(device)
 
 
239
  self.seed = -1
 
240
  self.a_prompt = 'best quality, extremely detailed'
241
  self.n_prompt = 'longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality'
242
 
 
244
  print("===>Starting hed2image Inference")
245
  image_path, instruct_text = inputs.split(",")[0], ','.join(inputs.split(',')[1:])
246
  image = Image.open(image_path)
 
 
 
 
247
  self.seed = random.randint(0, 65535)
248
  seed_everything(self.seed)
249
+ prompt = instruct_text + ', ' + self.a_prompt
250
+ image = self.pipe(prompt, image, num_inference_steps=20, eta=0.0, negative_prompt=self.n_prompt, guidance_scale=9.0).images[0]
 
 
 
 
251
  updated_image_path = get_new_image_name(image_path, func_name="hed2image")
252
  image.save(updated_image_path)
253
  return updated_image_path
254
 
255
+ class image2scribble:
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
256
  def __init__(self):
257
  print("Direct detect scribble.")
258
  self.detector = HEDdetector.from_pretrained('lllyasviel/ControlNet')
 
259
 
260
  def inference(self, inputs):
261
  print("===>Starting image2scribble Inference")
262
  image = Image.open(inputs)
 
 
 
 
263
  scribble = self.detector(image, scribble=True)
 
 
 
264
  updated_image_path = get_new_image_name(inputs, func_name="scribble")
265
  scribble.save(updated_image_path)
266
  return updated_image_path
267
 
268
+ class scribble2image:
269
  def __init__(self, device):
270
+ self.controlnet = ControlNetModel.from_pretrained("fusing/stable-diffusion-v1-5-controlnet-scribble", torch_dtype=torch.float16)
 
 
 
271
  self.pipe = StableDiffusionControlNetPipeline.from_pretrained(
272
+ "runwayml/stable-diffusion-v1-5", controlnet=self.controlnet, safety_checker=None, torch_dtype=torch.float16
273
  )
 
274
  self.pipe.scheduler = UniPCMultistepScheduler.from_config(self.pipe.scheduler.config)
275
  self.pipe.to(device)
 
 
276
  self.seed = -1
 
277
  self.a_prompt = 'best quality, extremely detailed'
278
  self.n_prompt = 'longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality'
279
 
 
281
  print("===>Starting scribble2image Inference")
282
  image_path, instruct_text = inputs.split(",")[0], ','.join(inputs.split(',')[1:])
283
  image = Image.open(image_path)
 
 
 
 
 
284
  self.seed = random.randint(0, 65535)
285
  seed_everything(self.seed)
286
+ prompt = instruct_text + ', ' + self.a_prompt
287
+ image = self.pipe(prompt, image, num_inference_steps=20, eta=0.0, negative_prompt=self.n_prompt,guidance_scale=9.0).images[0]
 
 
 
 
288
  updated_image_path = get_new_image_name(image_path, func_name="scribble2image")
289
  image.save(updated_image_path)
290
  return updated_image_path
291
 
292
+ class image2pose:
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
293
  def __init__(self):
294
  self.detector = OpenposeDetector.from_pretrained('lllyasviel/ControlNet')
 
295
 
296
  def inference(self, inputs):
297
  print("===>Starting image2pose Inference")
298
  image = Image.open(inputs)
 
 
 
 
299
  pose = self.detector(image)
 
300
  updated_image_path = get_new_image_name(inputs, func_name="human-pose")
301
  pose.save(updated_image_path)
302
  return updated_image_path
303
 
304
+ class pose2image:
305
  def __init__(self, device):
306
+ self.controlnet = ControlNetModel.from_pretrained("fusing/stable-diffusion-v1-5-controlnet-openpose", torch_dtype=torch.float16)
 
 
 
307
  self.pipe = StableDiffusionControlNetPipeline.from_pretrained(
308
+ "runwayml/stable-diffusion-v1-5", controlnet=self.controlnet, safety_checker=None, torch_dtype=torch.float16
309
  )
 
310
  self.pipe.scheduler = UniPCMultistepScheduler.from_config(self.pipe.scheduler.config)
311
  self.pipe.to(device)
 
312
  self.num_inference_steps = 20
313
  self.seed = -1
314
  self.unconditional_guidance_scale = 9.0
 
319
  print("===>Starting pose2image Inference")
320
  image_path, instruct_text = inputs.split(",")[0], ','.join(inputs.split(',')[1:])
321
  image = Image.open(image_path)
 
 
 
 
322
  self.seed = random.randint(0, 65535)
323
  seed_everything(self.seed)
324
+ prompt = instruct_text + ', ' + self.a_prompt
325
+ image = self.pipe(prompt, image, num_inference_steps=20, eta=0.0, negative_prompt=self.n_prompt, guidance_scale=9.0).images[0]
 
 
 
 
326
  updated_image_path = get_new_image_name(image_path, func_name="pose2image")
327
  image.save(updated_image_path)
328
  return updated_image_path
329
 
330
+ class image2seg:
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
331
  def __init__(self):
332
  print("Initialize image2segmentation Inference")
333
  self.image_processor = AutoImageProcessor.from_pretrained("openmmlab/upernet-convnext-small")
334
  self.image_segmentor = UperNetForSemanticSegmentation.from_pretrained("openmmlab/upernet-convnext-small")
335
+
336
+ self.ade_palette = [[120, 120, 120], [180, 120, 120], [6, 230, 230], [80, 50, 50],
337
+ [4, 200, 3], [120, 120, 80], [140, 140, 140], [204, 5, 255],
338
+ [230, 230, 230], [4, 250, 7], [224, 5, 255], [235, 255, 7],
339
+ [150, 5, 61], [120, 120, 70], [8, 255, 51], [255, 6, 82],
340
+ [143, 255, 140], [204, 255, 4], [255, 51, 7], [204, 70, 3],
341
+ [0, 102, 200], [61, 230, 250], [255, 6, 51], [11, 102, 255],
342
+ [255, 7, 71], [255, 9, 224], [9, 7, 230], [220, 220, 220],
343
+ [255, 9, 92], [112, 9, 255], [8, 255, 214], [7, 255, 224],
344
+ [255, 184, 6], [10, 255, 71], [255, 41, 10], [7, 255, 255],
345
+ [224, 255, 8], [102, 8, 255], [255, 61, 6], [255, 194, 7],
346
+ [255, 122, 8], [0, 255, 20], [255, 8, 41], [255, 5, 153],
347
+ [6, 51, 255], [235, 12, 255], [160, 150, 20], [0, 163, 255],
348
+ [140, 140, 140], [250, 10, 15], [20, 255, 0], [31, 255, 0],
349
+ [255, 31, 0], [255, 224, 0], [153, 255, 0], [0, 0, 255],
350
+ [255, 71, 0], [0, 235, 255], [0, 173, 255], [31, 0, 255],
351
+ [11, 200, 200], [255, 82, 0], [0, 255, 245], [0, 61, 255],
352
+ [0, 255, 112], [0, 255, 133], [255, 0, 0], [255, 163, 0],
353
+ [255, 102, 0], [194, 255, 0], [0, 143, 255], [51, 255, 0],
354
+ [0, 82, 255], [0, 255, 41], [0, 255, 173], [10, 0, 255],
355
+ [173, 255, 0], [0, 255, 153], [255, 92, 0], [255, 0, 255],
356
+ [255, 0, 245], [255, 0, 102], [255, 173, 0], [255, 0, 20],
357
+ [255, 184, 184], [0, 31, 255], [0, 255, 61], [0, 71, 255],
358
+ [255, 0, 204], [0, 255, 194], [0, 255, 82], [0, 10, 255],
359
+ [0, 112, 255], [51, 0, 255], [0, 194, 255], [0, 122, 255],
360
+ [0, 255, 163], [255, 153, 0], [0, 255, 10], [255, 112, 0],
361
+ [143, 255, 0], [82, 0, 255], [163, 255, 0], [255, 235, 0],
362
+ [8, 184, 170], [133, 0, 255], [0, 255, 92], [184, 0, 255],
363
+ [255, 0, 31], [0, 184, 255], [0, 214, 255], [255, 0, 112],
364
+ [92, 255, 0], [0, 224, 255], [112, 224, 255], [70, 184, 160],
365
+ [163, 0, 255], [153, 0, 255], [71, 255, 0], [255, 0, 163],
366
+ [255, 204, 0], [255, 0, 143], [0, 255, 235], [133, 255, 0],
367
+ [255, 0, 235], [245, 0, 255], [255, 0, 122], [255, 245, 0],
368
+ [10, 190, 212], [214, 255, 0], [0, 204, 255], [20, 0, 255],
369
+ [255, 255, 0], [0, 153, 255], [0, 41, 255], [0, 255, 204],
370
+ [41, 0, 255], [41, 255, 0], [173, 0, 255], [0, 245, 255],
371
+ [71, 0, 255], [122, 0, 255], [0, 255, 184], [0, 92, 255],
372
+ [184, 255, 0], [0, 133, 255], [255, 214, 0], [25, 194, 194],
373
+ [102, 255, 0], [92, 0, 255]]
374
 
375
  def inference(self, inputs):
376
  image = Image.open(inputs)
 
 
 
 
377
  pixel_values = self.image_processor(image, return_tensors="pt").pixel_values
 
378
  with torch.no_grad():
379
  outputs = self.image_segmentor(pixel_values)
 
380
  seg = self.image_processor.post_process_semantic_segmentation(outputs, target_sizes=[image.size[::-1]])[0]
 
381
  color_seg = np.zeros((seg.shape[0], seg.shape[1], 3), dtype=np.uint8) # height, width, 3
382
+ palette = np.array(self.ade_palette)
 
 
383
  for label, color in enumerate(palette):
384
  color_seg[seg == label, :] = color
 
385
  color_seg = color_seg.astype(np.uint8)
 
386
  segmentation = Image.fromarray(color_seg)
387
  updated_image_path = get_new_image_name(inputs, func_name="segmentation")
388
  segmentation.save(updated_image_path)
389
  return updated_image_path
390
 
391
+ class seg2image:
392
  def __init__(self, device):
393
+ self.controlnet = ControlNetModel.from_pretrained("fusing/stable-diffusion-v1-5-controlnet-seg", torch_dtype=torch.float16)
 
 
 
394
  self.pipe = StableDiffusionControlNetPipeline.from_pretrained(
395
+ "runwayml/stable-diffusion-v1-5", controlnet=self.controlnet, safety_checker=None, torch_dtype=torch.float16
396
  )
 
397
  self.pipe.scheduler = UniPCMultistepScheduler.from_config(self.pipe.scheduler.config)
398
  self.pipe.to(device)
 
 
399
  self.seed = -1
 
400
  self.a_prompt = 'best quality, extremely detailed'
401
  self.n_prompt = 'longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality'
402
 
 
404
  print("===>Starting seg2image Inference")
405
  image_path, instruct_text = inputs.split(",")[0], ','.join(inputs.split(',')[1:])
406
  image = Image.open(image_path)
 
 
 
 
407
  self.seed = random.randint(0, 65535)
408
  seed_everything(self.seed)
409
+ prompt = instruct_text + ', ' + self.a_prompt
410
+ image = self.pipe(prompt, image, num_inference_steps=20, eta=0.0, negative_prompt=self.n_prompt, guidance_scale=9.0).images[0]
 
 
 
 
411
  updated_image_path = get_new_image_name(image_path, func_name="segment2image")
412
  image.save(updated_image_path)
413
  return updated_image_path
414
 
415
+ class image2depth:
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
416
  def __init__(self):
417
  print("initialize depth estimation")
418
  self.depth_estimator = pipeline('depth-estimation')
 
419
 
420
  def inference(self, inputs):
421
  image = Image.open(inputs)
 
 
 
 
422
  depth = self.depth_estimator(image)['depth']
423
  depth = np.array(depth)
424
  depth = depth[:, :, None]
 
428
  depth.save(updated_image_path)
429
  return updated_image_path
430
 
431
+ class depth2image:
432
  def __init__(self, device):
433
+ self.controlnet = ControlNetModel.from_pretrained("fusing/stable-diffusion-v1-5-controlnet-depth", torch_dtype=torch.float16)
 
 
 
434
  self.pipe = StableDiffusionControlNetPipeline.from_pretrained(
435
+ "runwayml/stable-diffusion-v1-5", controlnet=self.controlnet, safety_checker=None, torch_dtype=torch.float16
436
  )
 
437
  self.pipe.scheduler = UniPCMultistepScheduler.from_config(self.pipe.scheduler.config)
438
  self.pipe.to(device)
 
 
439
  self.seed = -1
 
440
  self.a_prompt = 'best quality, extremely detailed'
441
  self.n_prompt = 'longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality'
442
 
 
444
  print("===>Starting depth2image Inference")
445
  image_path, instruct_text = inputs.split(",")[0], ','.join(inputs.split(',')[1:])
446
  image = Image.open(image_path)
 
 
 
 
447
  self.seed = random.randint(0, 65535)
448
  seed_everything(self.seed)
449
+ prompt = instruct_text + ', ' + self.a_prompt
450
+ image = self.pipe(prompt, image, num_inference_steps=20, eta=0.0, negative_prompt=self.n_prompt, guidance_scale=9.0).images[0]
 
 
 
 
451
  updated_image_path = get_new_image_name(image_path, func_name="depth2image")
452
  image.save(updated_image_path)
453
  return updated_image_path
454
 
455
+ class image2normal:
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
456
  def __init__(self):
457
  print("normal estimation")
458
  self.depth_estimator = pipeline("depth-estimation", model="Intel/dpt-hybrid-midas")
 
459
  self.bg_threhold = 0.4
460
 
461
  def inference(self, inputs):
462
  image = Image.open(inputs)
463
+ original_size = image.size
 
 
 
464
  image = self.depth_estimator(image)['predicted_depth'][0]
 
465
  image = image.numpy()
 
466
  image_depth = image.copy()
467
  image_depth -= np.min(image_depth)
468
  image_depth /= np.max(image_depth)
469
 
 
 
470
  x = cv2.Sobel(image, cv2.CV_32F, 1, 0, ksize=3)
471
+ x[image_depth < self.bg_threhold] = 0
472
 
473
  y = cv2.Sobel(image, cv2.CV_32F, 0, 1, ksize=3)
474
+ y[image_depth < self.bg_threhold] = 0
475
 
476
  z = np.ones_like(x) * np.pi * 2.0
 
477
  image = np.stack([x, y, z], axis=2)
478
  image /= np.sum(image ** 2.0, axis=2, keepdims=True) ** 0.5
479
  image = (image * 127.5 + 127.5).clip(0, 255).astype(np.uint8)
480
  image = Image.fromarray(image)
481
+ image = image.resize(original_size)
482
  updated_image_path = get_new_image_name(inputs, func_name="normal-map")
483
  image.save(updated_image_path)
484
  return updated_image_path
485
 
486
+ class normal2image:
487
  def __init__(self, device):
488
+ self.controlnet = ControlNetModel.from_pretrained("fusing/stable-diffusion-v1-5-controlnet-normal", torch_dtype=torch.float16)
 
 
 
489
  self.pipe = StableDiffusionControlNetPipeline.from_pretrained(
490
+ "runwayml/stable-diffusion-v1-5", controlnet=self.controlnet, safety_checker=None, torch_dtype=torch.float16
491
  )
 
492
  self.pipe.scheduler = UniPCMultistepScheduler.from_config(self.pipe.scheduler.config)
493
  self.pipe.to(device)
 
 
494
  self.seed = -1
 
495
  self.a_prompt = 'best quality, extremely detailed'
496
  self.n_prompt = 'longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality'
497
 
 
499
  print("===>Starting normal2image Inference")
500
  image_path, instruct_text = inputs.split(",")[0], ','.join(inputs.split(',')[1:])
501
  image = Image.open(image_path)
 
 
 
 
502
  self.seed = random.randint(0, 65535)
503
  seed_everything(self.seed)
504
+ prompt = instruct_text + ', ' + self.a_prompt
505
+ image = self.pipe(prompt, image, num_inference_steps=20, eta=0.0, negative_prompt=self.n_prompt, guidance_scale=9.0).images[0]
 
 
 
 
506
  updated_image_path = get_new_image_name(image_path, func_name="normal2image")
507
  image.save(updated_image_path)
508
  return updated_image_path
509
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
510
  class BLIPVQA:
511
  def __init__(self, device):
512
  print("Initializing BLIP VQA to %s" % device)
513
  self.device = device
514
+ self.processor = BlipProcessor.from_pretrained("Salesforce/blip-vqa-base", torch_dtype=torch.float16)
515
+ self.model = BlipForQuestionAnswering.from_pretrained("Salesforce/blip-vqa-base", torch_dtype=torch.float16).to(self.device)
516
 
517
  def get_answer_from_question_and_image(self, inputs):
518
  image_path, question = inputs.split(",")