Better results.

app.py

@@ -10,10 +10,11 @@ from shapely.geometry.polygon import Polygon
 from shapely.affinity import scale
 from PIL import Image, ImageDraw, ImageOps, ImageFilter, ImageFont, ImageColor
 
-os.system('pip3 install gradio==2.7.5')
+os.system('pip3 install gradio==3.14.0')
 import gradio as gr
 
 from transformers import AutoTokenizer, AutoConfig, AutoModelForCausalLM
+from transformers.models.gptj.modeling_gptj import apply_rotary_pos_emb as apply_rotary_pos_emb_pt
 
 tokenizer = AutoTokenizer.from_pretrained("architext/gptj-162M")
 finetuned = AutoModelForCausalLM.from_pretrained("architext/gptj-162M")
@@ -44,7 +45,7 @@ def creativity(intensity):
 
 housegan_labels = {"living_room": 1, "kitchen": 2, "bedroom": 3, "bathroom": 4, "missing": 5, "closet": 6, 
                          "balcony": 7, "corridor": 8, "dining_room": 9, "laundry_room": 10}
-
+default_color = [128, 128, 128]
 architext_colors = [[0, 0, 0], [249, 222, 182], [195, 209, 217], [250, 120, 128], [126, 202, 234], [190, 0, 198], [255, 255, 255], 
                    [6, 53, 17], [17, 33, 58], [132, 151, 246], [197, 203, 159], [6, 53, 17],]
 
@@ -116,8 +117,8 @@ def prompt_to_layout(user_prompt, intensity, fpath=None):
     geom = []
     for poly in polygons:
         scaled_poly = scale(Polygon(np.array(poly, dtype=int)), xfact=2, yfact=2, origin=(0,0))
-        geom.append(scaled_poly)  
-    colors = [architext_colors[housegan_labels[space]] for space in spaces]
+        geom.append(scaled_poly)
+    colors = [architext_colors[housegan_labels[space]] if space in housegan_labels else default_color for space in spaces]
     _, im = draw_polygons(geom, colors, fpath=fpath)
     html = '<img class="labels" src="images/labels.png" />'
     legend = Image.open("labels.png")
@@ -253,7 +254,7 @@ custom_css="""
 creative_slider = gr.inputs.Radio(["Low", "Medium", "High"], default="Low", label='Creativity')
 textbox = gr.inputs.Textbox(placeholder='An apartment with two bedrooms and one bathroom', lines="3", 
                             label="DESCRIBE YOUR IDEAL APARTMENT")
-generated = gr.outputs.Image(label='Generated Layout')
+generated = gr.outputs.Image(label='Generated Layout', type='numpy')
 layout = gr.outputs.Textbox(label='Layout Coordinates')
 
 iface = gr.Interface(fn=prompt_to_layout, inputs=[textbox, creative_slider], 

requirements.txt

@@ -3,7 +3,7 @@ numpy
 torch
 shapely
 pillow
-gradio==2.7.5
+gradio==3.14.0
 jinja2
-git+https://github.com/StellaAthena/transformers 
+git+https://github.com/huggingface/transformers@423d087ffd591158c7d8eecda211041bb907c658
 accelerate>=0.20.3