Update app.py

app.py

@@ -1,11 +1,10 @@
-# Import necessary libraries
 import streamlit as st
-from tungsten import TungstenAI
+from tungstenkit import models
 
 # Load Tungsten AI model
-model = TungstenAI()
+model = models.get("text-to-image")
 
-# Create Streamlit app
+# Streamlit app title
 st.title("Text to Image Generation with Tungsten AI")
 
 # User input for text
@@ -14,7 +13,10 @@ text_input = st.text_area("Enter text for image generation:")
 # Generate Image button
 if st.button("Generate Image"):
     # Generate image based on input text
-    generated_image = model.generate(text_input)
+    result = model.predict({"prompt": text_input})
+    
+    # Get the image path from the result
+    image_path = result["image"]
 
     # Display the generated image
-    st.image(generated_image, caption="Generated Image", use_column_width=True)
+    st.image(image_path, caption="Generated Image", use_column_width=True)