import gradio as gr
import requests
from PIL import Image
import io
import os

# Define the API URL for the Craiyon model (lightweight text-to-image generator)
API_URL = "https://api-inference.huggingface.co/models/dalle-mini"

# Function to call the model and generate images
def generate_comic(prompt):
    api_token = os.environ.get("API_TOKEN")  # Securely access the token
    
    headers = {"Authorization": f"Bearer {api_token}"}
    
    # Sending request to the model with user prompt
    response = requests.post(API_URL, headers=headers, json={"inputs": prompt})
    
    # Check if the response is successful
    if response.status_code != 200:
        return f"Error: {response.status_code}, {response.text}"
    
    # Extracting the response data
    images = response.json().get("generated_images", [])

    if not images:
        return "No images were generated, please try again with a different prompt."
    
    # Process the images (assuming they are base64 encoded or URL)
    pil_images = []
    for img in images:
        image_data = base64.b64decode(img)
        image = Image.open(io.BytesIO(image_data))
        pil_images.append(image)
    
    return pil_images  # Return the list of PIL images

# Gradio interface setup
def gradio_interface():
    with gr.Blocks() as demo:
        gr.Markdown("## GenArt Narrative - Turn Your Story into Comic Panels!")
        prompt = gr.Textbox(label="Enter your short story description", placeholder="Once upon a time...")
        output_gallery = gr.Gallery(label="Generated Comic Panels", columns=3, height=300)
        submit_button = gr.Button("Generate Comic")
        submit_button.click(fn=generate_comic, inputs=prompt, outputs=output_gallery)
    return demo

# Run the Gradio app
if __name__ == "__main__":
    app = gradio_interface()
    app.launch()