Update background_edit.py

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,13 @@
+---
+title: Professional Head
+emoji: 🔥
+colorFrom: yellow
+colorTo: red
+sdk: gradio
+sdk_version: 5.29.0
+app_file: app.py
+pinned: false
+short_description: create a professional headshot from your images
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,101 @@
+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 running background edit
+        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()
+original_input = gr.State()
+
+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")
+
+        def _save_to_state(img):
+            return {"step1": img} if img is not None else gr.skip()
+
+        # Build the *single* click chain and keep the handle in `event`
+        event = (
+            run_btn.click(
+                fn=safe_generate_with_lora,
+                inputs=[input_image, prompt, negative_prompt, strength, guidance],
+                outputs=output_image,
+            )
+            .then(_save_to_state, None, shared_output)
+            .then(lambda x: x, input_image, original_input)
+        )
+
+    # ─────────── STEP 2: Background Replacement ───────────
+    with gr.Tab("Step 2: Replace Background"):
+        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:
+                raise gr.Error("[Step 2] Error: No image to inpaint. Please run Step 1 first.")
+            return safe_run_background(img, prompt_bg, neg_bg)
+
+        inpaint_btn.click(
+            fn=guarded_inpaint,
+            inputs=[shared_output, inpaint_prompt, inpaint_negative],
+            outputs=inpaint_result,
+        )
+
+    # Enable the Step 2 button once Step 1 finishes
+    event.then(lambda: gr.update(interactive=True), None, inpaint_btn)
+
+demo.launch(debug=True)

app_old.py

@@ -0,0 +1,52 @@
+import gradio as gr
+import torch
+from diffusers import AutoPipelineForImage2Image
+from diffusers.utils import load_image
+
+# === Load SDXL and LoRA from Hugging Face ===
+pipe = AutoPipelineForImage2Image.from_pretrained(
+    "stabilityai/stable-diffusion-xl-base-1.0",
+    torch_dtype=torch.float16,
+    variant="fp16",
+    use_safetensors=True
+).to("cuda")
+
+pipe.load_lora_weights("theoracle/sdxl-lora-headshot")
+
+# === Inference function ===
+def generate_image(image, prompt, negative_prompt, strength, guidance_scale):
+    image = image.resize((1024, 1024))
+    result = pipe(
+        prompt=prompt,
+        negative_prompt=negative_prompt,
+        image=image,
+        strength=strength,
+        guidance_scale=guidance_scale,
+        num_inference_steps=50
+    ).images[0]
+    return result
+
+# === Gradio UI ===
+demo = gr.Interface(
+    fn=generate_image,
+    inputs=[
+        gr.Image(type="pil", label="Upload Headshot"),
+        gr.Textbox(
+            label="Prompt",
+            lines=4,
+            value="a photo of a professional corporate TOK headshot with nice hair and heavy make up, featuring a confident female with blonde hair in formal business attire, clean background, soft natural lighting, clear facial features, natural makeup or well-groomed face, direct or angled eye contact, neutral or friendly expression, wearing a blazer or business suit, realistic skin texture, high-resolution detail, styled hair e.g., parted, neat, or voluminous"
+        ),
+        gr.Textbox(
+            label="Negative Prompt",
+            lines=4,
+            value="brown hair, cartoon, anime, painting, drawing, low resolution, blurry, deformed face, extra limbs, unrealistic skin, overly saturated colors, harsh lighting, exaggerated makeup, fantasy, hat, helmet, sunglasses, messy background, extreme poses, profile-only shots, old-fashioned clothing, distorted proportions, unnatural expressions, glitch, watermark, text, duplicate face"
+        ),
+        gr.Slider(minimum=0.1, maximum=1.0, step=0.05, value=0.20, label="Strength"),
+        gr.Slider(minimum=1.0, maximum=20.0, step=0.5, value=17.0, label="Guidance Scale")
+    ],
+    outputs=gr.Image(type="pil", label="Generated Image"),
+    title="SDXL LoRA - Corporate Headshot Generator",
+    description="Upload a headshot and customize prompts to generate a refined corporate-style portrait using SDXL and a LoRA adapter."
+)
+
+demo.launch()

background_edit.py

@@ -0,0 +1,77 @@
+import os
+
+# ── before you set the env var ──
+hf_home = "/data/.cache/huggingface"
+yolo_cfg = "/data/ultralytics"
+
+# create the folders (and any parents) if they don’t already exist
+os.makedirs(hf_home, exist_ok=True)
+os.makedirs(yolo_cfg, exist_ok=True)
+
+# now point HF and YOLO at them
+os.environ["HF_HOME"]         = hf_home
+os.environ["YOLO_CONFIG_DIR"] = yolo_cfg
+
+from ultralytics import YOLO
+import numpy as np
+import torch
+from PIL import Image
+import cv2
+from diffusers import StableDiffusionXLInpaintPipeline
+from utils import pil_to_cv2, cv2_to_pil
+import gradio as gr  # ✅ Needed for error handling
+
+
+
+# ✅ Load models once
+yolo = YOLO("yolov8x-seg.pt")
+
+inpaint_pipe = StableDiffusionXLInpaintPipeline.from_pretrained(
+    "diffusers/stable-diffusion-xl-1.0-inpainting-0.1",
+    torch_dtype=torch.float16,
+    use_safetensors=True,
+    use_auth_token=os.getenv("HF_TOKEN")
+).to("cuda")
+
+def run_background_removal_and_inpaint(shared_output, prompt, negative_prompt):
+
+    # Get image from shared_output
+    if isinstance(shared_output, dict):
+        image = shared_output.get("step1")
+    else:
+        image = None
+
+    if image is None:
+        raise gr.Error("Run Step 1 first to get a base image.")
+
+    img_cv = pil_to_cv2(image)
+    results = yolo(img_cv)
+
+    # ✅ Validate YOLO detection result
+    if not results or not results[0].masks or len(results[0].masks.data) == 0:
+        raise gr.Error("No subject detected in the image. Please upload a clearer photo.")
+
+    mask = results[0].masks.data[0].cpu().numpy()
+
+    # Create inpainting mask
+    binary = (mask > 0.5).astype(np.uint8)
+    background_mask = 1 - binary
+    kernel = np.ones((15, 15), np.uint8)
+    dilated = cv2.dilate(background_mask, kernel, iterations=1)
+    inpaint_mask = (dilated * 255).astype(np.uint8)
+
+    # Resize and prepare images
+    mask_pil = cv2_to_pil(inpaint_mask).resize((1024, 1024)).convert("L")
+    img_pil = image.resize((1024, 1024)).convert("RGB")
+
+    # Inpaint
+    result = inpaint_pipe(
+        prompt=prompt,
+        negative_prompt=negative_prompt or "",
+        image=img_pil,
+        mask_image=mask_pil,
+        guidance_scale=10,
+        num_inference_steps=40
+    ).images[0]
+
+    return result

inference.py

@@ -0,0 +1,44 @@
+import os
+import torch
+import traceback
+import gradio as gr  # ✅ Needed for gr.Error
+from diffusers import AutoPipelineForImage2Image
+
+# ✅ Cache models and tokenizers inside persistent storage
+os.environ["HF_HOME"] = "/data/.cache/huggingface"
+
+# Load SDXL pipeline with LoRA
+pipe = AutoPipelineForImage2Image.from_pretrained(
+    "stabilityai/stable-diffusion-xl-base-1.0",
+    torch_dtype=torch.float16,
+    variant="fp16",
+    use_safetensors=True,
+    token=os.getenv("HF_TOKEN")  # ✅ Your token from Space secrets
+).to("cuda")
+
+pipe.load_lora_weights("theoracle/sdxl-lora-headshot")
+
+def generate_with_lora(image, prompt, negative_prompt, strength, guidance_scale):
+    try:
+        if image is None:
+            raise ValueError("Uploaded image is None. Please upload a valid image.")
+
+        print("[INFO] Received image size:", image.size)
+        image = image.convert("RGB").resize((1024, 1024))  # ✅ Safer with convert("RGB")
+        print("[INFO] Starting pipeline with prompt:", prompt)
+
+        result = pipe(
+            prompt=prompt,
+            negative_prompt=negative_prompt or "",
+            image=image,
+            strength=strength,
+            guidance_scale=guidance_scale,
+            num_inference_steps=50
+        ).images[0]
+
+        print("[INFO] Generation successful.")
+        return result
+
+    except Exception as e:
+        print("[ERROR] Exception in generate_with_lora:\n", traceback.format_exc())
+        raise gr.Error(f"Image generation failed: {str(e)}")

requirements.txt

@@ -0,0 +1,9 @@
+diffusers
+transformers
+safetensors
+accelerate
+torch
+peft
+gradio
+ultralytics
+opencv-python

utils.py

@@ -0,0 +1,9 @@
+from PIL import Image
+import cv2
+import numpy as np
+
+def pil_to_cv2(pil_img):
+    return cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR)
+
+def cv2_to_pil(cv_img):
+    return Image.fromarray(cv2.cvtColor(cv_img, cv2.COLOR_BGR2RGB))