Spaces:
Runtime error
Runtime error
File size: 1,727 Bytes
2de3774 |
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 |
import modules.async_worker as worker
import os
import cv2
import imageio
import numpy as np
import rembg
import torch
import PIL.Image
from PIL import Image
from typing import Any
class pipeline:
def remove_background(
self,
image: PIL.Image.Image,
rembg_session: Any = None,
force: bool = False,
**rembg_kwargs,
) -> PIL.Image.Image:
do_remove = True
if image.mode == "RGBA" and image.getextrema()[3][0] < 255:
do_remove = False
do_remove = do_remove or force
if do_remove:
image = rembg.remove(image, session=rembg_session, **rembg_kwargs)
return image
pipeline_type = ["rembg"]
model_hash = ""
# Optional function
def parse_gen_data(self, gen_data):
gen_data["original_image_number"] = gen_data["image_number"]
gen_data["image_number"] = 1
gen_data["show_preview"] = False
return gen_data
def load_base_model(self, name):
return
def load_keywords(self, lora):
return " "
def load_loras(self, loras):
return
def refresh_controlnet(self, name=None):
return
def clean_prompt_cond_caches(self):
return
def process(
self,
gen_data=None,
callback=None,
):
worker.add_result(
gen_data["task_id"],
"preview",
(-1, f"Removing background ...", None)
)
rembg_session = rembg.new_session()
image = self.remove_background(gen_data["input_image"], rembg_session)
# Return finished image to preview
if callback is not None:
callback(1, 0, 0, 1, image)
return [image]
|