Update app.py

app.py

@@ -208,9 +208,7 @@ def resize_keep_aspect_ratio(pil_image, target_size=1024):
 def generate_image(
     prompt,   
     image_1, caption_1,
-    image_2, caption_2,
-    use_id_1 = True,
-    use_id_2 = True,
+    image_2 = None, caption_2 = None,
     cond_size = 256, 
     target_height = 768, 
     target_width = 768, 
@@ -231,9 +229,18 @@ def generate_image(
     torch.cuda.empty_cache()
     num_images = 1
 
-    images = [image_1, image_2]
-    captions = [caption_1, caption_2]
-    idips_checkboxes = [use_id_1, use_id_2]
+    images = []
+    captions = []
+    idips_checkboxes = []
+
+    if image_1 != None:
+        images.append(image_1)
+        captions.append(caption_1)
+        idips_checkboxes.append(True)
+    else image_2 != None:
+        images.append(image_2)
+        captions.append(caption_2)
+        idips_checkboxes.append(True)
 
     print(f"Length of images: {len(images)}")
     print(f"Length of captions: {len(captions)}")
@@ -640,8 +647,6 @@ if __name__ == "__main__":
                 prompt, 
                 images[0], captions[0], 
                 images[1], captions[1], 
-                idip_checkboxes[0],
-                idip_checkboxes[1],  
                 cond_size, 
                 target_height, 
                 target_width,