From d6d14859e331c0bd9bf79dcef0319304f30f2417 Mon Sep 17 00:00:00 2001 From: Artiprocher Date: Wed, 21 Aug 2024 16:57:56 +0800 Subject: [PATCH] update UI --- apps/gradio/DiffSynth_Studio.py | 235 ++++++++++++++++++ .../streamlit/DiffSynth_Studio.py | 0 .../streamlit/pages}/1_Image_Creator.py | 13 +- .../streamlit/pages}/2_Video_Creator.py | 0 mask.jpg | Bin 0 -> 17011 bytes 5 files changed, 247 insertions(+), 1 deletion(-) create mode 100644 apps/gradio/DiffSynth_Studio.py rename DiffSynth_Studio.py => apps/streamlit/DiffSynth_Studio.py (100%) rename {pages => apps/streamlit/pages}/1_Image_Creator.py (95%) rename {pages => apps/streamlit/pages}/2_Video_Creator.py (100%) create mode 100644 mask.jpg diff --git a/apps/gradio/DiffSynth_Studio.py b/apps/gradio/DiffSynth_Studio.py new file mode 100644 index 0000000..05dfba3 --- /dev/null +++ b/apps/gradio/DiffSynth_Studio.py @@ -0,0 +1,235 @@ +import gradio as gr +from diffsynth import ModelManager, SDImagePipeline, SDXLImagePipeline, SD3ImagePipeline, HunyuanDiTImagePipeline, FluxImagePipeline +import os, torch +from PIL import Image +import numpy as np + + +config = { + "Stable Diffusion": { + "model_folder": "models/stable_diffusion", + "pipeline_class": SDImagePipeline, + "default_parameters": { + "height": 512, + "width": 512, + } + }, + "Stable Diffusion XL": { + "model_folder": "models/stable_diffusion_xl", + "pipeline_class": SDXLImagePipeline, + "default_parameters": {} + }, + "Stable Diffusion 3": { + "model_folder": "models/stable_diffusion_3", + "pipeline_class": SD3ImagePipeline, + "default_parameters": {} + }, + "Stable Diffusion XL Turbo": { + "model_folder": "models/stable_diffusion_xl_turbo", + "pipeline_class": SDXLImagePipeline, + "default_parameters": { + "negative_prompt": "", + "cfg_scale": 1.0, + "num_inference_steps": 1, + "height": 512, + "width": 512, + } + }, + "Kolors": { + "model_folder": "models/kolors", + "pipeline_class": SDXLImagePipeline, + "default_parameters": {} + }, + "HunyuanDiT": { + "model_folder": "models/HunyuanDiT", + "pipeline_class": HunyuanDiTImagePipeline, + "default_parameters": {} + }, + "FLUX": { + "model_folder": "models/FLUX", + "pipeline_class": FluxImagePipeline, + "default_parameters": { + "cfg_scale": 1.0, + } + } +} +MAX_NUM_PAINTER_LAYERS = 8 + + +def load_model_list(model_type): + if model_type is None: + return [] + folder = config[model_type]["model_folder"] + file_list = [i for i in os.listdir(folder) if i.endswith(".safetensors")] + if model_type in ["HunyuanDiT", "Kolors", "FLUX"]: + file_list += [i for i in os.listdir(folder) if os.path.isdir(os.path.join(folder, i))] + file_list = sorted(file_list) + return file_list + + +def load_model(model_type, model_path): + model_path = os.path.join(config[model_type]["model_folder"], model_path) + model_manager = ModelManager() + if model_type == "HunyuanDiT": + model_manager.load_models([ + os.path.join(model_path, "clip_text_encoder/pytorch_model.bin"), + os.path.join(model_path, "mt5/pytorch_model.bin"), + os.path.join(model_path, "model/pytorch_model_ema.pt"), + os.path.join(model_path, "sdxl-vae-fp16-fix/diffusion_pytorch_model.bin"), + ]) + elif model_type == "Kolors": + model_manager.load_models([ + os.path.join(model_path, "text_encoder"), + os.path.join(model_path, "unet/diffusion_pytorch_model.safetensors"), + os.path.join(model_path, "vae/diffusion_pytorch_model.safetensors"), + ]) + elif model_type == "FLUX": + model_manager.torch_dtype = torch.bfloat16 + file_list = [ + os.path.join(model_path, "text_encoder/model.safetensors"), + os.path.join(model_path, "text_encoder_2"), + ] + for file_name in os.listdir(model_path): + if file_name.endswith(".safetensors"): + file_list.append(os.path.join(model_path, file_name)) + model_manager.load_models(file_list) + else: + model_manager.load_model(model_path) + pipe = config[model_type]["pipeline_class"].from_model_manager(model_manager) + return model_manager, pipe + + + +model_manager: ModelManager = None +pipe = None + +with gr.Blocks() as app: + gr.Markdown("# DiffSynth-Studio Painter") + with gr.Row(): + with gr.Column(scale=382, min_width=100): + + with gr.Accordion(label="Model"): + model_type = gr.Dropdown(choices=[i for i in config], label="Model type") + model_path = gr.Dropdown(choices=[], interactive=True, label="Model path") + + @gr.on(inputs=model_type, outputs=model_path, triggers=model_type.change) + def model_type_to_model_path(model_type): + return gr.Dropdown(choices=load_model_list(model_type)) + + with gr.Accordion(label="Prompt"): + prompt = gr.Textbox(label="Prompt", lines=3) + negative_prompt = gr.Textbox(label="Negative prompt", lines=1) + cfg_scale = gr.Slider(minimum=1.0, maximum=10.0, value=7.0, step=0.1, interactive=True, label="Classifier-free guidance scale") + embedded_guidance = gr.Slider(minimum=0.0, maximum=10.0, value=0.0, step=0.1, interactive=True, label="Embedded guidance scale (only for FLUX)") + + with gr.Accordion(label="Image"): + num_inference_steps = gr.Slider(minimum=1, maximum=100, value=20, step=1, interactive=True, label="Inference steps") + height = gr.Slider(minimum=64, maximum=2048, value=1024, step=64, interactive=True, label="Height") + width = gr.Slider(minimum=64, maximum=2048, value=1024, step=64, interactive=True, label="Width") + with gr.Column(): + use_fixed_seed = gr.Checkbox(value=True, interactive=False, label="Use fixed seed") + seed = gr.Number(minimum=0, maximum=10**9, value=0, interactive=True, label="Random seed", show_label=False) + + @gr.on( + inputs=[model_type, model_path, prompt, negative_prompt, cfg_scale, embedded_guidance, num_inference_steps, height, width], + outputs=[prompt, negative_prompt, cfg_scale, embedded_guidance, num_inference_steps, height, width], + triggers=model_path.change + ) + def model_path_to_default_params(model_type, model_path, prompt, negative_prompt, cfg_scale, embedded_guidance, num_inference_steps, height, width): + global model_manager, pipe + if isinstance(model_manager, ModelManager): + model_manager.to("cpu") + torch.cuda.empty_cache() + model_manager, pipe = load_model(model_type, model_path) + cfg_scale = config[model_type]["default_parameters"].get("cfg_scale", cfg_scale) + embedded_guidance = config[model_type]["default_parameters"].get("embedded_guidance", embedded_guidance) + num_inference_steps = config[model_type]["default_parameters"].get("num_inference_steps", num_inference_steps) + height = config[model_type]["default_parameters"].get("height", height) + width = config[model_type]["default_parameters"].get("width", width) + return prompt, negative_prompt, cfg_scale, embedded_guidance, num_inference_steps, height, width + + + with gr.Column(scale=618, min_width=100): + with gr.Accordion(label="Painter"): + enable_local_prompt_list = [] + local_prompt_list = [] + mask_scale_list = [] + canvas_list = [] + for painter_layer_id in range(MAX_NUM_PAINTER_LAYERS): + with gr.Tab(label=f"Layer {painter_layer_id}"): + enable_local_prompt = gr.Checkbox(label="Enable", value=False, key=f"enable_local_prompt_{painter_layer_id}") + local_prompt = gr.Textbox(label="Local prompt", key=f"local_prompt_{painter_layer_id}") + mask_scale = gr.Slider(minimum=0.0, maximum=5.0, value=1.0, step=0.1, interactive=True, label="Mask scale", key=f"mask_scale_{painter_layer_id}") + canvas = gr.ImageEditor(canvas_size=(512, 1), sources=None, layers=False, interactive=True, image_mode="RGBA", + brush=gr.Brush(default_size=100, default_color="#000000", colors=["#000000"]), + label="Painter", key=f"canvas_{painter_layer_id}") + @gr.on(inputs=[height, width, canvas], outputs=canvas, triggers=[height.change, width.change, canvas.clear, enable_local_prompt.change], show_progress="hidden") + def resize_canvas(height, width, canvas): + h, w = canvas["background"].shape[:2] + if h != height or width != w: + return np.ones((height, width, 3), dtype=np.uint8) * 255 + else: + return canvas + + enable_local_prompt_list.append(enable_local_prompt) + local_prompt_list.append(local_prompt) + mask_scale_list.append(mask_scale) + canvas_list.append(canvas) + with gr.Accordion(label="Results"): + run_button = gr.Button(value="Generate", variant="primary") + output_image = gr.Image(sources=None, show_label=False, interactive=False, type="pil") + with gr.Row(): + with gr.Column(): + output_to_painter_button = gr.Button(value="Set as painter's background") + with gr.Column(): + output_to_input_button = gr.Button(value="Set as input image") + painter_background = gr.State(None) + input_background = gr.State(None) + @gr.on( + inputs=[prompt, negative_prompt, cfg_scale, embedded_guidance, num_inference_steps, height, width, seed] + enable_local_prompt_list + local_prompt_list + mask_scale_list + canvas_list, + outputs=[output_image], + triggers=run_button.click + ) + def generate_image(prompt, negative_prompt, cfg_scale, embedded_guidance, num_inference_steps, height, width, seed, *args, progress=gr.Progress()): + global pipe + input_params = { + "prompt": prompt, + "negative_prompt": negative_prompt, + "cfg_scale": cfg_scale, + "num_inference_steps": num_inference_steps, + "height": height, + "width": width, + "progress_bar_cmd": progress.tqdm, + } + if isinstance(pipe, FluxImagePipeline): + input_params["embedded_guidance"] = embedded_guidance + enable_local_prompt_list, local_prompt_list, mask_scale_list, canvas_list = ( + args[0 * MAX_NUM_PAINTER_LAYERS: 1 * MAX_NUM_PAINTER_LAYERS], + args[1 * MAX_NUM_PAINTER_LAYERS: 2 * MAX_NUM_PAINTER_LAYERS], + args[2 * MAX_NUM_PAINTER_LAYERS: 3 * MAX_NUM_PAINTER_LAYERS], + args[3 * MAX_NUM_PAINTER_LAYERS: 4 * MAX_NUM_PAINTER_LAYERS] + ) + local_prompts, masks, mask_scales = [], [], [] + for enable_local_prompt, local_prompt, mask_scale, canvas in zip( + enable_local_prompt_list, local_prompt_list, mask_scale_list, canvas_list + ): + if enable_local_prompt: + local_prompts.append(local_prompt) + masks.append(Image.fromarray(canvas["layers"][0][:, :, -1]).convert("RGB")) + mask_scales.append(mask_scale) + input_params.update({ + "local_prompts": local_prompts, + "masks": masks, + "mask_scales": mask_scales, + }) + torch.manual_seed(seed) + image = pipe(**input_params) + return image + + @gr.on(inputs=[output_image] + canvas_list, outputs=canvas_list, triggers=output_to_painter_button.click) + def send_output_to_painter_background(output_image, *canvas_list): + for canvas in canvas_list: + h, w = canvas["background"].shape[:2] + canvas["background"] = output_image.resize((w, h)) + return tuple(canvas_list) +app.launch() diff --git a/DiffSynth_Studio.py b/apps/streamlit/DiffSynth_Studio.py similarity index 100% rename from DiffSynth_Studio.py rename to apps/streamlit/DiffSynth_Studio.py diff --git a/pages/1_Image_Creator.py b/apps/streamlit/pages/1_Image_Creator.py similarity index 95% rename from pages/1_Image_Creator.py rename to apps/streamlit/pages/1_Image_Creator.py index 3b8ad45..732d219 100644 --- a/pages/1_Image_Creator.py +++ b/apps/streamlit/pages/1_Image_Creator.py @@ -1,4 +1,4 @@ -import torch, os, io +import torch, os, io, json, time import numpy as np from PIL import Image import streamlit as st @@ -275,6 +275,7 @@ with column_input: num_painter_layer = st.number_input("Number of painter layers", min_value=0, max_value=10, step=1, value=0) local_prompts, masks, mask_scales = [], [], [] white_board = Image.fromarray(np.ones((512, 512, 3), dtype=np.uint8) * 255) + painter_layers_json_data = [] for painter_tab_id in range(num_painter_layer): with st.expander(f"Painter layer {painter_tab_id}", expanded=True): enable_local_prompt = st.checkbox(f"Enable prompt {painter_tab_id}", value=True) @@ -293,6 +294,9 @@ with column_input: drawing_mode="freedraw", key=f"canvas_{painter_tab_id}" ) + if canvas_result_local.json_data is not None: + painter_layers_json_data.append(canvas_result_local.json_data.copy()) + painter_layers_json_data[-1]["prompt"] = local_prompt if enable_local_prompt: local_prompts.append(local_prompt) if canvas_result_local.image_data is not None: @@ -302,6 +306,13 @@ with column_input: mask = Image.fromarray(255 - np.array(mask)) masks.append(mask) mask_scales.append(mask_scale) + save_painter_layers = st.button("Save painter layers") + if save_painter_layers: + os.makedirs("data/painter_layers", exist_ok=True) + json_file_path = f"data/painter_layers/{time.time_ns()}.json" + with open(json_file_path, "w") as f: + json.dump(painter_layers_json_data, f, indent=4) + st.markdown(f"Painter layers are saved in {json_file_path}.") with column_output: diff --git a/pages/2_Video_Creator.py b/apps/streamlit/pages/2_Video_Creator.py similarity index 100% rename from pages/2_Video_Creator.py rename to apps/streamlit/pages/2_Video_Creator.py diff --git a/mask.jpg b/mask.jpg new file mode 100644 index 0000000000000000000000000000000000000000..b003d9a16d1d4dbe630cae4962a8d844ec1d3310 GIT binary patch literal 17011 zcmeIwIZzW(9LMqZ_PxE5&A=uBLnatG0M{ zys<;Ys}jYt!xOL(Z>_)^3)c-^GgfM2^ZUPf)xO!;@4k=KYds>FRjJApp_CA6FJe6; z6-1FFSrQdlmQ_`8XkM4*bZR+nPlh+w9~_@dV~F(a9wnV5@nL7x>gR?dyS^nHYHNt(Un!xGZ;`}fPfBu+QykM+3enN3c>BNeu(^AuCR94NLHG9t7dGi-6T(o$}(q+q6tX#EvP2<{i z>o;!Nyk+aQ?JYZY?%KU)@4o#94jn#n^w{we?I%y2KGS*j-1!R^FJI}tdhPm+o40P? zxqI*agWkUWfyYmtK70P+<*V02Z{EIp|Ka1O;m_$VYIpvaZ(+Z6>2{aEIpcD=iwceQ zW;z!`3CUenE!Wn2!bX$gEpKh>=u#uen!&8Ph8{qV>+2ahQGAL4+_}$2QUBwFaQHE00S@p126ysFaQHE00S@p126ysFaQHE00S@p z126ysFaQHE00S@p126ysFaQHE00S@p126ysFaQHE00S@p126ysFaQHE00S@p126ys PFaQHE@LvsNSbbjsIC>6U literal 0 HcmV?d00001