Update app.py

app.py

@@ -5,7 +5,7 @@ import torch
 import spaces
 
 from PIL import Image
-from diffusers import FlowMatchEulerDiscreteScheduler
+from diffusers import FlowMatchEulerDiscreteScheduler,DiffusionPipeline
 from optimization import optimize_pipeline_
 from qwenimage.pipeline_qwenimage_edit_plus import QwenImageEditPlusPipeline
 from qwenimage.transformer_qwenimage import QwenImageTransformer2DModel
@@ -16,6 +16,11 @@ import math
 from huggingface_hub import hf_hub_download
 from safetensors.torch import load_file
 
+from basicsr.archs.rrdbnet_arch import RRDBNet
+from basicsr.utils.download_util import load_file_from_url
+from realesrgan import RealESRGANer
+from realesrgan.archs.srvgg_arch import SRVGGNetCompact
+
 import os
 import base64
 from io import BytesIO
@@ -28,6 +33,85 @@ from PIL import Image
 import os
 import gradio as gr
 
+
+# --- Upscaling ---
+MAX_SEED = np.iinfo(np.int32).max
+UPSAMPLER_CACHE = {}
+GFPGAN_FACE_ENHANCER = {}
+
+def rnd_string(x): return "".join(random.choice("abcdefghijklmnopqrstuvwxyz_0123456789") for _ in range(x))
+
+
+def get_model_and_paths(model_name, denoise_strength):
+     if model_name in ('RealESRGAN_x4plus', 'RealESRNet_x4plus'):
+        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
+        netscale = 4
+        file_url = ['https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth'] \
+            if model_name == 'RealESRGAN_x4plus' else \
+            ['https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.1/RealESRNet_x4plus.pth']
+    elif model_name == 'RealESRGAN_x4plus_anime_6B':
+        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=6, num_grow_ch=32, scale=4)
+        netscale = 4
+        file_url = ['https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth']
+    elif model_name == 'RealESRGAN_x2plus':
+        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=2)
+        netscale = 2
+        file_url = ['https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth']
+    elif model_name == 'realesr-general-x4v3':
+        model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=32, upscale=2, act_type='prelu')
+        netscale = 4
+        file_url = [
+            'https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-wdn-x4v3.pth',
+            'https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth'
+        ]
+    else:
+        raise ValueError(f"Unsupported model: {model_name}")
+
+    model_path = os.path.join("weights", model_name + ".pth")
+    if not os.path.isfile(model_path):
+        ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
+        for url in file_url:
+            model_path = load_file_from_url(url=url, model_dir=os.path.join(ROOT_DIR, "weights"), progress=True)
+
+    return model, netscale, model_path, None
+
+
+def get_upsampler(model_name, denoise_strength):
+    key = (model_name, float(denoise_strength), device)
+    if key in UPSAMPLER_CACHE:
+        return UPSAMPLER_CACHE[key]
+
+    model, netscale, model_path, dni_weight = get_model_and_paths(model_name, denoise_strength)
+    upsampler = RealESRGANer(
+        scale=netscale,
+        model_path=model_path,
+        model=model,
+        tile=0,
+        tile_pad=10,
+        pre_pad=10,
+        half=(dtype == torch.bfloat16),
+        gpu_id=0 if device == "cuda" else None,
+    )
+    UPSAMPLER_CACHE[key] = upsampler
+    return upsampler
+
+@spaces.GPU
+def realesrgan(img, model_name, denoise_strength, outscale=4, progress=gr.Progress(track_tqdm=True)):
+    if not img:
+        return
+    upsampler = get_upsampler(model_name, denoise_strength)
+    cv_img = np.array(img.convert("RGB"))
+    bgr = cv2.cvtColor(cv_img, cv2.COLOR_RGB2BGR)
+    try:
+        output, _ = upsampler.enhance(bgr, outscale=int(outscale))
+    except Exception as e:
+        print("Upscale error:", e)
+        return img
+    out_filename = f"output_{rnd_string(8)}.jpg"
+    cv2.imwrite(out_filename, output)
+    return out_filename
+
+
 def turn_into_video(input_images, output_images, prompt, progress=gr.Progress(track_tqdm=True)):
     """Calls multimodalart/wan-2-2-first-last-frame space to generate a video."""
     if not input_images or not output_images:
@@ -498,9 +582,9 @@ def infer(
         true_cfg_scale=true_guidance_scale,
         num_images_per_prompt=num_images_per_prompt,
     ).images
-
+    upscaled = realesrgan(image, "realesr-general-x4v3", 0.5)
     # Return images, seed, and make button visible
-    return image, seed, gr.update(visible=True), gr.update(visible=True)
+    return upscaled, seed, gr.update(visible=True), gr.update(visible=True)
 
 
 # --- Examples and UI Layout ---