Upload 2 files

app/upscale.py

@@ -0,0 +1,62 @@
+import replicate
+from PIL import Image 
+import asyncio
+import requests
+from io import BytesIO
+#from dotenv import load_dotenv
+import os
+
+#load_dotenv("../.env")
+
+class ImageUpscaler():
+    
+    
+    def pil2byte(self, image):
+        byte_obj = BytesIO()
+        try:
+            image.save(byte_obj, format='PNG')
+            byte_obj.seek(0)
+        except Exception as e:
+            print(f"Unable to process the image: {e}")
+        return byte_obj
+    
+            
+    def img2byte(self, img_path):
+        with Image.open(img_path) as img:
+
+            img_byte_arr = BytesIO()
+
+            ext = os.path.splitext(img_path)[1].replace('.','').upper()
+            if ext == 'JPG':
+                ext = 'JPEG'
+            img.save(img_byte_arr, format=ext)
+            img_byte_arr.seek(0) 
+            return img_byte_arr
+        
+
+    
+    async def upscaler(self, target_img=None, factor=4, byte_stream=True):
+        try:
+            model = replicate.models.get("nightmareai/real-esrgan")
+            version = model.versions.get("42fed1c4974146d4d2414e2be2c5277c7fcf05fcc3a73abf41610695738c1d7b")
+            
+            if not byte_stream:
+                inputs = {
+                            'image': open(target_img,'rb'),
+
+                        }
+            else:
+                inputs = {
+                            'image': target_img,
+
+                        }
+            inputs['scale'] = factor
+            inputs['face_enhance'] = True
+            img_output = version.predict(**inputs)
+            upscaled_response = requests.get(img_output)
+            upscaled_bin = BytesIO(upscaled_response.content)
+        
+        except Exception as e:
+            return f"An error occurred during face image generation: {str(e)}"
+        
+        return upscaled_bin