update app

updating for gpt-j model

app.py

@@ -62,14 +62,14 @@ def draw_polygons(polygons, colors, im_size=(256, 256), b_color="white", fpath=N
 
     return draw, image
 
-def prompt_to_layout(user_prompt, fpath=None):
+def prompt_to_layout(user_prompt, top_p, top_k, fpath=None):
     
     model_prompt = '[User prompt] {} [Layout]'.format(user_prompt)
     input_ids = tokenizer(model_prompt, return_tensors='pt')
-    output = finetuned.generate(**input_ids, do_sample=True, top_p=0.94, top_k=100, max_length=300)
+    output = finetuned.generate(**input_ids, do_sample=True, top_p=top_p, top_k=top_k, max_length=300)
     output = tokenizer.batch_decode(output, skip_special_tokens=True)
     
-    layout = output[0].split('[User prompt]')[1].split('[Layout]')[1].split(', ')
+    layout = output[0].split('[User prompt]')[1].split('[Layout] ')[1].split(', ')
     spaces = [txt.split(':')[0] for txt in layout]
 
     coordinates = [txt.split(':')[1] for txt in layout]
@@ -83,11 +83,11 @@ def prompt_to_layout(user_prompt, fpath=None):
     for poly in polygons:
         geom.append(Polygon(np.array(poly, dtype=int)))
         
-    colors = [architext_colors[housegan_labels[space]] for space in spaces]
+    colors = [architext_colors2[housegan_labels[space]] for space in spaces]
     
     _, im = draw_polygons(geom, colors, fpath=fpath)
     
-    legend = Image.open("legend.png")
+    legend = Image.open(r"F:\Personal\Startup\ArchiText\Media\legend.png")
     
     im = np.array(im)
     im[:40, :] = np.array(legend)
@@ -96,7 +96,7 @@ def prompt_to_layout(user_prompt, fpath=None):
     return im, layout, output
     
 def mut_txt2layout(mut_output):
-    output = mut_output[0].split('[User prompt]')[1].split('[Layout]')[1].split(', ')
+    output = mut_output[0].rstrip().split('[User prompt]')[1].split('[Layout]')[1].split(', ')
     spaces = [txt.split(':')[0].strip(' ') for txt in output]
     coordinates = [txt.split(':')[1] for txt in output]
     coordinates = [re.findall(regex, coord) for coord in coordinates]
@@ -109,10 +109,10 @@ def mut_txt2layout(mut_output):
     for poly in polygons:
         geom.append(Polygon(np.array(poly, dtype=int)))
 
-    colors = [architext_colors[housegan_labels[space]] for space in spaces]
+    colors = [architext_colors2[housegan_labels[space]] for space in spaces]
     _, im = draw_polygons(geom, colors, fpath=None)
 
-    legend = Image.open("legend.png")
+    legend = Image.open(r"F:\Personal\Startup\ArchiText\Media\legend.png")
 
     im = np.array(im)
     im[:40, :] = np.array(legend)
@@ -120,19 +120,20 @@ def mut_txt2layout(mut_output):
     
     return im
     
-def prompt_with_mutation(user_prompt, mut_rate, fpath=None):
+def prompt_with_mutation(user_prompt, top_p, top_k, mut_rate, fpath=None):
     
     #Create initial layout based on prompt
-    im, layout, output = prompt_to_layout(user_prompt)
+    im, layout, output = prompt_to_layout(user_prompt, top_p=top_p, top_k=top_k)
         
     #Create mutated layout based on initial
     mut_len = int((1-mut_rate)*len(layout))
     index1 = random.randrange(0,len(layout)-mut_len)
     rooms = layout[index1:index1+mut_len]
-    rooms = ', '.join(rooms) + ', '
+    rooms[-1] = rooms[-1].split(':')[0] + ':'
+    rooms = ', '.join(rooms)# + ', '
     new_prompt = '[User prompt] {} [Layout] {}'.format(user_prompt, rooms)
     input_ids = tokenizer(new_prompt, return_tensors='pt')
-    mut_output = finetuned.generate(**input_ids, do_sample=True, top_p=0.94, top_k=100, max_length=300)
+    mut_output = finetuned.generate(**input_ids, do_sample=True, top_p=top_p, top_k=top_k, max_length=400)
     mut_output = tokenizer.batch_decode(mut_output, skip_special_tokens=True)
     mut_im = mut_txt2layout(mut_output)
     
@@ -213,28 +214,30 @@ custom_css="""
   opacity: 1 !important;
 }"""
 
-def gen_and_mutate(user_prompt, mutate=False, mut_rate=0.2):    
+def gen_and_mutate(user_prompt, mutate=False, top_p=0.94, top_k=100, mut_rate=0.2):    
     if(mutate):
         im, mut_im = None, None
         while (mut_im is None):
             try:
-                im, mut_im = prompt_with_mutation(user_prompt, mut_rate)
+                im, mut_im = prompt_with_mutation(user_prompt, top_p, top_k, mut_rate)
             except:
                 pass
     else:
         mut_im=Image.open("empty.png")
-        im, _, _ = prompt_to_layout(user_prompt)
+        im, _, _ = prompt_to_layout(user_prompt, top_p, top_k)
         
     return im, mut_im
     
 checkbox =  gr.inputs.Checkbox(label='Mutate')
-slider = gr.inputs.Slider(0.2, 0.8, 0.1, label='Mutation rate')
-textbox = gr.inputs.Textbox(placeholder='house with two bedrooms and one bathroom', lines="3", label="DESCRIBE YOUR DESIGN")
-
+topp_slider = gr.inputs.Slider(0.1, 1.0, 0.01, default=0.94, label='top_p')
+topk_slider = gr.inputs.Slider(0, 100, 25, default=0, label='top_k')
+mut_slider = gr.inputs.Slider(0.2, 0.8, 0.1, default=0.3, label='Mutation rate')
+textbox = gr.inputs.Textbox(placeholder='house with two bedrooms and one bathroom', lines="2", 
+                            label="DESCRIBE YOUR DESIGN")
 generated = gr.outputs.Image(label='Generated Layout')
 mutated = gr.outputs.Image(label='Mutated Layout')
 
-iface = gr.Interface(fn=gen_and_mutate, inputs=[textbox, checkbox, slider], outputs=[generated, mutated],
+iface = gr.Interface(fn=gen_and_mutate, inputs=[textbox, checkbox, topp_slider, topk_slider, mut_slider], outputs=[generated, mutated],
                       css=custom_css,
                       thumbnail="thumbnail_gradio.PNG",
                       description='Demo of Semantic Generation of Residential Layouts \n',