import replicate


def image_generation(prompt: dict):
    output = replicate.run(
        "black-forest-labs/flux-schnell",
        input=prompt,
    )

    return output[0]