# -*- coding: utf-8 -*-
"""StorySprout.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1rPolQFgot4IyiE6Egh8ULu-9wx86Kb4J
"""

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

HF_TOKEN = os.environ.get("HF_TOKEN")
def llm_model_load_pipe(model_id, quant=False):

  bnb_config = transformers.BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type='nf4',
    bnb_4bit_use_double_quant=True,
    bnb_4bit_compute_dtype=torch.bfloat16
    )

  model = transformers.AutoModelForCausalLM.from_pretrained(
      model_id,
      torch_dtype=torch.bfloat16,
      quantization_config=bnb_config,
      device_map='auto',
  )

  tokenizer = transformers.AutoTokenizer.from_pretrained(
      model_id,
      torch_dtype=torch.bfloat16,
      quantization_config=bnb_config,
      device_map='auto',
  )

  generate_text = transformers.pipeline(
      model=model, tokenizer=tokenizer,
      return_full_text=False,  # langchain expects the full text
      task='text-generation'
  )

  return generate_text, model, tokenizer

t2t_pipe, t2t_model, t2t_token = llm_model_load_pipe(model_id="HuggingFaceH4/zephyr-7b-beta", quant=True)

# Function to query the Hugging Face Zephyr model for story generation
def generate_story(user_prompt):
    # pipe = pipeline("text-generation", model="HuggingFaceH4/zephyr-7b-beta", torch_dtype=torch.bfloat16, device_map="auto")

    # We use the tokenizer's chat template to format each message - see https://huggingface.co/docs/transformers/main/en/chat_templating
    messages = [
        {
            "role": "system",
            "content": "You are a story generator for kids. Using the keywords in user content generate a short and educational story for children. Use aminals as charachters. Keywords: ",
        },
        {"role": "user", "content": user_prompt},
    ]
    prompt = t2t_pipe.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    outputs = t2t_pipe(prompt, max_new_tokens=256, do_sample=True, temperature=0.1, top_k=50, top_p=0.95)

    return outputs[0]["generated_text"]


# Function to query the Facebook MMS TTS API
def text_to_speech(text):
    API_URL = "https://api-inference.huggingface.co/models/facebook/mms-tts-eng"
    headers = {"Authorization": f"Bearer {HF_TOKEN}"}
    payload = {"inputs": text}
    response = requests.post(API_URL, headers=headers, json=payload)
    return response.content


# Function to query the Hugging Face Stable AI model for image generation
def get_images(story):
    API_URL = "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-2-1"
    headers = {"Authorization": f"Bearer {HF_TOKEN}"}
    payload = {"inputs": story}
    response = requests.post(API_URL, headers=headers, json=payload)
    image_bytes = response.content
    image = Image.open(io.BytesIO(image_bytes))
    return image

def run_pipe(user_prompt):
    story =  generate_story(user_prompt)
    audio = text_to_speech(story)
    image = get_images(story)
    yield story, audio, image

# Create the Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("<h1>StorySprout</h1>")
    system_prompt_input = gr.Textbox(lines=3, label="Enter a keyword for the story")
    with gr.Row():
        with gr.Column():
            stream_as_file_btn = gr.Button("Submit")
            image_output = gr.Image(label="Generated Images")
            audio_output = gr.Audio( autoplay=True,label="Text-to-Speech Audio")
            story_output = gr.Textbox(label="Generated Story")
            stream_as_file_btn.click(run_pipe, system_prompt_input, outputs=[story_output,audio_output,image_output])
demo.launch()