Update app.py

app.py

@@ -20,6 +20,8 @@ from PIL import Image
 from functools import partial
 import io
 
+app = FastAPI()
+
 
 subprocess.run(shlex.split('pip install wheel/torchmcubes-0.1.0-cp310-cp310-linux_x86_64.whl'))
 
@@ -34,6 +36,9 @@ if torch.cuda.is_available():
 else:
     device = "cpu"
 
+torch.cuda.synchronize()
+
+
 model = TSR.from_pretrained(
     "stabilityai/TripoSR",
     config_name="config.yaml",
@@ -47,17 +52,30 @@ 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'])
+
+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')
     
-#     return Image.open(io.BytesIO(base64_image_data))
+    if object_name is None:
+        object_name = file_path
+    
+    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}.")
+    return True
 
 
 def gen_pos_prompt(text):
@@ -164,88 +182,47 @@ def generate(image, mc_resolution, formats=["obj", "glb"]):
     return mesh_path_obj.name, mesh_path_glb.name
 
 def run_example(text_prompt,seed ,do_remove_background, foreground_ratio, mc_resolution):
-    # Step 1: Generate the image from text prompt
     image_pil = generate_image_from_text(text_prompt, seed)
-    
-    # Step 2: Preprocess the image
+
     preprocessed = preprocess(image_pil, do_remove_background, foreground_ratio)
     
-    # Step 3: Generate the 3D model
     mesh_name_obj, mesh_name_glb = generate(preprocessed, mc_resolution, ["obj", "glb"])
     
     return preprocessed, mesh_name_obj, mesh_name_glb
 
-with gr.Blocks() as demo:
-    gr.Markdown(HEADER)
-    with gr.Row(variant="panel"):
-        with gr.Column():
-            with gr.Row():
-                text_prompt = gr.Textbox(
-                    label="Text Prompt",
-                    placeholder="Enter a text prompt for image generation"
-                )
-                input_image = gr.Image(
-                    label="Generated Image",
-                    image_mode="RGBA",
-                    sources="upload",
-                    type="pil",
-                    elem_id="content_image",
-                    visible=False  # Hidden since we generate the image from text
-                )
-                seed = gr.Number(value=0)
-                processed_image = gr.Image(label="Processed Image", interactive=False, visible=False)
-            with gr.Row():
-                with gr.Group():
-                    do_remove_background = gr.Checkbox(
-                        label="Remove Background", value=True
-                    )
-                    foreground_ratio = gr.Slider(
-                        label="Foreground Ratio",
-                        minimum=0.5,
-                        maximum=1.0,
-                        value=0.85,
-                        step=0.05,
-                    )
-                    mc_resolution = gr.Slider(
-                        label="Marching Cubes Resolution",
-                        minimum=32,
-                        maximum=320,
-                        value=256,
-                        step=32
-                     )
-            with gr.Row():
-                submit = gr.Button("Generate", elem_id="generate", variant="primary")
-        with gr.Column():
-            with gr.Tab("OBJ"):
-                output_model_obj = gr.Model3D(
-                    label="Output Model (OBJ Format)",
-                    interactive=False,
-                )
-                gr.Markdown("Note: Downloaded object will be flipped in case of .obj export. Export .glb instead or manually flip it before usage.")
-            with gr.Tab("GLB"):
-                output_model_glb = gr.Model3D(
-                    label="Output Model (GLB Format)",
-                    interactive=False,
-                )
-                gr.Markdown("Note: The model shown here has a darker appearance. Download to get correct results.")
-    # with gr.Row(variant="panel"):
-    #     gr.Examples(
-    #         examples=[
-    #             os.path.join("examples", img_name) for img_name in sorted(os.listdir("examples"))
-    #         ],
-    #         inputs=[text_prompt],
-    #         outputs=[processed_image, output_model_obj, output_model_glb],
-    #         cache_examples=True,
-    #         fn=partial(run_example, do_remove_background=True, foreground_ratio=0.85, mc_resolution=256),
-    #         label="Examples",
-    #         examples_per_page=20
-    #     )
-    submit.click(fn=check_input_image, inputs=[text_prompt]).success(
-        fn=run_example,
-        inputs=[text_prompt, seed, do_remove_background, foreground_ratio, mc_resolution],
-        outputs=[processed_image, output_model_obj, output_model_glb],
-        # outputs=[output_model_obj, output_model_glb],
-    )
-
-demo.queue(max_size=10)
-demo.launch(auth=(os.getenv('USERNAME'), os.getenv('PASSWORD')))
+
+@app.post("/process_text/")
+async def process_image(
+    text_prompt: str = Form(...),
+    seed: int = Form(...),
+    do_remove_background: bool = Form(...),
+    foreground_ratio: float = Form(...),
+    mc_resolution: int = Form(...),
+    auth: str = Form(...)
+):
+    
+    if auth == os.getenv("AUTHORIZE"):
+
+        preprocessed, mesh_name_obj, mesh_name_glb = run_example(text_prompt,seed ,do_remove_background, foreground_ratio, mc_resolution)
+        # 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 {"Authentication":"Failed"}
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)