Update app.py

app.py

@@ -75,7 +75,7 @@ def prompt_to_layout(user_prompt, top_p, top_k, fpath=None):
     output = tokenizer.batch_decode(output, skip_special_tokens=True)
     #print(output)
     
-    layout = output[0].lstrip().split('[User prompt]')[1].split('[Layout]')[1].split(', ')
+    layout = output[0].lstrip().split('[User prompt]')[1].split('[Layout] ')[1].split(', ')
     spaces = [txt.split(':')[0] for txt in layout]
 
     coordinates = [txt.split(':')[1] for txt in layout]