Update app.py

app.py

@@ -1,13 +1,12 @@
-
 import torch
 from PIL import Image
 from RealESRGAN import RealESRGAN
 import gradio as gr
 import numpy as np
-import tempfile
-import time
+import io
 import zipfile
 import os
+import time
 
 # Set the device to CUDA if available, otherwise CPU
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
@@ -20,7 +19,6 @@ def load_model(scale):
         print(f"Weights for scale {scale} loaded successfully.")
     except Exception as e:
         print(f"Error loading weights for scale {scale}: {e}")
-        model.load_weights(weights_path, download=False)
     return model
 
 # Load models for different scales
@@ -33,15 +31,8 @@ def enhance_image(image, scale):
         print(f"Enhancing image with scale {scale}...")
         start_time = time.time()
         image_np = np.array(image.convert('RGB'))
-        print(f"Image converted to numpy array: shape {image_np.shape}, dtype {image_np.dtype}")
-        
-        if scale == '2x':
-            result = model2.predict(image_np)
-        elif scale == '4x':
-            result = model4.predict(image_np)
-        else:
-            result = model8.predict(image_np)
-            
+        model = model2 if scale == '2x' else model4 if scale == '4x' else model8
+        result = model.predict(image_np)
         enhanced_image = Image.fromarray(np.uint8(result))
         print(f"Image enhanced in {time.time() - start_time:.2f} seconds")
         return enhanced_image
@@ -49,55 +40,48 @@ def enhance_image(image, scale):
         print(f"Error enhancing image: {e}")
         return image
 
-def muda_dpi(input_image, dpi):
-    dpi_tuple = (dpi, dpi)
-    image = Image.fromarray(input_image.astype('uint8'), 'RGB')
-    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.jpg')
-    image.save(temp_file, format='JPEG', dpi=dpi_tuple)
-    temp_file.close()
-    return Image.open(temp_file.name)
+def muda_dpi(image, dpi):
+    try:
+        with io.BytesIO() as output:
+            image.save(output, format='JPEG', dpi=(dpi, dpi))
+            return Image.open(output)
+    except Exception as e:
+        print(f"Error adjusting DPI: {e}")
+        return image
 
-def resize_image(input_image, width, height):
-    image = Image.fromarray(input_image.astype('uint8'), 'RGB')
-    resized_image = image.resize((width, height))
-    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.jpg')
-    resized_image.save(temp_file, format='JPEG')
-    temp_file.close()
-    return Image.open(temp_file.name)
+def resize_image(image, width, height):
+    try:
+        resized_image = image.resize((width, height))
+        return resized_image
+    except Exception as e:
+        print(f"Error resizing image: {e}")
+        return image
 
 def process_images(image_files, enhance, scale, adjust_dpi, dpi, resize, width, height):
     processed_images = []
-    temp_dir = tempfile.mkdtemp()
+    zip_buffer = io.BytesIO()
     
     for image_file in image_files:
-        input_image = np.array(Image.open(image_file).convert('RGB'))
-        original_image = Image.fromarray(input_image.astype('uint8'), 'RGB')
+        image = Image.open(image_file).convert('RGB')
         
         if enhance:
-            original_image = enhance_image(original_image, scale)
+            image = enhance_image(image, scale)
         
         if adjust_dpi:
-            original_image = muda_dpi(np.array(original_image), dpi)
-            
+            image = muda_dpi(image, dpi)
+        
         if resize:
-            original_image = resize_image(np.array(original_image), width, height)
+            image = resize_image(image, width, height)
         
-        # Save each image as JPEG, preserving the original filename
-        file_name = os.path.basename(image_file.name)
-        output_path = os.path.join(temp_dir, file_name)
-        original_image.save(output_path, format='JPEG')
-        processed_images.append(output_path)
-    
-    # Create a ZIP file with all processed images
-    zip_path = os.path.join(temp_dir, 'processed_images.zip')
-    with zipfile.ZipFile(zip_path, 'w') as zipf:
-        for file_path in processed_images:
-            zipf.write(file_path, os.path.basename(file_path))
+        # Save image to the in-memory ZIP buffer
+        buffer = io.BytesIO()
+        image.save(buffer, format='JPEG')
+        processed_images.append(Image.open(io.BytesIO(buffer.getvalue())))
+        with zipfile.ZipFile(zip_buffer, 'a') as zipf:
+            zipf.writestr(os.path.basename(image_file.name), buffer.getvalue())
     
-    # Load images for display in the gallery
-    display_images = [Image.open(img_path) for img_path in processed_images]
-    
-    return display_images, zip_path
+    zip_buffer.seek(0)
+    return processed_images, zip_buffer
 
 iface = gr.Interface(
     fn=process_images,
@@ -120,6 +104,3 @@ iface = gr.Interface(
 )
 
 iface.launch(debug=True)
-
-
-