Update app.py

app.py

@@ -3,7 +3,6 @@ import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from fastapi import FastAPI
 from pydantic import BaseModel
-import uvicorn
 
 # Load model and tokenizer
 model_name = "roneneldan/TinyStories-33M"
@@ -29,12 +28,14 @@ def generate_story(prompt):
     
     return generated_text
 
-# Set up FastAPI
+# Set up FastAPI app
 app = FastAPI()
 
+# FastAPI request model
 class StoryRequest(BaseModel):
     prompt: str
 
+# FastAPI route to generate story
 @app.post("/generate")
 async def generate(request: StoryRequest):
     generated_text = generate_story(request.prompt)
@@ -49,9 +50,10 @@ iface = gr.Interface(
     description="Generate short stories using the TinyStories-33M model."
 )
 
-# Mount Gradio interface on FastAPI
+# Mount Gradio app to FastAPI
 app = gr.mount_gradio_app(app, iface, path="/")
 
-# Ensure Uvicorn runs the FastAPI app
+# For local testing
 if __name__ == "__main__":
+    import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=7860)