Update app.py

app.py

@@ -22,9 +22,57 @@ DEFAULT_NEGATIVE_PROMPT = (
 
 @spaces.GPU
 def run(*args):
-    id_image = args[0]
-    supp_images = args[1:4]
-    prompt, neg_prompt, scale, n_samples, seed, steps, H, W, id_scale, mode, id_mix = args[4:]
+
+    aws_access_key_id = 'AKIA2NMAMYX4K55CZ7HR'
+    BUCKET = 'syntheticai-headshots'
+    s3_client = boto3.client(
+            's3',
+            aws_access_key_id=aws_access_key_id,
+            aws_secret_access_key=os.getenv('AMAZON_SECRET_KEY')
+        )
+    
+    INPUT_BUCKET_FOLDER = bucket_folder #user_id/request_id/input/'
+    local_dir = req_id
+    os.makedirs(req_id, exist_ok=True)
+
+    # try:
+    response = s3_client.list_objects_v2(Bucket=BUCKET, Prefix=INPUT_BUCKET_FOLDER)
+
+    if 'Contents' in response:
+        for obj in response['Contents']:
+            s3_key = obj['Key']
+
+            if s3_key.endswith('/'):
+                continue
+
+            file_name = os.path.basename(s3_key)
+            local_path = os.path.join(local_dir, file_name)
+
+            s3_client.download_file(BUCKET, s3_key, local_path)
+
+    else:
+        print("No files found in that folder.")
+    
+    # Get a list of image file extensions you want to include
+    image_extensions = ('.jpg', '.jpeg', '.png', '.bmp', '.gif')
+    # Read image paths into a list
+    image_paths = [
+        os.path.join(local_dir, file)
+        for file in os.listdir(local_dir)
+        if file.lower().endswith(image_extensions)
+    ]
+
+    prompt = 'Professional LinkedIn-style headshot, symmetrical full face and upper body visible including shoulders and chest, centered composition with a small space above the head, wearing a formal suit and white shirt, neutral expression, captured from a short distance, realistic skin texture, exact face preserved, plain white or gray background, sharp focus, studio lighting, high-resolution, suitable for CV or resume'
+    neg_prompt = 'Wrong face, flaws in the eyes, flaws in the face, lowres, artifacts noise, text, deformed, partially rendered objects,  deformed or partially rendered eyes, deformed eyeballs, cross-eyed, blurry'
+    scale = 1.2
+    n_samples = 5
+    seed = 0
+    steps = 1
+    H = 1024
+    W = 768
+    id_scale = 0.8
+    mode = 'fidelity'
+    id_mix = False 
 
     pipeline.debug_img_list = []
     if mode == 'fidelity':
@@ -38,16 +86,17 @@ def run(*args):
     else:
         raise ValueError
 
-    if id_image is not None:
-        id_image = resize_numpy_image_long(id_image, 1024)
+    id_image = image_paths[0]
+
+    if image_paths is not None:
+        id_image = resize_numpy_image_long(image_paths[0], 1024)
         id_embeddings = pipeline.get_id_embedding(id_image)
-        for supp_id_image in supp_images:
-            if supp_id_image is not None:
-                supp_id_image = resize_numpy_image_long(supp_id_image, 1024)
-                supp_id_embeddings = pipeline.get_id_embedding(supp_id_image)
-                id_embeddings = torch.cat(
-                    (id_embeddings, supp_id_embeddings if id_mix else supp_id_embeddings[:, :5]), dim=1
-                )
+        for i in range(1,len(image_paths)):
+            supp_id_image = resize_numpy_image_long(image_paths[i], 1024)
+            supp_id_embeddings = pipeline.get_id_embedding(supp_id_image)
+            id_embeddings = torch.cat(
+                (id_embeddings, supp_id_embeddings if id_mix else supp_id_embeddings[:, :5]), dim=1
+            )
     else:
         id_embeddings = None
 
@@ -57,6 +106,36 @@ def run(*args):
         img = pipeline.inference(prompt, (1, H, W), neg_prompt, id_embeddings, id_scale, scale, steps)[0]
         ims.append(np.array(img))
 
+    file_paths = []
+    for i, img in enumerate(ims):
+        if isinstance(img, torch.Tensor):
+            img = img.detach().cpu()
+            img = transforms.ToPILImage()(img)
+        elif not isinstance(img, Image.Image):
+            continue  # skip unknown formats
+        temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
+        img.save(temp_file.name)
+        file_paths.append(temp_file.name)
+
+    # Upload images to S3
+    OUTPUT_BUCKET_FOLDER = os.path.join(os.path.dirname(os.path.dirname(bucket_folder)), 'output') #'user_id/request_id/output/'
+    os.makedirs(OUTPUT_BUCKET_FOLDER, exist_ok=True)
+    
+    try:
+        for local_file in file_paths:
+            file_name = os.path.basename(local_file)
+            s3_key = os.path.join(OUTPUT_BUCKET_FOLDER, file_name)
+            s3_client.upload_file(local_file, BUCKET, s3_key, ExtraArgs={
+        'ContentType': 'image/jpeg'
+    })
+    except ClientError as e:
+        print('e', e)
+        print('ERROR OCCURRED while uploading data to S3')
+        return False
+
+    shutil.rmtree(req_id)
+    shutil.rmtree(OUTPUT_BUCKET_FOLDER)
+
     return ims, pipeline.debug_img_list
 
 
@@ -99,127 +178,11 @@ If you have any questions, feel free to open a discussion or contact us at <b>wu
 """  # noqa E501
 
 
-with gr.Blocks(title="PuLID", css=".gr-box {border-color: #8136e2}") as demo:
-    gr.Markdown(_HEADER_)
-    with gr.Row():
-        with gr.Column():
-            with gr.Row():
-                face_image = gr.Image(label="ID image (main)", sources="upload", type="numpy", height=256)
-                supp_image1 = gr.Image(
-                    label="Additional ID image (auxiliary)", sources="upload", type="numpy", height=256
-                )
-                supp_image2 = gr.Image(
-                    label="Additional ID image (auxiliary)", sources="upload", type="numpy", height=256
-                )
-                supp_image3 = gr.Image(
-                    label="Additional ID image (auxiliary)", sources="upload", type="numpy", height=256
-                )
-            prompt = gr.Textbox(label="Prompt", value='portrait,cinematic,wolf ears,white hair')
-            submit = gr.Button("Generate")
-            neg_prompt = gr.Textbox(label="Negative Prompt", value=DEFAULT_NEGATIVE_PROMPT)
-            scale = gr.Slider(
-                label="CFG, recommend value range [1, 1.5], 1 will be faster ",
-                value=1.2,
-                minimum=1,
-                maximum=1.5,
-                step=0.1,
-            )
-            n_samples = gr.Slider(label="Num samples", value=4, minimum=1, maximum=4, step=1)
-            seed = gr.Slider(
-                label="Seed", value=42, minimum=np.iinfo(np.uint32).min, maximum=np.iinfo(np.uint32).max, step=1
-            )
-            steps = gr.Slider(label="Steps", value=4, minimum=1, maximum=8, step=1)
-            with gr.Row():
-                H = gr.Slider(label="Height", value=1024, minimum=512, maximum=1280, step=64)
-                W = gr.Slider(label="Width", value=768, minimum=512, maximum=1280, step=64)
-            with gr.Row():
-                id_scale = gr.Slider(label="ID scale", minimum=0, maximum=5, step=0.05, value=0.8, interactive=True)
-                mode = gr.Dropdown(label="mode", choices=['fidelity', 'extremely style'], value='fidelity')
-                id_mix = gr.Checkbox(
-                    label="ID Mix (if you want to mix two ID image, please turn this on, otherwise, turn this off)",
-                    value=False,
-                )
-
-            gr.Markdown("## Examples")
-            example_inps = [
-                [
-                    'portrait,cinematic,wolf ears,white hair',
-                    'example_inputs/liuyifei.png',
-                    'fidelity',
-                ]
-            ]
-            gr.Examples(examples=example_inps, inputs=[prompt, face_image, mode], label='realistic')
-
-            example_inps = [
-                [
-                    'portrait, impressionist painting, loose brushwork, vibrant color, light and shadow play',
-                    'example_inputs/zcy.webp',
-                    'fidelity',
-                ]
-            ]
-            gr.Examples(examples=example_inps, inputs=[prompt, face_image, mode], label='painting style')
-
-            example_inps = [
-                [
-                    'portrait, flat papercut style, silhouette, clean cuts, paper, sharp edges, minimalist,color block,man',
-                    'example_inputs/lecun.jpg',
-                    'fidelity',
-                ]
-            ]
-            gr.Examples(examples=example_inps, inputs=[prompt, face_image, mode], label='papercut style')
-
-            example_inps = [
-                [
-                    'woman,cartoon,solo,Popmart Blind Box, Super Mario, 3d',
-                    'example_inputs/rihanna.webp',
-                    'fidelity',
-                ]
-            ]
-            gr.Examples(examples=example_inps, inputs=[prompt, face_image, mode], label='3d style')
-
-            example_inps = [
-                [
-                    'portrait, the legend of zelda, anime',
-                    'example_inputs/liuyifei.png',
-                    'extremely style',
-                ]
-            ]
-            gr.Examples(examples=example_inps, inputs=[prompt, face_image, mode], label='anime style')
-
-            example_inps = [
-                [
-                    'portrait, superman',
-                    'example_inputs/lecun.jpg',
-                    'example_inputs/lifeifei.jpg',
-                    'fidelity',
-                    True,
-                ]
-            ]
-            gr.Examples(examples=example_inps, inputs=[prompt, face_image, supp_image1, mode, id_mix], label='id mix')
-
-        with gr.Column():
-            output = gr.Gallery(label='Output', elem_id="gallery")
-            intermediate_output = gr.Gallery(label='DebugImage', elem_id="gallery", visible=False)
-            gr.Markdown(_CITE_)
-
-    inps = [
-        face_image,
-        supp_image1,
-        supp_image2,
-        supp_image3,
-        prompt,
-        neg_prompt,
-        scale,
-        n_samples,
-        seed,
-        steps,
-        H,
-        W,
-        id_scale,
-        mode,
-        id_mix,
-    ]
-    submit.click(fn=run, inputs=inps, outputs=[output, intermediate_output])
+with gr.Blocks(title="AI headshot Generator", css=".gr-box {border-color: #8136e2}") as demo:
+
+    output = gr.Gallery(label='Output', elem_id="gallery")
+    intermediate_output = gr.Gallery(label='DebugImage', elem_id="gallery", visible=False)
 
+    submit.click(fn=run, inputs=['textbox', 'textbox', 'textbox'], outputs=[output, intermediate_output])
 
 demo.launch()