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, trn_steps):
+def train(prompt, init_image, trn_text, trn_steps):
     init_image = Image.open(init_image).convert("RGB")
     init_image = init_image.resize((256, 256))
 
@@ -32,7 +32,7 @@ def infer(prompt, init_image, trn_steps):
         guidance_scale=7.5,
         num_inference_steps=50,
         generator=generator,
-        text_embedding_optimization_steps=500,
+        text_embedding_optimization_steps=trn_text,
         model_fine_tuning_optimization_steps=trn_steps)
     
     with torch.no_grad():
@@ -40,12 +40,35 @@ def infer(prompt, init_image, trn_steps):
     
     
     
-    res = pipe(alpha=1)
     
     
-    return res.images[0]
+    
+    return "Training is finished !"
    
+def generate(prompt, init_image):
+    init_image = Image.open(init_image).convert("RGB")
+    init_image = init_image.resize((256, 256))
 
+   
+    res = pipe.train(
+        prompt,
+        init_image,
+        guidance_scale=7.5,
+        num_inference_steps=50,
+        generator=generator,
+        text_embedding_optimization_steps=0,
+        model_fine_tuning_optimization_steps=0)
+    
+    with torch.no_grad():
+        torch.cuda.empty_cache()
+    
+    
+    
+    res = pipe(alpha=1)
+    
+    
+    return res.images[0]
+    
 title = """
     <div style="text-align: center; max-width: 650px; margin: 0 auto;">
         <div
@@ -117,17 +140,23 @@ 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, label="finetuning steps")
-        submit_btn = gr.Button("Train")
-        
+        with gr.Row():
+            trn_text = gr.Slider(100, 500, value=250, label="text embedding")
+            trn_steps = gr.Slider(250, 1000, value=500, label="finetuning steps")
+        with gr.Row():
+            train_btn = gr.Button("1.Train")
+            gen_btn = gr.Button("2.Generate")
+
+        training_status = gr.Textbox(label="training status")
         image_output = gr.Image(label="Edited image")
 
-        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,trn_steps], outputs=[image_output], cache_examples=False, run_on_click=False)
+        #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,trn_steps], outputs=[image_output], cache_examples=False, run_on_click=False)
         
         
         gr.HTML(article)
 
-    submit_btn.click(fn=infer, inputs=[prompt_input,image_init,trn_steps], outputs=[image_output])
+    train_btn.click(fn=train, inputs=[prompt_input,image_init,trn_text,trn_steps], outputs=[training_status])
+    gen_btn.click(fn=generate, inputs=[prompt_input,image_init], outputs=[image_output])
     
 block.queue(max_size=12).launch(show_api=False)