Update app.py

app.py

@@ -1,13 +1,22 @@
-import streamlit as st
-import torch
+from typing import Optional
+import uvicorn
+from fastapi import FastAPI, Form, File, UploadFile
+from transformers import AutoTokenizer, AutoModelForImageGeneration
+from PIL import Image
+import io
 
-st.title("Hugging Face Space GPU Info")
+app = FastAPI()
 
-# Check if GPU is available
-gpu_available = torch.cuda.is_available()
+tokenizer = AutoTokenizer.from_pretrained("EleutherAI/sd-1.5")
+model = AutoModelForImageGeneration.from_pretrained("EleutherAI/sd-1.5")
+
+@app.post("/generate_drawing/")
+async def generate_drawing(prompt: str = Form(...)):
+    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)
 
-if gpu_available:
-    gpu_name = torch.cuda.get_device_name(0)
-    st.write(f"GPU available: {gpu_name}")
-else:
-    st.write("No GPU available")