Spaces:
Running
on
Zero
Running
on
Zero
File size: 4,855 Bytes
94d7430 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 |
import gradio as gr
from inference import generate_with_lora
from background_edit import run_background_removal_and_inpaint
import traceback, torch, gc
# ───────────────────── Helpers ─────────────────────
def _print_trace():
traceback.print_exc()
def safe_generate_with_lora(*a, **kw):
try:
return generate_with_lora(*a, **kw)
except gr.Error:
_print_trace()
raise
except Exception as e:
_print_trace()
raise gr.Error(f"Image generation failed: {e}")
def unload_models():
torch.cuda.empty_cache()
gc.collect()
def safe_run_background(*args, **kwargs):
try:
unload_models() # free VRAM before loading the inpainting model
return run_background_removal_and_inpaint(*args, **kwargs)
except Exception as e:
_print_trace()
raise gr.Error(f"[Step 2] Background replacement failed: {type(e).__name__}: {e}")
# ───────────────────── UI ─────────────────────
shared_output = gr.State() # holds the Step 1 output image
original_input = gr.State() # holds the original upload (optional)
with gr.Blocks() as demo:
demo.queue() # enable batching / concurrency
# ─────────── STEP 1: Headshot Refinement ───────────
with gr.Tab("Step 1: Headshot Refinement"):
with gr.Row():
input_image = gr.Image(type="pil", label="Upload Headshot")
output_image = gr.Image(type="pil", label="Refined Output")
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"
)
with gr.Row():
strength = gr.Slider(0.1, 1.0, value=0.20, step=0.05, label="Strength")
guidance = gr.Slider(1, 20, value=17.0, step=0.5, label="Guidance Scale")
run_btn = gr.Button("Generate")
def _save_to_state(img):
return {"step1": img} if img is not None else gr.skip()
event = (
run_btn.click(
fn=safe_generate_with_lora,
inputs=[input_image, prompt, negative_prompt, strength, guidance],
outputs=output_image,
)
.then(_save_to_state, output_image, shared_output)
.then(lambda x: x, input_image, original_input)
)
# ─────────── STEP 2: Background Replacement ───────────
with gr.Tab("Step 2: Replace Background"):
# Show formatted error messages
error_box = gr.Markdown(value="", visible=True)
with gr.Row():
inpaint_prompt = gr.Textbox(
label="New Background Prompt",
value="modern open-plan startup office background, natural lighting, glass walls, clean design, minimalistic decor"
)
inpaint_negative = gr.Textbox(
label="Negative Prompt",
value="dark lighting, cluttered background, fantasy elements, cartoon, anime, painting, low quality, distorted shapes"
)
with gr.Row():
inpaint_result = gr.Image(type="pil", label="Inpainted Image")
with gr.Row():
inpaint_btn = gr.Button("Remove Background & Inpaint", interactive=False)
def guarded_inpaint(img, prompt_bg, neg_bg):
if img is None:
return None, "**🛑 Error:** No headshot found — please run Step 1 first."
try:
print("[DEBUG] Starting background removal and inpainting…", flush=True)
result = safe_run_background(img, prompt_bg, neg_bg)
return result, "" # Clear error on success
except gr.Error as e:
print(f"[Step 2 gr.Error] {e}", flush=True)
return None, f"**🛑 Step 2 Failed:** {str(e)}"
except Exception as e:
print(f"[Step 2 UNEXPECTED ERROR] {type(e).__name__}: {e}", flush=True)
return None, f"**❌ Unexpected Error:** {type(e).__name__}: {e}"
inpaint_btn.click(
fn=guarded_inpaint,
inputs=[shared_output, inpaint_prompt, inpaint_negative],
outputs=[inpaint_result, error_box],
)
# Enable Step 2 after Step 1 completes
event.then(lambda: gr.update(interactive=True), None, inpaint_btn)
demo.launch(debug=True)
|