update

app.py

@@ -1,3 +1,5 @@
+import cv2
+import torch
 import gradio as gr
 from langchain_core.messages import HumanMessage, AIMessage
 from llm import DeepSeekLLM, OpenRouterLLM, TongYiLLM
@@ -5,11 +7,8 @@ from config import settings
 from modelscope.outputs import OutputKeys
 from modelscope.pipelines import pipeline
 from modelscope.utils.constant import Tasks
-import cv2
-from diffusers import StableDiffusionXLPipeline
-import torch
 from diffusers import (
-    DiffusionPipeline,
+    StableDiffusionXLPipeline,
     DPMSolverMultistepScheduler,
     DDIMScheduler,
     HeunDiscreteScheduler,
@@ -71,14 +70,14 @@ def update_chat(_provider: str, _chat, _model: str, _temperature: float, _max_to
     return _chat
 
 
-def object_remove(_image, refined):
+def object_remove(_image, _refined: bool):
     mask = _image['layers'][0]
     mask = mask.convert('L')
     _input = {
         'img': _image['background'].convert('RGB'),
         'mask': mask,
     }
-    inpainting = pipeline(Tasks.image_inpainting, model='damo/cv_fft_inpainting_lama', refined=refined)
+    inpainting = pipeline(Tasks.image_inpainting, model='damo/cv_fft_inpainting_lama', refined=_refined)
     result = inpainting(_input)
     vis_img = result[OutputKeys.OUTPUT_IMG]
     vis_img = cv2.cvtColor(vis_img, cv2.COLOR_BGR2RGB)
@@ -119,6 +118,14 @@ def text_to_image(_prompt: str, _n_prompt: str, _scheduler: str, _inference_step
         return result
 
 
+def image_upscale(_image, _size: str):
+    sr = pipeline(Tasks.image_super_resolution, model='damo/cv_rrdb_image-super-resolution')
+    result = sr(_image['background'].convert('RGB'))
+    vis_img = result[OutputKeys.OUTPUT_IMG]
+    vis_img = cv2.cvtColor(vis_img, cv2.COLOR_BGR2RGB)
+    return vis_img
+
+
 with gr.Blocks() as app:
     with gr.Tab('聊天'):
         chat_engine = gr.State(value=None)
@@ -275,8 +282,12 @@ with gr.Blocks() as app:
                 with gr.Accordion(label="背景移除"):
                     bg_remove_type = gr.Radio(["人像", "通用"], label="类型", value='人像')
                     bg_remove_btn = gr.Button('背景移除', variant='primary')
+                with gr.Accordion(label="高清放大"):
+                    upscale_size = gr.Radio(["X2", "X4"], label="放大倍数", value='X2')
+                    upscale_btn = gr.Button('高清放大', variant='primary')
             object_remove_btn.click(fn=object_remove, inputs=[image, object_remove_refined], outputs=[image_preview, mask_preview])
             bg_remove_btn.click(fn=bg_remove, inputs=[image, bg_remove_type], outputs=[image_preview])
+            upscale_btn.click(fn=image_upscale, inputs=[image, upscale_size], outputs=[image_preview])
 
     with gr.Tab('画图(GPU)'):
         with gr.Row():