Update main.py

main.py

@@ -9,6 +9,7 @@ import requests
 import os
 import socket
 
+#--------------------------------------------------- Definizione Server FAST API ------------------------------------------------------
 app = FastAPI()
 client = InferenceClient("mistralai/Mixtral-8x7B-Instruct-v0.1")
 
@@ -34,20 +35,7 @@ class InputImage(BaseModel):
     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}"    
-    
-    for user_prompt, bot_response in history:
-        prompt += f"[INST] {user_prompt} [/INST]"
-        prompt += f" {bot_response}</s> "
-    now = datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f")
-    prompt += f"[{now}] [INST] {message} [/INST]"
-
-    return prompt
-
+#--------------------------------------------------- Generazione TESTO ------------------------------------------------------
 @app.post("/Genera")
 def read_root(request: Request, input_data: InputData):
     input_text = input_data.input
@@ -55,40 +43,15 @@ def read_root(request: Request, input_data: InputData):
     max_new_tokens = input_data.max_new_tokens
     top_p = input_data.top_p
     repetition_penalty = input_data.repetition_penalty
-
-    history = []  # Puoi definire la history se necessario
+    history = [] 
     generated_response = generate(input_text, history, temperature, max_new_tokens, top_p, repetition_penalty)
     return {"response": generated_response}
 
-@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(
-        input_data.input,
-        input_data.negativePrompt,
-        input_data.steps, 
-        input_data.cfg,
-        1024,
-        1024,
-        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("/")
-def read_general():
-    return {"response": "Benvenuto. Per maggiori info vai a /docs"}  # Restituisci la risposta generata come JSON
-
 def generate(prompt, history, temperature=0.2, max_new_tokens=30000, top_p=0.95, repetition_penalty=1.0):
     temperature = float(temperature)
     if temperature < 1e-2:
         temperature = 1e-2
     top_p = float(top_p)
-
     generate_kwargs = dict(
         temperature=temperature,
         max_new_tokens=max_new_tokens,
@@ -100,10 +63,50 @@ def generate(prompt, history, temperature=0.2, max_new_tokens=30000, top_p=0.95,
     formatted_prompt = format_prompt(prompt, history)
     output = client.text_generation(formatted_prompt, **generate_kwargs, stream=False, details=False)
     return output
+    
+def format_prompt(message, history):
+    prompt = "<s>"
+    for user_prompt, bot_response in history:
+        prompt += f"[INST] {user_prompt} [/INST]"
+        prompt += f" {bot_response}</s> "
+    now = datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f")
+    prompt += f"[{now}] [INST] {message} [/INST]"
+    return prompt
+    
+#--------------------------------------------------- Generazione IMMAGINE ------------------------------------------------------
+@app.post("/Immagine")
+def generate_image(request: Request, input_data: InputImage):
+    client = Client("https://openskyml-fast-sdxl-stable-diffusion-xl.hf.space/--replicas/545b5tw7n/")
+    max_attempts = 10
+    attempt = 0
+    while attempt < max_attempts:
+        try:
+            result = client.predict(
+                input_data.input,
+                input_data.negativePrompt,
+                input_data.steps, 
+                input_data.cfg,
+                1024,
+                1024,
+                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}
+        except requests.exceptions.HTTPError as e:
+            if e.response.status_code == 500:
+                attempt += 1
+                if attempt < max_attempts:
+                    continue
+                else:
+                    return {"error": "Errore interno del server persistente"}
+            else:
+                return {"error": "Errore diverso da 500"}
+    return {"error": "Numero massimo di tentativi raggiunto"}
 
-    #stream = client.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=False, return_full_text=False)
-    # Accumula l'output in una lista
-    #output_list = []
-    #for response in stream:
-    #    output_list.append(response.token.text)
-    #return iter(output_list)  # Restituisci la lista come un iteratore
+@app.get("/")
+def read_general():
+    return {"response": "Benvenuto. Per maggiori info: https://matteoscript-fastapi.hf.space/docs"}