Update app.py

app.py

@@ -30,6 +30,7 @@ global text_encoder
 global tokenizer
 global noise_scheduler
 global network
+global original_image
 device = "cuda:0"
 generator = torch.Generator(device=device)
 from gradio_imageslider import ImageSlider
@@ -110,9 +111,8 @@ def inference(prompt, negative_prompt, guidance_scale, ddim_steps, seed):
 
 
 
-# @spaces.GPU()
 @torch.no_grad()
-def edit_inference(input_image, prompt, negative_prompt, guidance_scale, ddim_steps, seed, start_noise, a1, a2, a3, a4):
+def edit_inference(prompt, negative_prompt, guidance_scale, ddim_steps, seed, start_noise, a1, a2, a3, a4):
 
     global device
     global generator 
@@ -125,9 +125,12 @@ def edit_inference(input_image, prompt, negative_prompt, guidance_scale, ddim_st
     global pointy
     global wavy
     global large
+    global original_image
+   
     
     original_weights = network.proj.clone()
     
+    
     #pad to same number of PCs
     pcs_original = original_weights.shape[1]
     pcs_edits = young.shape[1]
@@ -194,20 +197,21 @@ def edit_inference(input_image, prompt, negative_prompt, guidance_scale, ddim_st
     #reset weights back to original 
     network.proj = torch.nn.Parameter(original_weights)
     network.reset()
+   
+    return (original_image, image)
 
-    return [image]
-
-# @spaces.GPU()
 def sample_then_run():
+    global original_image
     sample_model()    
     prompt = "sks person"
     negative_prompt = "low quality, blurry, unfinished, nudity, weapon"
     seed = 5
     cfg = 3.0
     steps = 50
-    image = inference( prompt, negative_prompt, cfg, steps, seed)
+    original_image = inference( prompt, negative_prompt, cfg, steps, seed)
     torch.save(network.proj, "model.pt" )
-    return image, "model.pt"
+    
+    return (original_image, original_image), "model.pt"
 
 
 global young
@@ -406,7 +410,7 @@ def file_upload(file):
     
     
     
-    
+   
 
     
     
@@ -438,11 +442,11 @@ with gr.Blocks(css="style.css") as demo:
     with gr.Column():
             with gr.Row():
                 with gr.Column():
-                    input_image = gr.Image(label="Identity from Sampled or Uploaded Model", height=512, width=512)
-                    # 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")
+                    sample = gr.Button("🎲 Sample New Model")
+                    file_input = gr.File(label="Upload Model", container=True)
+                    file_output = gr.File(label="Download Sampled Model", container=True, interactive=False)
+
+
                         # invert_button = gr.Button("⏪ Invert")
                 with gr.Column():
                     image_slider = ImageSlider(position=0.5, type="pil", height=512, width=512)
@@ -507,10 +511,10 @@ with gr.Blocks(css="style.css") as demo:
      
 
 
-    gr.Markdown("""<div style="text-align: justify;"> After sampling a new model or inverting, you can download the model below.""")
+    # 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)
+    # with gr.Row():
+    #     file_output = gr.File(label="Download Sampled Model", container=True, interactive=False)
         
 
 
@@ -521,7 +525,7 @@ with gr.Blocks(css="style.css") as demo:
         #             outputs = [image_slider, file_output])
         
             
-        sample.click(fn=sample_then_run, outputs=[input_image, file_output])
+    sample.click(fn=sample_then_run, outputs=[image_slider, file_output])
 
         # submit1.click(fn=inference,
         #             inputs=[prompt1, negative_prompt1, cfg1, steps1, seed1],
@@ -529,12 +533,18 @@ with gr.Blocks(css="style.css") as demo:
         # 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)
+    submit1.click(
+    fn=edit_inference, inputs=[ 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(share=True)
+
+
+    
+    
+