Update model.py

model.py

@@ -114,10 +114,10 @@ class Model:
         condition_img = condition_img.resize((512,512))
         W, H = condition_img.size
 
-        condition_img = torch.from_numpy(np.array(condition_img)).unsqueeze(0).permute(0,3,1,2).repeat(2,1,1,1)
+        condition_img = torch.from_numpy(np.array(condition_img)).unsqueeze(0).permute(0,3,1,2).repeat(3,1,1,1)
         condition_img = condition_img.to(self.device)
         condition_img = 2*(condition_img/255 - 0.5)
-        prompts = [prompt] * 2
+        prompts = [prompt] * 3
         caption_embs, emb_masks = self.t5_model.get_text_embeddings(prompts)
 
         print(f"processing left-padding...")
@@ -160,7 +160,7 @@ class Model:
 
         samples = torch.cat((condition_img[0:1], samples), dim=0)
         samples = 255 * (samples * 0.5 + 0.5)
-        samples = [image] + [
+        samples = [
             Image.fromarray(
                 sample.permute(1, 2, 0).cpu().detach().numpy().clip(
                     0, 255).astype(np.uint8)) for sample in samples
@@ -204,10 +204,10 @@ class Model:
         condition_img = condition_img.resize((512,512))
         W, H = condition_img.size
 
-        condition_img = torch.from_numpy(np.array(condition_img)).unsqueeze(0).permute(0,3,1,2).repeat(2,1,1,1)
+        condition_img = torch.from_numpy(np.array(condition_img)).unsqueeze(0).permute(0,3,1,2).repeat(3,1,1,1)
         condition_img = condition_img.to(self.device)
         condition_img = 2*(condition_img/255 - 0.5)
-        prompts = [prompt] * 2
+        prompts = [prompt] * 3
         caption_embs, emb_masks = self.t5_model.get_text_embeddings(prompts)
 
         print(f"processing left-padding...")
@@ -250,7 +250,7 @@ class Model:
         samples = samples.cpu()
         samples = torch.cat((condition_img[0:1], samples), dim=0)
         samples = 255 * (samples * 0.5 + 0.5)
-        samples = [image] + [
+        samples = [
             Image.fromarray(
                 sample.permute(1, 2, 0).cpu().detach().numpy().clip(0, 255).astype(np.uint8))
             for sample in samples