Update app.py

app.py

@@ -120,6 +120,23 @@ with gr.Blocks() as app:
                 outputs=gallery
             )
 
+        with gr.Tab("Custom Defect"):
+            with gr.Row():
+                custom_prompt_input = gr.Textbox(label="Custom Defect")
+                number_input_custom = gr.Number(label="Number of images to generate", value=1, minimum=1, maximum=10)
+                submit_button_custom = gr.Button("Generate")
+            image_outputs_custom = gr.Gallery()
+
+            def on_submit_click_custom(custom_prompt, number_of_images):
+                images = process_railway_defects(custom_prompt, number_of_images)
+                return images
+
+            submit_button_custom.click(
+                fn=on_submit_click_custom,
+                inputs=[custom_prompt_input, number_input_custom],
+                outputs=image_outputs_custom
+            )
+
         with gr.Tab("Inpaint Defect"):
             image_input = gr.Image(label="Upload Image")
             inpaint_prompt_input = gr.Textbox(label="Describe the defect")