Spaces:
Running
on
Zero
Running
on
Zero
File size: 5,041 Bytes
94d7430 10bd531 e0eea92 10bd531 94d7430 79e7646 10bd531 79e7646 10bd531 94d7430 10bd531 94d7430 10bd531 79e7646 10bd531 |
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 121 122 123 124 125 126 127 128 129 130 131 |
import os
import traceback
from datetime import datetime
import torch, gc
from PIL import Image
import gradio as gr
from inference import generate_with_lora
from background_edit import run_background_removal_and_inpaint
# ───────────────────── 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(image_path, *args, **kwargs):
try:
unload_models() # free VRAM before loading the inpainting model
return run_background_removal_and_inpaint(image_path, *args, **kwargs)
except Exception as e:
_print_trace()
raise gr.Error(f"[Step 2] Background replacement failed: {type(e).__name__}: {e}")
def _save_to_disk(img):
if img is None:
return gr.skip()
os.makedirs("/tmp/gradio_outputs", exist_ok=True)
ts = datetime.now().strftime("%Y%m%d_%H%M%S")
path = f"/tmp/gradio_outputs/step1_result_{ts}.png"
img.save(path)
return path
# ───────────────────── UI ─────────────────────
shared_output_path = gr.State() # holds file path to Step 1 output
original_input = gr.State() # holds the original upload (if needed)
with gr.Blocks() as demo:
demo.queue()
# ─────────── 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")
event = (
run_btn.click(
fn=safe_generate_with_lora,
inputs=[input_image, prompt, negative_prompt, strength, guidance],
outputs=output_image,
)
.then(_save_to_disk, output_image, shared_output_path)
.then(lambda x: x, input_image, original_input)
)
# ─────────── STEP 2: Background Replacement ───────────
with gr.Tab("Step 2: Replace Background"):
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(image_path, prompt_bg, neg_bg):
if not image_path or not os.path.isfile(image_path):
return None, "**🛑 Error:** No valid headshot found — please run Step 1 first."
try:
print(f"[DEBUG] Loading image from: {image_path}", flush=True)
result = safe_run_background(image_path, prompt_bg, neg_bg)
return result, ""
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_path, inpaint_prompt, inpaint_negative],
outputs=[inpaint_result, error_box],
)
event.then(lambda: gr.update(interactive=True), None, inpaint_btn)
demo.launch(debug=True)
|