Add separate prompt and guidance control for both steps

app.py

@@ -16,31 +16,40 @@ def unload_models():
     torch.cuda.empty_cache()
     gc.collect()
 
-def safe_generate_and_inpaint(image, prompt, negative_prompt, strength, guidance_scale):
+def safe_generate_and_inpaint(
+    image,
+    prompt_1, neg_1, strength_1, guidance_1,
+    prompt_2, neg_2, guidance_2
+):
     try:
         if image is None:
             raise gr.Error("Please upload an image first.")
 
-        # Step 1: Refinement
-        print("[INFO] Step 1: Generating refined image...", flush=True)
+        # Step 1: Headshot Refinement
+        print("[INFO] Step 1: Refining headshot...", flush=True)
         refined = generate_with_lora(
             image=image,
-            prompt=prompt,
-            negative_prompt=negative_prompt,
-            strength=strength,
-            guidance_scale=guidance_scale,
+            prompt=prompt_1,
+            negative_prompt=neg_1,
+            strength=strength_1,
+            guidance_scale=guidance_1,
         )
 
-        # Save to disk
+        # Save intermediate result to disk
         os.makedirs("./outputs", exist_ok=True)
         ts = datetime.now().strftime("%Y%m%d_%H%M%S")
         path = f"./outputs/step1_result_{ts}.png"
         refined.save(path)
 
-        # Step 2: Background removal and inpainting
+        # Step 2: Background Inpainting
         print("[INFO] Step 2: Inpainting background...", flush=True)
         unload_models()
-        result = run_background_removal_and_inpaint(path, prompt, negative_prompt)
+        result = run_background_removal_and_inpaint(
+            image_path=path,
+            prompt=prompt_2,
+            negative_prompt=neg_2,
+            guidance_scale=guidance_2
+        )
 
         return refined, result, ""
 
@@ -52,37 +61,44 @@ def safe_generate_and_inpaint(image, prompt, negative_prompt, strength, guidance
 
 # ───────────────────── Gradio UI ─────────────────────
 with gr.Blocks() as demo:
-    gr.Markdown("## 🖼️ AI Headshot Enhancer + Background Replacer")
-    gr.Markdown("Upload a headshot, adjust the prompt, and click one button. We'll refine the image and replace the background automatically.")
+    gr.Markdown("## 🧠 Headshot + Background Generator (Full Prompt Control)")
 
     with gr.Row():
         input_image = gr.Image(type="pil", label="Upload Headshot")
 
+    gr.Markdown("### Step 1: Headshot Refinement (LoRA)")
+
     with gr.Row():
-        prompt = gr.Textbox(
-            label="Prompt",
-            value="a professional corporate headshot of a confident woman in her 30s with blonde hair"
-        )
-        negative_prompt = gr.Textbox(
-            label="Negative Prompt",
-            value="deformed, cartoon, anime, illustration, painting, drawing, sketch, low resolution, blurry, out of focus, pixelated"
-        )
+        prompt_1 = gr.Textbox(label="Headshot Prompt", value="a professional headshot of a confident woman in her 30s with blonde hair")
+        neg_1 = gr.Textbox(label="Headshot Negative Prompt", value="deformed, cartoon, anime, sketch, blurry, low quality")
+
+    with gr.Row():
+        strength_1 = gr.Slider(0.1, 1.0, value=0.2, step=0.05, label="Refinement Strength")
+        guidance_1 = gr.Slider(1, 20, value=17, step=0.5, label="Guidance Scale (Headshot)")
+
+    gr.Markdown("### Step 2: Background Inpainting (SDXL)")
+
+    with gr.Row():
+        prompt_2 = gr.Textbox(label="Background Prompt", value="modern hospital background, clean, soft lighting")
+        neg_2 = gr.Textbox(label="Background Negative Prompt", value="fantasy, cartoon, cluttered, sketch")
 
     with gr.Row():
-        strength = gr.Slider(0.1, 1.0, value=0.20, step=0.05, label="Refinement Strength")
-        guidance = gr.Slider(1, 20, value=17.0, step=0.5, label="Guidance Scale")
+        guidance_2 = gr.Slider(1, 20, value=10, step=0.5, label="Guidance Scale (Background)")
 
-    go_btn = gr.Button("✨ Run Full Process (Refine + Inpaint)")
+    go_btn = gr.Button("✨ Generate Refined Headshot + Background")
 
     with gr.Row():
         output_refined = gr.Image(type="pil", label="Step 1: Refined Headshot")
-        output_final   = gr.Image(type="pil", label="Step 2: Final Image with Background")
+        output_final = gr.Image(type="pil", label="Step 2: Final Image with Background")
 
     error_box = gr.Markdown(label="Error", value="", visible=True)
 
     go_btn.click(
         fn=safe_generate_and_inpaint,
-        inputs=[input_image, prompt, negative_prompt, strength, guidance],
+        inputs=[
+            input_image, prompt_1, neg_1, strength_1, guidance_1,
+            prompt_2, neg_2, guidance_2
+        ],
         outputs=[output_refined, output_final, error_box]
     )
 

background_edit.py

@@ -33,7 +33,7 @@ inpaint_pipe = StableDiffusionXLInpaintPipeline.from_pretrained(
     use_auth_token=os.getenv("HF_TOKEN")
 ).to("cuda")
 
-def run_background_removal_and_inpaint(image_path, prompt, negative_prompt):
+def run_background_removal_and_inpaint(image_path, prompt, negative_prompt, guidance_scale=10):
     if not image_path or not os.path.isfile(image_path):
         raise gr.Error("No valid image found. Please run Step 1 first.")
 
@@ -59,7 +59,7 @@ def run_background_removal_and_inpaint(image_path, prompt, negative_prompt):
         negative_prompt=negative_prompt or "",
         image=img_pil,
         mask_image=mask_pil,
-        guidance_scale=10,
+        guidance_scale=guidance_scale,
         num_inference_steps=40
     ).images[0]