""" ui.py — Gradio Blocks frontend for the Passport Photo Maker. Entry point for the HF Space (Gradio SDK auto-runs this file as `app.py`). """ from __future__ import annotations import gc import io import logging import os import tempfile import zipfile import gradio as gr from PIL import Image from engine import ( STANDARDS, PAPER_SIZES_MM, MAX_BATCH_SIZE, process_photo, process_batch, format_compliance_markdown, list_garments, warm_up, ) logging.basicConfig(level=logging.INFO) logger = logging.getLogger("passport-maker") STANDARD_CHOICES = list(STANDARDS.keys()) PAPER_CHOICES = ["None (single photo only)"] + list(PAPER_SIZES_MM.keys()) OUTFIT_CHOICES = ["None (keep original clothing)"] + list_garments() # --------------------------------------------------------------------------- # Single-photo handler # --------------------------------------------------------------------------- def run_pipeline( image: Image.Image | None, spec_choice: str, bg_color: str | None, paper_choice: str, zoom: float, x_offset: float, y_offset: float, auto_straighten: bool, outfit_choice: str, progress: gr.Progress = gr.Progress(), ): """Gradio click handler. Wraps engine.process_photo with progress updates and converts internal exceptions into user-facing gr.Error popups instead of raw tracebacks. """ if image is None: raise gr.Error("Please upload a photo first.") try: progress(0.15, desc="Preparing image...") progress(0.35, desc="Removing background (this takes a few seconds on CPU)...") paper_key = None if paper_choice.startswith("None") else paper_choice outfit_label = None if outfit_choice.startswith("None") else outfit_choice photo, sheet, checks, bg_removed_preview, face_thumb, straighten_angle, outfit_applied, outfit_error = process_photo( image=image, spec_key=spec_choice, bg_hex=bg_color, paper_key=paper_key, zoom=zoom, x_offset=x_offset, y_offset=y_offset, auto_straighten=auto_straighten, outfit_label=outfit_label, ) progress(0.9, desc="Finalizing...") gc.collect() progress(1.0, desc="Done") status = f"✅ **{spec_choice}** — {STANDARDS[spec_choice].width_mm}×{STANDARDS[spec_choice].height_mm}mm @ 300 DPI ({photo.width}×{photo.height}px)" if sheet is not None: status += f" · tiled on {paper_key}" if abs(straighten_angle) >= 0.5: status += f" · auto-straightened {abs(straighten_angle):.1f}°" if outfit_label: # Only claim the outfit was applied if it genuinely was — a # prior version of this code claimed success here whenever # outfit_label was merely selected, even when compositing had # silently failed and fallen back to the original photo. Now # honest either way. if outfit_applied: status += f" · outfit: {outfit_label}" else: status += ( f" · ⚠️ outfit '{outfit_label}' could not be applied " f"(showing original clothing)" ) if outfit_error: status += f" — {outfit_error}" compliance_md = format_compliance_markdown(checks) # Stage-by-stage gallery — mirrors cutout.pro's Original / BG-removed # / Face cutout / Result breakdown. Each tuple is (image, caption); # Gradio Gallery renders the caption under the thumbnail. stages = [ (image, "1. Original"), (bg_removed_preview, "2. Background Removed"), (face_thumb, "3. Face Cutout"), (photo, "4. Cropped / Final"), ] # ImageSlider needs a (before, after) pair at the same canvas. We # hand it (original upload, final result) — Gradio letterboxes/ # fits internally, so mismatched source aspect ratios display fine; # it's a visual compare, not a pixel-aligned overlay. return stages, (image, photo), sheet, status, compliance_md except ValueError as e: # Expected, user-facing errors (no face detected, bad spec, photo # doesn't fit paper, etc) — clean message, no traceback. raise gr.Error(str(e)) except Exception: logger.exception("Unhandled error in process_photo") raise gr.Error( "Something went wrong processing this photo. Try a different " "image, or a smaller file size." ) # --------------------------------------------------------------------------- # Batch handler # --------------------------------------------------------------------------- def run_batch( files: list | None, spec_choice: str, bg_color: str | None, zoom: float, x_offset: float, y_offset: float, auto_straighten: bool, outfit_choice: str, progress: gr.Progress = gr.Progress(), ): """Gradio click handler for the Batch tab. Accepts a list of uploaded file paths (from gr.File multiple), runs each through process_photo, and returns a gallery of successes + a status report + a zip download. Per-image failures never abort the batch — see engine.process_batch. """ if not files: raise gr.Error("Please upload at least one photo.") if len(files) > MAX_BATCH_SIZE: raise gr.Error( f"Batch limit is {MAX_BATCH_SIZE} photos. You uploaded {len(files)} " f"— please remove some and try again." ) images: list[tuple[str, Image.Image]] = [] for f in files: path = f.name if hasattr(f, "name") else f try: img = Image.open(path) img.load() images.append((os.path.basename(path), img)) except Exception: images.append((os.path.basename(path), None)) progress(0.1, desc=f"Processing {len(images)} photos...") # Filter out unreadable files up front with a clear per-file error, # rather than letting them crash into process_batch's image pipeline. valid = [(name, img) for name, img in images if img is not None] bad_names = [name for name, img in images if img is None] outfit_label = None if outfit_choice.startswith("None") else outfit_choice try: results = process_batch( valid, spec_key=spec_choice, bg_hex=bg_color, zoom=zoom, x_offset=x_offset, y_offset=y_offset, auto_straighten=auto_straighten, outfit_label=outfit_label, ) except ValueError as e: raise gr.Error(str(e)) except Exception: logger.exception("Unhandled error in process_batch") raise gr.Error("Something went wrong processing this batch.") progress(0.85, desc="Packaging results...") gallery_items = [] ok_count = 0 status_lines = [] zip_path = None good_results = [r for r in results if r.photo is not None] if good_results: tmp_dir = tempfile.mkdtemp(prefix="passport_batch_") zip_path = os.path.join(tmp_dir, "passport_photos.zip") with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zf: for r in good_results: buf = io.BytesIO() r.photo.save(buf, format="PNG") out_name = os.path.splitext(r.filename)[0] + "_passport.png" zf.writestr(out_name, buf.getvalue()) gallery_items.append((r.photo, r.filename)) ok_count += 1 for r in results: if r.photo is not None: status_lines.append(f"✅ {r.filename}") else: status_lines.append(f"❌ {r.filename} — {r.error}") for name in bad_names: status_lines.append(f"❌ {name} — Could not read this file as an image.") header = f"**{ok_count}/{len(images)} photos processed successfully.**" status_md = header + "\n\n" + "\n".join(status_lines) gc.collect() progress(1.0, desc="Done") return gallery_items, status_md, zip_path def build_interface() -> gr.Blocks: with gr.Blocks( title="Passport Photo Maker — Free, Instant, 300 DPI", ) as demo: gr.Markdown( """ # 🛂 Passport Photo Maker AI background removal + auto face-centering + exact 300 DPI passport, visa, and ID photo sizing. Free. No login. No watermark. """ ) with gr.Tabs(): # ----------------------------------------------------------- # TAB 1 — Single photo # ----------------------------------------------------------- with gr.Tab("Single Photo"): with gr.Row(): with gr.Column(scale=1): inp_image = gr.Image( type="pil", label="Upload your photo", sources=["upload", "webcam"], ) spec_dropdown = gr.Dropdown( choices=STANDARD_CHOICES, value=STANDARD_CHOICES[0], label="Photo Standard", ) bg_color = gr.ColorPicker( value="#1E3A8A", label="Background Color (overrides standard default)", ) paper_dropdown = gr.Dropdown( choices=PAPER_CHOICES, value=PAPER_CHOICES[0], label="Print Sheet Layout", ) outfit_dropdown = gr.Dropdown( choices=OUTFIT_CHOICES, value=OUTFIT_CHOICES[0], label="Outfit (replaces clothing below the neck)", ) with gr.Accordion("Adjust Crop (optional — auto-crop is usually correct)", open=False): zoom_slider = gr.Slider( minimum=0.5, maximum=2.0, value=1.0, step=0.05, label="Zoom (higher = tighter crop on face)", ) x_slider = gr.Slider( minimum=-0.5, maximum=0.5, value=0.0, step=0.02, label="Shift Left / Right", ) y_slider = gr.Slider( minimum=-0.5, maximum=0.5, value=0.0, step=0.02, label="Shift Up / Down", ) reset_crop_btn = gr.Button("Reset Crop", size="sm") straighten_checkbox = gr.Checkbox( value=True, label="Auto-Straighten (levels head + shoulders if the photo is tilted)", ) generate_btn = gr.Button("Generate Photo", variant="primary", size="lg") with gr.Column(scale=1): out_stages = gr.Gallery( label="Processing Stages", columns=4, object_fit="contain", height=200, ) out_compare = gr.ImageSlider( label="Before / After — drag to compare", type="pil", ) status_box = gr.Markdown() compliance_box = gr.Markdown() out_sheet = gr.Image(type="pil", label="Print Sheet") # Update the color picker default whenever the standard # changes, so the user sees the correct expected background # before they touch anything — manual override still wins. def _sync_bg_default(spec_choice: str): return gr.update(value=STANDARDS[spec_choice].bg_hex) spec_dropdown.change( fn=_sync_bg_default, inputs=spec_dropdown, outputs=bg_color ) generate_btn.click( fn=run_pipeline, inputs=[ inp_image, spec_dropdown, bg_color, paper_dropdown, zoom_slider, x_slider, y_slider, straighten_checkbox, outfit_dropdown, ], outputs=[out_stages, out_compare, out_sheet, status_box, compliance_box], ) reset_crop_btn.click( fn=lambda: (gr.update(value=1.0), gr.update(value=0.0), gr.update(value=0.0)), inputs=None, outputs=[zoom_slider, x_slider, y_slider], ) # ----------------------------------------------------------- # TAB 2 — Batch # ----------------------------------------------------------- with gr.Tab("Batch (up to 10 photos)"): gr.Markdown( "Upload multiple photos, apply one standard + background to all, " "download every result in a single ZIP. Each photo is processed " "independently — one bad photo won't stop the rest." ) with gr.Row(): with gr.Column(scale=1): batch_files = gr.File( label="Upload photos", file_count="multiple", file_types=["image"], ) batch_spec_dropdown = gr.Dropdown( choices=STANDARD_CHOICES, value=STANDARD_CHOICES[0], label="Photo Standard (applied to all)", ) batch_bg_color = gr.ColorPicker( value="#1E3A8A", label="Background Color (overrides standard default)", ) with gr.Accordion("Adjust Crop (applied to all photos)", open=False): batch_zoom_slider = gr.Slider( minimum=0.5, maximum=2.0, value=1.0, step=0.05, label="Zoom", ) batch_x_slider = gr.Slider( minimum=-0.5, maximum=0.5, value=0.0, step=0.02, label="Shift Left / Right", ) batch_y_slider = gr.Slider( minimum=-0.5, maximum=0.5, value=0.0, step=0.02, label="Shift Up / Down", ) def _sync_batch_bg_default(spec_choice: str): return gr.update(value=STANDARDS[spec_choice].bg_hex) batch_spec_dropdown.change( fn=_sync_batch_bg_default, inputs=batch_spec_dropdown, outputs=batch_bg_color, ) batch_straighten_checkbox = gr.Checkbox( value=True, label="Auto-Straighten (applied to all photos)", ) batch_outfit_dropdown = gr.Dropdown( choices=OUTFIT_CHOICES, value=OUTFIT_CHOICES[0], label="Outfit (applied to all photos)", ) batch_generate_btn = gr.Button( "Generate All", variant="primary", size="lg" ) with gr.Column(scale=1): batch_gallery = gr.Gallery( label="Results", columns=3, object_fit="contain", height="auto" ) batch_status = gr.Markdown() batch_zip_out = gr.File(label="Download all (ZIP)") batch_generate_btn.click( fn=run_batch, inputs=[ batch_files, batch_spec_dropdown, batch_bg_color, batch_zoom_slider, batch_x_slider, batch_y_slider, batch_straighten_checkbox, batch_outfit_dropdown, ], outputs=[batch_gallery, batch_status, batch_zip_out], ) gr.Markdown( """ --- ⚠️ Photos are processed in-memory and are not stored. Compliance checks are automated heuristics, not a guarantee — verify final prints against your destination country's exact requirements before submission. """ ) return demo # Concurrency is capped at 2 to protect the single CPU-basic worker from # being driven into OOM by parallel BiRefNet inferences — each inference # holds a 1024x1024 float32 activation stack in memory; more than a # couple concurrent requests on 16GB shared RAM risks the crash this # entire architecture is built to avoid. Batch requests still queue # through the same limit — a 10-photo batch is 10 sequential dispatches # on the caller's side (see engine.process_batch), not 10 parallel ones. demo = build_interface() demo.queue(max_size=20, default_concurrency_limit=2) if __name__ == "__main__": logger.info("Warming up segmentation model...") warm_up() logger.info("Warm-up complete. Launching Gradio.") demo.launch( server_name="0.0.0.0", server_port=7860, theme=gr.themes.Soft(primary_hue="blue"), )