| import subprocess |
| from PIL import Image |
|
|
| def download_file(url, output_filename): |
| command = ['wget', '-O', output_filename, '-q', url] |
| subprocess.run(command, check=True) |
|
|
| url1 = 'https://storage.googleapis.com/mediapipe-models/image_segmenter/selfie_multiclass_256x256/float32/latest/selfie_multiclass_256x256.tflite' |
| url2 = 'https://storage.googleapis.com/mediapipe-models/image_segmenter/selfie_segmenter/float16/latest/selfie_segmenter.tflite' |
|
|
| filename1 = 'selfie_multiclass_256x256.tflite' |
| filename2 = 'selfie_segmenter.tflite' |
|
|
| download_file(url1, filename1) |
| download_file(url2, filename2) |
|
|
| import cv2 |
| import mediapipe as mp |
| import numpy as np |
| from mediapipe.tasks import python |
| from mediapipe.tasks.python import vision |
| import random |
| import gradio as gr |
| import spaces |
| import torch |
| from diffusers import FluxInpaintPipeline |
| from diffusers import FlowMatchEulerDiscreteScheduler, AutoencoderKL |
| from diffusers.models.transformers.transformer_flux import FluxTransformer2DModel |
| from transformers import CLIPTextModel, CLIPTokenizer,T5EncoderModel, T5TokenizerFast |
|
|
| DEVICE = "cuda" if torch.cuda.is_available() else "cpu" |
|
|
| bfl_repo="black-forest-labs/FLUX.1-dev" |
|
|
| BG_COLOR = (255, 255, 255) |
| MASK_COLOR = (0, 0 , 0) |
|
|
| def maskPerson(input): |
| base_options = python.BaseOptions(model_asset_path='selfie_multiclass_256x256.tflite') |
| options = vision.ImageSegmenterOptions(base_options=base_options, |
| output_category_mask=True) |
|
|
| with vision.ImageSegmenter.create_from_options(options) as segmenter: |
| image = mp.Image.create_from_file(input) |
| segmentation_result = segmenter.segment(image) |
| person_mask = segmentation_result.confidence_masks[0] |
|
|
| image_data = image.numpy_view() |
| fg_image = np.zeros(image_data.shape, dtype=np.uint8) |
| fg_image[:] = MASK_COLOR |
| bg_image = np.zeros(image_data.shape, dtype=np.uint8) |
| bg_image[:] = BG_COLOR |
|
|
| condition = np.stack((person_mask.numpy_view(),) * 3, axis=-1) > 0.2 |
| output_image = np.where(condition, fg_image, bg_image) |
|
|
| return output_image |
|
|
| def maskHead(input): |
| base_options = python.BaseOptions(model_asset_path='selfie_multiclass_256x256.tflite') |
| options = vision.ImageSegmenterOptions(base_options=base_options, |
| output_category_mask=True) |
|
|
| with vision.ImageSegmenter.create_from_options(options) as segmenter: |
| image = mp.Image.create_from_file(input) |
|
|
| segmentation_result = segmenter.segment(image) |
|
|
| hairmask = segmentation_result.confidence_masks[1] |
| facemask = segmentation_result.confidence_masks[3] |
|
|
| image_data = image.numpy_view() |
| fg_image = np.zeros(image_data.shape, dtype=np.uint8) |
| fg_image[:] = MASK_COLOR |
| bg_image = np.zeros(image_data.shape, dtype=np.uint8) |
| bg_image[:] = BG_COLOR |
|
|
| combined_mask = np.maximum(hairmask.numpy_view(), facemask.numpy_view()) |
|
|
| condition = np.stack((combined_mask,) * 3, axis=-1) > 0.2 |
| output_image = np.where(condition, fg_image, bg_image) |
|
|
| return output_image |
|
|
| def random_positioning(input, output_size=(1024, 1024)): |
| background = cv2.imread("default.jpeg") |
| if background is None: |
| raise ValueError("Unable to load background image") |
|
|
| background = cv2.resize(background, output_size, interpolation=cv2.INTER_AREA) |
| |
| if input is None: |
| raise ValueError("Unable to load input image") |
|
|
| scale_factor = random.uniform(0.5, 1.0) |
| new_size = (int(input.shape[1] * scale_factor), int(input.shape[0] * scale_factor)) |
| |
| resized_image = cv2.resize(input, new_size, interpolation=cv2.INTER_AREA) |
|
|
| if background.shape[2] != resized_image.shape[2]: |
| raise ValueError("Input image and background image must have the same number of channels") |
|
|
| x_offset = random.randint(0, output_size[0] - new_size[0]) |
| y_offset = random.randint(0, output_size[1] - new_size[1]) |
| background[y_offset:y_offset+new_size[1], x_offset:x_offset+new_size[0]] = resized_image |
| |
| return background |
|
|
|
|
| def remove_background(image_path, mask): |
| image = cv2.imread(image_path) |
| inverted_mask = cv2.bitwise_not(mask) |
| |
| _, binary_mask = cv2.threshold(inverted_mask, 127, 255, cv2.THRESH_BINARY) |
| |
| result = np.zeros_like(image, dtype=np.uint8) |
| |
| result[binary_mask == 255] = image[binary_mask == 255] |
| |
| return result |
| |
| pipe = FluxInpaintPipeline.from_pretrained(bfl_repo, torch_dtype=torch.bfloat16).to(DEVICE) |
| MAX_SEED = np.iinfo(np.int32).max |
| TRIGGER = "a photo of TOK" |
|
|
|
|
| @spaces.GPU(duration=75) |
| def execute(image, prompt, debug=False): |
| if not prompt : |
| gr.Info("Please enter a text prompt.") |
| return None |
|
|
| if not image : |
| gr.Info("Please upload a image.") |
| return None |
|
|
| img = cv2.imread(image) |
| img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) |
|
|
| imgs = [ random_positioning(img)] |
|
|
| pipe.load_lora_weights("XLabs-AI/flux-RealismLora", weight_name='lora.safetensors') |
| response = [] |
|
|
| seed_slicer = random.randint(0, MAX_SEED) |
| generator = torch.Generator().manual_seed(seed_slicer) |
|
|
| for image in range(len(imgs)): |
| current_img = imgs[image] |
| cv2.imwrite('base_image.jpg', current_img) |
| cv2.imwrite("mask_person.jpg", maskPerson('base_image.jpg')) |
| cv2.imwrite("mask_face.jpg", maskHead('base_image.jpg')) |
| |
| |
| im = Image.open('base_image.jpg') |
| np_arr = np.array(im) |
| rgb_image = cv2.cvtColor(np_arr, cv2.COLOR_BGR2RGB) |
|
|
| im = Image.fromarray(rgb_image) |
| |
| person = np.array(Image.open('mask_person.jpg')) |
| face = np.array(Image.open('mask_face.jpg')) |
|
|
| person_gray = cv2.cvtColor(person, cv2.COLOR_BGR2GRAY) |
| face_gray = cv2.cvtColor(face, cv2.COLOR_BGR2GRAY) |
|
|
| _, mask = cv2.threshold(face_gray, 1, 255, cv2.THRESH_BINARY_INV) |
| mask_inv = cv2.bitwise_not(mask) |
| person_masked = cv2.bitwise_and(person_gray, person_gray, mask=mask_inv) |
| face_masked = cv2.bitwise_and(face_gray, face_gray, mask=mask) |
| result = cv2.add(person_masked, face_masked) |
| cv2.imwrite('join.jpg', result) |
|
|
| fund_mask = Image.open('join.jpg') |
| |
| result0 = pipe( |
| prompt=f"{prompt} {TRIGGER}", |
| image=im, |
| mask_image=fund_mask, |
| width=1024, |
| height=1024, |
| strength=0.85, |
| generator=generator, |
| num_inference_steps=28, |
| max_sequence_length=256, |
| joint_attention_kwargs={"scale": 0.9}, |
| ).images[0] |
|
|
| arr = np.array(result0) |
| rgb_image = cv2.cvtColor(arr, cv2.COLOR_BGR2RGB) |
| cv2.imwrite('person.jpg', rgb_image) |
| cv2.imwrite("mask.jpg", maskHead('person.jpg')) |
| mask = Image.open('mask.jpg') |
|
|
| result = pipe( |
| prompt=f"{prompt} {TRIGGER}", |
| image=result0, |
| mask_image=mask, |
| width=1024, |
| height=1024, |
| strength=0.85, |
| generator=generator, |
| num_inference_steps=28, |
| max_sequence_length=256, |
| joint_attention_kwargs={"scale": 0.9}, |
| ).images[0] |
|
|
| if debug: |
| response.append(im) |
| response.append(person) |
| response.append(face) |
| response.append(fund_mask) |
| response.append(result0) |
| response.append(mask) |
|
|
| response.append(result) |
| return response |
|
|
| |
| |
| description = "This is an unofficial implementation of the ip face adapter for FLUX DEV and does not explicitly follow the ip face model, I created a wrapper with inpaint and mediapipe, I like to call Fake IP Adapter" |
| title = "Flux IP Face Adapter" |
| iface = gr.Interface( |
| fn=execute, |
| description=description, |
| title=title, |
| inputs=[ |
| gr.Image(type="filepath"), |
| gr.Textbox(label="Prompt"), |
| gr.Checkbox(label="Debug Mode") |
| ], |
| outputs="gallery" |
| ) |
|
|
| iface.launch(share=True, debug=True) |
