Update main.py

main.py

@@ -27,9 +27,15 @@ class InputData(BaseModel):
     top_p: float = 0.95
     repetition_penalty: float = 1.0
 
+class InputImage(BaseModel):
+    input: str
+    negativePrompt: str = '' 
+    steps: int = 25
+    cfg: int = 5
+    seed: int = 453666937
+
 def format_prompt(message, history):
     prompt = "<s>"
-    
     #with open('Manuale.txt', 'r') as file:
     #    manual_content = file.read()
     #    prompt += f"Leggi questo manuale dopo ti farò delle domande: {manual_content}"    
@@ -52,31 +58,26 @@ def read_root(request: Request, input_data: InputData):
 
     history = []  # Puoi definire la history se necessario
     generated_response = generate(input_text, history, temperature, max_new_tokens, top_p, repetition_penalty)
-    return {"response": generated_response}
+    return generated_response
 
-@app.get("/Immagine")
-def generate_image():
+@app.post("/Immagine")
+def generate_image(request: Request, input_data: InputImage):
     client = Client("https://openskyml-fast-sdxl-stable-diffusion-xl.hf.space/--replicas/545b5tw7n/")
     result = client.predict(
-        "a giant monster with yellow face hybrid of dragon and spider, in dark dense foggy forest ",
-        "",
-        25,
-        5,
+        input_data.input,
+        input_data.negativePrompt,
+        input_data.steps, 
+        input_data.cfg,
         1024,
         1024,
-        453666937,
+        input_data.seed,
         fn_index=0
     )
     image_url = result
     with open(image_url, 'rb') as img_file:
         img_binary = img_file.read()
         img_base64 = base64.b64encode(img_binary).decode('utf-8')
-    return {"response": img_base64}
-
-@app.get("/Test")
-def generate_image():
-    result = "a giant monster hybrid of dragon and spider, in dark dense foggy forest "
-    return {"response": result}
+    return img_base64
 
 @app.get("/")
 def read_general():