import torch
from diffusers import StableDiffusionPipeline
import requests
# You can access the image with PIL.Image for example
import io
from PIL import Image
class CaesarAIART:
    def __init__(self,CURRENT_DIR=""):
        self.model_id = "CompVis/stable-diffusion-v1-4"
        self.device = "cuda"
        self.pipe = StableDiffusionPipeline.from_pretrained(self.model_id, torch_dtype=torch.float16)
        self.CURRENT_DIR = CURRENT_DIR
        self.pipe = self.pipe.to(self.device)
    def generate(self,prompt):
        image = self.pipe(prompt).images[0]  
        image.save(f"{self.CURRENT_DIR}/CaesarAIART/caesarart.png")
    @staticmethod
    def generate_api(prompt):


        API_URL = "https://api-inference.huggingface.co/models/CompVis/stable-diffusion-v1-4"
        headers = {"Authorization": "Bearer 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"}

        def query(payload):
            response = requests.post(API_URL, headers=headers, json=payload)
            return response.content
        image_bytes = query({"inputs": prompt})

        image = Image.open(io.BytesIO(image_bytes))
          # create a thumbnail image
        imgio = io.BytesIO()
        image.save(imgio, 'JPEG')
        imgio.seek(0)
        return imgio
        #return StreamingResponse(content=imgio, media_type="image/jpeg"

if __name__ == "__main__":
    prompt = "a photo of an astronaut riding a horse on mars"
    image = CaesarAIART.generate_api(prompt)


    #caesaraiart = CaesarAIART()
    #caesaraiart.generate()