Update app.py

app.py

@@ -95,44 +95,21 @@ def gen_pos_prompt(text):
   pos_prompt = json.loads(response.get('body').read())['results'][0]['outputText']
   return pos_prompt
 
-# def encode_image_to_base64(image_path):
-#     with open(image_path, "rb") as image_file:
-#         encoded_string = base64.b64encode(image_file.read()).decode('utf-8')
-#     return encoded_string
-
 def encode_image_to_base64(image):
-    buffered = BytesIO()
-    image.save(buffered, format="PNG")  # Save the image to the buffer
-    encoded_string = base64.b64encode(buffered.getvalue()).decode('utf-8')
-    return encoded_string
+    with io.BytesIO() as buffered:
+        image.save(buffered, format="PNG")
+        return base64.b64encode(buffered.getvalue()).decode('utf-8')
 
 
 def generate_image_from_text(encoded_image,pos_prompt=None):
     neg_prompt = '''Detailed, complex textures, intricate patterns, realistic lighting, high contrast, reflections, fuzzy surface, realistic proportions, photographic quality, vibrant colors, detailed background, shadows, disfigured, deformed, ugly, multiple, duplicate.'''
     encoded_str = encode_image_to_base64(encoded_image)
-    if pos_prompt:
-      new_prompt = gen_pos_prompt(pos_prompt)
-      parameters = {
-          'taskType': 'IMAGE_VARIATION',
-          'imageVariationParams': {
-              'images': [encoded_str],
-              'text': new_prompt,
-              'negativeText': neg_prompt,
-              'similarityStrength': 0.7
-          },
-          'imageGenerationConfig': {
-              "cfgScale": 8,
-              "seed": 0,
-              "width": 512,
-              "height": 512,
-              "numberOfImages": 1
-          }
-      }
-    else:
-      parameters = {
+    new_prompt = gen_pos_prompt(pos_prompt)
+    parameters = {
           'taskType': 'IMAGE_VARIATION',
           'imageVariationParams': {
               'images': [encoded_str],
+              'text': gen_pos_prompt(pos_prompt) if pos_prompt else None,
               'negativeText': neg_prompt,
               'similarityStrength': 0.7
           },
@@ -143,8 +120,7 @@ def generate_image_from_text(encoded_image,pos_prompt=None):
               "height": 512,
               "numberOfImages": 1
           }
-      }
-    
+      }  
     request_body = json.dumps(parameters)
     response = bedrock_runtime.invoke_model(body=request_body, modelId='amazon.titan-image-generator-v1')
     response_body = json.loads(response.get('body').read())
@@ -189,14 +165,12 @@ def generate(image, mc_resolution, formats=["obj", "glb"]):
     return mesh_path_obj.name, mesh_path_glb.name
 
 def run_example(image, do_remove_background, foreground_ratio, mc_resolution, text_prompt=None):
-    if text_prompt:
-        image_pil = generate_image_from_text(encoded_image=image, pos_prompt=text_prompt)
-    else:
-        image_pil = generate_image_from_text(encoded_image=image)
+    image_pil = generate_image_from_text(encoded_image=image, pos_prompt=text_prompt)
     preprocessed = preprocess(image_pil, False, 0.9)
     mesh_name_obj, mesh_name_glb = generate(preprocessed, 256, ["obj", "glb"])
     return preprocessed, mesh_name_obj, mesh_name_glb
 
+    
 with gr.Blocks() as demo:
     gr.Markdown(HEADER)
     with gr.Row(variant="panel"):
@@ -265,7 +239,7 @@ with gr.Blocks() as demo:
     submit.click(fn=check_input_image, inputs=[input_image]).success(
         fn=run_example,
         inputs=[input_image, do_remove_background, foreground_ratio, mc_resolution, text_prompt],
-        outputs=[processed_image, output_model_obj, output_model_glb],
+        outputs=[output_model_obj, output_model_glb],
         # outputs=[output_model_obj, output_model_glb],
     )