Fixed some missing operations to process batches

app.py

@@ -54,14 +54,17 @@ def text_to_image(image, prompt, duplications: float):
                 batch['pixel_values'].pop()
 
                 # transform lists to tensors
-                batch['input_ids'] = torch.cat(batch['input_ids'], dim=0)
-                batch['attention_mask'] = torch.cat(batch['input_ids'], dim=0)
-                batch['pixel_values'] = torch.cat(batch['input_ids'], dim=0)
+                batch['input_ids'] = torch.stack(batch['input_ids'], dim=0)
+                batch['attention_mask'] = torch.stack(batch['attention_mask'], dim=0)
+                batch['pixel_values'] = torch.stack(batch['pixel_values'], dim=0)
 
                 # Add to the batched_inputs
                 batched_inputs.append(batch)
+                batch = dict(input_ids=list(), attention_mask=list(), pixel_values=list())
             else:
                 i += 1
+    else:
+        batched_inputs.append(inputs)
 
     maurice_description = list()
     maurice_embeddings = list()