File size: 4,858 Bytes
5b344d3 ec86d75 5b344d3 f68c7aa 5b344d3 3c0e164 56db7c1 5b344d3 ec86d75 b44d446 ec86d75 d87c251 ec86d75 5b344d3 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 |
from typing import Optional
import numpy as np
import cv2
import streamlit as st
from PIL import Image
import os
import tempfile
from sdfile import PIPELINES, generate
DEFAULT_PROMPT = "belted shirt black belted portrait-collar wrap blouse with black prints"
DEAFULT_WIDTH, DEFAULT_HEIGHT = 512,512
OUTPUT_IMAGE_KEY = "output_img"
LOADED_IMAGE_KEY = "loaded_img"
def get_image(key: str) -> Optional[Image.Image]:
if key in st.session_state:
return st.session_state[key]
return None
def set_image(key:str, img: Image.Image):
st.session_state[key] = img
def prompt_and_generate_button(prefix, pipeline_name: PIPELINES, **kwargs):
prompt = st.text_area(
"Prompt",
value = DEFAULT_PROMPT,
key = f"{prefix}-prompt"
)
negative_prompt = st.text_area(
"Negative prompt",
value = "",
key =f"{prefix}-negative_prompt",
)
col1,col2 =st.columns(2)
with col1:
steps = st.slider(
"Number of inference steps",
min_value=1,
max_value=200,
value=30,
key=f"{prefix}-inference-steps",
)
with col2:
guidance_scale = st.slider(
"Guidance scale",
min_value=0.0,
max_value=20.0,
value= 7.5,
step = 0.5,
key=f"{prefix}-guidance-scale",
)
enable_cpu_offload = st.checkbox(
"Enable CPU offload if you run out of memory",
key =f"{prefix}-cpu-offload",
value= False,
)
if st.button("Generate Image", key = f"{prefix}-btn"):
with st.spinner("Generating image ..."):
image = generate(
prompt,
pipeline_name,
negative_prompt=negative_prompt,
num_inference_steps=steps,
guidance_scale=guidance_scale,
enable_cpu_offload=enable_cpu_offload,
**kwargs,
)
set_image(OUTPUT_IMAGE_KEY,image.copy())
st.image(image)
def width_and_height_sliders(prefix):
col1, col2 = st.columns(2)
with col1:
width = st.slider(
"Width",
min_value=64,
max_value=1600,
step=16,
value=512,
key=f"{prefix}-width",
)
with col2:
height = st.slider(
"Height",
min_value=64,
max_value=1600,
step=16,
value=512,
key=f"{prefix}-height",
)
return width, height
def image_uploader(prefix):
image = st.file_uploader("Image", ["jpg", "png"], key=f"{prefix}-uploader")
if image:
image = Image.open(image)
print(f"loaded input image of size ({image.width}, {image.height})")
return image
return get_image(LOADED_IMAGE_KEY)
def sketching():
image = image_uploader("sketch2img")
if not image:
return None,None
with tempfile.TemporaryDirectory() as temp_dir:
temp_image_path = os.path.join(temp_dir, "uploaded_image.jpg")
image.save(temp_image_path)
image = cv2.imread(temp_image_path)
image = cv2.cvtColor(image,cv2.COLOR_BGR2GRAY)
image_blur = cv2.GaussianBlur(image,(5,5),0)
sketch = cv2.adaptiveThreshold(image_blur, 255, cv2.ADAPTIVE_THRESH_MEAN_C,cv2.THRESH_BINARY,11,2)
sketch_pil = Image.fromarray(sketch)
return sketch_pil
def txt2img_tab():
prefix = "txt2img"
width, height = width_and_height_sliders(prefix)
prompt_and_generate_button(prefix,"txt2img",width=width,height=height)
def sketching_tab():
prefix = "sketch2img"
col1,col2 = st.columns(2)
with col1:
sketch_pil = sketching()
with col2:
if sketch_pil:
controlnet_conditioning_scale = st.slider(
"Strength or dependence on the input sketch",
min_value=0.0,
max_value= 1.0,
value = 0.5,
step = 0.05,
key=f"{prefix}-controlnet_conditioning_scale",
)
prompt_and_generate_button(
prefix,
"sketch2img",
sketch_pil=sketch_pil,
controlnet_conditioning_scale=controlnet_conditioning_scale,
)
def main():
st.set_page_config(layout="wide")
st.title("Fashion-SDX: Playground")
tab1,tab2 = st.tabs(
["Text to image", "Sketch to image"]
)
with tab1:
txt2img_tab()
with tab2:
sketching_tab()
with st.sidebar:
st.header("Most Recent Output Image")
output_image = get_image((OUTPUT_IMAGE_KEY))
if output_image:
st.image(output_image)
else:
st.markdown("no output generated yet")
if __name__ =="__main__":
main() |