negative prompt

app.py

@@ -73,7 +73,7 @@ def encode_target_images(images, pipeline):
 
 
 @spaces.GPU(duration=120)
-def generate_image(text, img1, img2, img3, guidance_scale, inference_steps, seed, rigid_object, enable_cpu_offload=False):
+def generate_image(text, img1, img2, img3, guidance_scale, inference_steps, seed, enable_cpu_offload=False, neg_prompt="", true_cfg=1.0, image_cfg=0.0):
     if enable_cpu_offload:
         pipeline.enable_sequential_cpu_offload()
     input_images = [img1, img2, img3]
@@ -104,6 +104,9 @@ def generate_image(text, img1, img2, img3, guidance_scale, inference_steps, seed
         generator = torch.Generator(device="cpu").manual_seed(seed),
         joint_attention_kwargs={'shared_attn': True, 'num': numref},
         return_dict=False,
+        negative_prompt=neg_prompt,
+        true_cfg_scale=true_cfg,
+        image_cfg_scale=image_cfg,
     )[0][0]
     output = rearrange(output, "b c h (n w) -> (b n) c h w", n=numref)[::numref]
     img = Image.fromarray( (( torch.clip(output[0].float(), -1., 1.).permute(1,2,0).cpu().numpy()*0.5+0.5)*255).astype(np.uint8) )
@@ -120,7 +123,10 @@ def get_example():
             "./imgs/test_cases/action_figure/2.jpg",
             3.5,
             42,
-            True,
+            False,
+            "",
+            1.0,
+            0.0,
         ],
         [
             "A penguin plushie wearing pink sunglasses is lounging on a beach. Realistic shot.",
@@ -129,7 +135,10 @@ def get_example():
             "./imgs/test_cases/penguin/2.jpg",
             3.5,
             42,
-            True,
+            False,
+            "",
+            1.0,
+            0.0,
         ],
         [
             "A toy on a beach. Waves in the background. Realistic shot.",
@@ -138,16 +147,19 @@ def get_example():
             "./imgs/test_cases/rc_car/04.jpg",
             3.5,
             42,
-            True,
+            False,
+            "",
+            1.0,
+            0.0,
         ],
     ]
     return case
 
-def run_for_examples(text, img1, img2, img3, guidance_scale, seed, rigid_object, enable_cpu_offload=False):    
+def run_for_examples(text, img1, img2, img3, guidance_scale, seed, enable_cpu_offload=False, neg_prompt="", true_cfg=1.0, image_cfg=0.0):    
     inference_steps = 30
     
     return generate_image(
-        text, img1, img2, img3, guidance_scale, inference_steps, seed, rigid_object, enable_cpu_offload
+        text, img1, img2, img3, guidance_scale, inference_steps, seed, enable_cpu_offload, neg_prompt, true_cfg, image_cfg
     )
 
 description = """
@@ -210,13 +222,17 @@ with gr.Blocks() as demo:
                 label="Seed", minimum=0, maximum=2147483647, value=42, step=1
             )
 
-            rigid_object = gr.Checkbox(
-                label="rigid_object", info="Whether its a rigid object or a deformable object like pet animals, wearable etc.", value=True,
-            )
             enable_cpu_offload = gr.Checkbox(
                 label="Enable CPU Offload", info="Enable CPU Offload to avoid memory issues", value=False,
             )
 
+            with gr.Accordion("Advanced Options (True CFG, true_cfg_scale=1 means use fake CFG, >1 means use true CFG", open=False):    # noqa E501
+                neg_prompt = gr.Textbox(
+                    label="Negative Prompt",
+                    value="")
+                true_cfg = gr.Slider(1.0, 10.0, 1.5, step=0.1, label="true CFG. Recommended to be 1.5")
+                image_cfg = gr.Slider(0.0, 10.0, 0.0, step=0.1, label="image CFG scale, will increase the image alignment but longer run time and lower text alignment. Recommended to be 1.0")
+
             # generate
             generate_button = gr.Button("Generate Image")
             
@@ -236,8 +252,10 @@ with gr.Blocks() as demo:
             guidance_scale_input,
             num_inference_steps,
             seed_input,
-            rigid_object,   
             enable_cpu_offload,
+            neg_prompt,
+            true_cfg,
+            image_cfg,
         ],
         outputs=output_image,
     )
@@ -252,7 +270,10 @@ with gr.Blocks() as demo:
             image_input_3,
             guidance_scale_input,
             seed_input,
-            rigid_object,
+            enable_cpu_offload,
+            neg_prompt,
+            true_cfg,
+            image_cfg,
         ],
         outputs=output_image,
     )

pipelines/flux_pipeline/pipeline.py

@@ -98,6 +98,19 @@ def retrieve_timesteps(
     return timesteps, num_inference_steps
 
 
+def normalized_guidance_image(neg_noise_pred, noise_pred, image_noise_pred, true_cfg_scale, image_cfg_scale):
+    diff_img = image_noise_pred - neg_noise_pred
+    diff_txt = noise_pred - image_noise_pred
+
+    diff_norm_txt = diff_txt.norm(p=2, dim=[-1, -2], keepdim=True) 
+    diff_norm_img = diff_img.norm(p=2, dim=[-1, -2], keepdim=True) 
+    min_norm = torch.minimum(diff_norm_img, diff_norm_txt)
+    diff_txt = diff_txt * torch.minimum(torch.ones_like(diff_txt), min_norm / diff_norm_txt)
+    diff_img = diff_img * torch.minimum(torch.ones_like(diff_txt), min_norm / diff_norm_img)
+    pred_guided = image_noise_pred + image_cfg_scale * diff_img + true_cfg_scale *  diff_txt
+    return pred_guided
+
+
 class SynCDFluxPipeline(FluxPipeline):
     
     model_cpu_offload_seq = "text_encoder->text_encoder_2->transformer->vae"
@@ -162,6 +175,7 @@ class SynCDFluxPipeline(FluxPipeline):
         latents_ref: Optional[torch.Tensor] = None,
         latents_mask: Optional[torch.Tensor] = None,
         return_latents: bool=False,
+        image_cfg_scale: float=0.0,
     ):
         r"""
         Function invoked when calling the pipeline for generation.
@@ -389,7 +403,7 @@ class SynCDFluxPipeline(FluxPipeline):
                     return_dict=False,
                 )[0]
 
-                if do_true_cfg:
+                if do_true_cfg and i>=1:
                     neg_noise_pred = self.transformer(
                         hidden_states=latents,
                         timestep=timestep / 1000,
@@ -398,10 +412,28 @@ class SynCDFluxPipeline(FluxPipeline):
                         encoder_hidden_states=negative_prompt_embeds,
                         txt_ids=text_ids,
                         img_ids=latent_image_ids,
-                        joint_attention_kwargs=self.joint_attention_kwargs,
+                        joint_attention_kwargs={**self.joint_attention_kwargs, 'neg_mode': True},
                         return_dict=False,
                     )[0]
-                    noise_pred = neg_noise_pred + true_cfg_scale * (noise_pred - neg_noise_pred)
+
+                    if image_cfg_scale > 0:
+                        image_noise_pred = self.transformer(
+                            hidden_states=latents,                            
+                            timestep=timestep / 1000,
+                            guidance=guidance,
+                            pooled_projections=negative_pooled_prompt_embeds,
+                            encoder_hidden_states=negative_prompt_embeds,
+                            txt_ids=text_ids,
+                            img_ids=latent_image_ids,
+                            joint_attention_kwargs=self.joint_attention_kwargs,
+                            return_dict=False,
+                        )[0]
+                        
+                    if image_cfg_scale == 0:
+                        noise_pred = neg_noise_pred + true_cfg_scale * (noise_pred - neg_noise_pred)
+                    else:
+                        noise_pred = normalized_guidance_image(neg_noise_pred, noise_pred, image_noise_pred, true_cfg_scale, image_cfg_scale)
+
 
                 # compute the previous noisy sample x_t -> x_t-1
                 latents_dtype = latents.dtype

pipelines/flux_pipeline/transformer.py

@@ -79,6 +79,7 @@ class FluxAttnProcessor2_0:
         scale: float = 1.0,
         timestep: float = 0,
         val: bool = False,
+        neg_mode: bool = False,
     ) -> torch.FloatTensor:
         if mode == 'w': # and single:
             ref_dict[self.name] = hidden_states.detach()
@@ -133,7 +134,19 @@ class FluxAttnProcessor2_0:
             query = apply_rotary_emb(query, image_rotary_emb)
             key = apply_rotary_emb(key, image_rotary_emb)
 
-        hidden_states = F.scaled_dot_product_attention(query, key, value, dropout_p=0.0, is_causal=False, attn_mask=mask if timestep < 1. else None)
+        if neg_mode:
+            res = int(math.sqrt((end_of_hidden_states-(text_seq if encoder_hidden_states is None else 0)) // num))
+            hw = res*res
+            mask_ = torch.ones(1, res, num*res, res, num*res).to(query.device)
+            for i in range(num):
+                mask_[:, :, i*res:(i+1)*res, :, i*res:(i+1)*res] = 1
+            mask_ = rearrange(mask_, "b h w h1 w1 -> b (h w) (h1 w1)")
+            mask = torch.ones(1, num*hw + 512, num*hw + 512, device=query.device, dtype=query.dtype)
+            mask[:, 512:, 512:] = mask_
+            mask = mask.bool()
+            mask = rearrange(mask.unsqueeze(0).expand(attn.heads, -1, -1, -1), "nh b ... -> b nh ...")
+            
+        hidden_states = F.scaled_dot_product_attention(query, key, value, dropout_p=0.0, is_causal=False, attn_mask=mask)
         hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
 
         hidden_states = hidden_states.to(query.dtype)