Update app.py

app.py

@@ -1,17 +1,34 @@
 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
 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_response(prompt):
-    inputs = model.prepare_inputs_for_generation(prompt, max_length=1024)
-    outputs = model.generate(**inputs)
-    return outputs[0]
+    return model.generate(prompt, max_length=1024)[0]
+
+# 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 SquanchNastyAI model
-interface = gr.Interface(fn=generate_response, inputs="text", outputs="text")
+interface = gr.Interface(
+    fn=generate_response,
+    inputs="text",
+    outputs=["text", "image"],
+    components={
+        "text": gr.Text(),
+        "image": gr.Image(),
+    },
+    layout="row",
+)
 
 # Launch the Gradio interface
 interface.launch()