Update app.py

app.py

@@ -24,14 +24,27 @@ pipe = FluxPipeline.from_pretrained(
 ).to("cuda")
 
 @spaces.GPU
-def infer(control_image, seed=42, randomize_seed=False, width=1024, height=1024, guidance_scale=3.5, num_inference_steps=28, progress=gr.Progress(track_tqdm=True)):
+def infer(control_image, prompt, reference_scale= 0.03 , seed=42, randomize_seed=False, width=1024, height=1024, guidance_scale=3.5, num_inference_steps=28, progress=gr.Progress(track_tqdm=True)):
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
-    pipe_prior_output = pipe_prior_redux(control_image)
+    pipe_prior_output = pipe_prior_redux(control_image, prompt=prompt)
+    cond_size = 729
+    hidden_size = 4096
+    max_sequence_length = 512
+    full_attention_size = max_sequence_length + hidden_size + cond_size
+    attention_mask = torch.zeros(
+        (full_attention_size, full_attention_size), device="cuda", dtype=torch.bfloat16
+    )
+    bias = torch.log(
+        torch.tensor(reference_scale, dtype=torch.bfloat16, device="cuda").clamp(min=1e-5, max=1)
+    )
+    attention_mask[:, max_sequence_length : max_sequence_length + cond_size] = bias
+    joint_attention_kwargs=dict(attention_mask=attention_mask)
     images = pipe(
         guidance_scale=guidance_scale,
         num_inference_steps=num_inference_steps,
         generator=torch.Generator("cpu").manual_seed(seed),
+        joint_attention_kwargs=joint_attention_kwargs,
         **pipe_prior_output,
     ).images[0]
     return images, seed
@@ -53,6 +66,20 @@ An adapter for FLUX [dev] to create image variations
         with gr.Row():
             with gr.Column():
                     input_image = gr.Image(label="Image to create variations", type="pil")
+                    prompt = gr.Text(
+                        label="Prompt",
+                        show_label=False,
+                        max_lines=1,
+                        placeholder="Enter your prompt",
+                        container=False,
+                    )
+                    reference_scale = gr.Slider(
+                    label="Masking Scale",
+                    minimum=0.01,
+                    maximum=0.08,
+                    step=0.001,
+                    value=0.03,
+                    )
                     run_button = gr.Button("Run")
             result = gr.Image(label="Result", show_label=False)
         
@@ -107,7 +134,7 @@ An adapter for FLUX [dev] to create image variations
     gr.on(
         triggers=[run_button.click],
         fn = infer,
-        inputs = [input_image, seed, randomize_seed, width, height, guidance_scale, num_inference_steps],
+        inputs = [input_image, prompt, reference_scale, seed, randomize_seed, width, height, guidance_scale, num_inference_steps],
         outputs = [result, seed]
     )