Allow using different schedulers and negative prompts

app.py

@@ -8,6 +8,7 @@ from utils.parse import filter_boxes
 from generation import run as run_ours
 from baseline import run as run_baseline
 import torch
+from shared import DEFAULT_SO_NEGATIVE_PROMPT, DEFAULT_OVERALL_NEGATIVE_PROMPT
 from examples import stage1_examples, stage2_examples
 
 print(f"Is CUDA available: {torch.cuda.is_available()}")
@@ -87,7 +88,7 @@ def get_layout_image(response):
 def get_layout_image_gallery(response):
     return [get_layout_image(response)]
 
-def get_ours_image(response, seed, num_inference_steps, fg_seed_start, fg_blending_ratio=0.1, frozen_step_ratio=0.4, gligen_scheduled_sampling_beta=0.3, show_so_imgs=False, scale_boxes=False):
+def get_ours_image(response, seed, num_inference_steps, dpm_scheduler, fg_seed_start, fg_blending_ratio=0.1, frozen_step_ratio=0.4, gligen_scheduled_sampling_beta=0.3, so_negative_prompt="", overall_negative_prompt="", show_so_imgs=False, scale_boxes=False):
     if response == "":
         response = layout_placeholder
     gen_boxes, bg_prompt = parse_input(response)
@@ -98,10 +99,18 @@ def get_ours_image(response, seed, num_inference_steps, fg_seed_start, fg_blendi
         'gen_boxes': gen_boxes,
         'bg_prompt': bg_prompt
     }
+    
+    if dpm_scheduler:
+        scheduler_key = "dpm_scheduler"
+    else:
+        scheduler_key = "scheduler"
+        
     image_np, so_img_list = run_ours(
         spec, bg_seed=seed, fg_seed_start=fg_seed_start, 
         fg_blending_ratio=fg_blending_ratio,frozen_step_ratio=frozen_step_ratio, 
-        gligen_scheduled_sampling_beta=gligen_scheduled_sampling_beta, num_inference_steps=num_inference_steps)
+        gligen_scheduled_sampling_beta=gligen_scheduled_sampling_beta, num_inference_steps=num_inference_steps, scheduler_key=scheduler_key,
+        so_negative_prompt=so_negative_prompt, overall_negative_prompt=overall_negative_prompt
+    )
     images = [image_np]
     if show_so_imgs:
         images.extend([np.asarray(so_img) for so_img in so_img_list])
@@ -110,7 +119,10 @@ def get_ours_image(response, seed, num_inference_steps, fg_seed_start, fg_blendi
 def get_baseline_image(prompt, seed):
     if prompt == "":
         prompt = prompt_placeholder
-    image_np = run_baseline(prompt, bg_seed=seed)
+    
+    scheduler_key = "dpm_scheduler"
+    
+    image_np = run_baseline(prompt, bg_seed=seed, scheduler_key=scheduler_key)
     return [image_np]
 
 def parse_input(text=None):
@@ -232,17 +244,20 @@ with gr.Blocks(
                 with gr.Accordion("Advanced options", open=False):
                     seed = gr.Slider(0, 10000, value=0, step=1, label="Seed")
                     num_inference_steps = gr.Slider(1, 50, value=20, step=1, label="Number of inference steps")
+                    dpm_scheduler = gr.Checkbox(label="Use DPM scheduler (unchecked: DDIM scheduler, may have better coherence, recommend 50 inference steps)", show_label=False, value=True)
                     fg_seed_start = gr.Slider(0, 10000, value=20, step=1, label="Seed for foreground variation")
                     fg_blending_ratio = gr.Slider(0, 1, value=0.1, step=0.01, label="Variations added to foreground for single object generation (0: no variation, 1: max variation)")
                     frozen_step_ratio = gr.Slider(0, 1, value=0.4, step=0.1, label="Foreground frozen steps ratio (higher: preserve object attributes; lower: higher coherence; set to 0: (almost) equivalent to vanilla GLIGEN except details)")
                     gligen_scheduled_sampling_beta = gr.Slider(0, 1, value=0.3, step=0.1, label="GLIGEN guidance steps ratio (the beta value)")
+                    so_negative_prompt = gr.Textbox(lines=1, label="Negative prompt for single object generation", value=DEFAULT_SO_NEGATIVE_PROMPT)
+                    overall_negative_prompt = gr.Textbox(lines=1, label="Negative prompt for overall generation", value=DEFAULT_OVERALL_NEGATIVE_PROMPT)
                     show_so_imgs = gr.Checkbox(label="Show annotated single object generations", show_label=False)
             with gr.Column(scale=1):
                 gallery = gr.Gallery(
                     label="Generated image", show_label=False, elem_id="gallery"
                 ).style(columns=[1], rows=[1], object_fit="contain", preview=True)
         visualize_btn.click(fn=get_layout_image_gallery, inputs=response, outputs=gallery, api_name="visualize-layout")
-        generate_btn.click(fn=get_ours_image, inputs=[response, seed, num_inference_steps, fg_seed_start, fg_blending_ratio, frozen_step_ratio, gligen_scheduled_sampling_beta, show_so_imgs], outputs=gallery, api_name="layout-to-image")
+        generate_btn.click(fn=get_ours_image, inputs=[response, seed, num_inference_steps, dpm_scheduler, fg_seed_start, fg_blending_ratio, frozen_step_ratio, gligen_scheduled_sampling_beta, so_negative_prompt, overall_negative_prompt, show_so_imgs], outputs=gallery, api_name="layout-to-image")
 
         gr.Examples(
             stage2_examples,

baseline.py

@@ -3,7 +3,7 @@
 import torch
 import models
 from models import pipelines
-from shared import model_dict
+from shared import model_dict, DEFAULT_OVERALL_NEGATIVE_PROMPT
 
 vae, tokenizer, text_encoder, unet, scheduler, dtype = model_dict.vae, model_dict.tokenizer, model_dict.text_encoder, model_dict.unet, model_dict.scheduler, model_dict.dtype
 
@@ -17,9 +17,10 @@ batch_size = 1
 # h, w
 image_scale = (512, 512)
 
-bg_negative = 'artifacts, blurry, smooth texture, bad quality, distortions, unrealistic, distorted image, bad proportions, duplicate'
+bg_negative = DEFAULT_OVERALL_NEGATIVE_PROMPT
 
-def run(prompt, bg_seed=1, num_inference_steps=20):
+# Using dpm scheduler by default
+def run(prompt, scheduler_key='dpm_scheduler', bg_seed=1, num_inference_steps=20):
     print(f"prompt: {prompt}")
     generator = torch.Generator(models.torch_device).manual_seed(bg_seed)
     
@@ -34,7 +35,7 @@ def run(prompt, bg_seed=1, num_inference_steps=20):
     pipelines.gligen_enable_fuser(model_dict['unet'], enabled=False)
     _, images = pipelines.generate(
         model_dict, latents, input_embeddings, num_inference_steps,  
-        guidance_scale=guidance_scale
+        guidance_scale=guidance_scale, scheduler_key=scheduler_key
     )
 
     return images[0]

generation.py

@@ -1,17 +1,15 @@
 version = "v3.0"
 
-from PIL import Image
 import torch
 import models
-from models import load_sd
 import utils
 from models import pipelines, sam
 from utils import parse, latents
-from shared import model_dict, sam_model_dict
+from shared import model_dict, sam_model_dict, DEFAULT_SO_NEGATIVE_PROMPT, DEFAULT_OVERALL_NEGATIVE_PROMPT
 
 verbose = False
 
-vae, tokenizer, text_encoder, unet, scheduler, dtype = model_dict.vae, model_dict.tokenizer, model_dict.text_encoder, model_dict.unet, model_dict.scheduler, model_dict.dtype
+vae, tokenizer, text_encoder, unet, dtype = model_dict.vae, model_dict.tokenizer, model_dict.text_encoder, model_dict.unet, model_dict.dtype
 
 model_dict.update(sam_model_dict)
 
@@ -37,14 +35,14 @@ run_ind = None
 
 def generate_single_object_with_box(prompt, box, phrase, word, input_latents, input_embeddings, 
                                     sam_refine_kwargs, num_inference_steps, gligen_scheduled_sampling_beta=0.3, 
-                                    verbose=False, visualize=True):
+                                    verbose=False, scheduler_key=None, visualize=True):
     
     bboxes, phrases, words = [box], [phrase], [word]
     
     latents, single_object_images, single_object_pil_images_box_ann, latents_all = pipelines.generate_gligen(
         model_dict, input_latents, input_embeddings, num_inference_steps, bboxes, phrases, gligen_scheduled_sampling_beta=gligen_scheduled_sampling_beta, 
         guidance_scale=guidance_scale, return_saved_cross_attn=False,
-        return_box_vis=True, save_all_latents=True
+        return_box_vis=True, save_all_latents=True, scheduler_key=scheduler_key
     )
     
     mask_selected, conf_score_selected = sam.sam_refine_box(sam_input_image=single_object_images[0], box=box, model_dict=model_dict, verbose=verbose, **sam_refine_kwargs)
@@ -78,7 +76,7 @@ def get_masked_latents_all_list(so_prompt_phrase_word_box_list, input_latents_li
 
 def run(
     spec, bg_seed = 1, fg_seed_start = 20, frozen_step_ratio=0.4, gligen_scheduled_sampling_beta = 0.3, num_inference_steps = 20,
-    so_center_box = False, fg_blending_ratio = 0.1, so_horizontal_center_only = True,
+    so_center_box = False, fg_blending_ratio = 0.1, scheduler_key='dpm_scheduler', so_negative_prompt = DEFAULT_SO_NEGATIVE_PROMPT, overall_negative_prompt = DEFAULT_OVERALL_NEGATIVE_PROMPT, so_horizontal_center_only = True,
     align_with_overall_bboxes = False, horizontal_shift_only = True
 ):
     """    
@@ -106,13 +104,6 @@ def run(
             print(f"centered so_prompt_phrase_word_box_list: {so_prompt_phrase_word_box_list}")
     so_boxes = [item[-1] for item in so_prompt_phrase_word_box_list]
 
-    if True:
-        so_negative_prompt = "artifacts, blurry, smooth texture, bad quality, distortions, unrealistic, distorted image, bad proportions, duplicate, two, many, group, occlusion, occluded, side, border, collate"
-        overall_negative_prompt = "artifacts, blurry, smooth texture, bad quality, distortions, unrealistic, distorted image, bad proportions, duplicate"
-    else:
-        so_negative_prompt = ""
-        overall_negative_prompt = ""
-
     sam_refine_kwargs = dict(
         discourage_mask_below_confidence=discourage_mask_below_confidence, discourage_mask_below_coarse_iou=discourage_mask_below_coarse_iou,
         height=height, width=width, H=H, W=W
@@ -139,7 +130,7 @@ def run(
     latents_all_list, mask_tensor_list, so_img_list = get_masked_latents_all_list(
         so_prompt_phrase_word_box_list, input_latents_list, 
         gligen_scheduled_sampling_beta=gligen_scheduled_sampling_beta,
-        sam_refine_kwargs=sam_refine_kwargs, so_input_embeddings=so_input_embeddings, num_inference_steps=num_inference_steps, verbose=verbose
+        sam_refine_kwargs=sam_refine_kwargs, so_input_embeddings=so_input_embeddings, num_inference_steps=num_inference_steps, scheduler_key=scheduler_key, verbose=verbose
     )
 
     
@@ -166,7 +157,7 @@ def run(
         model_dict, composed_latents, overall_input_embeddings, num_inference_steps, 
         overall_bboxes_flattened, overall_phrases_flattened, guidance_scale=guidance_scale,
         gligen_scheduled_sampling_beta=gligen_scheduled_sampling_beta,
-        frozen_steps=frozen_steps, frozen_mask=frozen_mask
+        frozen_steps=frozen_steps, frozen_mask=frozen_mask, scheduler_key=scheduler_key
     )
 
     print(f"Generation with spatial guidance from input latents and first {frozen_steps} steps frozen (directly from the composed latents input)")

models/models.py

@@ -8,7 +8,7 @@ import numpy as np
 from utils.latents import get_unscaled_latents, get_scaled_latents, blend_latents
 from utils import torch_device
 
-def load_sd(key="runwayml/stable-diffusion-v1-5", use_fp16=False, load_inverse_scheduler=True, use_dpm_multistep_scheduler=False):
+def load_sd(key="runwayml/stable-diffusion-v1-5", use_fp16=False, load_inverse_scheduler=True):
     """
     Keys:
      key = "CompVis/stable-diffusion-v1-4"
@@ -22,7 +22,6 @@ def load_sd(key="runwayml/stable-diffusion-v1-5", use_fp16=False, load_inverse_s
     ```
     
     use_fp16: fp16 might have degraded performance
-    use_dpm_multistep_scheduler: DPMSolverMultistepScheduler
     """
     
     # run final results in fp32
@@ -37,12 +36,10 @@ def load_sd(key="runwayml/stable-diffusion-v1-5", use_fp16=False, load_inverse_s
     tokenizer = CLIPTokenizer.from_pretrained(key, subfolder="tokenizer", revision=revision, torch_dtype=dtype)
     text_encoder = CLIPTextModel.from_pretrained(key, subfolder="text_encoder", revision=revision, torch_dtype=dtype).to(torch_device)
     unet = UNet2DConditionModel.from_pretrained(key, subfolder="unet", revision=revision, torch_dtype=dtype).to(torch_device)
-    if use_dpm_multistep_scheduler:
-        scheduler = DPMSolverMultistepScheduler.from_pretrained(key, subfolder="scheduler", revision=revision, torch_dtype=dtype)
-    else:
-        scheduler = DDIMScheduler.from_pretrained(key, subfolder="scheduler", revision=revision, torch_dtype=dtype)
+    dpm_scheduler = DPMSolverMultistepScheduler.from_pretrained(key, subfolder="scheduler", revision=revision, torch_dtype=dtype)
+    scheduler = DDIMScheduler.from_pretrained(key, subfolder="scheduler", revision=revision, torch_dtype=dtype)
 
-    model_dict = EasyDict(vae=vae, tokenizer=tokenizer, text_encoder=text_encoder, unet=unet, scheduler=scheduler, dtype=dtype)
+    model_dict = EasyDict(vae=vae, tokenizer=tokenizer, text_encoder=text_encoder, unet=unet, scheduler=scheduler, dpm_scheduler=dpm_scheduler, dtype=dtype)
     
     if load_inverse_scheduler:
         inverse_scheduler = DDIMInverseScheduler.from_config(scheduler.config)

models/pipelines.py

@@ -53,8 +53,8 @@ def decode(vae, latents):
     return images
 
 @torch.no_grad()
-def generate(model_dict, latents, input_embeddings, num_inference_steps, guidance_scale = 7.5, no_set_timesteps=False):
-    vae, tokenizer, text_encoder, unet, scheduler, dtype = model_dict.vae, model_dict.tokenizer, model_dict.text_encoder, model_dict.unet, model_dict.scheduler, model_dict.dtype
+def generate(model_dict, latents, input_embeddings, num_inference_steps, guidance_scale = 7.5, no_set_timesteps=False, scheduler_key='dpm_scheduler'):
+    vae, tokenizer, text_encoder, unet, scheduler, dtype = model_dict.vae, model_dict.tokenizer, model_dict.text_encoder, model_dict.unet, model_dict[scheduler_key], model_dict.dtype
     text_embeddings, uncond_embeddings, cond_embeddings = input_embeddings
     
     if not no_set_timesteps:
@@ -93,11 +93,11 @@ def generate_gligen(model_dict, latents, input_embeddings, num_inference_steps,
     frozen_steps=20, frozen_mask=None,
     return_saved_cross_attn=False, saved_cross_attn_keys=None, return_cond_ca_only=False, return_token_ca_only=None, 
     offload_cross_attn_to_cpu=False, offload_latents_to_cpu=True,
-    return_box_vis=False, show_progress=True, save_all_latents=False):
+    return_box_vis=False, show_progress=True, save_all_latents=False, scheduler_key='dpm_scheduler'):
     """
     The `bboxes` should be a list, rather than a list of lists (one box per phrase, we can have multiple duplicated phrases).
     """
-    vae, tokenizer, text_encoder, unet, scheduler, dtype = model_dict.vae, model_dict.tokenizer, model_dict.text_encoder, model_dict.unet, model_dict.scheduler, model_dict.dtype
+    vae, tokenizer, text_encoder, unet, scheduler, dtype = model_dict.vae, model_dict.tokenizer, model_dict.text_encoder, model_dict.unet, model_dict[scheduler_key], model_dict.dtype
     text_embeddings, uncond_embeddings, cond_embeddings = input_embeddings
     
     if latents.dim() == 5:

shared.py

@@ -1,11 +1,15 @@
 from models import load_sd, sam
 
+
+DEFAULT_SO_NEGATIVE_PROMPT = "artifacts, blurry, smooth texture, bad quality, distortions, unrealistic, distorted image, bad proportions, duplicate, two, many, group, occlusion, occluded, side, border, collate"
+DEFAULT_OVERALL_NEGATIVE_PROMPT = "artifacts, blurry, smooth texture, bad quality, distortions, unrealistic, distorted image, bad proportions, duplicate"
+
+
 use_fp16 = False
-use_dpm = True
 
 sd_key = "gligen/diffusers-generation-text-box"
 
 print(f"Using SD: {sd_key}")
-model_dict = load_sd(key=sd_key, use_fp16=use_fp16, use_dpm_multistep_scheduler=use_dpm, load_inverse_scheduler=False)
+model_dict = load_sd(key=sd_key, use_fp16=use_fp16, load_inverse_scheduler=False)
 
 sam_model_dict = sam.load_sam()