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 = (0, 0, 0) # black
MASK_COLOR = (255, 255, 255) # white

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)):
    if input is None:
        raise ValueError("Impossible to load 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)
    
    background = np.zeros((output_size[1], output_size[0], 3), dtype=np.uint8)
    
    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
    background = np.clip(background, 0, 255) 
    background = background.astype(np.uint8)

    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=200)
def execute(image, prompt):
  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.jpg", cv2.bitwise_not(maskHead('base_image.jpg')))
    
    im = Image.open('base_image.jpg')
    mask = Image.open('mask.jpg')

    result = pipe(
            prompt=f"{prompt} {TRIGGER}",
            image=im,
            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]
    response.append(result)


  return response

iface = gr.Interface(
    fn=execute,
    inputs=[
        gr.Image(type="filepath"),
        gr.Textbox(label="Prompt")
    ],
    outputs="gallery"
)

iface.launch(share=True, debug=True)