a

pipeline.py

@@ -12,7 +12,7 @@ from bert_vectorize import vectorize_text_with_bert
 from logger import save_grid_with_label
 
 class TextToImagePipeline(Pipeline):
-    def __init__(self, model_dir: str, device: str = "cpu"):
+    def __init__(self, model_dir: str = "runs/run_3_jxa", device: str = "cpu"):
         # Initialize model, diffusion manager, and set up environment
         self.device = device
         self.model_dir = model_dir
@@ -29,9 +29,9 @@ class TextToImagePipeline(Pipeline):
         self.wrapper = DiffusionManager(self.net, device=self.device, noise_steps=1000)
         self.wrapper.set_schedule(Schedule.LINEAR)
 
-    def __call__(self, prompt: str, amt: int = 1):
-        # Generate images based on the prompt
-        return self.generate_sample_save_images(prompt, amt)
+    def __call__(self, inputs):
+        
+        return self.generate_sample_save_images(inputs, 8)
 
     def generate_sample_save_images(self, prompt: str, amt: int = 1):
         # Prepare the output path