Working on it

requirements-fastapi.txt

@@ -13,4 +13,5 @@ torch
 scipy
 ftfy
 accelerate
-uuid
+uuid
+bitsandbytes

routers/training.py

@@ -10,12 +10,14 @@ from io import BytesIO
 from pydantic import BaseModel
 import base64
 import uuid
-from transformers import AutoTokenizer
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
 from diffusers import StableDiffusionImg2ImgPipeline
 
 tokenizer = AutoTokenizer.from_pretrained("openlm-research/open_llama_7b")
-
+model = AutoModelForCausalLM.from_pretrained(
+    "openlm-research/open_llama_7b", device_map="auto", load_in_4bit=True
+)
 
 model_id_or_path = "runwayml/stable-diffusion-v1-5"
 pipe = StableDiffusionImg2ImgPipeline.from_pretrained(model_id_or_path, torch_dtype=torch.float16)
@@ -50,8 +52,6 @@ async def performAction(actionBody: ActionBody):
     imgUUID = str(uuid.uuid4())
     images[0].save(imgUUID+".png")
 
-    
-
     return {
         "imageName" : imgUUID+".png",
         "image": "data:image/jpeg;base64,"+img_str.decode(),