73 lines
2.1 KiB
Python
73 lines
2.1 KiB
Python
from celery import shared_task
|
|
from .models import Track
|
|
from django.core.files import File
|
|
from django.core.files.uploadedfile import InMemoryUploadedFile
|
|
from django.core.files.base import ContentFile
|
|
|
|
import os
|
|
import dotenv
|
|
|
|
import json
|
|
from requests import get, post
|
|
from requests.exceptions import ConnectionError
|
|
from PIL import Image
|
|
from io import BytesIO
|
|
|
|
dotenv.load_dotenv()
|
|
|
|
@shared_task
|
|
def get_banter(id):
|
|
"""
|
|
Requests an AI model for some banter.
|
|
Currently expects an ollama instance to be running,
|
|
but might be possible to use OpenAI API in future.
|
|
"""
|
|
track = Track.objects.get(pk=id)
|
|
track.banter = ""
|
|
|
|
try:
|
|
get(os.getenv('AI_ENDPOINT'))
|
|
except ConnectionError:
|
|
track.banter = "Error"
|
|
track.banter_done = True
|
|
track.save()
|
|
return f"{track}: AI endpoint not available."
|
|
|
|
try:
|
|
r = post(
|
|
os.getenv('AI_ENDPOINT'),
|
|
json={
|
|
'model': os.getenv('AI_MODEL'),
|
|
'prompt': str(track)
|
|
},
|
|
stream=True
|
|
)
|
|
r.raise_for_status()
|
|
|
|
for line in r.iter_lines():
|
|
body = json.loads(line)
|
|
if 'error' in body:
|
|
raise Exception(body['error'])
|
|
track.banter += body.get('response', '')
|
|
|
|
if body.get('done', False):
|
|
track.banter_done = True
|
|
track.save()
|
|
return f"{track}: {body['context']}"
|
|
except Exception as e:
|
|
track.banter_done = True
|
|
track.save()
|
|
return f"{track}: {e}"
|
|
|
|
@shared_task
|
|
def get_and_dither_image(url, instance, attribute):
|
|
"""Expects a url, and a Django model cls with an ImageField attribute"""
|
|
print(f"Parsing {url} for {instance} with attribute {attribute}")
|
|
r = get(url)
|
|
i = Image.open(BytesIO(r.content))
|
|
i = i.resize(size=(256, 256)).convert("P").quantize(colors=32)
|
|
buffer = BytesIO()
|
|
i.save(fp=buffer, format='PNG')
|
|
file = InMemoryUploadedFile(ContentFile(buffer.getvalue()), None, instance.name + ".png", "image/png", buffer.tell, "utf-8")
|
|
instance.image = file
|
|
instance.save() |