import torch
from diffusers import FluxInpaintPipeline
from diffusers.utils import load_image

from PIL import Image
import sys
import numpy as np

import json
import os
import spaces

device = "cuda"
pipeline_device = 0 if torch.cuda.is_available() else -1 # TODO mix above
torch_dtype = torch.float16
debug = True
@spaces.GPU
def make_inpaint_condition(image, image_mask):
    image = np.array(image.convert("RGB")).astype(np.float32) / 255.0
    image_mask = np.array(image_mask.convert("L")).astype(np.float32) / 255.0

    if image.shape[0:1] != image_mask.shape[0:1]:
        print("error image and image_mask must have the same image size")
        return None
    
    image[image_mask > 0.5] = -1.0  # set as masked pixel
    image = np.expand_dims(image, 0).transpose(0, 3, 1, 2)
    image = torch.from_numpy(image)
    return image





@spaces.GPU
def process_image(image,mask_image,prompt="a girl",negative_prompt="",model_id="black-forest-labs/FLUX.1-schnell",strength=0.75,seed=0,num_inference_steps=4):
    
    

    #control_image=make_inpaint_condition(image,mask_image)
    #image.save("_control.jpg")
    if image == None:
        return None
    
    pipe = FluxInpaintPipeline.from_pretrained(model_id, torch_dtype=torch.bfloat16)

    #batch_size =1
    generators = []
    generator = torch.Generator(device).manual_seed(seed)
    generators.append(generator)
    
    output = pipe(prompt=prompt, image=image, mask_image=mask_image,generator=generator)

    return output.images[0]


if __name__ == "__main__":
    image = Image.open(sys.argv[1])
    mask  = Image.open(sys.argv[2])
    output = process_image(image,mask)
    output.save(sys.argv[3])