feat: Add option to force PNG output for best quality

README.md

@@ -4,7 +4,7 @@ emoji: 📈
 colorFrom: blue
 colorTo: gray
 sdk: gradio
-sdk_version: 5.35.0
+sdk_version: 5.39.0
 app_file: app.py
 pinned: true
 license: apache-2.0

app.py

@@ -790,7 +790,7 @@ class Upscale:
         self.face_enhancer = GFPGANer(model_path=model_path, upscale=self.scale, arch=arch, channel_multiplier=channel_multiplier, model_rootpath=model_rootpath, det_model=face_detection, resolution=resolution)
 
 
-    def inference(self, gallery, face_restoration, upscale_model, scale: float, face_detection, face_detection_threshold: any, face_detection_only_center: bool, outputWithModelName: bool, progress=gr.Progress()):
+    def inference(self, gallery, face_restoration, upscale_model, scale: float, face_detection, face_detection_threshold: any, face_detection_only_center: bool, outputWithModelName: bool, save_as_png: bool, progress=gr.Progress()):
         try:
             if not gallery or (not face_restoration and not upscale_model):
                 raise ValueError("Invalid parameter setting")
@@ -905,9 +905,14 @@ class Upscale:
                        print(f"Warning: Processing resulted in no image for '{img_path}'. Skipping output.")
                        continue
 
-                    if not extension:
-                        extension = ".png" if img_mode == "RGBA" else ".jpg" # RGBA images should be saved in png format
-                    save_path = f"output/{basename}{self.modelInUse}{extension}"
+                    # Determine the file extension for the output image based on user preference and image properties.
+                    if save_as_png:
+                        # Force PNG output for the best quality, as requested by the user.
+                        final_extension = ".png"
+                    else:
+                        # Use original logic: PNG for images with an alpha channel (RGBA), otherwise use the original extension or default to jpg.
+                        final_extension = ".png" if img_mode == "RGBA" else (extension if extension else ".jpg")
+                    save_path = f"output/{basename}{self.modelInUse}{final_extension}"
                     self.imwriteUTF8(save_path, restored_img)
                     zipf_restore.write(save_path, arcname=os.path.basename(save_path))
 
@@ -1145,7 +1150,11 @@ def remove_image_from_gallery(gallery: list, selected_image: str):
 
 def main():
     if torch.cuda.is_available():
-        torch.cuda.set_per_process_memory_fraction(0.975, device='cuda:0')
+        # Sets the VRAM limit for the GPU. Fine-tune VRAM usage.
+        # Higher values allow larger image tiles to be processed at once, which is faster. If you see the 'depth' in the console log increase frequently, try raising this value.
+        # Lower it for complex models that need more memory overhead.
+        # This setting is locked in once processing starts. To apply a new value, you must restart the entire application.
+        torch.cuda.set_per_process_memory_fraction(0.900, device='cuda:0')
         # set torch options to avoid get black image for RTX16xx card
         # https://github.com/CompVis/stable-diffusion/issues/69#issuecomment-1260722801
         torch.backends.cudnn.enabled = True
@@ -1213,6 +1222,8 @@ def main():
                 face_detection_threshold   = gr.Number(label="Face eye dist threshold", value=10, info="A threshold to filter out faces with too small an eye distance (e.g., side faces).")
                 face_detection_only_center = gr.Checkbox(value=False, label="Face detection only center", info="If set to True, only the face closest to the center of the image will be kept.")
                 with_model_name            = gr.Checkbox(label="Output image files name with model name", value=True)
+                # Add a checkbox to always save the output as a PNG file for the best quality.
+                save_as_png                = gr.Checkbox(label="Always save output as PNG", value=True, info="If enabled, all output images will be saved in PNG format to ensure the best quality. If disabled, the format will be determined automatically (PNG for images with transparency, otherwise JPG).")
 
                 # Define the event listener to add the uploaded image to the gallery
                 input_image.change(append_gallery, inputs=[input_gallery, input_image], outputs=[input_gallery, input_image])
@@ -1237,6 +1248,7 @@ def main():
                             face_detection_threshold,
                             face_detection_only_center,
                             with_model_name,
+                            save_as_png,
                         ], variant="secondary", size="lg",)
             with gr.Column(variant="panel"):
                 gallerys = gr.Gallery(type="filepath", label="Output (The whole image)", format="png")
@@ -1272,6 +1284,7 @@ def main():
                 face_detection_threshold,
                 face_detection_only_center,
                 with_model_name,
+                save_as_png,
             ],
             outputs=[gallerys, outputs],
         )

requirements.txt

@@ -1,6 +1,6 @@
 --extra-index-url https://download.pytorch.org/whl/cu124
 
-gradio==5.35.0
+gradio==5.39.0
 
 basicsr @ git+https://github.com/avan06/BasicSR
 facexlib @ git+https://github.com/avan06/facexlib

webui.bat

@@ -1,21 +1,37 @@
 @echo off
 
-:: The source of the webui.bat file is stable-diffusion-webui
+:: The original source of the webui.bat file is stable-diffusion-webui
+:: Modified and enhanced by Gemini with features for venv management and requirements handling.
+
+:: --------- Configuration ---------
 set COMMANDLINE_ARGS=--input_images_limit -1
+:: Define the name of the Launch application
+set APPLICATION_NAME=app.py
+:: Define the name of the virtual environment directory
+set VENV_NAME=venv
+:: Set to 1 to always attempt to update packages from requirements.txt on every launch
+set ALWAYS_UPDATE_REQS=0
+:: ---------------------------------
+
 
+:: Set PYTHON executable if not already defined
 if not defined PYTHON (set PYTHON=python)
-if not defined VENV_DIR (set "VENV_DIR=%~dp0%venv")
+:: Set VENV_DIR using VENV_NAME if not already defined
+if not defined VENV_DIR (set "VENV_DIR=%~dp0%VENV_NAME%")
 
 mkdir tmp 2>NUL
 
+:: Check if Python is callable
 %PYTHON% -c "" >tmp/stdout.txt 2>tmp/stderr.txt
 if %ERRORLEVEL% == 0 goto :check_pip
 echo Couldn't launch python
 goto :show_stdout_stderr
 
 :check_pip
+:: Check if pip is available
 %PYTHON% -mpip --help >tmp/stdout.txt 2>tmp/stderr.txt
 if %ERRORLEVEL% == 0 goto :start_venv
+:: If pip is not available and PIP_INSTALLER_LOCATION is set, try to install pip
 if "%PIP_INSTALLER_LOCATION%" == "" goto :show_stdout_stderr
 %PYTHON% "%PIP_INSTALLER_LOCATION%" >tmp/stdout.txt 2>tmp/stderr.txt
 if %ERRORLEVEL% == 0 goto :start_venv
@@ -23,33 +39,106 @@ echo Couldn't install pip
 goto :show_stdout_stderr
 
 :start_venv
-if ["%VENV_DIR%"] == ["-"] goto :skip_venv
-if ["%SKIP_VENV%"] == ["1"] goto :skip_venv
+:: Skip venv creation/activation if VENV_DIR is explicitly set to "-"
+if ["%VENV_DIR%"] == ["-"] goto :skip_venv_entirely
+:: Skip venv creation/activation if SKIP_VENV is set to "1"
+if ["%SKIP_VENV%"] == ["1"] goto :skip_venv_entirely
 
+:: Check if the venv already exists by looking for Python.exe in its Scripts directory
 dir "%VENV_DIR%\Scripts\Python.exe" >tmp/stdout.txt 2>tmp/stderr.txt
-if %ERRORLEVEL% == 0 goto :activate_venv
+if %ERRORLEVEL% == 0 goto :activate_venv_and_maybe_update
 
+:: Venv does not exist, create it
+echo Virtual environment not found in "%VENV_DIR%". Creating a new one.
 for /f "delims=" %%i in ('CALL %PYTHON% -c "import sys; print(sys.executable)"') do set PYTHON_FULLNAME="%%i"
 echo Creating venv in directory %VENV_DIR% using python %PYTHON_FULLNAME%
 %PYTHON_FULLNAME% -m venv "%VENV_DIR%" >tmp/stdout.txt 2>tmp/stderr.txt
-if %ERRORLEVEL% == 0 goto :activate_venv
-echo Unable to create venv in directory "%VENV_DIR%"
-goto :show_stdout_stderr
-
-:activate_venv
+if %ERRORLEVEL% NEQ 0 (
+    echo Unable to create venv in directory "%VENV_DIR%"
+    goto :show_stdout_stderr
+)
+echo Venv created.
+
+:: Install requirements for the first time if venv was just created
+:: This section handles the initial installation of packages from requirements.txt
+:: immediately after a new virtual environment is created.
+echo Checking for requirements.txt for initial setup in %~dp0
+if exist "%~dp0requirements.txt" (
+    echo Found requirements.txt, attempting to install for initial setup...
+    call "%VENV_DIR%\Scripts\activate.bat"
+    echo Installing packages from requirements.txt ^(initial setup^)...
+    "%VENV_DIR%\Scripts\python.exe" -m pip install -r "%~dp0requirements.txt"
+    if %ERRORLEVEL% NEQ 0 (
+        echo Failed to install requirements during initial setup. Please check the output above.
+        pause
+        goto :show_stdout_stderr_custom_pip_initial 
+    )
+    echo Initial requirements installed successfully.
+    call "%VENV_DIR%\Scripts\deactivate.bat"
+) else (
+    echo No requirements.txt found for initial setup, skipping package installation.
+)
+goto :activate_venv_and_maybe_update
+
+
+:activate_venv_and_maybe_update
+:: This label is reached if the venv exists or was just created.
+:: Set PYTHON to point to the venv's Python interpreter.
 set PYTHON="%VENV_DIR%\Scripts\Python.exe"
-echo venv %PYTHON%
+echo Activating venv: %PYTHON%
+
+:: Always update requirements if ALWAYS_UPDATE_REQS is 1
+:: This section allows for updating packages from requirements.txt on every launch
+:: if the ALWAYS_UPDATE_REQS variable is set to 1.
+if defined ALWAYS_UPDATE_REQS (
+    if "%ALWAYS_UPDATE_REQS%"=="1" (
+        echo ALWAYS_UPDATE_REQS is enabled.
+        if exist "%~dp0requirements.txt" (
+            echo Attempting to update packages from requirements.txt...
+            REM No need to call activate.bat here again, PYTHON is already set to the venv's python
+            %PYTHON% -m pip install -r "%~dp0requirements.txt"
+            if %ERRORLEVEL% NEQ 0 (
+                echo Failed to update requirements. Please check the output above.
+                pause
+                goto :endofscript 
+            )
+            echo Requirements updated successfully.
+        ) else (
+            echo ALWAYS_UPDATE_REQS is enabled, but no requirements.txt found. Skipping update.
+        )
+    ) else (
+        echo ALWAYS_UPDATE_REQS is not enabled or not set to 1. Skipping routine update.
+    )
+)
 
-:skip_venv
+goto :launch
+
+:skip_venv_entirely
+:: This label is reached if venv usage is explicitly skipped.
+echo Skipping venv.
 goto :launch
 
 :launch
-%PYTHON% app.py %COMMANDLINE_ARGS% %*
+:: Launch the main application
+echo Launching Web UI with arguments: %COMMANDLINE_ARGS% %*
+%PYTHON% %APPLICATION_NAME% %COMMANDLINE_ARGS% %*
+echo Launch finished.
 pause
 exit /b
 
-:show_stdout_stderr
+:show_stdout_stderr_custom_pip_initial
+:: Custom error handler for failures during the initial pip install process.
+echo.
+echo exit code ^(pip initial install^): %errorlevel%
+echo Errors during initial pip install. See output above.
+echo.
+echo Launch unsuccessful. Exiting.
+pause
+exit /b
 
+
+:show_stdout_stderr
+:: General error handler: displays stdout and stderr from the tmp directory.
 echo.
 echo exit code: %errorlevel%
 
@@ -61,13 +150,13 @@ type tmp\stdout.txt
 
 :show_stderr
 for /f %%i in ("tmp\stderr.txt") do set size=%%~zi
-if %size% equ 0 goto :show_stderr
+if %size% equ 0 goto :endofscript
 echo.
 echo stderr:
 type tmp\stderr.txt
 
 :endofscript
-
 echo.
 echo Launch unsuccessful. Exiting.
 pause
+exit /b