Update main.py

main.py

@@ -1,52 +1,49 @@
-from fastapi import FastAPI, File, UploadFile, Form
+from fastapi import FastAPI, File, UploadFile, HTTPException
 from fastapi.responses import JSONResponse
-from fastapi.middleware.cors import CORSMiddleware
 from gradio_client import Client, handle_file
 import os
-import tempfile
 import base64
 
+# Initialize FastAPI
 app = FastAPI()
 
-# Setup CORS
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],  # Adjust as needed
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
+# Load environment variables if needed
 HF_TOKEN = os.getenv("HF_TOKEN")
 
 # Initialize the Gradio client with the token
-client = Client("Makhinur/Image_Face_Upscale_Restoration-GFPGAN", hf_token=HF_TOKEN)
+client = Client("Makhinur/Its", hf_token=HF_TOKEN)
 
 @app.post("/upload/")
-async def upload_file(file: UploadFile = File(...), version: str = Form(...), scale: int = Form(...)):
-    # Save the uploaded file temporarily
-    with tempfile.NamedTemporaryFile(delete=False) as temp_file:
-        temp_file.write(await file.read())
-        temp_file_path = temp_file.name
-
+async def upload_image(file: UploadFile = File(...)):
     try:
-        # Use handle_file to prepare the file for the Gradio client
-        result = client.predict(handle_file(temp_file_path), version, scale, api_name="/predict")
-
-        # Check if the result is valid
-        if result and len(result) == 2:
-            # Convert the image data to a base64 string
-            with open(result[0], "rb") as image_file:
-                image_data = base64.b64encode(image_file.read()).decode("utf-8")
-
-            return JSONResponse({
-                "sketch_image_base64": f"data:image/png;base64,{image_data}",
-                "result_file": result[1]
-            })
-        else:
-            return JSONResponse({"error": "Invalid result from the prediction API."}, status_code=500)
-    except Exception as e:
-        return JSONResponse({"error": str(e)}, status_code=500)
-    finally:
+        # Save the uploaded file to a temporary location
+        file_location = f"temp_{file.filename}"
+        with open(file_location, "wb") as f:
+            f.write(await file.read())
+
+        # Use the Gradio client to send the image to the Gradio app
+        result = client.predict(
+            img=handle_file(file_location),
+            api_name="/predict"
+        )
+
         # Clean up the temporary file
-        if os.path.exists(temp_file_path):
-            os.unlink(temp_file_path)
+        os.remove(file_location)
+
+        # Read the sketch image from the result
+        sketch_image_path = result[0]  # Path to the sketch image
+
+        # Convert the sketch image to base64
+        with open(sketch_image_path, "rb") as img_file:
+            sketch_image_base64 = base64.b64encode(img_file.read()).decode('utf-8')
+
+        # Prepare the response
+        response = {
+            "sketch_image_base64": f"data:image/jpeg;base64,{sketch_image_base64}"
+        }
+
+        return JSONResponse(content=response)
+
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+