update'

main.py

@@ -1,34 +1,56 @@
 from fastapi import FastAPI, Query
-from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor
+from transformers import (
+    AutoProcessor,
+    AutoModelForCausalLM,
+    AutoTokenizer,
+)
 from qwen_vl_utils import process_vision_info
 import torch
+import logging
+
+logging.basicConfig(level=logging.INFO)
 
 app = FastAPI()
 
-checkpoint = "Qwen/Qwen2.5-VL-3B-Instruct"
-min_pixels = 256*28*28
-max_pixels = 1280*28*28
+# Qwen2.5-VL Model Setup
+qwen_checkpoint = "Qwen/Qwen2.5-VL-3B-Instruct"
+min_pixels = 256 * 28 * 28
+max_pixels = 1280 * 28 * 28
+
 processor = AutoProcessor.from_pretrained(
-    checkpoint,
+    qwen_checkpoint,
     min_pixels=min_pixels,
-    max_pixels=max_pixels
+    max_pixels=max_pixels,
 )
-model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
-    checkpoint,
+
+qwen_model = AutoModelForCausalLM.from_pretrained(
+    qwen_checkpoint,
     torch_dtype=torch.bfloat16,
     device_map="auto",
-    # attn_implementation="flash_attention_2",
+)
+
+# LLaMA Model Setup
+llama_model_name = "path/to/llama-uncensored-model"
+llama_tokenizer = AutoTokenizer.from_pretrained(llama_model_name)
+llama_model = AutoModelForCausalLM.from_pretrained(
+    llama_model_name, torch_dtype=torch.float16, device_map="auto"
 )
 
 @app.get("/")
 def read_root():
-    return {"message": "API is live. Use the /predict endpoint."}
+    return {"message": "API is live. Use the /predict, /chat, or /llama_chat endpoints."}
 
 @app.get("/predict")
 def predict(image_url: str = Query(...), prompt: str = Query(...)):
     messages = [
         {"role": "system", "content": "You are a helpful assistant with vision abilities."},
-        {"role": "user", "content": [{"type": "image", "image": image_url}, {"type": "text", "text": prompt}]},
+        {
+            "role": "user",
+            "content": [
+                {"type": "image", "image": image_url},
+                {"type": "text", "text": prompt},
+            ],
+        },
     ]
     text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
     image_inputs, video_inputs = process_vision_info(messages)
@@ -38,31 +60,43 @@ def predict(image_url: str = Query(...), prompt: str = Query(...)):
         videos=video_inputs,
         padding=True,
         return_tensors="pt",
-    ).to(model.device)
+    ).to(qwen_model.device)
     with torch.no_grad():
-        generated_ids = model.generate(**inputs, max_new_tokens=128)
-    generated_ids_trimmed = [out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)]
+        generated_ids = qwen_model.generate(**inputs, max_new_tokens=128)
+    generated_ids_trimmed = [
+        out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs["input_ids"], generated_ids)
+    ]
     output_texts = processor.batch_decode(
         generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
     )
     return {"response": output_texts[0]}
 
 @app.get("/chat")
-def chat( prompt: str = Query(...)):
+def chat(prompt: str = Query(...)):
     messages = [
-        {"role": "system", "content": "You are a helpful assistant with vision abilities."},
-        {"role": "user", "content": [ {"type": "text", "text": prompt}]},
+        {"role": "system", "content": "You are a helpful assistant."},
+        {"role": "user", "content": [{"type": "text", "text": prompt}]},
     ]
     text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
     inputs = processor(
-        text=[text], 
+        text=[text],
         padding=True,
         return_tensors="pt",
-    ).to(model.device)
+    ).to(qwen_model.device)
     with torch.no_grad():
-        generated_ids = model.generate(**inputs, max_new_tokens=128)
-    generated_ids_trimmed = [out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)]
+        generated_ids = qwen_model.generate(**inputs, max_new_tokens=128)
+    generated_ids_trimmed = [
+        out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs["input_ids"], generated_ids)
+    ]
     output_texts = processor.batch_decode(
         generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
     )
     return {"response": output_texts[0]}
+
+@app.get("/llama_chat")
+def llama_chat(prompt: str = Query(...)):
+    inputs = llama_tokenizer(prompt, return_tensors="pt").to(llama_model.device)
+    with torch.no_grad():
+        outputs = llama_model.generate(**inputs, max_new_tokens=128)
+    response = llama_tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return {"response": response}

main.py.old

@@ -0,0 +1,78 @@
+from fastapi import FastAPI, Query
+from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor
+from qwen_vl_utils import process_vision_info
+import torch
+import logging
+
+logging.basicConfig(level=logging.INFO)
+
+try:
+    # Code that may raise an exception
+    x = 1 / 0
+except ZeroDivisionError as e:
+    logging.error("Error occurred: %s", e)
+    # Take alternative action to recover from the exception
+
+app = FastAPI()
+
+checkpoint = "Qwen/Qwen2.5-VL-3B-Instruct"
+min_pixels = 256*28*28
+max_pixels = 1280*28*28
+processor = AutoProcessor.from_pretrained(
+    checkpoint,
+    min_pixels=min_pixels,
+    max_pixels=max_pixels
+)
+model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
+    checkpoint,
+    torch_dtype=torch.bfloat16,
+    device_map="auto",
+    # attn_implementation="flash_attention_2",
+)
+
+@app.get("/")
+def read_root():
+    return {"message": "API is live. Use the /predict endpoint."}
+
+@app.get("/predict")
+def predict(image_url: str = Query(...), prompt: str = Query(...)):
+    messages = [
+        {"role": "system", "content": "You are a helpful assistant with vision abilities."},
+        {"role": "user", "content": [{"type": "image", "image": image_url}, {"type": "text", "text": prompt}]},
+    ]
+    text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    image_inputs, video_inputs = process_vision_info(messages)
+    inputs = processor(
+        text=[text],
+        images=image_inputs,
+        videos=video_inputs,
+        padding=True,
+        return_tensors="pt",
+    ).to(model.device)
+    with torch.no_grad():
+        generated_ids = model.generate(**inputs, max_new_tokens=128)
+    generated_ids_trimmed = [out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)]
+    output_texts = processor.batch_decode(
+        generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
+    )
+    return {"response": output_texts[0]}
+
+@app.get("/chat")
+def chat( prompt: str = Query(...)):
+    messages = [
+        {"role": "system", "content": "You are a helpful assistant with vision abilities."},
+        {"role": "user", "content": [ {"type": "text", "text": prompt}]},
+    ]
+    text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    inputs = processor(
+        text=[text], 
+        padding=True,
+        return_tensors="pt",
+    ).to(model.device)
+    with torch.no_grad():
+        generated_ids = model.generate(**inputs, max_new_tokens=128)
+    generated_ids_trimmed = [out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)]
+    output_texts = processor.batch_decode(
+        generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
+    )
+    return {"response": output_texts[0]}