Update app.py

app.py

@@ -23,17 +23,13 @@ from io import BytesIO
 from botocore.exceptions import NoCredentialsError, PartialCredentialsError
 import datetime
 
-
-
 app = FastAPI()
 
-
 subprocess.run(shlex.split('pip install wheel/torchmcubes-0.1.0-cp310-cp310-linux_x86_64.whl'))
 
 from tsr.system import TSR
 from tsr.utils import remove_background, resize_foreground, to_gradio_3d_orientation
 
-# os.environ["CUDA_VISIBLE_DEVICES"] = "1"
 
 HEADER = """FRAME AI"""
 
@@ -57,17 +53,6 @@ ACCESS = os.getenv("ACCESS")
 SECRET = os.getenv("SECRET")
 bedrock = boto3.client(service_name='bedrock', aws_access_key_id = ACCESS, aws_secret_access_key = SECRET, region_name='us-east-1')
 bedrock_runtime = boto3.client(service_name='bedrock-runtime', aws_access_key_id = ACCESS, aws_secret_access_key = SECRET, region_name='us-east-1')
-# def generate_image_from_text(pos_prompt):
-#     # bedrock_runtime = boto3.client(region_name = 'us-east-1', service_name='bedrock-runtime')
-#     parameters = {'text_prompts': [{'text': pos_prompt , 'weight':1},
-#                                     {'text': """Blurry, out of frame, out of focus, Detailed, dull, duplicate, bad quality, low resolution, cropped""", 'weight': -1}],
-#                   'cfg_scale': 7, 'seed': 0, 'samples': 1}
-#     request_body = json.dumps(parameters)
-#     response = bedrock_runtime.invoke_model(body=request_body,modelId = 'stability.stable-diffusion-xl-v1')
-#     response_body = json.loads(response.get('body').read())
-#     base64_image_data = base64.b64decode(response_body['artifacts'][0]['base64'])
-    
-#     return Image.open(io.BytesIO(base64_image_data))
 
 
 def gen_pos_prompt(text):
@@ -231,31 +216,36 @@ async def process_image(
     do_remove_background: bool = Form(...),
     foreground_ratio: float = Form(...),
     mc_resolution: int = Form(...),
+    auth: str = Form(...),
     text_prompt: Optional[str] = Form(None)
 ):
-    image_bytes = await file.read()
-    input_image = Image.open(BytesIO(image_bytes))
-
-    if enhance_image:
-        image_pil = generate_image_from_text(encoded_image=input_image, seed=seed, pos_prompt=text_prompt)
-    else:
-        image_pil = input_image
-
-    preprocessed = preprocess(image_pil, do_remove_background, foreground_ratio)
-    mesh_name_obj, mesh_name_glb = generate(preprocessed, mc_resolution)
-    timestamp = datetime.datetime.now().strftime('%Y%m%d%H%M%S%f')
-    object_name = f'object_{timestamp}_1.obj'
-    object_name_2 = f'object_{timestamp}_2.glb'
-
-    if upload_file_to_s3(mesh_name_obj, 'framebucket3d',object_name) and upload_file_to_s3(mesh_name_glb, 'framebucket3d',object_name_2):
-
-        return {
-            "obj_path": f"https://framebucket3d.s3.amazonaws.com/{object_name}",
-            "glb_path": f"https://framebucket3d.s3.amazonaws.com/{object_name_2}"
-        }
-
+    
+    if auth == os.getenv("AUTHORIZE"):
+        image_bytes = await file.read()
+        input_image = Image.open(BytesIO(image_bytes))
+    
+        if enhance_image:
+            image_pil = generate_image_from_text(encoded_image=input_image, seed=seed, pos_prompt=text_prompt)
+        else:
+            image_pil = input_image
+    
+        preprocessed = preprocess(image_pil, do_remove_background, foreground_ratio)
+        mesh_name_obj, mesh_name_glb = generate(preprocessed, mc_resolution)
+        timestamp = datetime.datetime.now().strftime('%Y%m%d%H%M%S%f')
+        object_name = f'object_{timestamp}_1.obj'
+        object_name_2 = f'object_{timestamp}_2.glb'
+    
+        if upload_file_to_s3(mesh_name_obj, 'framebucket3d',object_name) and upload_file_to_s3(mesh_name_glb, 'framebucket3d',object_name_2):
+    
+            return {
+                "obj_path": f"https://framebucket3d.s3.amazonaws.com/{object_name}",
+                "glb_path": f"https://framebucket3d.s3.amazonaws.com/{object_name_2}"
+            }
+    
+        else:
+            return {"Internal Server Error": False}
     else:
-        return {"Internal Server Error": False}
+        return {"Authentication":"Failed"}
 
 if __name__ == "__main__":
     import uvicorn