Update app.py

app.py

@@ -33,6 +33,26 @@ def enable_lora(lora_add):
     else:
         return lora_add
 
+def get_upscale_finegrain(prompt, img_path, upscale_factor):
+    client = Client("finegrain/finegrain-image-enhancer")
+    result = client.predict(
+        input_image=handle_file(img_path),
+        prompt=prompt,
+        negative_prompt="",
+        seed=42,
+        upscale_factor=upscale_factor,
+        controlnet_scale=0.6,
+        controlnet_decay=1,
+        condition_scale=6,
+        tile_width=112,
+        tile_height=144,
+        denoise_strength=0.35,
+        num_inference_steps=18,
+        solver="DDIM",
+        api_name="/process"
+    )
+    return result[1]
+
 async def generate_image(
     prompt:str,
     model:str,
@@ -41,7 +61,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)
@@ -74,19 +95,28 @@ async def gen(
     scales:float=3.5,
     steps:int=24,
     seed:int=-1,
-    progress=gr.Progress(track_tqdm=True)
+    progress=gr.Progress(track_tqdm=True),
+    upscale_factor:int=0
 ):
     model = enable_lora(lora_add)
     print(model)
     image, seed = await generate_image(prompt,model,lora_word,width,height,scales,steps,seed)
-    return image, seed
-     
+    if upscale_factor != 0:
+        image = get_upscale_finegrain(prompt, image, upscale_factor)
+    return image, seed, image
+
+def upscale_image(img_path, upscale_factor, prompt):
+    if upscale_factor == 0:
+        return img_path  
+    else:
+        return get_upscale_finegrain(prompt, img_path, upscale_factor)
+
 with gr.Blocks(css=CSS, js=JS, theme="Nymbo/Nymbo_Theme") as demo:
     gr.HTML("<h1><center>Flux Lab Light</center></h1>")
     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 Generated 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')
@@ -120,44 +150,57 @@ with gr.Blocks(css=CSS, js=JS, theme="Nymbo/Nymbo_Theme") as demo:
                     step=1,
                     value=24,
                 )
-                seed = gr.Slider(
-                    label="Seeds",
-                    minimum=-1,
-                    maximum=MAX_SEED,
-                    step=1,
-                    value=-1,
-                )
-                lora_add = gr.Textbox(
-                    label="Add Flux LoRA",
-                    info="Copy the HF LoRA model name here",
-                    lines=1,
-                    placeholder="Please use Warm status model",
-                )
-                lora_word = gr.Textbox(
-                    label="Add Flux LoRA Trigger Word",
-                    info="Add the Trigger Word",
-                    lines=1,
-                    value="",
-                )
+            seed = gr.Slider(
+                label="Seeds",
+                minimum=-1,
+                maximum=MAX_SEED,
+                step=1,
+                value=-1,
+            )
+            lora_add = gr.Textbox(
+                label="Add Flux LoRA",
+                info="Copy the HF LoRA model name here",
+                lines=1,
+                placeholder="Please use Warm status model",
+            )
+            lora_word = gr.Textbox(
+                label="Add Flux LoRA Trigger Word",
+                info="Add the Trigger Word",
+                lines=1,
+                value="",
+            )
+            upscale_factor = gr.Radio(
+                label="UpScale Factor",
+                choices=[
+                    0, 
+                    2, 
+                    3, 
+                    4
+                ],
+                value=0,
+                scale=2
+            )
+            output_res = gr.Image(label="Upscaled Image")
 
-    gr.on(
-        triggers=[
-            prompt.submit,
-            sendBtn.click,
-        ],
-        fn=gen,
-        inputs=[
-            prompt,
-            lora_add,
-            lora_word,
-            width, 
-            height, 
-            scales, 
-            steps, 
-            seed
-        ],
-        outputs=[img, seed]
-    )
-    
-if __name__ == "__main__":
-    demo.queue(api_open=False).launch(show_api=False, share=False)
+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, output_res]
+)
+
+if name == "main":
+demo.queue(api_open=False).launch(show_api=False, share=False)