from openai import OpenAI


def load_openai_model(model_name):
    client = OpenAI()

    if model_name == "Dalle-3":
        response = client.images.generate(
            model="dall-e-3",
            prompt="a white siamese cat",
            size="1024x1024",
            quality="standard",
            n=1,
            )
    elif model_name == "Dalle-2":
        response = client.images.generate(
            model="dall-e-2",
            prompt="a white siamese cat",
            size="512x512",
            quality="standard",
            n=1,
            )
    else:
        raise NotImplementedError

    image_url = response.data[0].url

    return image_url


if __name__ == "__main__":
    image_url = load_openai_model('Dalle-3')
    print(image_url)