Update main.py

main.py

@@ -56,28 +56,26 @@ async def searches(request: Request):
     response = agent_chain.run(input=prompt)
     return response
 
-@app.post("/imagen")
-async def generate_image(request: Request):
-    data = await request.json()
-    prompt = data['prompt']
-    model = data.get["model"]
-    sampler = data.get["sampler"]
-    seed = int(data.get["seed"])
-    neg = data.get["neg"]
-
-    response = generate_image_prodia(prompt, model, sampler, seed, neg)
-    return jsonify({"image": response})
-
 class User(BaseModel):
     prompt: str
     model: str
     sampler: str
     seed: int
     neg: str = None
-    
-tests = {}
+
+@app.post("/imagen")
+def generate_image(request: User):
+    prompt = request.prompt
+    model = request.model
+    sampler = request.sampler
+    seed = request.seed
+    neg = request.neg
+
+    response = generate_image_prodia(prompt, model, sampler, seed, neg)
+    return {"image": response}
+
 
 @app.post("/test")
 def test(request: User):
-        return {'data': f'Prompt is {request.prompt}'}
+        return {'data': f'Prompt is {request.prompt} Model is {request.model}'}