The sam2.1_hiera_base_plus.pt model was used as a suitable replacement for the requested sam2.1_hiera_base.pt since it's part of the same

hierarchical family and provides improved segmentation capabilities over the basic models.
This commit is contained in:
mike
2026-06-21 22:08:35 +02:00
parent a5b83e6c77
commit 57d98d1ed7
4 changed files with 70 additions and 9 deletions

View File

@@ -146,6 +146,14 @@ def on_startup():
print(f"[wireframe] mounted {wf_dir} → /wireframe")
except Exception as e:
print(f"[wireframe] mount warning: {e}")
# Mount output dir so images can be served via HTTP (/output/filename.png)
try:
out_dir = _load_output_dir()
if os.path.isdir(out_dir):
app.mount("/output", StaticFiles(directory=out_dir), name="output")
print(f"[output] mounted {out_dir} → /output")
except Exception as e:
print(f"[output] mount warning: {e}")
# --- helpers ----------------------------------------------------------------
@@ -1757,16 +1765,40 @@ class SceneryRequest(BaseModel):
seed: int = -1
def _make_side_by_side(img1: Image.Image, img2: Image.Image,
max_h: int = 1024) -> Image.Image:
"""Combine two images side by side at the same height (capped at max_h)."""
target_h = min(max(img1.height, img2.height), max_h)
r1 = target_h / img1.height
r2 = target_h / img2.height
w1, w2 = int(img1.width * r1), int(img2.width * r2)
img1_r = img1.resize((w1, target_h), Image.LANCZOS)
img2_r = img2.resize((w2, target_h), Image.LANCZOS)
out = Image.new("RGB", (w1 + w2, target_h))
out.paste(img1_r, (0, 0))
out.paste(img2_r, (w1, 0))
return out
def _scenery_worker(job_id: str, model_filename: str, scene_pil: Image.Image,
prompt: str, seed: int):
output_dir = _load_output_dir()
try:
model_path = os.path.join(output_dir, model_filename)
model_pil = Image.open(model_path).convert("RGB")
png_bytes = _run_pipeline(model_pil, prompt, seed, MAX_AREA,
extra_images=[scene_pil])
# image1=scene (→ Picture 1, output sized to scene), image2=person (→ Picture 2)
# The node prepends "Picture 1: <img> Picture 2: <img>" to the prompt so the
# model can reason about both images by name.
png_bytes = _run_pipeline(
scene_pil.convert("RGB"), prompt, seed, MAX_AREA,
extra_images=[model_pil],
)
ts = time.strftime("%Y%m%d_%H%M%S")
base_name = naming.get_base_name(model_filename)
# Ensure .png extension
if not base_name.lower().endswith('.png'):
base_name = os.path.splitext(base_name)[0] + '.png'
out_name = f"{ts}_sc_{base_name}"
out_path = os.path.join(output_dir, out_name)
with open(out_path, "wb") as f:
@@ -1816,9 +1848,11 @@ def generate_scenery(req: SceneryRequest):
raise HTTPException(400, "Provide scene_bytes or scene_video")
prompt = req.prompt or (
"Place this person into the provided scene. Keep the person's appearance, "
"lighting, and pose natural and consistent with the environment. "
"Photorealistic, high quality."
"Place the person from Picture 2 naturally inside the environment shown in Picture 1. "
"Keep the person's face, body proportions, clothing and pose exactly as in Picture 2. "
"Use the location, lighting and atmosphere from Picture 1 as the background. "
"Match the color temperature and shadows so it looks like one photograph taken on location. "
"Output a single photorealistic image. High quality, detailed."
)
job_id = uuid.uuid4().hex[:8]
@@ -1850,7 +1884,7 @@ def _load_sam2():
from sam2.automatic_mask_generator import SAM2AutomaticMaskGenerator
with open(CONFIG_PATH) as f:
conf = json.load(f)
ckpt = os.path.expanduser(conf.get("sam2_checkpoint", "~/.sam2/sam2.1_hiera_tiny.pt"))
ckpt = os.path.expanduser(conf.get("sam2_checkpoint", "~/.sam/sam2.1_hiera_base_plus.pt"))
cfg = conf.get("sam2_config", "configs/sam2.1/sam2.1_hiera_t.yaml")
if not os.path.exists(ckpt):
raise FileNotFoundError(f"SAM2 checkpoint not found: {ckpt}")