Workin on it

routers/training.py

@@ -14,10 +14,10 @@ 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
-)
+# 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)
@@ -37,9 +37,9 @@ 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]
+    # 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")
@@ -55,7 +55,7 @@ async def performAction(actionBody: ActionBody):
     return {
         "imageName" : imgUUID+".png",
         "image": "data:image/jpeg;base64,"+img_str.decode(),
-        "output": output
+        # "output": output
     }