Update app.py

app.py

@@ -1,34 +1,18 @@
-import torch
-from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
-from PIL import Image
 import gradio as gr
+from transformers import pipeline
 
-# Load text-to-image model
-text_to_image_model = torch.hub.load("ProPerNounpYK/texttoimage", "generator", map_location=torch.device("cuda" if torch.cuda.is_available() else "cpu"))
+# Text-to-Image model
+text_to_image = pipeline("text-to-image", model="ProPerNounpYK/texttoimage")
 
-# Load chat model
-chat_model = AutoModelForSeq2SeqLM.from_pretrained("ProPerNounpYK/chat")
-chat_tokenizer = AutoTokenizer.from_pretrained("ProPerNounpYK/chat")
+# Chat model
+chat = pipeline("conversational", model="ProPerNounpYK/chat")
 
-# Create multimodal interface
 interface = gr.Interface(
-    fn=lambda input_text, input_image: generate_response(input_text, input_image),
-    inputs=["text", "image"],
+    fn=lambda input: chat(input, text_to_image(input)),
+    inputs="text",
     outputs=["text", "image"],
-    title="Multimodal Conversational AI",
-    description="Talk to me, and I'll respond with images!"
+    title="Text-to-Image Chat",
+    description="Type something and get a response with an image!"
 )
 
-def generate_response(input_text, input_image):
-    # Process input text using chat model
-    chat_output = chat_model(input_text)
-    chat_response = chat_tokenizer.decode(chat_output, skip_special_tokens=True)
-
-    # Process input image using text-to-image model
-    generated_image = text_to_image_model(input_text, input_image)
-
-    # Return response as a tuple of text and image
-    return chat_response, generated_image
-
-# Launch the interface
 interface.launch()