Fixed missing output for last prediction

app.py

@@ -7,6 +7,8 @@ from transformers import BitsAndBytesConfig
 
 from sentence_transformers import SentenceTransformer, util
 
+from transformers import PretrainedConfig
+
 quantization_config = BitsAndBytesConfig(
     load_in_4bit=True,
     bnb_4bit_compute_dtype=torch.float16
@@ -20,7 +22,8 @@ model = LlavaForConditionalGeneration.from_pretrained(
     quantization_config=quantization_config,
     device_map="auto",
     # use_flash_attention_2=True,
-    low_cpu_mem_usage=True
+    low_cpu_mem_usage=True,
+    # config=PretrainedConfig(do_sample=True)
 )
 
 MAXIMUM_PIXEL_VALUES = 3725568
@@ -63,6 +66,14 @@ def text_to_image(image, prompt, duplications: float):
                 batch = dict(input_ids=list(), attention_mask=list(), pixel_values=list())
             else:
                 i += 1
+        if i >= len(inputs['pixel_values']) and len(batch['input_ids']) > 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:
         batched_inputs.append(inputs)
 
@@ -73,8 +84,8 @@ def text_to_image(image, prompt, duplications: float):
         batch['input_ids'] = batch['input_ids'].to(model.device)
         batch['attention_mask'] = batch['attention_mask'].to(model.device)
         batch['pixel_values'] = batch['pixel_values'].to(model.device)
-        output = model.generate(**batch, max_new_tokens=500, temperature=0.3)
-
+        # output = model.generate(**batch, max_new_tokens=500, temperature=0.3)
+        output = model.generate(**batch, max_new_tokens=500)
         # Unload GPU
         batch['input_ids'].to('cpu')
         batch['attention_mask'].to('cpu')