Update app.py

app.py

@@ -9,14 +9,14 @@ import torch
 MODEL_NAME = "runwayml/stable-diffusion-v1-5"
 OUTPUT_DIR = "/content/stable_diffusion_weights/custom_model"
 
-def fine_tune(instance_prompt, image1, image2=None):
+def fine_tune(instance_prompt, image1, image2):
     instance_data_dir = "/content/instance_images"
     if os.path.exists(instance_data_dir):
         shutil.rmtree(instance_data_dir)
     os.makedirs(instance_data_dir, exist_ok=True)
 
     image1.save(os.path.join(instance_data_dir, "instance_0.png"))
-    if image2:
+    if image2 is not None:
         image2.save(os.path.join(instance_data_dir, "instance_1.png"))
 
     fine_tune_model(instance_data_dir, instance_prompt, MODEL_NAME, OUTPUT_DIR)
@@ -42,7 +42,7 @@ def gradio_app():
                 with gr.Column():
                     instance_prompt = gr.Textbox(label="Instance Prompt")
                     image1 = gr.Image(label="Upload Image 1", type="pil")
-                    image2 = gr.Image(label="Upload Image 2 (Optional)", type="pil", optional=True)
+                    image2 = gr.Image(label="Upload Image 2 (Optional)", type="pil")
                     fine_tune_button = gr.Button("Fine-Tune Model")
                     output_text = gr.Textbox(label="Output")
                 fine_tune_button.click(fine_tune, inputs=[instance_prompt, image1, image2], outputs=output_text)