Switch background removal to RMBG-2.0

[AmitMY]

Summary

• Replace briaai/RMBG-1.4 (pipeline) with briaai/RMBG-2.0 (AutoModelForImageSegmentation + manual transforms / sigmoid / mask resize, since RMBG-2.0 is not pipeline-compatible).
• Add kornia as a runtime dependency (required by the RMBG-2.0 trust-remote-code modeling file).
• Regenerate assets/examples/{amit,flux,stock}/masked.png and pose.pose with the new model.

Test plan

• Run python -m human_avatar.image_to_avatar end-to-end and visually inspect the regenerated masks.
• Confirm pip install -e . resolves cleanly with the new kornia dep.
• Spot-check the README image previews still render.

🤖 Generated with Claude Code

---

Note

Medium Risk
Moderate risk because it changes the core background-removal pipeline and introduces a new runtime dependency plus trust_remote_code model loading, which can affect output quality and runtime/environment compatibility.

Overview
Switches background removal from a Hugging Face pipeline (RMBG-1.4) to direct AutoModelForImageSegmentation inference with cached model loading, explicit preprocessing (Resize/Normalize), sigmoid mask generation, and alpha compositing.

Updates runtime dependencies by adding kornia to support the RMBG-2.0 trust_remote_code model implementation.

^{Reviewed by Cursor Bugbot for commit 29a5cca. Bugbot is set up for automated code reviews on this repo. Configure here.}

Summary by CodeRabbit

• Improvements
  • Upgraded background removal functionality with an enhanced segmentation model that delivers more accurate detection and higher quality image processing results compared to previous version
  • Background removal now outputs images with fully transparent backgrounds in RGBA format, enabling improved compositing and blending capabilities for professional-grade image editing, design workflows, and creative applications

[coderabbitai]

Caution

Review failed

Pull request was closed or merged during review

📝 Walkthrough

Walkthrough

The background removal feature is upgraded from Hugging Face pipelines to direct RMBG-2.0 model inference. A new cached model loader with device management is introduced, and the image processing pipeline shifts to manual torchvision preprocessing, torch-based inference, and PIL-based postprocessing to return RGBA images. Dependencies are updated to include kornia.

Changes

Background Removal Implementation Upgrade

Layer / File(s)	Summary
Dependencies pyproject.toml	Added kornia as a required dependency for RMBG-2.0 model support.
Model Loading human_avatar/image_to_avatar.py	New cached load_rmbg_model() function loads briaai/RMBG-2.0 from Hugging Face, configures float32 matmul precision to "high", selects CUDA device when available, switches to eval mode, and returns the model and device.
Image Processing human_avatar/image_to_avatar.py	remove_image_background() reimplemented to preprocess RGB images with torchvision (resize to 1024×1024, tensor conversion, ImageNet normalization), run segmentation model inference under torch.no_grad(), apply sigmoid to logits, convert the mask to PIL format, and return an RGBA image with the mask as the alpha channel.
Imports human_avatar/image_to_avatar.py	Updated to include torch, torchvision.transforms, and AutoModelForImageSegmentation for direct model handling instead of Hugging Face pipelines.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~22 minutes

Poem

🐰 A rabbit hops through pixel dreams,
Background whisked away by RMBG streams,
Direct inference now—no pipeline detour,
With CUDA and masks we refine the core.
RGBA magic makes images pure! ✨

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title directly and clearly summarizes the main change: switching the background removal implementation from a previous version to RMBG-2.0.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch rmbg-2.0

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Review rate limit: 0/1 reviews remaining, refill in 60 minutes.}

_{Comment @coderabbitai help to get the list of available commands and usage tips.}