app.py

from fastapi import FastAPI, UploadFile, File
from transformers import pipeline
from fastai.vision.all import *
from PIL import Image
import os
import io
import json

access_token = os.getenv("access_token")

# NOTE - we configure docs_url to serve the interactive Docs at the root path
# of the app. This way, we can use the docs as a landing page for the app on Spaces.
app = FastAPI(docs_url="/")

pipe = pipeline("text2text-generation", model="google/flan-t5-small")
categories = ('Heart', 'Oblong', 'Oval', 'Round', 'Square')
learn = load_learner('model.pkl')

# Initialize the Code Llama Instruct pipeline (example with 7B model)
llama_model_id = "meta-llama/CodeLlama-7b-Instruct-hf"
llama_pipeline = pipeline(
    "text-generation",
    model=llama_model_id,
    model_kwargs={"torch_dtype": torch.bfloat16},
    device_map="auto",
    use_auth_token=access_token  # Use the access token for authentication
)

@app.get("/generate")
def generate(text: str):
    """
    Using the text2text-generation pipeline from `transformers`, generate text
    from the given input text. The model used is `google/flan-t5-small`, which
    can be found [here](https://huggingface.co/google/flan-t5-small).
    """
    output = pipe(text)
    return {"output": output[0]["generated_text"]}

@app.post("/face-analyse")
async def face_analyse(file: UploadFile = File(...)):
    # Read the uploaded file content
    request_object_content = await file.read()
    
    try:
        # Attempt to open the image
        img = Image.open(io.BytesIO(request_object_content))
    except Exception as e:
        return {"error": "Failed to open the image file. Make sure it is a valid image file."}

    # Check if img is None or not
    if img is None:
        return {"error": "Failed to open the image file."}
    
    try:
        # Resize the image to 300x300 pixels
        img = img.resize((300, 300))
    except Exception as e:
        return {"error": "Failed to resize the image."}

    try:
        # Assuming 'learn' is your image classifier model
        pred, idx, probs = learn.predict(img)
    except Exception as e:
        return {"error": "Failed to make predictions."}

    # Assuming categories is a list of category labels
    return dict(zip(categories, map(float, probs)))

@app.post("/extract-frame-details")
def extract_frame_details(text: str):
    """
    Using the Code Llama Instruct pipeline from `transformers`, extract frame
    details from the given input text. The model used is `meta-llama/CodeLlama-7b-Instruct-hf`.
    """
    messages = [
        {"role": "system", "content": "Please provide details about frames in JSON format."},
        {"role": "user", "content": text},
    ]

    terminators = [
        llama_pipeline.tokenizer.eos_token_id,
        llama_pipeline.tokenizer.convert_tokens_to_ids("")
    ]

    outputs = llama_pipeline(
        messages,
        max_new_tokens=256,
        eos_token_id=terminators,
        do_sample=True,
        temperature=0.6,
        top_p=0.9,
    )
    
    generated_text = outputs[0]["generated_text"]
    
    try:
        extracted_info = json.loads(generated_text)
    except json.JSONDecodeError:
        return {"error": "Failed to parse the generated text as JSON."}

    return extracted_info