High-level OSInsert interface. Modes ----- - ``aggressive``: ObjectStitch + SAM + InsertAnything - ``conservative``: background + bbox -> mask -> InsertAnything Notes ----- - InsertAnything is initialized when the model instance is created. - ObjectStitch and SAM ar
| 42 | |
| 43 | |
| 44 | class OSInsertModel: |
| 45 | """High-level OSInsert interface. |
| 46 | |
| 47 | Modes |
| 48 | ----- |
| 49 | - ``aggressive``: ObjectStitch + SAM + InsertAnything |
| 50 | - ``conservative``: background + bbox -> mask -> InsertAnything |
| 51 | |
| 52 | Notes |
| 53 | ----- |
| 54 | - InsertAnything is initialized when the model instance is created. |
| 55 | - ObjectStitch and SAM are lazily initialized on the first aggressive call, |
| 56 | and then cached/reused for subsequent calls. |
| 57 | """ |
| 58 | |
| 59 | def __init__( |
| 60 | self, |
| 61 | model_dir: str | Path, |
| 62 | device: str = "cuda:0", |
| 63 | *, |
| 64 | eager_aggressive_init: bool = False, |
| 65 | objectstitch_ckpt_path: str | Path | None = None, |
| 66 | objectstitch_config_path: str | Path | None = None, |
| 67 | objectstitch_clip_dir: str | Path | None = None, |
| 68 | sam_checkpoint: str | Path | None = None, |
| 69 | flux_fill_path: str | Path | None = None, |
| 70 | flux_redux_path: str | Path | None = None, |
| 71 | ia_lora_path: str | Path | None = None, |
| 72 | ) -> None: |
| 73 | self.config = OSInsertConfig( |
| 74 | model_dir=Path(model_dir), |
| 75 | device=device, |
| 76 | objectstitch_ckpt_path=Path(objectstitch_ckpt_path) if objectstitch_ckpt_path is not None else None, |
| 77 | objectstitch_config_path=Path(objectstitch_config_path) if objectstitch_config_path is not None else None, |
| 78 | objectstitch_clip_dir=Path(objectstitch_clip_dir) if objectstitch_clip_dir is not None else None, |
| 79 | sam_checkpoint=Path(sam_checkpoint) if sam_checkpoint is not None else None, |
| 80 | flux_fill_path=Path(flux_fill_path) if flux_fill_path is not None else None, |
| 81 | flux_redux_path=Path(flux_redux_path) if flux_redux_path is not None else None, |
| 82 | ia_lora_path=Path(ia_lora_path) if ia_lora_path is not None else None, |
| 83 | ) |
| 84 | |
| 85 | self._ia_net = InsertAnythingModel( |
| 86 | model_dir=self.config.model_dir, |
| 87 | flux_fill_path=self.config.flux_fill_path, |
| 88 | flux_redux_path=self.config.flux_redux_path, |
| 89 | ia_lora_path=self.config.ia_lora_path, |
| 90 | device=self.config.device, |
| 91 | ) |
| 92 | |
| 93 | self._os_model = None |
| 94 | self._os_sampler = None |
| 95 | self._sam_predictor = None |
| 96 | |
| 97 | if eager_aggressive_init: |
| 98 | self._get_objectstitch_model_and_sampler() |
| 99 | self._get_sam_predictor() |
| 100 | |
| 101 | def _get_objectstitch_model_and_sampler(self): |