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 = image.cpu().numpy()  # torch.Tensor を NumPy 配列に変換
    image = (image * 255).astype(np.uint8)  # 0~1 の範囲を 0~255 に戻し、uint8 型にキャスト
    image = image.transpose(0, 2, 3, 1)  # (1, 3, H, W) から (1, H, W, 3) に形状を変更

    return image[0]





@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):
    
    
    if image == None:
        return None
    
    #control_image=make_inpaint_condition(image,mask_image)
    #mask_image.save("_mask_image.jpg")
    #image = control_image
    
    pipe = FluxInpaintPipeline.from_pretrained(model_id, torch_dtype=torch.bfloat16)
    pipe.to(device)
    #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,strength=strength)

    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])