from groq import Groq
import os
import requests
import os
import IPython.display as display
from PIL import Image
from io import BytesIO
import streamlit as st

client = Groq()

def enhance_prompt(user_prompt):
    response = client.chat.completions.create(
        model="qwen-2.5-32b",
        messages=[{"role": "system", "content": "Enhance the given logo description to make it more detailed and creative, don't make it larger than 5 lines, don't add any extra details such as name if not provided in the orignal prompt and make the prompt humane in simple words. Try to make the text in the image more detailed."},
                {"role": "user", "content": user_prompt}],
        temperature=0.6,
        max_completion_tokens=4096,
        top_p=0.95,
        stream=False,
        stop=None,
    )
    return response.choices[0].message.content

HF_TOKEN = os.environ.get("HF_TOKEN")

API_URL = "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-xl-base-1.0"

HEADERS = {"Authorization": f"Bearer {HF_TOKEN}"}

def generate(prompt):
    """Generate and display an image using Hugging Face's SDXL API."""
    payload = {"inputs": prompt}
    response = requests.post(API_URL, headers=HEADERS, json=payload)

    if response.status_code == 200:
        # try:
        #     image = Image.open(BytesIO(response.content))
        #     display.display(image)
        # except Exception:
        #     print("⚠️ The response is not a valid image. Here is the raw response:", response.text)
        img = response.content
        return img
    else:
        print("❌ Error:", response.json())
st.title("AI-Powered-Design Generator")
prompt = st.text_input("",placeholder="Enter your prompt")
if st.button("Run"):
    
    user_prompt = enhance_prompt(prompt)
    st.markdown("### Enhanced Prompt")
    st.write(user_prompt)
    st.markdown("### Generated Image")
    generated_image = generate(user_prompt)
    try:
        st.image(generated_image)
    except Exception:
        print("⚠️ The response is not a valid image. Here is the raw response:", generated_image)