import torch
from diffusers import FluxImg2ImgPipeline

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):
    print("start process image process_image")
    if image == None:
        print("empty input image returned")
        return None
    
    pipe = FluxImg2ImgPipeline.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
    print(prompt)
    output = pipe(prompt=prompt, image=image,generator=generator,strength=strength
                  ,guidance_scale=0,num_inference_steps=num_inference_steps,max_sequence_length=256)

    # TODO support mask
    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])