Update app.py

app.py

@@ -406,6 +406,10 @@ def file_upload(file):
     
     
     
+    
+
+    
+    
 intro = """
 <div style="display: flex;align-items: center;justify-content: center">
     <h1 style="margin-left: 12px;text-align: center;margin-bottom: 7px;display: inline-block">weights2weights</h1>
@@ -418,51 +422,122 @@ intro = """
         display: inline-block;
     ">
     <img style="margin-top: -1em;margin-bottom: 0em;position: absolute;" src="https://bit.ly/3CWLGkA" alt="Duplicate Space"></a>
-</p>"""
+</p>
+"""
+
+
 
 with gr.Blocks(css="style.css") as demo:
     gr.HTML(intro)
-    with gr.Tab("Sampling Models + Editing"):
-        with gr.Column():
+    
+    gr.Markdown("""
+    Click sample (to sample an identity) *or* upload an image & click `invert` to get started ✨
+    > 💡 When inverting, draw a mask over the face for improved results. 
+    > To use a model previously downloaded from this demo see `Uplaoding a model` in the `Advanced options`
+    """)
+    with gr.Column():
             with gr.Row():
                 with gr.Column():
-                    input_image = gr.Image(label="Identity from Sampled or Uploaded Model", interactive=False)
-                    sample = gr.Button("🎲 Sample New Model")
-                    file_input = gr.File(label="Upload Model", container=True)
-                    #file_input.change(fn=file_upload, inputs=file_input, outputs = gallery1)
-        
-                    #with gr.Row():
-                    #    sample = gr.Button("🎲 Sample New Model")
-                        #invert_button = gr.Button("⏪ Invert")
+                    # input_image = gr.Image(source='upload', elem_id="image_upload", tool='sketch', type='pil', label="Upload image and draw to define mask",
+                    #                         height=512, width=512, brush_color='#00FFFF', mask_opacity=0.6)
+                    input_image = gr.ImageEditor(elem_id="image_upload", type='pil', label="Upload image and draw to define mask",
+                                            height=512, width=512, brush=gr.Brush(), layers=False) 
+                    with gr.Row():
+                        sample = gr.Button("🎲 Sample New Model")
+                        invert_button = gr.Button("⏪ Invert")
                 with gr.Column():
                     image_slider = ImageSlider(position=0.5, type="pil", height=512, width=512)
-                    prompt = gr.Textbox(label="Prompt",
-                                    info="Make sure to include 'sks person'" ,
-                                    placeholder="sks person", 
-                                    value="sks person")
-                    seed = gr.Number(value=5, label="Seed", precision=0, interactive=True)
+                    # gallery1 = gr.Image(label="Identity from Original Model",height=512, width=512, interactive=False)
+                    
+            prompt1 = gr.Textbox(label="Prompt",
+                                            info="Make sure to include 'sks person'" ,
+                                            placeholder="sks person", 
+                                            value="sks person")
+                
+            # Editing 
+            with gr.Column():
+                #gallery2 = gr.Image(label="Identity from Edited Model", interactive=False, visible=False )
+                with gr.Row():
+                    a1 = gr.Slider(label="- Young +", value=0, step=0.001, minimum=-1, maximum=1, interactive=True)
+                    a2 = gr.Slider(label="- Pointy Nose +", value=0, step=0.001, minimum=-1, maximum=1, interactive=True)       
+                with gr.Row():        
+                    a3 = gr.Slider(label="- Curly Hair +", value=0, step=0.001, minimum=-1, maximum=1, interactive=True)
+                    a4 = gr.Slider(label="- Thick Eyebrows +", value=0, step=0.001, minimum=-1, maximum=1, interactive=True)
+#                 prompt2 = gr.Textbox(label="Prompt",
+#                                         info="Make sure to include 'sks person'" ,
+#                                         placeholder="sks person", 
+#                                         value="sks person", visible=False)
+#                 seed2 = gr.Number(value=5, label="Seed", precision=0, interactive=True, visible=False)
+
+
+
+#                 submit2 = gr.Button("Generate", visible=False)
+
+            
+            
+            with gr.Accordion("Advanced Options", open=False):
+                with gr.Tab("Inversion"):
+                    with gr.Row():     
+                        lr = gr.Number(value=1e-1, label="Learning Rate", interactive=True)
+                        pcs = gr.Slider(label="# Principal Components", value=10000, step=1, minimum=1, maximum=10000, interactive=True)
+                    with gr.Row():
+                        epochs = gr.Slider(label="Epochs", value=400, step=1, minimum=1, maximum=2000, interactive=True)
+                        weight_decay = gr.Number(value=1e-10, label="Weight Decay", interactive=True)
+                with gr.Tab("Sampling"):
+                    with gr.Row():     
+                            cfg1= gr.Slider(label="CFG", value=3.0, step=0.1, minimum=0, maximum=10, interactive=True)
+                            steps1 = gr.Slider(label="Inference Steps",  value=50, step=1, minimum=0, maximum=100, interactive=True)
+                            seed1 = gr.Number(value=5, label="Seed", precision=0, interactive=True)
                     with gr.Row():
-                        a1 = gr.Slider(label="- Young +", value=0, step=0.001, minimum=-1, maximum=1, interactive=True)
-                        a2 = gr.Slider(label="- Pointy Nose +", value=0, step=0.001, minimum=-1, maximum=1, interactive=True)       
-                    with gr.Row():        
-                        a3 = gr.Slider(label="- Curly Hair +", value=0, step=0.001, minimum=-1, maximum=1, interactive=True)
-                        a4 = gr.Slider(label="- placeholder +", value=0, step=0.001, minimum=-1, maximum=1, interactive=True)
+                            negative_prompt1 = gr.Textbox(label="Negative Prompt", placeholder="low quality, blurry, unfinished, nudity, weapon", value="low quality, blurry, unfinished, nudity, weapon")
+                            injection_step = gr.Slider(label="Injection Step",  value=800, step=1, minimum=0, maximum=1000, interactive=True)
+                            
+                            
+                # with gr.Tab("Editing"):
+                #     with gr.Column():     
+                        # cfg2 = gr.Slider(label="CFG", value=3.0, step=0.1, minimum=0, maximum=10, interactive=True)
+                        # steps2 = gr.Slider(label="Inference Steps",  value=50, step=1, minimum=0, maximum=100, interactive=True)
+                        # injection_step = gr.Slider(label="Injection Step",  value=800, step=1, minimum=0, maximum=1000, interactive=True)
+                        # negative_prompt2 = gr.Textbox(label="Negative Prompt", placeholder="low quality, blurry, unfinished, nudity, weapon", value="low quality, blurry, unfinished, nudity, weapon")
+                with gr.Tab("Uploading a model"):
+                    gr.Markdown("""<div style="text-align: justify;">Upload a model below downloaded from this demo.""")
 
-                    
-                    with gr.Accordion("Advanced Options", open=False):
-                        with gr.Column():     
-                                cfg = gr.Slider(label="CFG", value=3.0, step=0.1, minimum=0, maximum=10, interactive=True)
-                                steps = gr.Slider(label="Inference Steps",  value=50, step=1, minimum=0, maximum=100, interactive=True)
-                                injection_step = gr.Slider(label="Injection Step",  value=800, step=1, minimum=0, maximum=1000, interactive=True)
-                                negative_prompt = gr.Textbox(label="Negative Prompt", placeholder="low quality, blurry, unfinished, nudity", value="low quality, blurry, unfinished, nudity")
-
-                    submit = gr.Button("Generate")
-                    # gallery1 = gr.Image(label="Identity from Original Model",height=512, width=512, interactive=False)
+                    file_input = gr.File(label="Upload Model", container=True)
+
+            submit1 = gr.Button("Generate")
+     
+
+
+    gr.Markdown("""<div style="text-align: justify;"> After sampling a new model or inverting, you can download the model below.""")
+
+    with gr.Row():
+        file_output = gr.File(label="Download Sampled Model", container=True, interactive=False)
+        
 
 
 
+
+        invert_button.click(fn=run_inversion,
+                    inputs=[input_image, pcs, epochs, weight_decay,lr], 
+                    outputs = [image_slider, file_output])
+        
+            
+        sample.click(fn=sample_then_run, outputs=[input_image, file_output])
+
+        # submit1.click(fn=inference,
+        #             inputs=[prompt1, negative_prompt1, cfg1, steps1, seed1],
+        #             outputs=gallery1)
+        # submit1.click(fn=edit_inference,
+        #             inputs=[input_image, prompt1, negative_prompt1, cfg1, steps1, seed1, injection_step, a1, a2, a3, a4],
+        #             outputs=image_slider)
+        submit1.click(
+        fn=edit_inference, inputs=[input_image, prompt1, negative_prompt1, cfg1, steps1, seed1, injection_step, a1, a2, a3, a4], outputs=[image_slider]
+    )
+        file_input.change(fn=file_upload, inputs=file_input, outputs = input_image)
+
+
     
     
-
+        
             
 demo.queue().launch()