Add debugging for image generation

app.py

@@ -15,6 +15,9 @@ from transformers import GPT2TokenizerFast
 # update requirements.txt with:
 # C:\Users\Grant\PycharmProjects\test_space\venv\Scripts\pip3.exe freeze > requirements.txt
 
+# Huggingface Spaces have 16GB RAM and 8 CPU cores
+# See https://huggingface.co/docs/hub/spaces-overview#hardware-resources
+
 pretrained_weights = 'gpt2'
 tokenizer = GPT2TokenizerFast.from_pretrained(pretrained_weights)
 
@@ -47,6 +50,7 @@ model = MinDalle(
 
 
 def gen_image(prompt):
+    print(f'RUNNING gen_image with pronpt: {prompt}')
     images = model.generate_images(
         text=prompt,
         seed=-1,
@@ -57,28 +61,29 @@ def gen_image(prompt):
         supercondition_factor=16,
         is_verbose=False
     )
+    print('COMPLETED GENERATION')
     images = images.to('cpu').numpy()
     images = images.astype(np.uint8)
     return Image.fromarray(images[0])
 
 
-gpu = False
-# init only once
-learner = load_learner('export.pkl',
-                       cpu=not gpu)  # cpu=False uses GPU; make sure installed torch is GPU e.g. `pip3 install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu116`
-
-
-def run_model(name):
-    prompt = f"Name: {name}\r\n"
-    prompt_ids = tokenizer.encode(prompt)
-    if gpu:
-        inp = tensor(prompt_ids)[None].cuda()  # Use .cuda() for torch GPU
-    else:
-        inp = tensor(prompt_ids)[None]
-    preds = learner.model.generate(inp, max_length=1024, num_beams=5, temperature=1.5, do_sample=True)
-    result = tokenizer.decode(preds[0].cpu().numpy())
-    result = result.split('###')[0].replace(r'\r\n', '\n')
-    return result
+# gpu = False
+# # init only once
+# learner = load_learner('export.pkl',
+#                        cpu=not gpu)  # cpu=False uses GPU; make sure installed torch is GPU e.g. `pip3 install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu116`
+#
+#
+# def run_model(name):
+#     prompt = f"Name: {name}\r\n"
+#     prompt_ids = tokenizer.encode(prompt)
+#     if gpu:
+#         inp = tensor(prompt_ids)[None].cuda()  # Use .cuda() for torch GPU
+#     else:
+#         inp = tensor(prompt_ids)[None]
+#     preds = learner.model.generate(inp, max_length=1024, num_beams=5, temperature=1.5, do_sample=True)
+#     result = tokenizer.decode(preds[0].cpu().numpy())
+#     result = result.split('###')[0].replace(r'\r\n', '\n')
+#     return result
 
 
 iface = gr.Interface(fn=gen_image, inputs="text", outputs="pil")