import gradio as gr
from huggingface_hub import InferenceClient
from gradio_client import Client
from PIL import Image
import requests
from io import BytesIO
import streamlit as st

# Initialize the HuggingFace Inference Client with the specified model
client_mistral = InferenceClient("mistralai/Mistral-7B-Instruct-v0.2")
# Initialize the Playground AI client
client_playground = Client("https://playgroundai-playground-v2-5.hf.space/--replicas/c9ozb/")

def format_prompt(logo_request):
    system_prompt = """
    You are an advanced language model designed to create detailed and creative image prompts for logo generation. Based on the user's input, generate an elaborate and descriptive image prompt that can be used to create a high-quality logo. Ensure that the prompt is clear, imaginative, and provides specific details that will guide the logo creation process effectively.
    """
    prompt = f"<s>[SYS] {system_prompt} [/SYS][INST] {logo_request} [/INST]</s>"
    return prompt

def generate_improved_prompt(logo_request, temperature=0.9, max_new_tokens=512, top_p=0.95, repetition_penalty=1.0):
    temperature = float(temperature)
    if temperature < 1e-2:
        temperature = 1e-2
    top_p = float(top_p)

    generate_kwargs = {
        "temperature": temperature,
        "max_new_tokens": max_new_tokens,
        "top_p": top_p,
        "repetition_penalty": repetition_penalty,
        "do_sample": True,
        "seed": 42,
    }

    formatted_prompt = format_prompt(logo_request)
    stream = client_mistral.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False)
    output = ""

    for response in stream:
        output += response.token.text
        yield output

def generate_image(prompt, negative_prompt="", use_negative_prompt=False, seed=0, width=1024, height=1024, guidance_scale=7.5, randomize_seed=True):
    result = client_playground.predict(
        prompt,
        negative_prompt,
        use_negative_prompt,
        seed,
        width,
        height,
        guidance_scale,
        randomize_seed,
        api_name="/run"
    )

    # Extract the image URL from the result
    image_path = result[0][0]["image"]
    image_url = "https://playgroundai-playground-v2-5.hf.space/--replicas/c9ozb/file=" + image_path
    
    # Fetch and display the result image
    response = requests.get(image_url)
    
    if response.headers['Content-Type'].startswith('image'):
        img = Image.open(BytesIO(response.content))
        return img
    else:
        return None

css = """
#mkd {
    height: 500px; 
    overflow: auto; 
    border: 1px solid #ccc; 
}
"""

with gr.Blocks(css=css) as gpt:
    with gr.Row():
        with gr.Column(scale=2):
            gr.HTML("<h1>Settings</h1>")
            logo_input = gr.Textbox(label="Input your logo request", placeholder="Describe the logo you want...")
        with gr.Column(scale=3):
            gr.HTML("<h1><center>Logo Prompt Generator<h1><center>")
            generate_button = gr.Button("Generate")
            output_area = gr.Textbox(label="AI Response", interactive=False, lines=10)
            generate_button.click(
                fn=generate_improved_prompt,
                inputs=[logo_input],
                outputs=output_area
            )

    gr.Markdown("""
    ---
    ### Meta Information
    **Project Title**: Magic AI Website Logo Creator
    
    **Github**: [https://github.com/pacnimo/gpt-prompt-generator](https://github.com/pacnimo/)
    
    **Description**: Logo Prompt Generator is Free and Easy to Use. Create a GPT Prompt Based on the Logo Request. 1 Click Prompt Generator.
    
    **Footer**: © 2024 by [pacnimo](https://github.com/pacnimo/). All rights reserved.
    """) # Meta, project description, and footer added here

gpt.launch(debug=True)

st.title("Image Generation using API")

# Streamlit app input fields
prompt = st.text_input("Enter prompt:", "")
negative_prompt = st.text_input("Enter negative prompt:", "")
use_negative_prompt = st.checkbox("Use negative prompt", False)
seed = st.number_input("Seed (0-2147483647):", value=0, min_value=0, max_value=2147483647)
width = st.slider("Width:", 1024, 1536)
height = st.slider("Height:", 1024, 1536)
guidance_scale = st.slider("Guidance Scale:", 3.0, 20.0, 0.1)
randomize_seed = st.checkbox("Randomize seed", True)

if st.button("Generate Image"):
    with st.spinner("Generating image..."):
        img = generate_image(prompt, negative_prompt, use_negative_prompt, seed, width, height, guidance_scale, randomize_seed)
        if img:
            st.image(img, caption="Generated Image")
        else:
            st.error("The content retrieved is not an image.")