Include advanced options to modify guidance scale, seed

app.py

@@ -82,7 +82,7 @@ def pgd_inpaint(X, target, model, criterion, eps=0.1, step_size=0.015, iters=40,
             
     return X_adv
     
-def process_image_img2img(raw_image,prompt):
+def process_image_img2img(raw_image,prompt, scale, num_steps, seed):
     resize = T.transforms.Resize(512)
     center_crop = T.transforms.CenterCrop(512)
     init_image = center_crop(resize(raw_image))
@@ -103,13 +103,13 @@ def process_image_img2img(raw_image,prompt):
     adv_image = to_pil(adv_X[0]).convert("RGB")
     
     # a good seed (uncomment the line below to generate new images)
-    SEED = 9222
+    SEED = seed# Default is 9222
     # SEED = np.random.randint(low=0, high=10000)
     
     # Play with these for improving generated image quality
     STRENGTH = 0.5
-    GUIDANCE = 7.5
-    NUM_STEPS = 50
+    GUIDANCE = scale # Default is 7.5
+    NUM_STEPS = num_steps # Default is 50
     
     with torch.autocast('cuda'):
         torch.manual_seed(SEED)
@@ -119,7 +119,7 @@ def process_image_img2img(raw_image,prompt):
 
     return [(init_image,"Source Image"), (adv_image, "Adv Image"), (image_nat,"Gen. Image Nat"), (image_adv, "Gen. Image Adv")]
 
-def process_image_inpaint(raw_image,mask, prompt):
+def process_image_inpaint(raw_image,mask, prompt,scale, num_steps, seed):
     init_image = raw_image.convert('RGB').resize((512,512))
     mask_image = mask.convert('RGB')
     mask_image = ImageOps.invert(mask_image).resize((512,512))
@@ -156,7 +156,7 @@ def process_image_inpaint(raw_image,mask, prompt):
     adv_image = recover_image(adv_image, init_image, mask_image, background=True)
 
     # A good seed
-    SEED = 9209
+    SEED = seed #Default is 9209
     
     # Uncomment the below to generated other images
     # SEED = np.random.randint(low=0, high=100000)
@@ -165,8 +165,8 @@ def process_image_inpaint(raw_image,mask, prompt):
     print(SEED)
     
     #strength = 0.7
-    guidance_scale = 7.5
-    num_inference_steps = 100
+    guidance_scale = scale# Default is 7.5
+    num_inference_steps = num_steps # Default is 100
     
     image_nat = pipe_inpaint(prompt=prompt, 
                          image=init_image, 
@@ -192,7 +192,7 @@ def process_image_inpaint(raw_image,mask, prompt):
     return [(init_image,"Source Image"), (adv_image, "Adv Image"), (image_nat,"Gen. Image Nat"), (image_adv, "Gen. Image Adv")]
 
 
-examples_list = [["dog.png", "dog under heavy rain and muddy ground real"]]
+examples_list = [["dog.png", "dog under heavy rain and muddy ground real", 7.5, 50, 9222]]
 
 
 with gr.Blocks() as demo:
@@ -203,7 +203,7 @@ with gr.Blocks() as demo:
      <p style="margin-bottom: 10px; font-size: 94%">This is an unofficial demo for Photoguard, which is an approach to safeguarding images against manipulation by ML-powered photo-editing models such as stable diffusion through immunization of images. The demo is based on the <a href='https://github.com/MadryLab/photoguard' style='text-decoration: underline;' target='_blank'> Github </a> implementation provided by the authors.</p>
           ''')
     gr.HTML('''
-   <p align="center"><img src="https://raw.githubusercontent.com/MadryLab/photoguard/main/assets/hero_fig.PNG" style="width:40%"/></p>
+   <p align="center"><img src="https://raw.githubusercontent.com/MadryLab/photoguard/main/assets/hero_fig.PNG" style="width:60%"/></p>
     ''')
     gr.HTML('''
     <p style="margin-bottom: 10px; font-size: 94%"> A malevolent actor might download
@@ -224,15 +224,20 @@ we disrupt their ability to successfully perform such edits forcing them to gene
             mask_image_inpaint = gr.Image(type="pil", label = "Mask")
             input_prompt_inpaint = gr.Textbox(label="Prompt")
             run_btn_inpaint = gr.Button('Run')
-        
+
+        with gr.Accordion("Advanced options", open=False):
+                scale = gr.Slider(label="Guidance Scale", minimum=0.1, maximum=30.0, value=7.5, step=0.1)
+                num_steps = gr.Slider(label="Number of Inference Steps", minimum=5, maximum=125, value=100, step=5)
+                seed = gr.Slider(label="Seed", minimum=0, maximum=2147483647, step=1, randomize=True)
+            
         with gr.Row():
             result_gallery = gr.Gallery(
             label="Generated images", show_label=False, elem_id="gallery"
         ).style(grid=[2], height="auto")
             
-    run_btn_img2img.click(process_image_img2img, inputs = [input_image_img2img,input_prompt_img2img], outputs = [result_gallery])
-    examples = gr.Examples(examples=examples_list,inputs = [input_image_img2img,input_prompt_img2img],  outputs = [result_gallery], cache_examples = True, fn = process_image_img2img)
-    run_btn_inpaint.click(process_image_inpaint, inputs = [input_image_inpaint,mask_image_inpaint,input_prompt_inpaint], outputs = [result_gallery])   
+    run_btn_img2img.click(process_image_img2img, inputs = [input_image_img2img,input_prompt_img2img, scale, num_steps, seed], outputs = [result_gallery])
+    examples = gr.Examples(examples=examples_list,inputs = [input_image_img2img,input_prompt_img2img,scale, num_steps, seed],  outputs = [result_gallery], cache_examples = True, fn = process_image_img2img)
+    run_btn_inpaint.click(process_image_inpaint, inputs = [input_image_inpaint,mask_image_inpaint,input_prompt_inpaint,scale, num_steps, seed], outputs = [result_gallery])   
     
                      
 demo.launch(debug=True)