Update app.py

app.py

@@ -23,15 +23,16 @@ device="cuda"
 pipe = pipe.to(device)
 
 @spaces.GPU
-def run(prompt, negative_prompt="", guidance_scale=7.0, pag_scale=3.0, randomize_seed=True, seed=42, progress=gr.Progress(track_tqdm=True)):
+def run(prompt, negative_prompt="", guidance_scale=7.0, pag_scale=3.0, pag_layers=["mid"], randomize_seed=True, seed=42, progress=gr.Progress(track_tqdm=True)):
     prompt = prompt.strip()
+    negative_prompt = negative_prompt.strip()
     if(randomize_seed):
         seed = random.randint(0, sys.maxsize)
-    if(prompt == ""):
+    if(prompt == "" and negative_prompt == ""):
         guidance_scale = 0.0
         
     generator = torch.Generator(device="cuda").manual_seed(seed)
-    image_pag = pipe(prompt, guidance_scale=guidance_scale, pag_scale=3.0, pag_applied_layers=['mid'], generator=generator, num_inference_steps=25).images[0]    
+    image_pag = pipe(prompt, guidance_scale=guidance_scale, pag_scale=pag_scale, pag_applied_layers=pag_layers, generator=generator, num_inference_steps=25).images[0]    
     
     generator = torch.Generator(device="cuda").manual_seed(seed)
     image_normal = pipe(prompt, guidance_scale=guidance_scale, generator=generator, num_inference_steps=25).images[0]
@@ -57,6 +58,7 @@ with gr.Blocks(css=css, theme=theme) as demo:
         guidance_scale = gr.Number(label="Guidance Scale", value=7.0)
         negative_prompt = gr.Textbox(label="Negative prompt", info="Is only applied for the CFG part, leave blank for unconditional generation")
         pag_scale = gr.Number(label="Pag Scale", value=3.0)
+        pag_layers = gr.Dropdown(label="Model layers to apply Pag to", info="mid is the one used on the paper, up and down blocks seem unstable", choices=["up", "mid", "down"], multiselect=True, value="mid")
         randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
         seed = gr.Slider(minimum=1, maximum=18446744073709551615, step=1, randomize=True)
     gr.Examples(fn=run, examples=[" ", "an insect robot preparing a delicious meal, anime style", "a photo of a group of friends at an amusement park"], inputs=prompt, outputs=[output, seed], cache_examples=True)
@@ -66,7 +68,7 @@ with gr.Blocks(css=css, theme=theme) as demo:
             prompt.submit
         ],
         fn=run,
-        inputs=[prompt, negative_prompt, guidance_scale, pag_scale, randomize_seed, seed],
+        inputs=[prompt, negative_prompt, guidance_scale, pag_scale, pag_layers, randomize_seed, seed],
         outputs=[output, seed],
     )
 if __name__ == "__main__":