Up

app.py

@@ -31,7 +31,7 @@ st.set_page_config(
 
 save_memory = False
 
-@st.experimental_singleton
+@st.experimental_memo
 def load_model():
     model_path = hf_hub_download('lllyasviel/ControlNet', 'models/control_sd15_scribble.pth')
     model = create_model('./models/cldm_v15.yaml').cpu()
@@ -39,7 +39,6 @@ def load_model():
     model = model.cuda()
     return model
 
-@st.experimental_singleton
 def process(input_image, prompt, a_prompt, n_prompt, num_samples, image_resolution, detect_resolution, ddim_steps, guess_mode, strength, scale, seed, eta):
     with torch.no_grad():
    
@@ -60,7 +59,7 @@ def process(input_image, prompt, a_prompt, n_prompt, num_samples, image_resoluti
         control = einops.rearrange(control, 'b h w c -> b c h w').clone()
 
         if seed == -1:
-            seed = random.randint(0, 65535)
+            seed = random.randint(0, 2147483647)
         seed_everything(seed)
 
         if save_memory:
@@ -105,8 +104,8 @@ def main():
     st.header("Generate image with ControllNet")
     with st.sidebar:
         st_lottie(lottie_penguin, height=200)
-        choose = option_menu("Generate image", ["Upload", "Canvas"],
-                            icons=['collection', 'file-plus'],
+        choose = option_menu("Generate image", ["Upload", "Canvas", "Image Gallery"],
+                            icons=['cloud-upload', 'file-plus', 'collection'],
                             menu_icon="infinity", default_index=0,
                             styles={
                                 "container": {"padding": ".0rem", "font-size": "14px"},
@@ -158,7 +157,7 @@ def main():
                 # file_bytes = np.asarray(bytearray(upload_file.read()), dtype=np.uint8)
                 # imageBGR = cv2.imdecode(file_bytes, cv2.IMREAD_COLOR)
                 # input_image = cv2.cvtColor(imageBGR , cv2.COLOR_BGR2RGB)
-                input_image = np.asarray(Image.open(upload_file))
+                input_image = np.asarray(Image.open(upload_file).convert("RGB"))
                 print("input_image", input_image.shape)             
 
             if generate_button:
@@ -167,20 +166,17 @@ def main():
                     print("input_image", input_image.shape)
                     print("results", results[0].shape)
                     H, W, C = input_image.shape
-                    # output_image = cv2.resize(results[0], (W, H), interpolation=cv2.INTER_AREA)
+                    output_image = cv2.resize(results[0], (W, H), interpolation=cv2.INTER_AREA)
                     col11.image(input_image, channels='RGB', width=None, clamp=False, caption='Input image')
-                    col12.image(results[0], channels='RGB', width=None, clamp=False, caption='Generated image')
+                    col12.image(output_image, channels='RGB', width=None, clamp=False, caption='Generated image')
 
     elif choose == 'Canvas':
-        with st.form(key='canvas_form'):
+        with st.form(key='canvas_generate_form'):
             # Specify canvas parameters in application
             stroke_width = st.sidebar.slider("Stroke width: ", 1, 25, 3)
             stroke_color = st.sidebar.color_picker("Stroke color hex: ")
             bg_color = st.sidebar.color_picker("Background color hex: ", "#eee")
-            bg_height = st.sidebar.slider("Canvas height", min_value=256, max_value=512, value=512, step=64)
-            bg_width = st.sidebar.slider("Canvas width", min_value=256, max_value=512, value=512, step=64)
             realtime_update = st.sidebar.checkbox("Update in realtime", True)
-
             # Create a canvas component
             col31, col32 = st.columns(2)
             with col31:
@@ -191,22 +187,25 @@ def main():
                     background_color=bg_color,
                     background_image=None,
                     update_streamlit=realtime_update,
-                    height=bg_height,
-                    width=bg_width,
+                    height=512,
+                    width=512,
                     drawing_mode="freedraw",
                     point_display_radius=0,
                     key="canvas",
                 )
-            prompt = st.text_input(label="Prompt", placeholder='Type your instruction')
-            
+
+            prompt = st.text_input(label="Prompt", placeholder='Type your instruction')   
+
             with st.expander('Advanced option', expanded=False):
                 col41, col42 = st.columns(2)
+                
                 with col41:
                     image_resolution = st.slider(label="Image Resolution", min_value=256, max_value=512, value=512, step=256)
                     strength = st.slider(label="Control Strength", min_value=0.0, max_value=2.0, value=1.0, step=0.01)
                     guess_mode = st.checkbox(label='Guess Mode', value=False)
                     detect_resolution = st.slider(label="HED Resolution", min_value=128, max_value=1024, value=512, step=1)
                     ddim_steps = st.slider(label="Steps", min_value=1, max_value=100, value=20, step=1)
+                
                 with col42:
                     scale = st.slider(label="Guidance Scale", min_value=0.1, max_value=30.0, value=9.0, step=0.1)
                     seed = st.number_input(label="Seed", min_value=-1, value=-1)
@@ -217,13 +216,27 @@ def main():
 
             # Do something interesting with the image data and paths
             generate_button = st.form_submit_button(label='Generate Image')
-            if canvas_result.image_data is not None:
-                input_image = canvas_result.image_data
-                with st.spinner(text=f"It may take up to 1 minute under high load. Generating images..."):
-                    results = process(input_image, prompt, a_prompt, n_prompt, 1, image_resolution, detect_resolution, ddim_steps, guess_mode, strength, scale, seed, eta)
-                    H, W, C = input_image.shape
-                    output_image = cv2.resize(results[0], (W, H), interpolation=cv2.INTER_AREA)
-                    col32.image(output_image, channels='RGB', width=384, clamp=True, caption='Generated image')
+            if generate_button:
+                if canvas_result.image_data is not None:
+                    input_image = canvas_result.image_data
+                    with st.spinner(text=f"It may take up to 1 minute under high load. Generating images..."):
+                        results = process(input_image, prompt, a_prompt, n_prompt, 1, image_resolution, detect_resolution, ddim_steps, guess_mode, strength, scale, seed, eta)
+                        H, W, C = input_image.shape
+                        output_image = cv2.resize(results[0], (W, H), interpolation=cv2.INTER_AREA)
+                        col32.image(output_image, channels='RGB', width=None, clamp=True, caption='Generated image')
+    
+    elif choose == "Image Gallery":
+        with st.expander('Image gallery', expanded=True):
+            col01, col02, = st.columns(2)
+            with col01:
+                st.image('demo/example_1.jpg', caption="Sport car")
+                st.image('demo/example_2.jpg', caption="Dog house")
+                st.image('demo/example_3.jpg', caption="Guitar")
+            with col02:
+                st.image('demo/example_4.jpg', caption="Sport car")
+                st.image('demo/example_5.jpg', caption="Dog house")
+                st.image('demo/example_6.jpg', caption="Guitar")
+
         
 if __name__ == '__main__':
     main()