CaesarAIART test

CaesarAIART/caesaraiart.py

@@ -1,6 +1,9 @@
 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"
@@ -11,8 +14,28 @@ class CaesarAIART:
     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))
+        filtered_image = io.BytesIO()
+        image.save(filtered_image, "PNG")
+        filtered_image.seek(0)
+        return filtered_image
 
 if __name__ == "__main__":
     prompt = "a photo of an astronaut riding a horse on mars"
-    caesaraiart = CaesarAIART()
-    caesaraiart.generate()
+    image = CaesarAIART.generate_api(prompt)
+
+
+    #caesaraiart = CaesarAIART()
+    #caesaraiart.generate()

main.py

@@ -15,7 +15,7 @@ import speech_recognition as sr
 import uvicorn
 from fastapi import FastAPI, File, UploadFile,Depends, WebSocket, WebSocketDisconnect
 from fastapi.middleware.cors import CORSMiddleware
-from fastapi.responses import FileResponse
+from fastapi.responses import FileResponse,StreamingResponse
 from tqdm import tqdm
 from transformers import pipeline
 
@@ -30,6 +30,7 @@ from RequestModels import *
 from CaesarFaceRecognition.caesardeepface import CaesarDeepFace
 from CaesarAIART.caesaraiart import CaesarAIART
 #from CaesarAIMusicLoad.caesaraimusicload import CaesarAITelegramBOT
+
 importcsv  = ImportCSV("CaesarAI")
 caesaryolo = CaesarYolo()
 caesarfacedetectmodel = CaesarFaceDetection()
@@ -38,7 +39,7 @@ caesarfacedetectmodel = CaesarFaceDetection()
 app = FastAPI()
 
 CURRENT_DIR = os.path.realpath(__file__).replace(f"/main.py","")
-caesaraiart = CaesarAIART(CURRENT_DIR)
+#caesaraiart = CaesarAIART(CURRENT_DIR)
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],  # can alter with time
@@ -134,8 +135,8 @@ def caesarfacesnap(frames: CaesarOCRHTTPModel):
 def caesarart(promptjson: CaesarAIARTModel):
     try:
         promptjson = dict(promptjson)
-        caesaraiart.generate(promptjson["prompt"])
-        return FileResponse(f"{CURRENT_DIR}/CaesarAIART/caesarart.png")
+        image = CaesarAIART.generate_api(promptjson["prompt"])
+        return StreamingResponse(image, media_type="image/png")#FileResponse(f"{CURRENT_DIR}/CaesarAIART/caesarart.png")
 
     except Exception as ex:
         return {"error":f"{type(ex)},{ex}"}