Update app.py

app.py

@@ -21,7 +21,7 @@ pipe = ImagicStableDiffusionPipeline.from_pretrained(
 
 generator = torch.Generator("cuda").manual_seed(0)
 
-def infer(prompt, init_image):
+def infer(prompt, init_image, trn_steps):
     init_image = Image.open(init_image).convert("RGB")
     init_image = init_image.resize((256, 256))
 
@@ -33,7 +33,7 @@ def infer(prompt, init_image):
         num_inference_steps=50,
         generator=generator,
         text_embedding_optimization_steps=500,
-        model_fine_tuning_optimization_steps=480)
+        model_fine_tuning_optimization_steps=trn_steps)
     
    # with torch.no_grad():
    #     torch.cuda.empty_cache()
@@ -118,17 +118,17 @@ with gr.Blocks(css=css) as block:
 
         prompt_input = gr.Textbox(label="Target text", placeholder="Describe the image with what you want to change about the subject")
         image_init = gr.Image(source="upload", type="filepath",label="Input Image")
-        
+        trn_steps = gr.Slider(250, 1000, value=500)
         submit_btn = gr.Button("Train")
         
         image_output = gr.Image(label="Edited image")
 
-        examples=[['a sitting dog','imagic-dog.png'], ['a photo of a bird spreading wings','imagic-bird.png']]
+        examples=[['a sitting dog','imagic-dog.png', 250], ['a photo of a bird spreading wings','imagic-bird.png',250]]
         ex = gr.Examples(examples=examples, fn=infer, inputs=[prompt_input,image_init], outputs=[image_output], cache_examples=False, run_on_click=False)
         
         
         gr.HTML(article)
 
-    submit_btn.click(fn=infer, inputs=[prompt_input,image_init], outputs=[image_output])
+    submit_btn.click(fn=infer, inputs=[prompt_input,image_init,trn_steps], outputs=[image_output])
     
 block.queue(max_size=12).launch(show_api=False)