# https://github.com/sayakpaul/diffusers-torchao
#6.22s
import os
from typing import Any, Dict
from PIL import Image
import torch
import torch.distributed as dist
from huggingface_inference_toolkit.logging import logger

dist.init_process_group()
torch.cuda.set_device(dist.get_rank())

from para_attn.context_parallel import init_context_parallel_mesh
from para_attn.context_parallel.diffusers_adapters import parallelize_pipe
from para_attn.parallel_vae.diffusers_adapters import parallelize_vae

class EndpointHandler:
    def __init__(self,path=""):



    def __call__(self, data: Dict[str, Any]) -> Image.Image:
        logger.info(f"Received incoming request with {data=}")

        if "inputs" in data and isinstance(data["inputs"], str):
            prompt = data.pop("inputs")
        elif "prompt" in data and isinstance(data["prompt"], str):
            prompt = data.pop("prompt")
        else:
            raise ValueError(
                "Provided input body must contain either the key `inputs` or `prompt` with the"
                " prompt to use for the image generation, and it needs to be a non-empty string."
            )

        return "1"