from fastapi import FastAPI, UploadFile, Form
from fastapi.responses import StreamingResponse
#import torch
from PIL import Image
#from diffusers import StableDiffusionDepth2ImgPipeline
import numpy as np
from io import BytesIO

app = FastAPI()

"""
pipe = StableDiffusionDepth2ImgPipeline.from_pretrained(
   "stabilityai/stable-diffusion-2-depth",
   torch_dtype=torch.float16,
).to("cuda")
"""

def pad_image(input_image):
    pad_w, pad_h = np.max(((2, 2), np.ceil(
        np.array(input_image.size) / 64).astype(int)), axis=0) * 64 - input_image.size
    im_padded = Image.fromarray(
        np.pad(np.array(input_image), ((0, pad_h), (0, pad_w), (0, 0)), mode='edge'))
    w, h = im_padded.size
    if w == h:
        return im_padded
    elif w > h:
        new_image = Image.new(im_padded.mode, (w, w), (0, 0, 0))
        new_image.paste(im_padded, (0, (w - h) // 2))
        return new_image
    else:
        new_image = Image.new(im_padded.mode, (h, h), (0, 0, 0))
        new_image.paste(im_padded, ((h - w) // 2, 0))
        return new_image

def predict(input_image, prompt, steps, scale, seed, strength, depth_image=None):
    depth = None
    if depth_image is not None:
        depth_image = pad_image(depth_image)
        depth_image = depth_image.resize((512, 512))
        depth = np.array(depth_image.convert("L"))
        depth = depth.astype(np.float32) / 255.0
        depth = depth[None, None]
        depth = torch.from_numpy(depth)
    init_image = input_image.convert("RGB")
    image = pad_image(init_image)  # resize to integer multiple of 32
    image = image.resize((512, 512))
    result = pipe(prompt=prompt, image=image, strength=strength)

    return result['images']

def grayscale(image,
              prompt,
              steps,
              scale,
              seed,
              strength):
    image = image.convert('L') #convert to grayscale

    return image

@app.post("/convert_ifc_img/")
async def convert_ifc_img(file: UploadFile, 
                          prompt: str = Form(default=""), 
                          steps: int = Form(default=50), 
                          scale: float = Form(default=9), 
                          seed: int = Form(default=178106186), 
                          strength: float = Form(default=0.9)
                        ):
    
    image = Image.open(file.file)                        
    image_result = grayscale(image, prompt, steps, scale, seed, strength)
    buffer = BytesIO()
    image_result.save(buffer, format="PNG")
    buffer.seek(0)
    return StreamingResponse(buffer, media_type="image/png")