Update app.py

app.py

@@ -1,38 +1,51 @@
 import gradio as gr
-from huggingface_hub import from_pretrained_keras
-import tensorflow as tf
-import numpy as np
+from transformers import pipeline
+import torch
 from PIL import Image
 import io
 
-# Load the model from Hugging Face Hub using from_pretrained_keras
-model = from_pretrained_keras("google/maxim-s3-deblurring-reds")
+# Check for CUDA availability and set device
+if torch.cuda.is_available():
+    device = torch.device("cuda")
+    print(f"Using CUDA device: {torch.cuda.get_device_name(0)}")
+else:
+    device = torch.device("cpu")
+    print("Using CPU. CUDA is not available.")
+
+try:
+    # Initialize the deblurring pipeline with the specified model and device
+    deblurrer = pipeline("image-to-image", model="google/maxim-s3-deblurring-reds", device=device)
+except Exception as e:
+    print(f"Error loading the model: {e}")
+    exit()  # Exit if model loading fails
 
 def deblur_image(input_image):
-    # Preprocess the input image
-    image = np.array(input_image)
-    image = tf.convert_to_tensor(image)
-    image = tf.image.resize(image, (256, 256))
-    
-    # Make predictions with the model
-    predictions = model.predict(tf.expand_dims(image, 0))
-    
-    # Convert the prediction back to an image
-    output_image = predictions[0].numpy().astype(np.uint8)
-    output_image = Image.fromarray(output_image)
-    
-    # Save the result in memory as a byte array
-    byte_arr = io.BytesIO()
-    output_image.save(byte_arr, format='PNG')
-    byte_arr.seek(0)
-    
-    return byte_arr
+    try:
+        output = deblurrer(input_image)
+        deblurred_image = output[0]
+
+        # Convert PIL Image to Bytes for download
+        img_byte_arr = io.BytesIO()
+        deblurred_image.save(img_byte_arr, format='PNG')  # or JPEG, etc.
+        img_byte_arr = img_byte_arr.getvalue()
+
+        return deblurred_image, img_byte_arr  # Return both image and bytes
+    except Exception as e:
+        print(f"Error during deblurring: {e}")
+        return None, None
 
-# Set up the Gradio interface
-iface = gr.Interface(fn=deblur_image,
-                     inputs=gr.inputs.Image(type="pil"),
-                     outputs=gr.outputs.Image(type="file"),
-                     live=True)
+# Create the Gradio interface
+iface = gr.Interface(
+    fn=deblur_image,
+    inputs=gr.Image(type="pil", label="Upload Blurred Image"),
+    outputs=[
+        gr.Image(type="pil", label="Deblurred Image"),
+        gr.File(label="Download Deblurred Image", file_types=[".png", ".jpg", ".jpeg"]) # Added File output
+    ],
+    title="Deblurring App",
+    description="Deblur your images using the google/maxim-s3-deblurring-reds model.",
+    examples=[["blurred_image.jpg"]],
+)
 
-# Launch the Gradio interface
+# Launch the Gradio app
 iface.launch()