Update app.py

app.py

@@ -1,21 +1,17 @@
-from typing import Optional
-import uvicorn
-from fastapi import FastAPI, Form, File, UploadFile
+import streamlit as st
+import torch
 from transformers import AutoTokenizer, GPTNeoForCausalLM
 from PIL import Image
 import io
 
-app = FastAPI()
-
 tokenizer = AutoTokenizer.from_pretrained("EleutherAI/sd-1.5")
 model = GPTNeoForCausalLM.from_pretrained("EleutherAI/sd-1.5")
 
-@app.post("/generate_drawing/")
-async def generate_drawing(prompt: str = Form(...)):
+st.title("Scribble Drawing Generator")
+
+prompt = st.text_input("Enter a prompt:")
+if prompt:
     inputs = tokenizer(prompt, return_tensors="pt")
     outputs = model.generate(inputs['input_ids'], max_length=256, do_sample=True)
     image = Image.open(io.BytesIO(outputs[0].cpu().numpy()))
-    return {"image": image}
-
-if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=8000)
+    st.image(image, caption="Generated Drawing")