Update app.py

app.py

@@ -1,4 +1,5 @@
 from pathlib import Path
+from num2words import num2words
 import numpy as np
 import random
 import re
@@ -21,6 +22,12 @@ finetuned = finetuned.to(device)
 
 # Utility functions
 
+def containsNumber(value):
+    for character in value:
+        if character.isdigit():
+            return True
+    return False
+
 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}
 
@@ -67,7 +74,12 @@ def draw_polygons(polygons, colors, im_size=(256, 256), b_color="white", fpath=N
     return draw, image
 
 def prompt_to_layout(user_prompt, top_p, top_k, fpath=None):
-    
+
+    if(containsNumber(user_prompt) == True):
+        spaced_prompt = user_prompt.split(' ')
+        new_prompt = ' '.join([word if word.isdigit() == False else num2words(int(word)).lower() for word in spaced_prompt])
+        model_prompt = '[User prompt] {} [Layout]'.format(new_prompt)
+        
     model_prompt = '[User prompt] {} [Layout]'.format(user_prompt)
     #print(model_prompt)
     input_ids = tokenizer(model_prompt, return_tensors='pt').to(device)