CaesarAIART test

CaesarAIART/caesaraiart.py

@@ -0,0 +1,18 @@
+import torch
+from diffusers import StableDiffusionPipeline
+
+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")
+
+if __name__ == "__main__":
+    prompt = "a photo of an astronaut riding a horse on mars"
+    caesaraiart = CaesarAIART()
+    caesaraiart.generate()

RequestModels.py

@@ -33,7 +33,8 @@ class CaesarStockInfoModel(BaseModel):
 class CaesarTelegramMessageModel(BaseModel):
     chat_id: str
     message: str
-
+class CaesarAIARTModel(BaseModel):
+    prompt: str
 class CaesarVoiceModel(BaseModel):
     filename: Union[str,None]
     language: Union[str,None]

main.py

@@ -28,14 +28,17 @@ from CaesarVoice import CaesarVoice
 from csv_to_db import ImportCSV
 from RequestModels import *
 from CaesarFaceRecognition.caesardeepface import CaesarDeepFace
-from CaesarAIMusicLoad.caesaraimusicload import CaesarAITelegramBOT
+from CaesarAIART.caesaraiart import CaesarAIART
+#from CaesarAIMusicLoad.caesaraimusicload import CaesarAITelegramBOT
 importcsv  = ImportCSV("CaesarAI")
 caesaryolo = CaesarYolo()
 caesarfacedetectmodel = CaesarFaceDetection()
-caesartelgegrambot = CaesarAITelegramBOT()
+#caesartelgegrambot = CaesarAITelegramBOT()
+
 app = FastAPI()
 
 CURRENT_DIR = os.path.realpath(__file__).replace(f"/main.py","")
+caesaraiart = CaesarAIART(CURRENT_DIR)
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],  # can alter with time
@@ -92,6 +95,7 @@ async def caesarfacedetect(websocket: WebSocket):
 
     except WebSocketDisconnect:
         print("Client disconnected")
+"""
 @app.get("/getcaesartelegramchatids")
 def getcaesartelegramchatid():
     try:
@@ -109,7 +113,7 @@ def caesartelegramsendmessage(messageinfo: CaesarTelegramMessageModel):
         return chatids
     except Exception as ex:
         return {"error":f"{type(ex)},{ex}"}
-
+"""
 @app.post("/caesarfacesnap")
 def caesarfacesnap(frames: CaesarOCRHTTPModel):
     try:
@@ -126,6 +130,16 @@ def caesarfacesnap(frames: CaesarOCRHTTPModel):
     except Exception as ex:
         return {"error":f"{type(ex)},{ex}"}
 
+@app.post("/caesaraiart")
+def caesarart(promptjson: CaesarAIARTModel):
+    try:
+        promptjson = dict(promptjson)
+        caesaraiart.generate(promptjson["prompt"])
+        return FileResponse(f"{CURRENT_DIR}/CaesarAIART/caesarart.png")
+
+    except Exception as ex:
+        return {"error":f"{type(ex)},{ex}"}
+
 @app.post("/faceauthorizeuser")
 def faceauthorizeuser(recogperson: CaesarFaceAuth = Depends(),file: UploadFile = File(...)):
 

requirements.txt

@@ -103,6 +103,9 @@ tensorboard-plugin-wit==1.8.1
 tensorflow==2.11.0
 tensorflow-estimator==2.11.0
 tensorflow-intel
+diffusers
+scipy 
+accelerate
 tensorflow-io-gcs-filesystem==0.31.0
 termcolor==2.2.0
 thinc==8.1.7