Update app.py

app.py

@@ -1,9 +1,12 @@
 import gradio as gr
 import os
-import subprocess
 import torch
 from PIL import Image
-import time
+import logging
+from NAFNet.demo import run_model  # Import the model directly instead of using subprocess
+
+# Setup logging
+logging.basicConfig(level=logging.INFO)
 
 UPLOAD_FOLDER = 'uploads'
 OUTPUT_FOLDER = 'outputs'
@@ -19,44 +22,28 @@ def gradio_interface(image):
     output_path = os.path.join(OUTPUT_FOLDER, "output.png")
 
     image.save(input_path)
-    print(f"Input image saved at: {input_path}")
+    logging.info(f"Input image saved at: {input_path}")
 
     try:
         # Ensure CUDA memory is freed before running inference
         torch.cuda.empty_cache()
 
-        # Run the NAFNet model with controlled memory usage
-        command = [
-            "python", "NAFNet/demo.py",
-            "-opt", "NAFNet/options/test/REDS/NAFNet-width64.yml",
-            "--input_path", input_path,
-            "--output_path", output_path
-        ]
+        logging.info("Running model...")
         
-        print("Running model...")
-        result = subprocess.run(command, capture_output=True, text=True)
-        print("Model execution completed.")
-
-        # Log the output
-        print("STDOUT:", result.stdout)
-        print("STDERR:", result.stderr)
-
-        if result.returncode != 0:
-            return f"Error: {result.stderr}"
-
-        # Wait for output file to be generated
-        time.sleep(2)  # Give time for file system updates
-        if not os.path.exists(output_path):
-            return "Error: Output image not generated."
+        # Run NAFNet directly instead of subprocess
+        run_model(
+            opt_path="NAFNet/options/test/REDS/NAFNet-width64.yml",
+            input_path=input_path,
+            output_path=output_path
+        )
 
-        print(f"Output image generated: {output_path}")
+        logging.info("Model execution completed.")
 
-        # Explicitly reload the image to avoid caching issues
-        return Image.open(output_path).copy()
+        return Image.open(output_path)
 
     except Exception as e:
-        print(f"Exception occurred: {str(e)}")
-        return f"Exception: {str(e)}"
+        logging.error(f"Exception: {str(e)}")
+        return f"Error: {str(e)}"
 
 # Launch Gradio
 iface = gr.Interface(
@@ -66,5 +53,4 @@ iface = gr.Interface(
     title="Image Restoration with NAFNet"
 )
 
-print("Starting Gradio interface...")
-iface.launch()
+iface.launch()  # No `share=True` for Hugging Face Spaces