background_edit.py

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