Distribución de poder de procesamiento

app.py

@@ -1,12 +1,8 @@
+from fastapi import FastAPI, Form
 from fastapi import FastAPI, File, UploadFile
-from fastapi.responses import StreamingResponse
-
-import io
+from fastapi.responses import StreamingResponse, FileResponse
 from io import BytesIO
-
-from fastapi import FastAPI, Form
-
-import funciones
+import funciones, globales
 
 app = FastAPI()
 
@@ -18,15 +14,14 @@ async def echo_image(image: UploadFile = File(...)):
     contents = await image.read()
     return StreamingResponse(BytesIO(contents), media_type=image.content_type)
 
-@app.post("/get-platillo/")
-async def get_platillo_image(prompt: str = Form(...)):
-    
-    imagen_pil = funciones.genera_platillo(prompt)
-
-    img_io = io.BytesIO()
-    imagen_pil.save(img_io, "PNG")
-    img_io.seek(0)
-    
-    #Ver cual es el mejor resultado para backend, si forzar la disposición de archivo o permitir que el navegador decida (puede que sea irrelevante para un consumo de la api directo)
-    #return StreamingResponse(content=img_io, media_type="image/png", headers={"Content-Disposition": "attachment; filename=platillo.png"})
-    return StreamingResponse(content=img_io, media_type="image/png")
+@app.post("/genera-imagen/")
+async def genera_imagen(platillo: str = Form(...)):
+
+    if globales.seconds_available > 25:
+        print("GPU...")
+        resultado = funciones.genera_platillo_gpu(platillo)
+        return FileResponse(resultado, media_type="image/png", filename="imagen.png")
+    else: 
+        print("Inference...")
+        resultado = funciones.genera_platillo_inference(platillo)
+        return StreamingResponse(content=resultado, media_type="image/png")

funciones.py

@@ -1,22 +1,58 @@
 import bridges
 from huggingface_hub import InferenceClient
-from PIL import Image
+import gradio_client
+import io
+import globales
 
-def genera_platillo(prompt): 
-    enlace = "black-forest-labs/FLUX.1-dev"
-    proveedor =  "hf-inference"
+
+previo = "Una fotografía de un plato blanco con "
+
+def genera_platillo_gpu(platillo):
+
+    client = gradio_client.Client(globales.espacio, hf_token=bridges.hug)
+
+    prompt = previo + platillo
+
+    print("Eso es el prompt final:", prompt)
+
+    kwargs = {
+            "prompt": prompt,
+            "api_name": "/infer"
+        }
+    
+
+    try:
+        result = client.predict(**kwargs
+        # prompt=prompt,
+        # negative_prompt="",
+        # seed=42,
+        # randomize_seed=True,
+        # width=1024,
+        # height=1024,
+        # guidance_scale=3.5,
+        # num_inference_steps=28,
+        # api_name="/infer"
+        )
+        
+        return result[0]
+
+    except Exception as e:
+        print("Excepción es: ", e)
+
+
+def genera_platillo_inference(platillo):
 
     client = InferenceClient(
-            provider= proveedor,
+            provider= globales.proveedor,
             api_key=bridges.hug
-        ) 
-    
-    #Testing change,
+        )
+
+    prompt = previo + platillo
 
     try: 
         image = client.text_to_image(
         prompt,
-        model=enlace,
+        model=globales.inferencia,
         #seed=42, #default varía pero el default es que siempre sea la misma.
         #guidance_scale=7.5,
         #num_inference_steps=50,
@@ -24,7 +60,11 @@ def genera_platillo(prompt):
         #height=1024 #El límite de replicate es 1024.
         )
 
-        return image        
+        img_io = io.BytesIO()
+        image.save(img_io, "PNG")
+        img_io.seek(0)
+
+        return img_io         
         
     except Exception as e:
         print("Excepción es: ", e)

globales.py

@@ -0,0 +1,7 @@
+seconds_available = 0
+
+#espacio = "black-forest-labs/FLUX.1-schnell"
+espacio = "black-forest-labs/FLUX.1-dev"
+inferencia = "black-forest-labs/FLUX.1-dev"
+
+proveedor =  "hf-inference"

requirements.txt

@@ -2,5 +2,6 @@ fastapi
 fastapi[standard]
 
 huggingface_hub
+gradio_client
 
 Pillow

tester.py

@@ -1,3 +0,0 @@
-import funciones
-
-funciones.genera_platillo()