Update app.py

app.py

@@ -23,6 +23,8 @@ import io
 from io import BytesIO
 from botocore.exceptions import NoCredentialsError, PartialCredentialsError
 import datetime
+from transformers.utils import move_cache
+move_cache()
 
 app = FastAPI()
 
@@ -54,6 +56,7 @@ 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')
+s3_client = boto3.client('s3',aws_access_key_id = ACCESS, aws_secret_access_key = SECRET, region_name='us-east-1')
 
 
 def gen_pos_prompt(text):
@@ -114,7 +117,7 @@ def generate_image_from_text(encoded_image, seed, pos_prompt=None):
               },
               'imageGenerationConfig': {
                   "cfgScale": 8,
-                  "seed": int(seed),
+                  "seed": seed,
                   "width": 512,
                   "height": 512,
                   "numberOfImages": 1
@@ -130,7 +133,7 @@ def generate_image_from_text(encoded_image, seed, pos_prompt=None):
           },
           'imageGenerationConfig': {
               "cfgScale": 8,
-              "seed": int(seed),
+              "seed": seed,
               "width": 512,
               "height": 512,
               "numberOfImages": 1
@@ -147,20 +150,25 @@ def check_input_image(input_image):
         raise gr.Error("No image uploaded!")
 
 def preprocess(input_image, do_remove_background, foreground_ratio):
-    torch.cuda.synchronize()
     def fill_background(image):
+        torch.cuda.synchronize()
+        torch.cuda.empty_cache()
         image = np.array(image).astype(np.float32) / 255.0
         image = image[:, :, :3] * image[:, :, 3:4] + (1 - image[:, :, 3:4]) * 0.5
         image = Image.fromarray((image * 255.0).astype(np.uint8))
         return image
 
     if do_remove_background:
+        torch.cuda.synchronize()
+        torch.cuda.empty_cache()
         image = input_image.convert("RGB")
         image = remove_background(image, rembg_session)
         image = resize_foreground(image, foreground_ratio)
         image = fill_background(image)
-        print("do_remove_back")
+        
+        torch.cuda.synchronize()
         torch.cuda.empty_cache()
+        print("Background Removed")
     else:
         image = input_image
         if image.mode == "RGBA":
@@ -175,40 +183,29 @@ def generate(image, mc_resolution, formats=["obj", "glb"]):
     scene_codes = model(image, device=device)
     torch.cuda.synchronize()
     mesh = model.extract_mesh(scene_codes, resolution=mc_resolution)[0]
+    torch.cuda.synchronize()
     mesh = to_gradio_3d_orientation(mesh)
+    torch.cuda.synchronize()
 
     mesh_path_glb = tempfile.NamedTemporaryFile(suffix=f".glb", delete=False)
+    torch.cuda.synchronize()
     mesh.export(mesh_path_glb.name)
+    torch.cuda.synchronize()
 
     mesh_path_obj = tempfile.NamedTemporaryFile(suffix=f".obj", delete=False)
+    torch.cuda.synchronize()
     mesh.apply_scale([-1, 1, 1])  # Otherwise the visualized .obj will be flipped
     mesh.export(mesh_path_obj.name)
     torch.cuda.synchronize()  # Ensure all CUDA operations are complete before clearing cache
     torch.cuda.empty_cache()
     return mesh_path_obj.name, mesh_path_glb.name
 
-def upload_file_to_s3(file_path, bucket_name, object_name=None):
-    s3_client = boto3.client('s3',aws_access_key_id = ACCESS, aws_secret_access_key = SECRET, region_name='us-east-1')
-    
-    if object_name is None:
-        object_name = file_path
+def upload_file_to_s3(file_path, bucket_name, object_name=None):    
+    s3_client.upload_file(file_path, bucket_name, object_name)
     
-    try:
-        s3_client.upload_file(file_path, bucket_name, object_name)
-    except FileNotFoundError:
-        print(f"The file {file_path} was not found.")
-        return False
-    except NoCredentialsError:
-        print("Credentials not available.")
-        return False
-    except PartialCredentialsError:
-        print("Incomplete credentials provided.")
-        return False
-    except Exception as e:
-        print(f"An error occurred: {e}")
-        return False
-    
-    print(f"File {file_path} uploaded successfully to {bucket_name}/{object_name}.")
+    # print(f"File {file_path} uploaded successfully to {bucket_name}/{object_name}.")
+    torch.cuda.synchronize()  # Wait for all CUDA operations to complete
+    torch.cuda.empty_cache()
     return True
 
     
@@ -220,8 +217,8 @@ async def process_image(
     enhance_image: bool = Form(False),  # Default enhance_image value
     do_remove_background: bool = Form(True),  # Default do_remove_background value
     foreground_ratio: float = Form(0.85, gt=0.0, lt=1.0),  # Ratio must be between 0.0 and 1.0 (exclusive)
-    mc_resolution: int = Form(285, ge=32, le=320),  # Resolution must be between 256 and 4096
-    auth: str = Form("default_auth_token"),  # Default auth token
+    mc_resolution: int = Form(286, ge=32, le=320),  # Resolution must be between 256 and 4096
+    auth: str = Form(...),  
     text_prompt: Optional[str] = Form(None)
 ):