Update app.py

app.py

@@ -9,14 +9,15 @@ import torch
 MODEL_NAME = "runwayml/stable-diffusion-v1-5"
 OUTPUT_DIR = "/content/stable_diffusion_weights/custom_model"
 
-def fine_tune(instance_prompt, images):
+def fine_tune(instance_prompt, image1, image2=None):
     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)
 
-    for i, img in enumerate(images):
-        img.save(os.path.join(instance_data_dir, f"instance_{i}.png"))
+    image1.save(os.path.join(instance_data_dir, "instance_0.png"))
+    if image2:
+        image2.save(os.path.join(instance_data_dir, "instance_1.png"))
 
     fine_tune_model(instance_data_dir, instance_prompt, MODEL_NAME, OUTPUT_DIR)
     return "Model fine-tuning complete."
@@ -40,10 +41,11 @@ def gradio_app():
             with gr.Row():
                 with gr.Column():
                     instance_prompt = gr.Textbox(label="Instance Prompt")
-                    image_input = gr.Image(label="Upload Images", type="pil", multiple=True)
+                    image1 = gr.Image(label="Upload Image 1", type="pil")
+                    image2 = gr.Image(label="Upload Image 2 (Optional)", type="pil", optional=True)
                     fine_tune_button = gr.Button("Fine-Tune Model")
                     output_text = gr.Textbox(label="Output")
-                fine_tune_button.click(fine_tune, inputs=[instance_prompt, image_input], outputs=output_text)
+                fine_tune_button.click(fine_tune, inputs=[instance_prompt, image1, image2], outputs=output_text)
         
         with gr.Tab("Generate Images"):
             with gr.Row():
@@ -62,4 +64,4 @@ def gradio_app():
     demo.launch()
 
 if __name__ == "__main__":
-    gradio_app()
+    gradio_app()