Update app.py

app.py

@@ -2,39 +2,11 @@ import gradio as gr
 import numpy as np
 import cv2
 from PIL import Image
-
+from region_control import MultiDiffusion, get_views, preprocess_mask
+from sketch_helper import get_high_freq_colors, color_quantization, create_binary_matrix
 MAX_COLORS = 12
 
-def get_high_freq_colors(image):
-  im = image.getcolors(maxcolors=1024*1024)
-  sorted_colors = sorted(im, key=lambda x: x[0], reverse=True)
-  
-  freqs = [c[0] for c in sorted_colors]
-  mean_freq = sum(freqs) / len(freqs)
-
-  high_freq_colors = [c for c in sorted_colors if c[0] > max(2, mean_freq/3)]  # Ignore colors that occur very few times (less than 2) or less than half the average frequency
-  return high_freq_colors
-
-def color_quantization(image, n_colors):
-    # Get color histogram
-    hist, _ = np.histogramdd(image.reshape(-1, 3), bins=(256, 256, 256), range=((0, 256), (0, 256), (0, 256)))
-    # Get most frequent colors
-    colors = np.argwhere(hist > 0)
-    colors = colors[np.argsort(hist[colors[:, 0], colors[:, 1], colors[:, 2]])[::-1]]
-    colors = colors[:n_colors]
-    # Replace each pixel with the closest color
-    dists = np.sum((image.reshape(-1, 1, 3) - colors.reshape(1, -1, 3))**2, axis=2)
-    labels = np.argmin(dists, axis=1)
-    return colors[labels].reshape((image.shape[0], image.shape[1], 3)).astype(np.uint8)
-
-def create_binary_matrix(img_arr, target_color):
-    print(target_color)
-    # Create mask of pixels with target color
-    mask = np.all(img_arr == target_color, axis=-1)
-    
-    # Convert mask to binary matrix
-    binary_matrix = mask.astype(int)
-    return binary_matrix
+sd = MultiDiffusion("cuda", "2.1")
 
 def process_sketch(image, binary_matrixes):
   high_freq_colors = get_high_freq_colors(image)
@@ -43,13 +15,12 @@ def process_sketch(image, binary_matrixes):
   im2arr = color_quantization(im2arr, n_colors=how_many_colors)
   
   colors_fixed = []
-  for color in high_freq_colors[1:]:
-    r = color[1][0]
-    g = color[1][1]
-    b = color[1][2]
-    binary_matrix = create_binary_matrix(im2arr, (r,g,b))
-    binary_matrixes.append(binary_matrix)
-    colors_fixed.append(gr.update(value=f'<div class="color-bg-item" style="background-color: rgb({r},{g},{b})"></div>'))
+  for color in high_freq_colors:
+    r, g, b = color[1]
+    if any(c != 255 for c in (r, g, b)):
+      binary_matrix = create_binary_matrix(im2arr, (r,g,b))
+      binary_matrixes.append(binary_matrix)
+      colors_fixed.append(gr.update(value=f'<div class="color-bg-item" style="background-color: rgb({r},{g},{b})"></div>'))
   visibilities = []
   colors = []
   for n in range(MAX_COLORS):
@@ -62,8 +33,15 @@ def process_sketch(image, binary_matrixes):
 
 def process_generation(binary_matrixes, master_prompt, *prompts):
     clipped_prompts = prompts[:len(binary_matrixes)]
-    #Now: master_prompt can be used as the main prompt, and binary_matrixes and clipped_prompts can be used as the masked inputs
-    pass
+    prompts = [master_prompt] + list(clipped_prompts)
+    neg_prompts = [""] * len(prompts)
+    fg_masks = torch.cat([preprocess_mask(mask_path, 512 // 8, 512 // 8, "cuda") for mask_path in binary_matrixes])
+    bg_mask = 1 - torch.sum(fg_masks, dim=0, keepdim=True)
+    bg_mask[bg_mask < 0] = 0
+    masks = torch.cat([bg_mask, fg_masks])
+    print(masks.size())
+    image = sd.generate(masks, prompts, neg_prompts, 512, 512, 50, bootstrapping=20)
+    return(image)
 
 css = '''
 #color-bg{display:flex;justify-content: center;align-items: center;}
@@ -72,15 +50,11 @@ css = '''
 '''
 def update_css(aspect):
   if(aspect=='Square'):
-    width = 512
-    height = 512
+    return [gr.update(visible=True), gr.update(visible=False), gr.update(visible=False)]
   elif(aspect == 'Horizontal'):
-    width = 768
-    height = 512
+   return [gr.update(visible=False), gr.update(visible=True), gr.update(visible=False)]
   elif(aspect=='Vertical'):
-    width = 512
-    height = 768
-  return gr.update(value=f"<style>#main-image{{width: {width}px}} .fixed-height{{height: {height}px !important}}</style>")
+    return [gr.update(visible=False), gr.update(visible=False), gr.update(visible=True)]
 
 with gr.Blocks(css=css) as demo:
   binary_matrixes = gr.State([])
@@ -89,11 +63,13 @@ with gr.Blocks(css=css) as demo:
   ''')
   with gr.Row():
     with gr.Box(elem_id="main-image"):
-      with gr.Row():
-          image = gr.Image(interactive=True, tool="color-sketch", source="canvas", type="pil")
-      with gr.Row():
-          aspect = gr.Radio(["Square", "Horizontal", "Vertical"], value="Square", label="Aspect Ratio")
-      button_run = gr.Button("I've finished my sketch",elem_id="main_button")
+      #with gr.Row():
+      image = gr.Image(interactive=True, tool="color-sketch", source="canvas", type="pil", shape=(512,512), brush_radius=45)
+      #image_horizontal = gr.Image(interactive=True, tool="color-sketch", source="canvas", type="pil", shape=(768,512), visible=False, brush_radius=45)
+      #image_vertical = gr.Image(interactive=True, tool="color-sketch", source="canvas", type="pil", shape=(512, 768), visible=False, brush_radius=45)
+      #with gr.Row():
+      #    aspect = gr.Radio(["Square", "Horizontal", "Vertical"], value="Square", label="Aspect Ratio")
+      button_run = gr.Button("I've finished my sketch",elem_id="main_button", interactive=True)
       
       prompts = []
       colors = []
@@ -111,9 +87,8 @@ with gr.Blocks(css=css) as demo:
   gr.Markdown('''
   ![Examples](https://multidiffusion.github.io/pics/tight.jpg)
   ''')
-  css_height = gr.HTML("<style>#main-image{width: 512px} .fixed-height{height: 512px !important}</style>")
-  
-  aspect.change(update_css, inputs=aspect, outputs=css_height)
+  #css_height = gr.HTML("<style>#main-image{width: 512px} .fixed-height{height: 512px !important}</style>")
+  #aspect.change(update_css, inputs=aspect, outputs=[image, image_horizontal, image_vertical])
   button_run.click(process_sketch, inputs=[image, binary_matrixes], outputs=[post_sketch, binary_matrixes, *color_row, *colors])
   final_run_btn.click(process_generation, inputs=[binary_matrixes, general_prompt, *prompts], outputs=out_image)
-demo.launch()
+demo.launch(debug=True)