import gradio as gr
from transformers import AutoModelForCausalLM
from PIL import Image
from huggingface_hub import pipeline


# Load the SquanchNastyAI model from Hugging Face Spaces
text_model = AutoModelForCausalLM.from_pretrained("or4cl3ai/SquanchNastyAI")
# Initialize the pipeline for image generation
image_pipeline = pipeline("image-generation", model="google/vit-base-patch16-384")


# Define a function to generate a text response to a prompt
def generate_text(prompt):
    response = text_model.generate(prompt, max_length=1024)[0]
    return response


# Define a function to generate an image from a prompt
def generate_image(prompt):
    image = image_pipeline(prompt)
    return image


# Create a Gradio interface for the AI model
def ai_interface(prompt):
    text_response = generate_text(prompt)
    image_response = generate_image(prompt)
    return text_response, image_response


inputs = gr.inputs.Textbox(label="Enter a prompt")
outputs = [
    gr.outputs.Textbox(label="Text Response"),
    gr.outputs.Image(label="Image Response")
]

interface = gr.Interface(fn=ai_interface, inputs=inputs, outputs=outputs)

interface.launch()