uppdated app.py

app.py

@@ -1,56 +1,84 @@
 # app.py
 
+# ── Monkey‐patch missing torchvision module ──
 import sys
 import torchvision.transforms.functional as F
 sys.modules['torchvision.transforms.functional_tensor'] = F
 
+import os
 import gradio as gr
 import torch
-from diffusers import StableDiffusionInpaintPipeline
-from realesrgan import RealESRGAN
+import numpy as np
+import cv2
 from PIL import Image
+from diffusers import StableDiffusionInpaintPipeline
 
+# Import the RealESRGANer helper and architecture
+from basicsr.archs.rrdbnet_arch import RRDBNet  # RRDB backbone :contentReference[oaicite:0]{index=0}
+from realesrgan.utils import RealESRGANer      # RealESRGANer class :contentReference[oaicite:1]{index=1}
 
-# 1. Initialize the inpainting pipeline on CPU (fp32)
+# 1. Initialize Stable Diffusion InpaintPipeline on CPU
 pipe = StableDiffusionInpaintPipeline.from_pretrained(
     "runwayml/stable-diffusion-inpainting",
     torch_dtype=torch.float32,
 )
 pipe.to("cpu")
 
-# 2. Initialize Real-ESRGAN (4×) on CPU, download weights if needed
-sr = RealESRGAN(device="cpu", scale=4)
-sr.load_weights("RealESRGAN_x4plus.pth", download=True)
+# 2. Build the RRDBNet model and RealESRGANer (4×) on CPU
+device = torch.device("cpu")
+rrdb = RRDBNet(
+    num_in_ch=3, num_out_ch=3,
+    num_feat=64, num_block=23,
+    num_grow_ch=32, scale=4
+)
+# Pass a GitHub URL so it downloads under-the-hood
+esrgan = RealESRGANer(
+    scale=4,
+    model_path="https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth",
+    model=rrdb,
+    tile=0, tile_pad=10, pre_pad=10,
+    half=False,
+    device=device,
+)
 
 def fill_and_upscale(input_img: Image.Image,
-                     mask_img: Image.Image,
-                     prompt: str):
-    # Ensure correct modes
+                     mask_img:  Image.Image,
+                     prompt:    str):
+    # Inpaint
     init = input_img.convert("RGB")
     mask = mask_img.convert("RGB")
+    filled: Image.Image = pipe(
+        prompt=prompt, image=init, mask_image=mask
+    ).images[0]
 
-    # 3. Generative fill (inpainting)
-    filled = pipe(prompt=prompt, image=init, mask_image=mask).images[0]
+    # Prepare for Real-ESRGANer (expects BGR numpy)
+    arr = np.array(filled)
+    bgr = cv2.cvtColor(arr, cv2.COLOR_RGB2BGR)
 
-    # 4. Upscale with Real-ESRGAN
-    upscaled = sr.predict(filled)
+    # Upscale
+    out_bgr, _ = esrgan.enhance(bgr, outscale=None)
+    out_rgb = cv2.cvtColor(out_bgr, cv2.COLOR_BGR2RGB)
+    upscaled = Image.fromarray(out_rgb)
 
     return filled, upscaled
 
-# 5. Build Gradio UI
+# 3. Gradio UI
 with gr.Blocks() as demo:
-    gr.Markdown("## Stable-Diffusion Inpainting + Real-ESRGAN Upscale (CPU)")
+    gr.Markdown("## Inpaint + 4× Upscale (CPU Only)")
     with gr.Row():
         with gr.Column():
-            inp     = gr.Image(type="pil", label="Input Image")
-            msk     = gr.Image(type="pil", label="Mask (white=fill)")
-            prompt  = gr.Textbox(label="Prompt", placeholder="e.g. A misty mountain lake")
-            run_btn = gr.Button("Generate & Upscale")
+            inp    = gr.Image(type="pil", label="Input Image")
+            msk    = gr.Image(type="pil", label="Mask (white=fill)")
+            prompt = gr.Textbox(
+                label="Prompt",
+                placeholder="e.g. A serene waterfall at dawn"
+            )
+            btn    = gr.Button("Run")
         with gr.Column():
             out1 = gr.Image(type="pil", label="Inpainted")
             out2 = gr.Image(type="pil", label="Upscaled")
 
-    run_btn.click(fill_and_upscale, inputs=[inp, msk, prompt], outputs=[out1, out2])
+    btn.click(fill_and_upscale, [inp, msk, prompt], [out1, out2])
 
-# 6. Launch (HuggingFace Spaces picks up port automatically)
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()