# -------------------------------------------------------------------------------- NOTEBOOK-CELL: CODE
# -*- coding: utf-8 -*-
import dataiku
import io
from PIL import Image

df = dataiku.Dataset("image_descriptions").get_dataframe()
folder = dataiku.Folder("pvFOY2UN")

MODEL_ID = dataiku.get_custom_variables()["image_generation_model_id"]
model = dataiku.api_client().get_default_project().get_llm(MODEL_ID)

def get_image(folder, image_path):
    """
    Get an image from a managed folder.
    """
    with folder.get_download_stream(image_path) as stream:
        return stream.read()

# -------------------------------------------------------------------------------- NOTEBOOK-CELL: CODE
for i in df.index:
    prompt = df.at[i, "caption"]
    
    generation = model.new_images_generation()
    generation.images_to_generate = 1
    generation.with_prompt(prompt)
    result = generation.execute()
    
    image = Image.open(io.BytesIO(result.images[0]))
    buf = io.BytesIO()
    image.save(buf, "JPEG")
    folder.upload_data(df.at[i, "filepath"], buf.getvalue())
