Working on it

routers/training.py

@@ -10,9 +10,13 @@ from io import BytesIO
 from pydantic import BaseModel
 import base64
 import uuid
+from transformers import AutoTokenizer
 
 from diffusers import StableDiffusionImg2ImgPipeline
 
+tokenizer = AutoTokenizer.from_pretrained("openlm-research/open_llama_7b")
+
+
 model_id_or_path = "runwayml/stable-diffusion-v1-5"
 pipe = StableDiffusionImg2ImgPipeline.from_pretrained(model_id_or_path, torch_dtype=torch.float16)
 pipe = pipe.to("cuda")
@@ -31,6 +35,10 @@ class ActionBody(BaseModel):
 @router.post("/perform-action")
 async def performAction(actionBody: ActionBody):
     
+    model_inputs = tokenizer(["A list of colors: red, blue"], return_tensors="pt").to("cuda")
+    generated_ids = model.generate(**model_inputs)
+    output = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
+
     response = requests.get(actionBody.url)
     init_image = Image.open(BytesIO(response.content)).convert("RGB")
     init_image = init_image.resize((actionBody.resizeW, actionBody.resizeH))
@@ -42,9 +50,12 @@ async def performAction(actionBody: ActionBody):
     imgUUID = str(uuid.uuid4())
     images[0].save(imgUUID+".png")
 
+    
+
     return {
         "imageName" : imgUUID+".png",
-        "image": "data:image/jpeg;base64,"+str(img_str)
+        "image": "data:image/jpeg;base64,"+img_str.decode(),
+        "output": output
     }
 
 
@@ -67,5 +78,5 @@ async def hifunction():
     # images[0].save("fantasy_landscape.png")
 
     return {
-        "image": "data:image/jpeg;base64,"+str(img_str)
+        "image": "data:image/jpeg;base64,"+img_str.decode()
     }