Raiff1982's picture
Create app/app.py
0462aae verified
# This Gradio app allows users to interact with a chatbot that can generate text and images based on user prompts.
import gradio as gr
import numpy as np
from transformers_js import pipeline # Corrected import to use transformers_js instead of transformers_js_py
# Define the available models
AVAILABLE_MODELS = {
"GPT-2": "gpt2",
"DALL-E": "dalle-mini/dalle-mini-1.3B"
}
# Initialize the text generation pipeline
text_generator = pipeline("text-generation", model=AVAILABLE_MODELS["GPT-2"])
# Initialize the image generation pipeline
image_generator = pipeline("image-generation", model=AVAILABLE_MODELS["DALL-E"])
# Function to generate text
def generate_text(prompt, model):
np.random.seed(42) # Set a seed for reproducibility
if model == "GPT-2":
return text_generator(prompt, max_length=50, num_return_sequences=1)[0]['generated_text']
else:
return "Model not supported for text generation"
# Function to generate images
def generate_image(prompt, model):
if model == "DALL-E":
image = image_generator(prompt, num_inference_steps=50, guidance_scale=7.5).images[0]
return image
else:
return "Model not supported for image generation"
# Create the Gradio interface
with gr.Blocks() as demo:
gr.Markdown("# Chatbot with Text and Image Generation")
with gr.Tab("Text Generation"):
text_prompt = gr.Textbox(label="Enter your text prompt")
text_model = gr.Radio(choices=list(AVAILABLE_MODELS.keys()), label="Choose a model", value="GPT-2")
text_output = gr.Textbox(label="Generated Text")
text_button = gr.Button("Generate Text")
text_button.click(generate_text, inputs=[text_prompt, text_model], outputs=text_output)
with gr.Tab("Image Generation"):
image_prompt = gr.Textbox(label="Enter your image prompt")
image_model = gr.Radio(choices=list(AVAILABLE_MODELS.keys()), label="Choose a model", value="DALL-E")
image_output = gr.Image(label="Generated Image")
image_button = gr.Button("Generate Image")
image_button.click(generate_image, inputs=[image_prompt, image_model], outputs=image_output)
# Launch the interface
demo.launch(show_error=True)