Update app.py

app.py

@@ -75,18 +75,21 @@ vae_hybridpony = AutoencoderKL.from_pretrained(os.path.join(ckpt_dir_hybridpony,
 pipe_realpony = StableDiffusionXLPipeline.from_pretrained(
     ckpt_dir_realpony,
     vae=vae_realpony,
+    custom_pipeline="multimodalart/sdxl_perturbed_attention_guidance",
     torch_dtype=torch.float16,
     use_safetensors=True,
 )
 pipe_ultpony = StableDiffusionXLPipeline.from_pretrained(
     ckpt_dir_ultpony,
     vae=vae_ultpony,
+    custom_pipeline="multimodalart/sdxl_perturbed_attention_guidance",
     torch_dtype=torch.float16,
     use_safetensors=True,
 )
 pipe_hybridpony = StableDiffusionXLPipeline.from_pretrained(
     ckpt_dir_hybridpony,
     vae=vae_hybridpony,
+    custom_pipeline="multimodalart/sdxl_perturbed_attention_guidance",
     torch_dtype=torch.float16,
     use_safetensors=True,
 )
@@ -204,7 +207,7 @@ def upscale_image(image, scale):
 
 @spaces.GPU(duration=120)
 def generate_image(model_choice, additional_positive_prompt, additional_negative_prompt, height, width, num_inference_steps,
-                   guidance_scale, num_images_per_prompt, use_random_seed, seed, sampler, clip_skip, 
+                   guidance_scale, pag_scale, pag_layers, num_images_per_prompt, use_random_seed, seed, sampler, clip_skip, 
                    use_florence2, use_medium_enhancer, use_long_enhancer,
                    use_positive_prefix, use_positive_suffix, use_negative_prefix, use_negative_suffix,
                    use_upscaler, upscale_factor,
@@ -271,6 +274,8 @@ def generate_image(model_choice, additional_positive_prompt, additional_negative
             width=width,
             num_inference_steps=num_inference_steps,
             guidance_scale=guidance_scale,
+            pag_scale=pag_scale,
+            pag_layers=pag_layers,
             num_images_per_prompt=num_images_per_prompt,
             generator=torch.Generator(pipe.device).manual_seed(seed)
         ).images
@@ -323,6 +328,8 @@ with gr.Blocks(theme='bethecloud/storj_theme') as demo:
                 width = gr.Slider(512, 2048, 1024, step=64, label="Width")
                 num_inference_steps = gr.Slider(20, 100, 30, step=1, label="Number of Inference Steps")
                 guidance_scale = gr.Slider(1, 20, 6, step=0.1, label="Guidance Scale")
+                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")
                 num_images_per_prompt = gr.Slider(1, 4, 1, step=1, label="Number of images per prompt")
                 use_random_seed = gr.Checkbox(label="Use Random Seed", value=True)
                 seed = gr.Number(label="Seed", value=0, precision=0)
@@ -374,7 +381,7 @@ with gr.Blocks(theme='bethecloud/storj_theme') as demo:
         inputs=[
             model_choice,  # Add this new input
             positive_prompt, negative_prompt, height, width, num_inference_steps,
-            guidance_scale, num_images_per_prompt, use_random_seed, seed, sampler,
+            guidance_scale, pag_scale, pag_layers, num_images_per_prompt, use_random_seed, seed, sampler,
             clip_skip, use_florence2, use_medium_enhancer, use_long_enhancer,
             use_positive_prefix, use_positive_suffix, use_negative_prefix, use_negative_suffix,
             use_upscaler, upscale_factor,