Update app.py

app.py

@@ -11,7 +11,7 @@ from PIL import Image
 
 translator = Translator()
 HF_TOKEN = os.environ.get("HF_TOKEN", None)
-basemodel = "black-forest-labs/FLUX.1-dev"
+basemodel = "XLabs-AI/flux-RealismLora"
 MAX_SEED = np.iinfo(np.int32).max
 
 CSS = """
@@ -27,12 +27,11 @@ JS = """function () {
   }
 }"""
 
-lora_add = gr.Textbox(
-    label="Add Flux LoRA",
-    info="Copy the HF LoRA model name here",
-    lines=1,
-    value="XLabs-AI/flux-RealismLora"
-)
+def enable_lora(lora_add):
+    if not lora_add:
+        return basemodel
+    else:
+        return lora_add
 
 async def generate_image(
     prompt:str,
@@ -42,7 +41,8 @@ async def generate_image(
     height:int=1024,
     scales:float=3.5,
     steps:int=24,
-    seed:int=-1):
+    seed:int=-1
+):
 
     if seed == -1:
         seed = random.randint(0, MAX_SEED)
@@ -66,6 +66,42 @@ async def generate_image(
     
     return image, seed
 
+async def upscale_image(
+    prompt:str,
+    img_path:str,
+    upscale_factor:int=2,
+    controlnet_scale:float=0.6,
+    controlnet_decay:float=1,
+    condition_scale:int=6,
+    tile_width:int=112,
+    tile_height:int=144,
+    denoise_strength:float=0.35,
+    num_inference_steps:int=18,
+    solver:str="DDIM"
+):
+    client = AsyncInferenceClient()
+    try:
+        result = await client.image_to_image(
+            prompt=prompt,
+            input_image=img_path,
+            negative_prompt="",
+            seed=42,
+            upscale_factor=upscale_factor,
+            controlnet_scale=controlnet_scale,
+            controlnet_decay=controlnet_decay,
+            condition_scale=condition_scale,
+            tile_width=tile_width,
+            tile_height=tile_height,
+            denoise_strength=denoise_strength,
+            num_inference_steps=num_inference_steps,
+            solver=solver,
+            model="finegrain/finegrain-image-enhancer",
+        )
+    except Exception as e:
+        raise gr.Error(f"Error in {e}")
+    
+    return result[0]
+
 async def gen(
     prompt:str,
     lora_add:str="",
@@ -75,20 +111,31 @@ async def gen(
     scales:float=3.5,
     steps:int=24,
     seed:int=-1,
+    upscale_factor:int=2,
+    controlnet_scale:float=0.6,
+    controlnet_decay:float=1,
+    condition_scale:int=6,
+    tile_width:int=112,
+    tile_height:int=144,
+    denoise_strength:float=0.35,
+    num_inference_steps:int=18,
+    solver:str="DDIM",
     progress=gr.Progress(track_tqdm=True)
 ):
     model = enable_lora(lora_add)
     print(model)
-    image, seed = await generate_image(prompt,model,lora_word,width,height,scales,steps,seed)
-    return image, seed
-     
+    image, seed = await generate_image(prompt, model, lora_word, width, height, scales, steps, seed)
+    upscale_img = await upscale_image(prompt, image, upscale_factor, controlnet_scale, controlnet_decay, condition_scale, tile_width, tile_height, denoise_strength, num_inference_steps, solver)
+    return image, upscale_img, seed
+
 with gr.Blocks(css=CSS, js=JS, theme="Nymbo/Nymbo_Theme") as demo:
     gr.HTML("<h1><center>Flux Lab Light</center></h1>")
     gr.HTML("<p><center>Powered By HF Inference API</center></p>")
     with gr.Row():
         with gr.Column(scale=4):
             with gr.Row():
-                img = gr.Image(type="filepath", label='flux Generated Image', height=600)
+                img = gr.Image(type="filepath", label='Flux Image', height=600)
+                upscale_img = gr.Image(type="filepath", label='Upscale Image', height=600)
             with gr.Row():
                 prompt = gr.Textbox(label='Enter Your Prompt (Multi-Languages)', placeholder="Enter prompt...", scale=6)
                 sendBtn = gr.Button(scale=1, variant='primary')
@@ -133,7 +180,7 @@ with gr.Blocks(css=CSS, js=JS, theme="Nymbo/Nymbo_Theme") as demo:
                     label="Add Flux LoRA",
                     info="Copy the HF LoRA model name here",
                     lines=1,
-                    placeholder="Please use Warm status model",
+                    value="XLabs-AI/flux-RealismLora"
                 )
                 lora_word = gr.Textbox(
                     label="Add Flux LoRA Trigger Word",
@@ -141,6 +188,71 @@ with gr.Blocks(css=CSS, js=JS, theme="Nymbo/Nymbo_Theme") as demo:
                     lines=1,
                     value="",
                 )
+                upscale_factor = gr.Radio(
+                    label="UpScale Factor",
+                    choices=[
+                        2, 3, 4
+                    ],
+                    value=2,
+                    scale=2
+                )
+                controlnet_scale = gr.Slider(
+                    label="ControlNet Scale",
+                    minimum=0.1,
+                    maximum=1.0,
+                    step=0.1,
+                    value=0.6
+                )
+                controlnet_decay = gr.Slider(
+                    label="ControlNet Decay",
+                    minimum=0.1,
+                    maximum=1.0,
+                    step=0.1,
+                    value=1
+                )
+                condition_scale = gr.Slider(
+                    label="Condition Scale",
+                    minimum=1,
+                    maximum=10,
+                    step=1,
+                    value=6
+                )
+                tile_width = gr.Slider(
+                    label="Tile Width",
+                    minimum=64,
+                    maximum=256,
+                    step=16,
+                    value=112
+                )
+                tile_height = gr.Slider(
+                    label="Tile Height",
+                    minimum=64,
+                    maximum=256,
+                    step=16,
+                    value=144
+                )
+                denoise_strength = gr.Slider(
+                    label="Denoise Strength",
+                    minimum=0.1,
+                    maximum=1.0,
+                    step=0.1,
+                    value=0.35
+                )
+                num_inference_steps = gr.Slider(
+                    label="Num Inference Steps",
+                    minimum=1,
+                    maximum=50,
+                    step=1,
+                    value=18
+                )
+                solver = gr.Radio(
+                    label="Solver",
+                    choices=[
+                        "DDIM", "DPM"
+                    ],
+                    value="DDIM",
+                    scale=2
+                )
 
     gr.on(
         triggers=[
@@ -156,9 +268,18 @@ with gr.Blocks(css=CSS, js=JS, theme="Nymbo/Nymbo_Theme") as demo:
             height, 
             scales, 
             steps, 
-            seed
+            seed,
+            upscale_factor,
+            controlnet_scale,
+            controlnet_decay,
+            condition_scale,
+            tile_width,
+            tile_height,
+            denoise_strength,
+            num_inference_steps,
+            solver
         ],
-        outputs=[img, seed]
+        outputs=[img, upscale_img, seed]
     )
     
 if __name__ == "__main__":