Update app.py

app.py

@@ -39,10 +39,10 @@ model.renderer.set_chunk_size(131072)
 model.to(device)
 
 rembg_session = rembg.new_session()
-# my_aws_access_key_id = os.getenv("ACCESS")
-# my_aws_secret_access_key = os.getenv("SECRET")
-# bedrock_runtime = boto3.client(service_name='bedrock-runtime', aws_access_key_id = my_aws_access_key_id, aws_secret_access_key = my_aws_secret_access_key, region_name='us-east-1')
-
+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},
@@ -55,11 +55,99 @@ rembg_session = rembg.new_session()
     
 #     return Image.open(io.BytesIO(base64_image_data))
 
+
+def gen_pos_prompt(text):
+  instruction = f'''Your task is to create a positive prompt for image generation.
+    Objective: Generate images that prioritize structural integrity and accurate shapes. The focus should be on the correct form and basic contours of objects, with minimal concern for colors.
+    
+    Guidelines:
+    
+    Complex Objects (e.g., animals, vehicles): For these, the image should resemble a toy object, emphasizing the correct shape and structure while minimizing details and color complexity.
+    
+    Example Input: A sports bike
+    Example Positive Prompt: Simple sports bike with accurate shape and structure, minimal details, digital painting, concept art style, basic contours, soft lighting, clean lines, neutral or muted colors, toy-like appearance, low contrast.
+    
+    Example Input: A lion
+    Example Positive Prompt: Toy-like depiction of a lion with a focus on structural accuracy, minimal details, digital painting, concept art style, basic contours, soft lighting, clean lines, neutral or muted colors, simplified features, low contrast.
+    
+    Simple Objects (e.g., a tennis ball): For these, the prompt should specify a realistic depiction, focusing on the accurate shape and structure.
+    
+    Example Input: A tennis ball
+    Example Positive Prompt: Realistic depiction of a tennis ball with accurate shape and texture, digital painting, clean lines, minimal additional details, soft lighting, neutral or muted colors, focus on structural integrity.
+    
+    Prompt Structure:
+    
+    Subject: Clearly describe the object and its essential shape and structure.
+    Medium: Specify the art style (e.g., digital painting, concept art).
+    Style: Include relevant style terms (e.g., simplified, toy-like for complex objects; realistic for simple objects).
+    Resolution: Mention resolution if necessary (e.g., basic resolution).
+    Lighting: Indicate the type of lighting (e.g., soft lighting).
+    Color: Use neutral or muted colors with minimal emphasis on color details.
+    Additional Details: Keep additional details minimal or specify if not desired.
+    Input: {text}
+    Positive Prompt: 
+    '''
+
+  body = json.dumps({'inputText': instruction,
+                     'textGenerationConfig': {'temperature': 0.1, 'topP': 0.01, 'maxTokenCount':512}})
+  response = bedrock_runtime.invoke_model(body=body, modelId='amazon.titan-text-express-v1')
+  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 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 = {
+          'taskType': 'IMAGE_VARIATION',
+          'imageVariationParams': {
+              'images': [encoded_str],
+              'negativeText': neg_prompt,
+              'similarityStrength': 0.7
+          },
+          'imageGenerationConfig': {
+              "cfgScale": 8,
+              "seed": 0,
+              "width": 512,
+              "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())
+    base64_image_data = base64.b64decode(response_body['images'][0])
+    return Image.open(io.BytesIO(base64_image_data))
+
 def check_input_image(input_image):
     if input_image is None:
         raise gr.Error("No image uploaded!")
 
-
 def preprocess(input_image, do_remove_background, foreground_ratio):
     def fill_background(image):
         image = np.array(image).astype(np.float32) / 255.0
@@ -78,8 +166,7 @@ def preprocess(input_image, do_remove_background, foreground_ratio):
             image = fill_background(image)
     return image
 
-
-@spaces.GPU
+# @spaces.GPU
 def generate(image, mc_resolution, formats=["obj", "glb"]):
     scene_codes = model(image, device=device)
     mesh = model.extract_mesh(scene_codes, resolution=mc_resolution)[0]
@@ -94,9 +181,10 @@ def generate(image, mc_resolution, formats=["obj", "glb"]):
     
     return mesh_path_obj.name, mesh_path_glb.name
 
-def run_example(image_pil):
-    preprocessed = preprocess(image_pil, False, 0.9)
-    mesh_name_obj, mesh_name_glb = generate(preprocessed, 256, ["obj", "glb"])
+def run_example(text_prompt, do_remove_background, foreground_ratio, mc_resolution):
+    image_pil = generate_image_from_text(text_prompt)  
+    preprocessed = preprocess(image_pil, do_remove_background, foreground_ratio)
+    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:
@@ -104,14 +192,20 @@ with gr.Blocks() as demo:
     with gr.Row(variant="panel"):
         with gr.Column():
             with gr.Row():
+                
                 input_image = gr.Image(
-                    label="Input Image",
+                    label="Generated Image",
                     image_mode="RGBA",
                     sources="upload",
                     type="pil",
                     elem_id="content_image",
+                    visible=True  # Hidden since we generate the image from text
+                )
+                text_prompt = gr.Textbox(
+                    label="Text Prompt",
+                    placeholder="Enter Positive Prompt"
                 )
-                processed_image = gr.Image(label="Processed Image", interactive=False)
+                processed_image = gr.Image(label="Processed Image", interactive=False, visible=False)
             with gr.Row():
                 with gr.Group():
                     do_remove_background = gr.Checkbox(
@@ -151,22 +245,19 @@ with gr.Blocks() as demo:
             examples=[
                 os.path.join("examples", img_name) for img_name in sorted(os.listdir("examples"))
             ],
-            inputs=[input_image],
+            inputs=[text_prompt],
             outputs=[processed_image, output_model_obj, output_model_glb],
             cache_examples=True,
-            fn=partial(run_example),
+            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=[input_image]).success(
-        fn=preprocess,
-        inputs=[input_image, do_remove_background, foreground_ratio],
-        outputs=[processed_image],
-    ).success(
-        fn=generate,
-        inputs=[processed_image, mc_resolution],
-        outputs=[output_model_obj, output_model_glb],
+    submit.click(fn=check_input_image, inputs=[input_image, text_prompt]).success(
+        fn=run_example,
+        inputs=[input_image,text_prompt, 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()
+demo.launch(auth=(os.getenv("USERNAME"), os.getenv("PASSWORD")))