Fix image generation function.

app.py

@@ -1,9 +1,11 @@
 # system
 
 import os
+from pathlib import Path
 
-os.system("gdown https://drive.google.com/uc?id=1--27E5dk8GzgvpVL0ofr-m631iymBpUH")
-os.system("gdown https://drive.google.com/uc?id=191a5lTsUPQ1hXaeo6kVNbo_W3WYuXsmF")
+if not Path('./Text2Punk-final-7.pt').exists() and not Path('./clip-final.pt').exists():
+    os.system("gdown https://drive.google.com/uc?id=1--27E5dk8GzgvpVL0ofr-m631iymBpUH")
+    os.system("gdown https://drive.google.com/uc?id=191a5lTsUPQ1hXaeo6kVNbo_W3WYuXsmF")
 
 # plot
 
@@ -17,12 +19,7 @@ import gradio as gr
 
 # text2punks utils
 
-from text2punks.utils import to_pil_image, model_loader, generate_image
-
-
-batch_size = 32
-num_images = 32
-top_prediction = 8
+from text2punks.utils import resize, to_pil_image, model_loader, generate_image
 
 # nobs to tune
 
@@ -34,21 +31,24 @@ temperature = 1.25
 def compose_predictions(images):
 
     increased_h = 0
-    h, w = images[0].shape[0], images[0].shape[1]
-    image_grid = Image.new("RGB", (len(images)*w, h))
+    b, c, h, w = *images.shape,
+    image_grid = Image.new("RGB", (b*w*4, h*4), color=0)
 
-    for i, img_ in enumerate(images):
-        image_grid.paste(to_pil_image(img_), (i*w, increased_h))
+    for i in range(b):
+        # resize(images[i], 96)
+        print(images[i].shape)
+        img_ = to_pil_image(images[i])
+        image_grid.paste(img_, (i*w*4, increased_h))
 
-    return img
+    return image_grid
 
 
-def run_inference(prompt, num_images=32, num_preds=8):
+def run_inference(prompt, num_images=32, batch_size=32, num_preds=8):
 
     t2p_path, clip_path = './Text2Punk-final-7.pt', './clip-final.pt'
     text2punk, clip = model_loader(t2p_path, clip_path)
 
-    images = generate_image(prompt_text=prompt, top_k=top_k, temperature=temperature, num_images=num_images, batch_size=batch_size, top_prediction=top_prediction, text2punk_model=text2punk, clip_model=clip)
+    images, _ = generate_image(prompt_text=prompt, top_k=top_k, temperature=temperature, num_images=num_images, batch_size=batch_size, top_prediction=num_preds, text2punk_model=text2punk, clip_model=clip)
     predictions = compose_predictions(images)
 
     output_title = f"""
@@ -69,7 +69,7 @@ Text2Cryptopunks is an AI model that generates Cryptopunks images from text prom
 
 gr.Interface(run_inference, 
     inputs=[gr.inputs.Textbox(label='type somthing like this : "An Ape CryptoPunk that has 2 Attributes, a Pigtails and a Medical Mask."')],
-    outputs=outputs, 
+    outputs=outputs,
     title='Text2Cryptopunks',
     description=description,
     article="<p style='text-align: center'> Created by kTonpa | <a href='https://github.com/kTonpa/Text2CryptoPunks'>GitHub</a>",

text2punks/utils.py

@@ -26,9 +26,11 @@ codebook = torch.load('./text2punks/data/codebook.pt')
 def exists(val):
     return val is not None
 
+def resize(image_tensor, size):
+    return F.resize(image_tensor, (size, size), F.InterpolationMode.NEAREST)
 
 def to_pil_image(image_tensor):
-    return F.to_pil_image(image_tensor)
+    return F.to_pil_image(image_tensor.type(torch.uint8))
 
 
 def model_loader(text2punk_path, clip_path):