def generate_image(prompt: str):
        safe, err = is_prompt_safe(prompt)
        if not safe:
            return None, err

        try:
            image = img_client.text_to_image(prompt=prompt, model=IMG_MODEL)
            return image, None
        except Exception as e:
            return None, f"Error during image generation: {e}"

    with st.form("img_form"):
        p = st.text_area("Image description:", height=120)
        ok = st.form_submit_button("Generate Image")

    if ok:
        if not p.strip():
            st.warning("⚠️ Enter a description.")
        else:
            with st.spinner("Generating image..."):
                im, err = generate_image(p.strip())

            if err:
                st.error(err)
            else:
                st.image(im, use_container_width=True)
                st.session_state.generated_image = im

    im = st.session_state.get("generated_image")
    if im:
        buf = BytesIO()
        im.save(buf, format="PNG")
        st.download_button(
            "📥 Download",
            buf.getvalue(),
            "ai_generated_image.png",
            "image/png"
        )