Spaces:
Running
on
Zero
Running
on
Zero
File size: 3,680 Bytes
79e7646 800f45f 79e7646 eb20c24 79e7646 eb20c24 79e7646 eb20c24 79e7646 3a30bfb 93be815 10bd531 79e7646 10bd531 79e7646 eb20c24 79e7646 93be815 79e7646 eb20c24 3a30bfb eb20c24 |
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 |
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
import spaces
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
INPAINT_SIZE = 1024
# Load clothing model
clothing_model = YOLO("deepfashion2_yolov8s-seg.pt")
# ✅ 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")
@spaces.GPU(duration=30)
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.")
image = Image.open(image_path).convert("RGB")
img_cv = pil_to_cv2(image)
results = yolo(img_cv)
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()
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)
mask_pil = cv2_to_pil(inpaint_mask).resize((INPAINT_SIZE, INPAINT_SIZE)).convert("L")
img_pil = image.resize((INPAINT_SIZE, INPAINT_SIZE)).convert("RGB")
result = inpaint_pipe(
prompt=prompt,
negative_prompt=negative_prompt or "",
image=img_pil,
mask_image=mask_pil,
guidance_scale=guidance_scale,
num_inference_steps=40
).images[0]
return result
@spaces.GPU(duration=30)
def run_clothing_inpaint(image, prompt, negative_prompt, guidance):
try:
print("[INFO] Step 3: Clothing segmentation and inpainting...", flush=True)
img_cv = np.array(image.convert("RGB"))[..., ::-1] # PIL → OpenCV BGR
h, w = img_cv.shape[:2]
# Segment clothing
results = clothing_model(img_cv)
masks = results[0].masks.data.cpu().numpy()
if len(masks) == 0:
raise gr.Error("No clothing detected. Try a different image.")
mask = masks[0]
resized_mask = cv2.resize(mask, (w, h), interpolation=cv2.INTER_NEAREST)
binary_mask = (resized_mask > 0.5).astype(np.uint8) * 255
mask_pil = Image.fromarray(binary_mask).convert("L").resize((INPAINT_SIZE, INPAINT_SIZE))
# Resize input image
resized_image = image.convert("RGB").resize((INPAINT_SIZE, INPAINT_SIZE))
# Inpaint clothing
result = inpaint_pipe(
prompt=prompt,
negative_prompt=negative_prompt,
image=resized_image,
mask_image=mask_pil,
guidance_scale=guidance,
num_inference_steps=50
).images[0]
return result, ""
except gr.Error as e:
return None, f"🛑 {str(e)}"
except Exception as e:
traceback.print_exc()
return None, f"❌ Unexpected Error: {type(e).__name__}: {str(e)}"
|