Update main.py

main.py

@@ -10,7 +10,7 @@ from stability_sdk import client
 import stability_sdk.interfaces.gooseai.generation.generation_pb2 as generation
 import replicate
 
-def generate_and_upscale_image(text_prompt, clipdrop_api_key, stability_api_key, replicate_api_token, download_path):
+def generate_and_upscale_image(text_prompt, clipdrop_api_key, stability_api_key, replicate_api_token):
 
     headers = {'x-api-key': clipdrop_api_key}
     body_params = {'prompt': (None, text_prompt, 'text/plain')}
@@ -68,4 +68,8 @@ def generate_and_upscale_image(text_prompt, clipdrop_api_key, stability_api_key,
 
     response = requests.get(output)
     final_img = Image.open(io.BytesIO(response.content))
-    final_img.save(os.path.join(download_path, "gfpgan_upscaled_image.png"))
+    img_byte_arr = io.BytesIO()
+    final_img.save(img_byte_arr, format='PNG')
+    img_byte_arr = img_byte_arr.getvalue()
+
+    return img_byte_arr