mirror of
https://github.com/modelscope/DiffSynth-Studio.git
synced 2026-03-18 22:08:13 +00:00
265 lines
11 KiB
Python
265 lines
11 KiB
Python
import torch, os, io
|
|
import numpy as np
|
|
from PIL import Image
|
|
import streamlit as st
|
|
st.set_page_config(layout="wide")
|
|
from streamlit_drawable_canvas import st_canvas
|
|
from diffsynth.models import ModelManager
|
|
from diffsynth.prompts import SDXLPrompter, SDPrompter
|
|
from diffsynth.pipelines import SDXLPipeline, SDPipeline
|
|
|
|
|
|
torch.cuda.set_per_process_memory_fraction(0.999, 0)
|
|
|
|
|
|
@st.cache_data
|
|
def load_model_list(folder):
|
|
file_list = os.listdir(folder)
|
|
file_list = [i for i in file_list if i.endswith(".safetensors")]
|
|
file_list = sorted(file_list)
|
|
return file_list
|
|
|
|
|
|
def detect_model_path(sd_model_path, sdxl_model_path):
|
|
if sd_model_path != "None":
|
|
model_path = os.path.join("models/stable_diffusion", sd_model_path)
|
|
elif sdxl_model_path != "None":
|
|
model_path = os.path.join("models/stable_diffusion_xl", sdxl_model_path)
|
|
else:
|
|
model_path = None
|
|
return model_path
|
|
|
|
|
|
def load_model(sd_model_path, sdxl_model_path):
|
|
if sd_model_path != "None":
|
|
model_path = os.path.join("models/stable_diffusion", sd_model_path)
|
|
model_manager = ModelManager()
|
|
model_manager.load_from_safetensors(model_path)
|
|
prompter = SDPrompter()
|
|
pipeline = SDPipeline()
|
|
elif sdxl_model_path != "None":
|
|
model_path = os.path.join("models/stable_diffusion_xl", sdxl_model_path)
|
|
model_manager = ModelManager()
|
|
model_manager.load_from_safetensors(model_path)
|
|
prompter = SDXLPrompter()
|
|
pipeline = SDXLPipeline()
|
|
else:
|
|
return None, None, None, None
|
|
return model_path, model_manager, prompter, pipeline
|
|
|
|
|
|
def release_model():
|
|
if "model_manager" in st.session_state:
|
|
st.session_state["model_manager"].to("cpu")
|
|
del st.session_state["loaded_model_path"]
|
|
del st.session_state["model_manager"]
|
|
del st.session_state["prompter"]
|
|
del st.session_state["pipeline"]
|
|
torch.cuda.empty_cache()
|
|
|
|
|
|
def use_output_image_as_input():
|
|
# Search for input image
|
|
output_image_id = 0
|
|
selected_output_image = None
|
|
while True:
|
|
if f"use_output_as_input_{output_image_id}" not in st.session_state:
|
|
break
|
|
if st.session_state[f"use_output_as_input_{output_image_id}"]:
|
|
selected_output_image = st.session_state["output_images"][output_image_id]
|
|
break
|
|
output_image_id += 1
|
|
if selected_output_image is not None:
|
|
st.session_state["input_image"] = selected_output_image
|
|
|
|
|
|
def apply_stroke_to_image(stroke_image, image):
|
|
image = np.array(image.convert("RGB")).astype(np.float32)
|
|
height, width, _ = image.shape
|
|
|
|
stroke_image = np.array(Image.fromarray(stroke_image).resize((width, height))).astype(np.float32)
|
|
weight = stroke_image[:, :, -1:] / 255
|
|
stroke_image = stroke_image[:, :, :-1]
|
|
|
|
image = stroke_image * weight + image * (1 - weight)
|
|
image = np.clip(image, 0, 255).astype(np.uint8)
|
|
image = Image.fromarray(image)
|
|
return image
|
|
|
|
|
|
@st.cache_data
|
|
def image2bits(image):
|
|
image_byte = io.BytesIO()
|
|
image.save(image_byte, format="PNG")
|
|
image_byte = image_byte.getvalue()
|
|
return image_byte
|
|
|
|
|
|
def show_output_image(image):
|
|
st.image(image, use_column_width="always")
|
|
st.button("Use it as input image", key=f"use_output_as_input_{image_id}")
|
|
st.download_button("Download", data=image2bits(image), file_name="image.png", mime="image/png", key=f"download_output_{image_id}")
|
|
|
|
|
|
column_input, column_output = st.columns(2)
|
|
|
|
# with column_input:
|
|
with st.sidebar:
|
|
# Select a model
|
|
with st.expander("Model", expanded=True):
|
|
sd_model_list = ["None"] + load_model_list("models/stable_diffusion")
|
|
sd_model_path = st.selectbox(
|
|
"Stable Diffusion", sd_model_list
|
|
)
|
|
sdxl_model_list = ["None"] + load_model_list("models/stable_diffusion_xl")
|
|
sdxl_model_path = st.selectbox(
|
|
"Stable Diffusion XL", sdxl_model_list
|
|
)
|
|
|
|
# Load the model
|
|
model_path = detect_model_path(sd_model_path, sdxl_model_path)
|
|
if model_path is None:
|
|
st.markdown("No models selected.")
|
|
release_model()
|
|
elif st.session_state.get("loaded_model_path", "") != model_path:
|
|
st.markdown(f"Using model at {model_path}.")
|
|
release_model()
|
|
model_path, model_manager, prompter, pipeline = load_model(sd_model_path, sdxl_model_path)
|
|
st.session_state.loaded_model_path = model_path
|
|
st.session_state.model_manager = model_manager
|
|
st.session_state.prompter = prompter
|
|
st.session_state.pipeline = pipeline
|
|
else:
|
|
st.markdown(f"Using model at {model_path}.")
|
|
model_path, model_manager, prompter, pipeline = (
|
|
st.session_state.loaded_model_path,
|
|
st.session_state.model_manager,
|
|
st.session_state.prompter,
|
|
st.session_state.pipeline,
|
|
)
|
|
|
|
# Show parameters
|
|
with st.expander("Prompt", expanded=True):
|
|
column_positive, column_negative = st.columns(2)
|
|
prompt = st.text_area("Positive prompt")
|
|
negative_prompt = st.text_area("Negative prompt")
|
|
with st.expander("Classifier-free guidance", expanded=True):
|
|
use_cfg = st.checkbox("Use classifier-free guidance", value=True)
|
|
if use_cfg:
|
|
cfg_scale = st.slider("Classifier-free guidance scale", min_value=1.0, max_value=10.0, step=0.1, value=7.5)
|
|
else:
|
|
cfg_scale = 1.0
|
|
with st.expander("Inference steps", expanded=True):
|
|
num_inference_steps = st.slider("Inference steps", min_value=1, max_value=100, value=20, label_visibility="hidden")
|
|
with st.expander("Image size", expanded=True):
|
|
height = st.select_slider("Height", options=[256, 512, 768, 1024, 2048], value=512)
|
|
width = st.select_slider("Width", options=[256, 512, 768, 1024, 2048], value=512)
|
|
with st.expander("Seed", expanded=True):
|
|
use_fixed_seed = st.checkbox("Use fixed seed", value=False)
|
|
if use_fixed_seed:
|
|
seed = st.number_input("Random seed", value=0, label_visibility="hidden")
|
|
with st.expander("Number of images", expanded=True):
|
|
num_images = st.number_input("Number of images", value=4, label_visibility="hidden")
|
|
with st.expander("Tile (for high resolution)", expanded=True):
|
|
tiled = st.checkbox("Use tile", value=False)
|
|
tile_size = st.select_slider("Tile size", options=[64, 128], value=64)
|
|
tile_stride = st.select_slider("Tile stride", options=[8, 16, 32, 64], value=32)
|
|
|
|
|
|
# Show input image
|
|
with column_input:
|
|
with st.expander("Input image (Optional)", expanded=True):
|
|
with st.container(border=True):
|
|
column_white_board, column_upload_image = st.columns([1, 2])
|
|
with column_white_board:
|
|
create_white_board = st.button("Create white board")
|
|
delete_input_image = st.button("Delete input image")
|
|
with column_upload_image:
|
|
upload_image = st.file_uploader("Upload image", type=["png", "jpg"], key="upload_image")
|
|
|
|
if upload_image is not None:
|
|
st.session_state["input_image"] = Image.open(upload_image)
|
|
elif create_white_board:
|
|
st.session_state["input_image"] = Image.fromarray(np.ones((1024, 1024, 3), dtype=np.uint8) * 255)
|
|
else:
|
|
use_output_image_as_input()
|
|
|
|
if delete_input_image and "input_image" in st.session_state:
|
|
del st.session_state.input_image
|
|
if delete_input_image and "upload_image" in st.session_state:
|
|
del st.session_state.upload_image
|
|
|
|
input_image = st.session_state.get("input_image", None)
|
|
if input_image is not None:
|
|
with st.container(border=True):
|
|
column_drawing_mode, column_color_1, column_color_2 = st.columns([4, 1, 1])
|
|
with column_drawing_mode:
|
|
drawing_mode = st.radio("Drawing tool", ["transform", "freedraw", "line", "rect"], horizontal=True, index=1)
|
|
with column_color_1:
|
|
stroke_color = st.color_picker("Stroke color")
|
|
with column_color_2:
|
|
fill_color = st.color_picker("Fill color")
|
|
stroke_width = st.slider("Stroke width", min_value=1, max_value=50, value=10)
|
|
with st.container(border=True):
|
|
denoising_strength = st.slider("Denoising strength", min_value=0.0, max_value=1.0, value=0.7)
|
|
with st.container(border=True):
|
|
input_width, input_height = input_image.size
|
|
canvas_result = st_canvas(
|
|
fill_color=fill_color,
|
|
stroke_width=stroke_width,
|
|
stroke_color=stroke_color,
|
|
background_color="rgba(255, 255, 255, 0)",
|
|
background_image=input_image,
|
|
update_streamlit=True,
|
|
height=int(512 / input_width * input_height),
|
|
width=512,
|
|
drawing_mode=drawing_mode,
|
|
key="canvas"
|
|
)
|
|
|
|
|
|
with column_output:
|
|
run_button = st.button("Generate image", type="primary")
|
|
auto_update = st.checkbox("Auto update", value=False)
|
|
num_image_columns = st.slider("Columns", min_value=1, max_value=8, value=2)
|
|
image_columns = st.columns(num_image_columns)
|
|
|
|
# Run
|
|
if (run_button or auto_update) and model_path is not None:
|
|
|
|
if not use_fixed_seed:
|
|
torch.manual_seed(np.random.randint(0, 10**9))
|
|
|
|
output_images = []
|
|
for image_id in range(num_images):
|
|
if use_fixed_seed:
|
|
torch.manual_seed(seed + image_id)
|
|
if input_image is not None:
|
|
input_image = input_image.resize((width, height))
|
|
if canvas_result.image_data is not None:
|
|
input_image = apply_stroke_to_image(canvas_result.image_data, input_image)
|
|
else:
|
|
denoising_strength = 1.0
|
|
with image_columns[image_id % num_image_columns]:
|
|
progress_bar = st.progress(0.0)
|
|
image = pipeline(
|
|
model_manager, prompter,
|
|
prompt, negative_prompt=negative_prompt, cfg_scale=cfg_scale,
|
|
num_inference_steps=num_inference_steps,
|
|
height=height, width=width,
|
|
init_image=input_image, denoising_strength=denoising_strength,
|
|
tiled=tiled, tile_size=tile_size, tile_stride=tile_stride,
|
|
progress_bar_st=progress_bar
|
|
)
|
|
output_images.append(image)
|
|
progress_bar.progress(1.0)
|
|
show_output_image(image)
|
|
st.session_state["output_images"] = output_images
|
|
|
|
elif "output_images" in st.session_state:
|
|
for image_id in range(len(st.session_state.output_images)):
|
|
with image_columns[image_id % num_image_columns]:
|
|
image = st.session_state.output_images[image_id]
|
|
progress_bar = st.progress(1.0)
|
|
show_output_image(image)
|