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

def generate_text(prompt, maxi, cohere_api_key):
    co = cohere.Client(cohere_api_key)

    response = co.generate(
        model='xlarge',  # Specify the model you are using
        prompt=prompt,
        temperature=0.7,
        max_tokens=maxi,
        stop_sequences=["--END--"]
    )
    return response.generations[0].text

def generate_image(prompt, gradio_api_key):
    r = requests.post('https://clipdrop-api.co/text-to-image/v1',
                      files={
                          'prompt': (None, prompt, 'text/plain')
                      },
                      headers={'x-api-key': gradio_api_key}
    )

    if r.ok:
        images = Image.open(io.BytesIO(r.content))
        return images
    else:
        raise ValueError("Failed to generate image")

def text_and_image_generator(prompt, maxi, cohere_api_key, gradio_api_key):
    text = generate_text(f"Generate a story with {prompt}. Conclude the generated story properly.", maxi, cohere_api_key)
    image = generate_image(text, gradio_api_key)
    title = generate_text(f"title for the story {prompt} within 5 words", 5, cohere_api_key)

    return title, text, image

app = gr.Interface(
    fn=text_and_image_generator,
    inputs=[
        gr.Textbox(label="Enter your prompt to generate a story", lines=2, placeholder="Type your prompt here..."),
        gr.Slider(320, 1000, step=1, label="Story length"),
        gr.Textbox(label="Cohere API Key", type="password"),
        gr.Textbox(label="Gradio API Key", type="password")
    ],
    outputs=[
        gr.Textbox(label="Story title"),
        gr.Textbox(label="Story"),
        gr.Image(type="pil", label="Image based on the Generated story")
    ],
    title="Story and Image Generator",
    theme="dark"
)

app.launch()