src/api/image_regeneration_api.py

"""
project @ NTO-TCP-HF
created @ 2024-10-29
author  @ github.com/ishworrsubedii
"""
import base64
import time
from io import BytesIO
from typing import Optional

import replicate
import requests
from PIL import Image
from fastapi import APIRouter, UploadFile, File, Form
from fastapi.responses import JSONResponse
from src.utils.logger import logger

image_regeneration_router = APIRouter()


def image_regeneration_replicate(input):
    output = replicate.run(
        "konieshadow/fooocus-api:fda927242b1db6affa1ece4f54c37f19b964666bf23b0d06ae2439067cd344a4",
        input=input
    )
    return output


@image_regeneration_router.post("/image_redesign")
async def image_re_gen(
        prompt: str = Form(...),
        negative_prompt: str = Form(""),
        image: UploadFile = File(...),
        mask_image: Optional[UploadFile] = File(default=None),
        reference_image_c1: Optional[UploadFile] = File(default=None),
        reference_image_c1_type: Optional[str] = Form(default=""),
        reference_image_c1_weight: Optional[float] = Form(default=0.0),
        reference_image_c1_stop: Optional[float] = Form(default=0.0),
        reference_image_c2: Optional[UploadFile] = File(default=None),
        reference_image_c2_type: Optional[str] = Form(default=""),
        reference_image_c2_weight: Optional[float] = Form(default=0.0),
        reference_image_c2_stop: Optional[float] = Form(default=0.0),
        reference_image_c3: Optional[UploadFile] = File(default=None),
        reference_image_c3_type: Optional[str] = Form(default=""),
        reference_image_c3_weight: Optional[float] = Form(default=0.0),
        reference_image_c3_stop: Optional[float] = Form(default=0.0),
        reference_image_c4: Optional[UploadFile] = File(default=None),
        reference_image_c4_type: Optional[str] = Form(default=""),
        reference_image_c4_weight: Optional[float] = Form(default=0.0),
        reference_image_c4_stop: Optional[float] = Form(default=0.0),
):
    logger.info("-" * 50)
    logger.info(">>> IMAGE REDESIGN STARTED <<<")
    start_time = time.time()

    try:
        async def process_reference_image(reference_image: Optional[UploadFile]) -> Optional[str]:
            if reference_image is not None:
                reference_image_bytes = await reference_image.read()
                reference_image = Image.open(BytesIO(reference_image_bytes)).convert("RGB")
                ref_img_base64 = BytesIO()
                reference_image.save(ref_img_base64, format="WEBP")
                reference_image_b64 = base64.b64encode(ref_img_base64.getvalue()).decode("utf-8")
                return f"data:image/WEBP;base64,{reference_image_b64}"
            return None
        logger.info(">>> REFERENCE IMAGE PROCESSING FUNCTION INITIALIZED <<<")
    except Exception as e:
        logger.error(f">>> REFERENCE IMAGE PROCESSING ERROR: {str(e)} <<<")
        return JSONResponse(status_code=500, 
                          content={"error": f"Error processing reference image: {str(e)}", "code": 500})

    try:
        image_bytes = await image.read()
        image = Image.open(BytesIO(image_bytes)).convert("RGB")
        img_base64 = BytesIO()
        image.save(img_base64, format="WEBP")
        image_data_uri = f"data:image/WEBP;base64,{base64.b64encode(img_base64.getvalue()).decode('utf-8')}"
        logger.info(">>> MAIN IMAGE PROCESSED SUCCESSFULLY <<<")
    except Exception as e:
        logger.error(f">>> MAIN IMAGE PROCESSING ERROR: {str(e)} <<<")
        return JSONResponse(status_code=500,
                          content={"error": f"Error processing main image: {str(e)}", "code": 500})

    try:
        reference_images = {
            'c1': await process_reference_image(reference_image_c1),
            'c2': await process_reference_image(reference_image_c2),
            'c3': await process_reference_image(reference_image_c3),
            'c4': await process_reference_image(reference_image_c4)
        }
        logger.info(">>> REFERENCE IMAGES PROCESSED SUCCESSFULLY <<<")
    except Exception as e:
        logger.error(f">>> REFERENCE IMAGES PROCESSING ERROR: {str(e)} <<<")
        return JSONResponse(status_code=500,
                          content={"error": f"Error processing reference images: {str(e)}", "code": 500})

    try:
        input_data = {
            "prompt": prompt,
            "inpaint_input_image": image_data_uri,
            "sharpness": 2,
            "guidance_scale": 4,
            "refiner_switch": 0.5,
            "performance_selection": "Quality",
            "aspect_ratios_selection": "1024*1024"
        }

        if negative_prompt:
            input_data["negative_prompt"] = negative_prompt

        if mask_image is not None:
            mask_image_bytes = await mask_image.read()
            mask_image = Image.open(BytesIO(mask_image_bytes)).convert("RGB")
            mask_base64 = BytesIO()
            mask_image.save(mask_base64, format="WEBP")
            mask_image_data_uri = f"data:image/WEBP;base64,{base64.b64encode(mask_base64.getvalue()).decode('utf-8')}"
            input_data["inpaint_input_mask"] = mask_image_data_uri
        logger.info(">>> INPUT DATA PREPARED SUCCESSFULLY <<<")
    except Exception as e:
        logger.error(f">>> INPUT DATA PREPARATION ERROR: {str(e)} <<<")
        return JSONResponse(status_code=500,
                          content={"error": f"Error preparing input data: {str(e)}", "code": 500})

    try:
        for i in range(1, 5):
            c = f'c{i}'
            if reference_images[c] is not None:
                input_data[f"cn_img{i}"] = reference_images[c]

                type_value = locals()[f'reference_image_{c}_type']
                if type_value:
                    input_data[f"cn_type{i}"] = type_value

                weight_value = locals()[f'reference_image_{c}_weight']
                if weight_value != 0.0:
                    input_data[f"cn_weight{i}"] = weight_value

                stop_value = locals()[f'reference_image_{c}_stop']
                if stop_value != 0.0 or stop_value != 0:
                    input_data[f"cn_stop{i}"] = stop_value
        logger.info(">>> REFERENCE IMAGE PARAMETERS PROCESSED <<<")
    except Exception as e:
        logger.error(f">>> REFERENCE IMAGE PARAMETERS ERROR: {str(e)} <<<")
        return JSONResponse(status_code=500,
                          content={"error": f"Error processing reference image parameters: {str(e)}", "code": 500})

    try:
        output = image_regeneration_replicate(input_data)
        response = requests.get(output[0])
        output_base64 = base64.b64encode(response.content).decode('utf-8')
        base64_prefix = image_data_uri.split(",")[0] + ","
        logger.info(">>> IMAGE REGENERATION COMPLETED <<<")
    except Exception as e:
        logger.error(f">>> IMAGE REGENERATION ERROR: {str(e)} <<<")
        return JSONResponse(status_code=500,
                          content={"error": f"Error generating image: {str(e)}", "code": 500})

    try:
        inference_time = round(time.time() - start_time, 2)
        response = {
            "output": f"{base64_prefix}{output_base64}",
            "inference_time": inference_time,
            "code": 200,
        }
        logger.info(f">>> TOTAL INFERENCE TIME: {inference_time}s <<<")
        logger.info(">>> REQUEST COMPLETED SUCCESSFULLY <<<")
        logger.info("-" * 50)
        return JSONResponse(content=response, status_code=200)
    except Exception as e:
        logger.error(f">>> RESPONSE CREATION ERROR: {str(e)} <<<")
        return JSONResponse(status_code=500,
                          content={"error": f"Error creating response: {str(e)}", "code": 500})