FelixLuoX/stable_diffusion_test
2
1import gradio as gr2import io 3from PIL import Image, PngImagePlugin4import base645import requests6import json 7import ui_functions as uifn8 9txt2img_defaults = {10 'prompt': '',11 'ddim_steps': 50,12 'toggles': [1, 2, 3],13 'sampler_name': 'k_lms',14 'ddim_eta': 0.0,15 'n_iter': 1,16 'batch_size': 1,17 'cfg_scale': 7.5,18 'seed': '',19 'height': 512,20 'width': 512,21 'fp': None,22 'variant_amount': 0.0,23 'variant_seed': '',24 'submit_on_enter': 'Yes',25}26 27img2img_defaults = {28 'prompt': '',29 'ddim_steps': 50,30 'toggles': [1, 4, 5],31 'sampler_name': 'k_lms',32 'ddim_eta': 0.0,33 'n_iter': 1,34 'batch_size': 1,35 'cfg_scale': 5.0,36 'denoising_strength': 0.75,37 'mask_mode': 1,38 'resize_mode': 0,39 'seed': '',40 'height': 512,41 'width': 512,42 'fp': None,43}44sample_img2img = None45job_manager = None46RealESRGAN = True47show_embeddings = False48 49img2img_resize_modes = [50 "Just resize",51 "Crop and resize",52 "Resize and fill",53]54 55img2img_toggles = [56 'Create prompt matrix (separate multiple prompts using |, and get all combinations of them)',57 'Normalize Prompt Weights (ensure sum of weights add up to 1.0)',58 'Loopback (use images from previous batch when creating next batch)',59 'Random loopback seed',60 'Save individual images',61 'Save grid',62 'Sort samples by prompt',63 'Write sample info files',64 'Write sample info to one file',65 'jpg samples',66]67 68img2img_toggle_defaults = [img2img_toggles[i] for i in img2img_defaults['toggles']]69 70def read_content(file_path: str) -> str:71 """read the content of target file72 """73 with open(file_path, 'r', encoding='utf-8') as f:74 content = f.read()75 76 return content77 78def base2picture(resbase64):79 res=resbase64.split(',')[1]80 img_b64decode = base64.b64decode(res)81 image = io.BytesIO(img_b64decode)82 img = Image.open(image)83 return img 84 85def filter_content(raw_style: str):86 if "(" in raw_style:87 i = raw_style.index("(")88 else :89 i = -190 91 if i == -1:92 return raw_style93 else :94 return raw_style[:i]95 96def request_images(raw_text, class_draw, style_draw, batch_size):97 if filter_content(class_draw) != "国画":98 if filter_content(class_draw) != "通用":99 raw_text = raw_text + f"，{filter_content(class_draw)}"100 101 for sty in style_draw:102 raw_text = raw_text + f"，{filter_content(sty)}"103 print(f"raw text is {raw_text}")104 url = "http://flagart.baai.ac.cn/api/general/"105 elif filter_content(class_draw) == "国画":106 if raw_text.endswith("国画"):107 pass 108 else :109 raw_text = raw_text + "，国画"110 url = "http://flagart.baai.ac.cn/api/guohua/"111 112 d = {"data":[raw_text, batch_size]}113 r = requests.post(url, json=d, headers={"Content-Type": "application/json", "Accept": "*/*", "Accept-Encoding": "gzip, deflate, br", "Connection": "keep-alive"})114 result_text = r.text115 content = json.loads(result_text)["data"][0]116 images = []117 for i in range(batch_size):118 # print(content[i])119 images.append(base2picture(content[i])) 120 121 return images 122 123def call_JS(sd_method, **kwargs):124 param_str = json.dumps(kwargs)125 return f"async (...x) => {{ return await SD.{sd_method}({{ x, ...{param_str} }}) ?? []; }}"126 127def encode_pil_to_base64(pil_image):128 with io.BytesIO() as output_bytes:129 130 # Copy any text-only metadata131 use_metadata = False132 metadata = PngImagePlugin.PngInfo()133 for key, value in pil_image.info.items():134 if isinstance(key, str) and isinstance(value, str):135 metadata.add_text(key, value)136 use_metadata = True137 138 pil_image.save(139 output_bytes, "PNG", pnginfo=(metadata if use_metadata else None)140 )141 bytes_data = output_bytes.getvalue()142 base64_str = str(base64.b64encode(bytes_data), "utf-8")143 return "data:image/png;base64," + base64_str144 145def img2img(*args):146 147 # 处理image148 for i, item in enumerate(args):149 # print(type(item))150 if type(item) == dict:151 args[i]['image'] = encode_pil_to_base64(item['image'])152 args[i]['mask'] = encode_pil_to_base64(item['mask'])153 else:154 print(i,type(item))155 print(item)156 157 batch_size = args[8]158 159 url = "http://flagart.baai.ac.cn/api/img2img/"160 d = {"data":args}161 r = requests.post(url, json=d, headers={"Content-Type": "application/json", "Accept": "*/*", "Accept-Encoding": "gzip, deflate, br", "Connection": "keep-alive"})162 print(r)163 result_text = r.text164 content = json.loads(result_text)["data"][0]165 images = []166 for i in range(batch_size):167 # print(content[i])168 images.append(base2picture(content[i])) 169 # content = json.loads(result_text)170 # print(result_text)171 # print("服务器已经把东西返回来啦！！！！！！！乌拉乌拉！！！！！")172 return images173 174 175examples = [176 '水墨蝴蝶和牡丹花，国画',177 '苍劲有力的墨竹，国画',178 '暴风雨中的灯塔',179 '机械小松鼠，科学幻想',180 '中国水墨山水画，国画',181 "Lighthouse in the storm",182 "A dog",183 "Landscape by 张大千",184 "A tiger 长了兔子耳朵",185 "A baby bird 铅笔素描",186 187]188 189if __name__ == "__main__":190 block = gr.Blocks(css=read_content('style.css'))191 192 with block:193 # gr.HTML(read_content("header.html"))194 195 with gr.Tabs() as tabs:196 197 with gr.TabItem("文生图(Text-to-img)"):198 199 with gr.Group():200 with gr.Box():201 with gr.Row().style(mobile_collapse=False, equal_height=True):202 text = gr.Textbox(203 label="Prompt",204 show_label=False,205 max_lines=1,206 placeholder="Input text(输入文字)",207 interactive=True,208 ).style(209 border=(True, False, True, True),210 rounded=(True, False, False, True),211 container=False,212 )213 214 btn = gr.Button("Generate image").style(215 margin=False,216 rounded=(True, True, True, True),217 )218 with gr.Row().style(mobile_collapse=False, equal_height=True):219 class_draw = gr.Dropdown(["通用(general)", "国画(traditional Chinese painting)", 220 "照片，摄影(picture photography)", "油画(oil painting)", 221 "铅笔素描(pencil sketch)", "CG", 222 "水彩画(watercolor painting)", "水墨画(ink and wash)",223 "插画(illustrations)", "3D", "图生图(img2img)"],224 label="生成类型(type)",225 show_label=True,226 value="通用(general)")227 with gr.Row().style(mobile_collapse=False, equal_height=True):228 style_draw = gr.CheckboxGroup(["蒸汽朋克(steampunk)", "电影摄影风格(film photography)", 229 "概念艺术(concept art)", "Warming lighting", 230 "Dramatic lighting", "Natural lighting", 231 "虚幻引擎(unreal engine)", "4k", "8k",232 "充满细节(full details)"],233 label="画面风格(style)",234 show_label=True,235 )236 with gr.Row().style(mobile_collapse=False, equal_height=True):237 sample_size = gr.Slider(minimum=1,238 maximum=4,239 step=1,240 label="生成数量(number)",241 show_label=True,242 interactive=True,243 )244 245 gallery = gr.Gallery(246 label="Generated images", show_label=False, elem_id="gallery"247 ).style(grid=[2], height="auto")248 249 gr.Examples(examples=examples, fn=request_images, inputs=text, outputs=gallery, examples_per_page=100)250 text.submit(request_images, inputs=[text, class_draw, style_draw, sample_size], outputs=gallery)251 btn.click(request_images, inputs=[text, class_draw, style_draw, sample_size], outputs=gallery)252 253 with gr.TabItem("图生图(Img-to-Img)", id="img2img_tab"):254 with gr.Row(elem_id="prompt_row"):255 img2img_prompt = gr.Textbox(label="Prompt",256 elem_id='img2img_prompt_input',257 placeholder="神奇的森林，流淌的河流.",258 lines=1,259 max_lines=1 if txt2img_defaults['submit_on_enter'] == 'Yes' else 25,260 value=img2img_defaults['prompt'],261 show_label=False).style()262 263 img2img_btn_mask = gr.Button("Generate", variant="primary", visible=False,264 elem_id="img2img_mask_btn")265 img2img_btn_editor = gr.Button("Generate", variant="primary", elem_id="img2img_edit_btn")266 with gr.Row().style(equal_height=False):267 with gr.Column():268 gr.Markdown('#### 输入图像')269 img2img_image_mask = gr.Image(270 value=sample_img2img,271 source="upload",272 interactive=True,273 tool="sketch",274 type='pil',275 elem_id="img2img_mask",276 image_mode="RGBA"277 )278 img2img_image_editor = gr.Image(279 value=sample_img2img,280 source="upload",281 interactive=True,282 tool="select",283 type='pil',284 visible=False,285 image_mode="RGBA",286 elem_id="img2img_editor"287 )288 289 with gr.Tabs():290 with gr.TabItem("编辑设置"):291 with gr.Row():292 # disable Uncrop for now293 choices=["Mask", "Crop", "Uncrop"]294 img2img_image_editor_mode = gr.Radio(choices=["Mask"],295 label="编辑模式",296 value="Mask", elem_id='edit_mode_select',297 visible=True)298 img2img_mask = gr.Radio(choices=["保留mask区域", "生成mask区域"],299 label="Mask 方式",300 #value=img2img_mask_modes[img2img_defaults['mask_mode']],301 value = "生成mask区域",302 visible=True)303 304 img2img_mask_blur_strength = gr.Slider(minimum=1, maximum=10, step=1,305 label="How much blurry should the mask be? (to avoid hard edges)",306 value=3, visible=False)307 308 img2img_resize = gr.Radio(label="Resize mode",309 choices=["Just resize", "Crop and resize",310 "Resize and fill"],311 value=img2img_resize_modes[312 img2img_defaults['resize_mode']], visible=False)313 314 img2img_painterro_btn = gr.Button("Advanced Editor",visible=False)315 # with gr.TabItem("Hints",visible=False):316 # img2img_help = gr.Markdown(visible=False, value=uifn.help_text)317 318 with gr.Column():319 gr.Markdown('#### 编辑后的图片')320 output_img2img_gallery = gr.Gallery(label="Images", elem_id="img2img_gallery_output").style(321 grid=[4, 4, 4])322 img2img_job_ui = job_manager.draw_gradio_ui() if job_manager else None323 with gr.Column(visible=False):324 with gr.Tabs(visible=False):325 with gr.TabItem("", id="img2img_actions_tab",visible=False):326 gr.Markdown("Select an image, then press one of the buttons below")327 with gr.Row():328 output_img2img_copy_to_clipboard_btn = gr.Button("Copy to clipboard")329 output_img2img_copy_to_input_btn = gr.Button("Push to img2img input")330 output_img2img_copy_to_mask_btn = gr.Button("Push to img2img input mask")331 332 gr.Markdown("Warning: This will clear your current image and mask settings!")333 with gr.TabItem("", id="img2img_output_info_tab",visible=False):334 output_img2img_params = gr.Textbox(label="Generation parameters")335 with gr.Row():336 output_img2img_copy_params = gr.Button("Copy full parameters").click(337 inputs=output_img2img_params, outputs=[],338 _js='(x) => {navigator.clipboard.writeText(x.replace(": ",":"))}', fn=None,339 show_progress=False)340 output_img2img_seed = gr.Number(label='Seed', interactive=False, visible=False)341 output_img2img_copy_seed = gr.Button("Copy only seed").click(342 inputs=output_img2img_seed, outputs=[],343 _js=call_JS("gradioInputToClipboard"), fn=None, show_progress=False)344 output_img2img_stats = gr.HTML(label='Stats')345 346 gr.Markdown('# 编辑设置')347 348 with gr.Row():349 with gr.Column():350 img2img_width = gr.Slider(minimum=64, maximum=2048, step=64, label="图片宽度",351 value=img2img_defaults["width"])352 img2img_height = gr.Slider(minimum=64, maximum=2048, step=64, label="图片高度",353 value=img2img_defaults["height"])354 img2img_cfg = gr.Slider(minimum=-40.0, maximum=30.0, step=0.5,355 label='文本引导强度',356 value=img2img_defaults['cfg_scale'], elem_id='cfg_slider')357 img2img_seed = gr.Textbox(label="随机种子", lines=1, max_lines=1,358 value=img2img_defaults["seed"])359 img2img_batch_count = gr.Slider(minimum=1, maximum=50, step=1,360 label='生成数量',361 value=img2img_defaults['n_iter'])362 img2img_dimensions_info_text_box = gr.Textbox(363 label="长宽比设置")364 with gr.Column():365 img2img_steps = gr.Slider(minimum=1, maximum=250, step=1, label="采样步数",366 value=img2img_defaults['ddim_steps'])367 368 img2img_sampling = gr.Dropdown(label='采样方式',369 choices=["DDIM", 'k_dpm_2_a', 'k_dpm_2', 'k_euler_a', 'k_euler',370 'k_heun', 'k_lms'],371 value=img2img_defaults['sampler_name'])372 373 img2img_denoising = gr.Slider(minimum=0.0, maximum=1.0, step=0.01, label='Denoising Strength',374 value=img2img_defaults['denoising_strength'],visible=False)375 376 img2img_toggles = gr.CheckboxGroup(label='', choices=img2img_toggles,377 value=img2img_toggle_defaults,visible=False)378 379 img2img_realesrgan_model_name = gr.Dropdown(label='RealESRGAN model',380 choices=['RealESRGAN_x4plus',381 'RealESRGAN_x4plus_anime_6B'],382 value='RealESRGAN_x4plus',383 visible=RealESRGAN is not None) # TODO: Feels like I shouldnt slot it in here.384 385 img2img_embeddings = gr.File(label="Embeddings file for textual inversion",386 visible=show_embeddings)387 388 img2img_image_editor_mode.change(389 uifn.change_image_editor_mode,390 [img2img_image_editor_mode,391 img2img_image_editor,392 img2img_image_mask,393 img2img_resize,394 img2img_width,395 img2img_height396 ],397 [img2img_image_editor, img2img_image_mask, img2img_btn_editor, img2img_btn_mask,398 img2img_painterro_btn, img2img_mask, img2img_mask_blur_strength]399 )400 401 # img2img_image_editor_mode.change(402 # uifn.update_image_mask,403 # [img2img_image_editor, img2img_resize, img2img_width, img2img_height],404 # img2img_image_mask405 # )406 407 # output_txt2img_copy_to_input_btn.click(408 # uifn.copy_img_to_input,409 # [output_txt2img_gallery],410 # [img2img_image_editor, img2img_image_mask, tabs],411 # _js=call_JS("moveImageFromGallery",412 # fromId="txt2img_gallery_output",413 # toId="img2img_editor")414 # )415 416 output_img2img_copy_to_input_btn.click(417 uifn.copy_img_to_edit,418 [output_img2img_gallery],419 [img2img_image_editor, tabs, img2img_image_editor_mode],420 _js=call_JS("moveImageFromGallery",421 fromId="img2img_gallery_output",422 toId="img2img_editor")423 )424 output_img2img_copy_to_mask_btn.click(425 uifn.copy_img_to_mask,426 [output_img2img_gallery],427 [img2img_image_mask, tabs, img2img_image_editor_mode],428 _js=call_JS("moveImageFromGallery",429 fromId="img2img_gallery_output",430 toId="img2img_editor")431 )432 433 output_img2img_copy_to_clipboard_btn.click(fn=None, inputs=output_img2img_gallery, outputs=[],434 _js=call_JS("copyImageFromGalleryToClipboard",435 fromId="img2img_gallery_output")436 )437 438 img2img_func = img2img439 img2img_inputs = [img2img_prompt, img2img_image_editor_mode, img2img_mask,440 img2img_mask_blur_strength, img2img_steps, img2img_sampling, img2img_toggles,441 img2img_realesrgan_model_name, img2img_batch_count, img2img_cfg,442 img2img_denoising, img2img_seed, img2img_height, img2img_width, img2img_resize,443 img2img_image_mask]444 # img2img_outputs = [output_img2img_gallery, output_img2img_seed, output_img2img_params,445 # output_img2img_stats]446 447 img2img_outputs = [output_img2img_gallery]448 449 # If a JobManager was passed in then wrap the Generate functions450 if img2img_job_ui:451 img2img_func, img2img_inputs, img2img_outputs = img2img_job_ui.wrap_func(452 func=img2img_func,453 inputs=img2img_inputs,454 outputs=img2img_outputs,455 )456 457 img2img_btn_mask.click(458 img2img_func,459 img2img_inputs,460 img2img_outputs461 )462 463 def img2img_submit_params():464 # print([img2img_prompt, img2img_image_editor_mode, img2img_mask,465 # img2img_mask_blur_strength, img2img_steps, img2img_sampling, img2img_toggles,466 # img2img_realesrgan_model_name, img2img_batch_count, img2img_cfg,467 # img2img_denoising, img2img_seed, img2img_height, img2img_width, img2img_resize,468 # img2img_image_editor, img2img_image_mask, img2img_embeddings])469 return (img2img_func,470 img2img_inputs,471 img2img_outputs)472 473 img2img_btn_editor.click(*img2img_submit_params())474 475 # GENERATE ON ENTER476 img2img_prompt.submit(None, None, None,477 _js=call_JS("clickFirstVisibleButton",478 rowId="prompt_row"))479 480 img2img_painterro_btn.click(None,481 [img2img_image_editor, img2img_image_mask, img2img_image_editor_mode],482 [img2img_image_editor, img2img_image_mask],483 _js=call_JS("Painterro.init", toId="img2img_editor")484 )485 486 img2img_width.change(fn=uifn.update_dimensions_info, inputs=[img2img_width, img2img_height],487 outputs=img2img_dimensions_info_text_box)488 img2img_height.change(fn=uifn.update_dimensions_info, inputs=[img2img_width, img2img_height],489 outputs=img2img_dimensions_info_text_box)490 491 # gr.HTML(read_content("footer.html"))492 # gr.Image('./contributors.png')493 494 block.queue(max_size=50, concurrency_count=20).launch()