Update main.py

main.py

@@ -58,19 +58,46 @@ def upscale_image_esrgan(image_bytes):
     return upscaled_img_bytes
 
 def further_upscale_image(image_bytes):
+    # Ensure environment variable is set correctly
+    print("Replicate API token: ", os.environ['REPLICATE_API_TOKEN'])
+
+    # Save the image bytes to a temporary file
+    temp_file_name = "temp.png"
+    with open(temp_file_name, 'wb') as temp_file:
+        temp_file.write(image_bytes)
+
     # Run the GFPGAN model
-    output = replicate.run(
-        "tencentarc/gfpgan:9283608cc6b7be6b65a8e44983db012355fde4132009bf99d976b2f0896856a3",
-        input={"img": BytesIO(image_bytes), "version": "v1.4", "scale": 16}
-    )
-    
-    # The output is a URI of the processed image
-    # We will retrieve the image data and save it
-    response = requests.get(output)
-    img = Image.open(BytesIO(response.content))
-    img.save("upscaled.png")  # Save the upscaled image
+    try:
+        print("Running GFPGAN model...")
+        output = replicate.run(
+            "tencentarc/gfpgan:9283608cc6b7be6b65a8e44983db012355fde4132009bf99d976b2f0896856a3",
+            input={"img": open(temp_file_name, "rb"), "version": "v1.4", "scale": 16}
+        )
+        print("Model output: ", output)
+    except Exception as e:
+        print("Error running GFPGAN model: ", e)
+        raise e
+
+    # Get the image data from the output URI
+    try:
+        print("Fetching image data from output URI...")
+        response = requests.get(output)
+    except Exception as e:
+        print("Error fetching image data from output URI: ", e)
+        raise e
+
+    # Open and save the image
+    try:
+        print("Saving upscaled image...")
+        img = Image.open(BytesIO(response.content))
+        img.save("upscaled.png")  # Save the upscaled image
+    except Exception as e:
+        print("Error saving upscaled image: ", e)
+        raise e
+
     return img
 
+
 def main():
     st.title("Image Generation and Upscaling")
     st.write("Enter a text prompt and an image will be generated and upscaled.")