import torch
from diffusers import FluxInpaintPipeline

from PIL import Image
import sys
import spaces

# I only test with FLUX.1-schnell

#@spaces.GPU
def process_image(image,mask_image,prompt="a person",model_id="black-forest-labs/FLUX.1-schnell",strength=0.75,seed=0,num_inference_steps=4):
    if image == None:
        return None
    
    pipe = FluxInpaintPipeline.from_pretrained(model_id, torch_dtype=torch.bfloat16)
    pipe.to("cuda")

    generators = []
    generator = torch.Generator("cuda").manual_seed(seed)
    generators.append(generator)
    # more parameter see https://huggingface.co/docs/diffusers/api/pipelines/flux#diffusers.FluxInpaintPipeline
    output = pipe(prompt=prompt, image=image, mask_image=mask_image,generator=generator,strength=strength)

    return output.images[0]

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