Update app.py

app.py

@@ -1,48 +1,38 @@
+import gradio as gr
 import torch
-from transformers import AutoModelForSequenceClassification, AutoTokenizer, AutoModelForImageGeneration
+from transformers import AutoModelForSequenceToSequence, AutoTokenizer
 from PIL import Image
-import gradio as gr
+from torchvision import transforms
 
-# Load the language model
-language_model = AutoModelForSequenceClassification.from_pretrained("mistralai/Mixtral-8x7B-Instruct-v0.1")
-language_tokenizer = AutoTokenizer.from_pretrained("mistralai/Mixtral-8x7B-Instruct-v0.1")
-
-# Load the image generation model
-image_model = AutoModelForImageGeneration.from_pretrained("artificialguybr/CuteCartoonRedmond-V2")
+# Load the Text-to-Image model
+image_model = AutoModelForSequenceToSequence.from_pretrained("artificialguybr/CuteCartoonRedmond-V2")
 image_tokenizer = AutoTokenizer.from_pretrained("artificialguybr/CuteCartoonRedmond-V2")
 
-# Define a function to generate an image based on a prompt
-def generate_image(prompt):
-    input_ids = image_tokenizer.encode(prompt, return_tensors="pt")
-    output = image_model.generate(input_ids)
+# Load the Text Generation model
+text_model = AutoModelForSequenceToSequence.from_pretrained("mistralai/Mixtral-8x7B-Instruct-v0.1")
+text_tokenizer = AutoTokenizer.from_pretrained("mistralai/Mixtral-8x7B-Instruct-v0.1")
+
+# Define a function to generate an image from text
+def generate_image(text):
+    inputs = image_tokenizer(text, return_tensors="pt")
+    output = image_model.generate(inputs["input_ids"], attention_mask=inputs["attention_mask"])
     image = Image.fromarray(output[0].detach().numpy())
     return image
 
-# Define a function to have a conversation
-def have_conversation(input_text):
-    input_ids = language_tokenizer.encode(input_text, return_tensors="pt")
-    output = language_model.generate(input_ids)
-    response = language_tokenizer.decode(output[0], skip_special_tokens=True)
-    return response
+# Define a function to generate text from text
+def generate_text(text):
+    inputs = text_tokenizer(text, return_tensors="pt")
+    output = text_model.generate(inputs["input_ids"], attention_mask=inputs["attention_mask"])
+    return text_tokenizer.decode(output[0], skip_special_tokens=True)
 
 # Create a Gradio interface
-iface = gr.Interface(
-    fn=have_conversation,
-    inputs="text",
-    outputs="text",
-    title="Converse with AI",
-    description="Talk to the AI and see its response!"
-)
-
-# Add an image generation feature to the interface
-iface.add_component(
-    gr.Image(type="pil"),
+demo = gr.Interface(
+    fn=lambda text: {"image": generate_image(text), "text": generate_text(text)},
     inputs="text",
-    outputs="image",
-    fn=generate_image,
-    title="Generate Image",
-    description="Enter a prompt and see the generated image!"
+    outputs=["image", "text"],
+    title="Text-to-Image and Text Generation",
+    description="Enter a prompt to generate both an image and text!"
 )
 
-# Launch the interface
-iface.launch()
+# Launch the Gradio app
+demo.launch()