Update app.py

app.py

@@ -1,115 +1,164 @@
 import os
-
 import sys
+import tempfile
+import cv2
+import torch
+import gradio as gr
 from torchvision.transforms import functional
-sys.modules["torchvision.transforms.functional_tensor"] = functional 
-# //sequntila NotImplemented
 
+# --- PATCH FOR COMPATIBILITY ---
+sys.modules["torchvision.transforms.functional_tensor"] = functional
+
+# --- EMBEDDED CSS FOR STYLING ---
+CSS_STYLING = """
+:root {
+    --primary: #6a35ee; --primary-dark: #4a1dcc; --secondary: #00c9ff;
+    --accent: #ff6b6b; --light: #f8f9ff; --dark: #1a1f36; --text: #4a5568;
+    --input-background-fill: var(--light) !important;
+    --input-border-color: #e0e0e0 !important;
+    --input-label-color: var(--text) !important;
+}
+.gradio-container { background: var(--light); font-family: 'Inter', sans-serif; }
+#main-title { color: var(--dark); text-align: center; font-size: 2.5rem !important; font-weight: 900; }
+#main-subtitle { color: var(--text); text-align: center; font-size: 1rem !important; margin-top: -15px; margin-bottom: 20px; }
+#submit-button { background: var(--primary); color: white; font-weight: bold; border-radius: 8px !important; transition: all 0.3s ease; }
+#submit-button:hover { background: var(--primary-dark); box-shadow: 0px 4px 15px rgba(106, 53, 238, 0.4); transform: translateY(-2px); }
+.gr-image { border: 1px dashed var(--input-border-color) !important; border-radius: 12px !important; min-height: 300px; }
+input[type="range"]::-webkit-slider-thumb { background: var(--primary) !important; }
+input[type="range"]::-moz-range-thumb { background: var(--primary) !important; }
+"""
+
+# --- DOWNLOAD HELPER FUNCTIONS ---
+def download_file(url, dir_path, file_name):
+    """Downloads a file if it doesn't exist."""
+    os.makedirs(dir_path, exist_ok=True)
+    file_path = os.path.join(dir_path, file_name)
+    if not os.path.exists(file_path):
+        print(f"Downloading {file_name}...")
+        try:
+            os.system(f"wget {url} -O {file_path}")
+            print("Download complete.")
+        except Exception as e:
+            print(f"Error downloading {file_name}: {e}")
+            # In case wget is not available, you might need to use requests or urllib
+            # import requests
+            # with open(file_path, 'wb') as f:
+            #     f.write(requests.get(url).content)
+    return file_path
+
+# --- DOWNLOAD MODELS AND EXAMPLES ---
+print("Checking for required files...")
+# Models
+models_dir = 'models'
+download_file('https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth', models_dir, 'realesr-general-x4v3.pth')
+download_file('https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth', models_dir, 'GFPGANv1.4.pth')
+download_file('https://github.com/TencentARC/GFPGAN/releases/download/v1.3.4/RestoreFormer.pth', models_dir, 'RestoreFormer.pth')
+
+# Example Images
+examples_dir = 'examples'
+example1_path = download_file('https://raw.githubusercontent.com/TencentARC/GFPGAN/master/inputs/whole_imgs/10045.png', examples_dir, 'example1.png')
+example2_path = download_file('https://raw.githubusercontent.com/TencentARC/GFPGAN/master/inputs/whole_imgs/Blake_Lively.jpg', examples_dir, 'example2.jpg')
+
+# --- LOAD MODELS INTO MEMORY ---
 from basicsr.archs.srvgg_arch import SRVGGNetCompact
 from gfpgan.utils import GFPGANer
 from realesrgan.utils import RealESRGANer
 
-import torch
-import cv2
-import gradio as gr
-
-
-#Download Required Models
-if not os.path.exists('realesr-general-x4v3.pth'):
-    os.system("wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth -P .")
-if not os.path.exists('GFPGANv1.2.pth'):
-    os.system("wget https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.2.pth -P .")
-if not os.path.exists('GFPGANv1.3.pth'):
-    os.system("wget https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.3.pth -P .")
-if not os.path.exists('GFPGANv1.4.pth'):
-    os.system("wget https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth -P .")
-if not os.path.exists('RestoreFormer.pth'):
-    os.system("wget https://github.com/TencentARC/GFPGAN/releases/download/v1.3.4/RestoreFormer.pth -P .")
-
-
-model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=32, upscale=4, act_type='prelu')
-model_path = 'realesr-general-x4v3.pth'
-half = True if torch.cuda.is_available() else False
-upsampler = RealESRGANer(scale=4, model_path=model_path, model=model, tile=0, tile_pad=10, pre_pad=0, half=half)
-
-
-# Save Image to the Directory
-# os.makedirs('output', exist_ok=True)
-
-def upscaler(img, version, scale):
+bg_upsampler = None
+try:
+    model_path = os.path.join(models_dir, 'realesr-general-x4v3.pth')
+    model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=32, upscale=4, act_type='prelu')
+    half = torch.cuda.is_available()
+    bg_upsampler = RealESRGANer(scale=4, model_path=model_path, model=model, tile=0, tile_pad=10, pre_pad=0, half=half)
+    print("Background Upsampler (Real-ESRGAN) loaded.")
+except Exception as e:
+    print(f"Error loading background upsampler: {e}. The app may not work correctly.")
+
+# --- CORE IMAGE PROCESSING FUNCTION ---
+def upscale_image(img_path, version, scale):
+    """Enhance an image using GFPGAN and Real-ESRGAN."""
+    if not img_path:
+        raise gr.Error("Please upload an image.")
+    if not bg_upsampler:
+        raise gr.Error("Background upsampler not loaded. Cannot proceed.")
 
     try:
+        img = cv2.imread(img_path, cv2.IMREAD_UNCHANGED)
+        if img is None: raise RuntimeError("Failed to read image.")
         
-        img = cv2.imread(img, cv2.IMREAD_UNCHANGED)
-        if len(img.shape) == 3 and img.shape[2] == 4:
-            img_mode = 'RGBA'
-        elif len(img.shape) == 2:
-            img_mode = None
-            img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
-        else:
-            img_mode = None
+        has_alpha = img.shape[2] == 4
 
-
-        h, w = img.shape[0:2]
-        if h < 300:
-            img = cv2.resize(img, (w * 2, h * 2), interpolation=cv2.INTER_LANCZOS4)
-
-        
         face_enhancer = GFPGANer(
-            model_path=f'{version}.pth', 
-            upscale=2, 
-            arch='RestoreFormer' if version=='RestoreFormer' else 'clean',
+            model_path=os.path.join(models_dir, f'{version}.pth'),
+            upscale=2, # GFPGAN native upscale factor
+            arch='RestoreFormer' if version == 'RestoreFormer' else 'clean',
             channel_multiplier=2,
-            bg_upsampler=upsampler
+            bg_upsampler=bg_upsampler
         )
+        
+        _, _, output = face_enhancer.enhance(img, has_aligned=False, only_center_face=False, paste_back=True)
+
+        if scale != 2:
+            h, w = output.shape[0:2]
+            target_w, target_h = int(w * scale / 2), int(h * scale / 2)
+            if target_w > 8000 or target_h > 8000:
+                raise gr.Error(f"Target size is too large. Please choose a smaller scale.")
+            interpolation = cv2.INTER_LANCZOS4 if scale > 2 else cv2.INTER_AREA
+            output = cv2.resize(output, (target_w, target_h), interpolation=interpolation)
+
+        output_rgb = cv2.cvtColor(output, cv2.COLOR_BGR2RGB)
+        ext = 'png' if has_alpha else 'jpg'
+        
+        # Save to a temporary file for download
+        with tempfile.NamedTemporaryFile(delete=False, suffix=f'.{ext}') as temp_file:
+            cv2.imwrite(temp_file.name, cv2.cvtColor(output_rgb, cv2.COLOR_RGB2BGR))
+            return output_rgb, temp_file.name
 
-
-        try:
-            _, _, output = face_enhancer.enhance(img, has_aligned=False, only_center_face=False, paste_back=True)
-        except RuntimeError as error:
-            print('Error', error)
-
-
-        try:
-            if scale != 2:
-                interpolation = cv2.INTER_AREA if scale < 2 else cv2.INTER_LANCZOS4
-                h, w = img.shape[0:2]
-                output = cv2.resize(output, (int(w * scale / 2), int(h * scale / 2)), interpolation=interpolation)
-        except Exception as error:
-            print('wrong scale input.', error)
-
-        # Save Image to the Directory
-        # ext = os.path.splitext(os.path.basename(str(img)))[1]
-        # if img_mode == 'RGBA':
-        #     ext = 'png'
-        # else:
-        #     ext = 'jpg'
-        #
-        # save_path = f'output/out.{ext}'
-        # cv2.imwrite(save_path, output)
-        # return output, save_path
-
-        output = cv2.cvtColor(output, cv2.COLOR_BGR2RGB)
-        return output
     except Exception as error:
-        print('global exception', error)
-        return None, None
+        print(f"Error processing image: {error}")
+        raise gr.Error(f"An error occurred: {error}")
+
+# --- GRADIO UI LAYOUT ---
+with gr.Blocks(css=CSS_STYLING, theme=gr.themes.Base()) as demo:
+    gr.Markdown("<h1 id='main-title'>NeuraVision AI Image Upscaler</h1>", elem_id="main-title")
+    gr.Markdown("<p id='main-subtitle'>Enhance old, blurry, and low-resolution photos with AI.</p>", elem_id="main-subtitle")
+
+    with gr.Row(variant="panel"):
+        # LEFT COLUMN (INPUT & CONTROLS)
+        with gr.Column(scale=1):
+            input_image = gr.Image(type="filepath", label="Upload Image")
+            
+            version = gr.Radio(
+                ['GFPGANv1.4', 'RestoreFormer'], value='GFPGANv1.4', 
+                label='AI Model', info="v1.4 is best for general use. RestoreFormer for very old photos."
+            )
+            
+            scale = gr.Slider(
+                minimum=1, maximum=8, step=0.5, value=4, 
+                label="Upscale Factor", info="How many times larger to make the image."
+            )
+            
+            submit_btn = gr.Button("Enhance Image", variant="primary", elem_id="submit-button")
+            
+            gr.Examples(
+                examples=[[example1_path, "RestoreFormer", 4], [example2_path, "GFPGANv1.4", 4]],
+                inputs=[input_image, version, scale],
+                label="Click an example to start"
+            )
+        
+        # RIGHT COLUMN (OUTPUT)
+        with gr.Column(scale=1):
+            output_image = gr.Image(type="numpy", label="Enhanced Result", interactive=False)
+            download_button = gr.File(label="Download Image", interactive=False)
+
+    # --- BUTTON & EVENT HANDLING ---
+    submit_btn.click(
+        fn=upscale_image, 
+        inputs=[input_image, version, scale], 
+        outputs=[output_image, download_button]
+    )
+    input_image.clear(lambda: (None, None), None, [output_image, download_button])
 
 if __name__ == "__main__":
-
-    title = "NeuraVision ai image upscale"
-
-    demo = gr.Interface(
-            upscaler, [
-                gr.Image(type="filepath", label="Input"),
-                gr.Radio(['GFPGANv1.2', 'GFPGANv1.3', 'GFPGANv1.4', 'RestoreFormer'], type="value", label='version'),
-                gr.Number(label="Rescaling factor"),
-            ], [
-                gr.Image(type="numpy", label="Output"),
-            ],
-            title=title,
-            allow_flagging="never"
-        )
-
     demo.queue()
-    demo.launch()
+    demo.launch(share=True) # Set share=False if you don't need a public link