Update app.py

app.py

@@ -1,16 +1,18 @@
-import os
 import gradio as gr
+import os
+from gradio_client import Client, handle_file
+from huggingface_hub import login
+from PIL import Image
 import numpy as np
 import random
-from huggingface_hub import AsyncInferenceClient
 from translatepy import Translator
 import requests
 import re
 import asyncio
-from PIL import Image
+
+login(token=os.environ.get("HF_TOKEN", None), username=os.environ.get("HF_USERNAME", None))
 
 translator = Translator()
-HF_TOKEN = os.environ.get("HF_TOKEN", None)
 basemodel = "black-forest-labs/FLUX.1-dev"
 MAX_SEED = np.iinfo(np.int32).max
 
@@ -33,40 +35,12 @@ def enable_lora(lora_add):
     else:
         return lora_add
 
-client = AsyncInferenceClient()
-
-async def generate_image(
-    prompt:str,
-    model:str,
-    lora_word:str,
-    width:int=768,
-    height:int=1024,
-    scales:float=3.5,
-    steps:int=24,
-    seed:int=-1):
-
-    if seed == -1:
-        seed = random.randint(0, MAX_SEED)
-    seed = int(seed)
-    print(f'prompt:{prompt}')
-    
-    text = str(translator.translate(prompt, 'English')) + "," + lora_word
-
-    try:
-        image = await client.text_to_image(
-            prompt=text,
-            height=height,
-            width=width,
-            guidance_scale=scales,
-            num_inference_steps=steps,
-            model=model,
-        )
-    except Exception as e:
-        raise gr.Error(f"Error in {e}")
-    
-    return image, seed
+def handle_file(img_path):
+    return Image.open(img_path)
 
 def get_upscale_finegrain(prompt, img_path, upscale_factor):
+    if upscale_factor == 0:
+        return handle_file(img_path)
     client = Client("finegrain/finegrain-image-enhancer")
     result = client.predict(
         input_image=handle_file(img_path),
@@ -99,6 +73,30 @@ async def upscale_image(image, upscale_factor):
     
     return result
 
+async def generate_image(
+    prompt:str,
+    model:str,
+    lora_word:str,
+    width:int=768,
+    height:int=1024,
+    scales:float=3.5,
+    steps:int=24,
+    seed:int=-1
+):
+    if seed == -1:
+        seed = random.randint(0, MAX_SEED)
+    seed = int(seed)
+    print(f'prompt:{prompt}')
+    
+    text = str(translator.translate(prompt, 'English')) + "," + lora_word
+
+    try:
+        image = gr.Image(type="pil", image=gr.processing_utils.encode_pil_image(text_to_image(text, height=height, width=width, guidance_scale=scales, num_inference_steps=steps, model=model)))
+    except Exception as e:
+        raise gr.Error(f"Error in {e}")
+    
+    return image, seed
+
 async def gen(
     prompt:str,
     lora_add:str="XLabs-AI/flux-RealismLora",
@@ -108,7 +106,7 @@ async def gen(
     scales:float=3.5,
     steps:int=24,
     seed:int=-1,
-    upscale_factor:int=2
+    upscale_factor:int=0
 ):
     model = enable_lora(lora_add)
     image, seed = await generate_image(prompt,model,lora_word,width,height,scales,steps,seed)
@@ -127,78 +125,11 @@ with gr.Blocks(css=CSS, js=JS, theme="Nymbo/Nymbo_Theme") as demo:
                 sendBtn = gr.Button(scale=1, variant='primary')
         with gr.Accordion("Opciones avanzadas", open=True):
             with gr.Column(scale=1):
-                width = gr.Slider(
-                    label="Ancho",
-                    minimum=512,
-                    maximum=1280,
-                    step=8,
-                    value=768,
-                )
-                height = gr.Slider(
-                    label="Alto",
-                    minimum=512,
-                    maximum=1280,
-                    step=8,
-                    value=1024,
-                )
-                scales = gr.Slider(
-                    label="Guía",
-                    minimum=3.5,
-                    maximum=7,
-                    step=0.1,
-                    value=3.5,
-                )
-                steps = gr.Slider(
-                    label="Pasos",
-                    minimum=1,
-                    maximum=100,
-                    step=1,
-                    value=24,
-                )
-                seed = gr.Slider(
-                    label="Semillas",
-                    minimum=-1,
-                    maximum=MAX_SEED,
-                    step=1,
-                    value=-1,
-                )
-                lora_add = gr.Textbox(
-                    label="Agregar Flux LoRA",
-                    info="Modelo de LoRA a agregar",
-                    lines=1,
-                    value="XLabs-AI/flux-RealismLora",
-                )
-                lora_word = gr.Textbox(
-                    label="Palabra clave de LoRA",
-                    info="Palabra clave para activar el modelo de LoRA",
-                    lines=1,
-                    value="",
-                )
-                upscale_factor = gr.Radio(
-                    label="Factor de escalado",
-                    choices=[2, 3, 4],
-                    value=2,
-                )
-
-    gr.on(
-        triggers=[
-            prompt.submit,
-            sendBtn.click,
-        ],
-        fn=gen,
-        inputs=[
-            prompt,
-            lora_add,
-            lora_word,
-            width, 
-            height, 
-            scales, 
-            steps, 
-            seed,
-            upscale_factor
-        ],
-        outputs=[img, seed]
-    )
-    
-if __name__ == "__main__":
-    demo.queue(api_open=False).launch(show_api=False, share=False)
+                width = gr.Slider(label="Ancho", minimum=512, maximum=1280, step=8, value=768)
+                height = gr.Slider(label="Alto", minimum=512, maximum=1280, step=8, value=1024)
+                scales = gr.Slider(label="Guía", minimum=3.5, maximum=7, step=0.1, value=3.5)
+                steps = gr.Slider(label="Pasos", minimum=1, maximum=50, step=1)
+                upscale_factor = gr.Slider(label="Factor de escala", minimum=0, maximum=4, step=1, value=0)
+                seed = gr.Number(label="Semilla", value=-1)
+                sendBtn.click(gen, inputs=[prompt, lora_add, lora_word, width, height, scales, steps, seed, upscale_factor], outputs=[img])
+demo.launch()