Update app.py

app.py

@@ -7,6 +7,7 @@ from io import BytesIO
 # from IPython.display import display
 import base64
 import time
+import random
 
 
 
@@ -21,6 +22,10 @@ def display_image(image=None,width=500,height=500):
     img = image.resize((width, height))
     return img
 
+def pretty_print(messages):
+    for message in messages:
+        return f"{message['role']}: {message['content']}"
+
 # API Gateway endpoint URL
 api_url = 'https://a02q342s5b.execute-api.us-east-2.amazonaws.com/reinvent-demo-inf2-sm-20231114'
 
@@ -45,18 +50,50 @@ api_url = 'https://a02q342s5b.execute-api.us-east-2.amazonaws.com/reinvent-demo-
 
 
 # Creating Tabs
-tab1, tab2, tab3 = st.tabs(["Image Generation", "Architecture", "Code"])
+tab1, tab2, tab3 = st.tabs(["Image Generation",  "Architecture", "Code"])
 
 with tab1:
     # Create two columns for layout
     left_column, right_column = st.columns(2)
+
+    with right_column:
+        cont = st.container()
+
     # ===========
     with left_column:
         # Define Streamlit UI elements
-        st.title('Stable Diffusion XL Image Generation with AWS Inferentia')
+        st.title('Stable Diffusion XL Image Generation with AWS Inferentia 2')
+
+        sample_prompts = [
+            "A futuristic cityscape at sunset, cyberpunk",
+            "A serene landscape with mountains and a river, photorealistic style",
+            "An astronaut riding a horse, artistic and surreal",
+            "A robot playing chess in a medieval setting, high detail",
+            "An underwater scene with colorful coral reefs and fish, vibrant colors",
+            "Raccoon astronaut in space, sci-fi, future, cold color palette, muted colors, detailed, 8k",
+            "A lost city rediscovered in the Amazon jungle, overgrown with plants, in the style of a vintage travel poster",
+            "A steampunk train emitting clouds of steam as it races through a mountain pass, digital art",
+            "An enchanted forest with bioluminescent trees and fairies dancing, in a Studio Ghibli style",
+            "A portrait of an elegant alien empress with a detailed headdress, reminiscent of Art Nouveau",
+            "A post-apocalyptic Tokyo with nature reclaiming skyscrapers, in the style of a concept art",
+            "A mythical phoenix rising from ashes, vibrant colors, with a nebula in the background",
+            "A cybernetic wolf in a neon-lit city, cyberpunk theme, rain-drenched streets",
+            "A high fantasy battle scene with dragons in the sky and knights on the ground, epic scale",
+            "An ice castle on a lonely mountain peak, under the northern lights, fantasy illustration",
+            "A surreal landscape where giant flowers bloom in the desert, with a distant thunderstorm, hyperrealism"
+        ]
+
+        def set_random_prompt():
+            # This function will be called when the button is clicked
+            random_prompt = random.choice(sample_prompts)
+            # Update the session state for the input field
+            st.session_state.prompt_one = random_prompt
 
         prompt_one = st.text_area("Enter your prompt:", 
-                            f"Raccoon astronaut in space, sci-fi, future, cold color palette, muted colors, detailed, 8k")
+                                  
+                                  key="prompt_one")
+        
+        st.button('Random Prompt', on_click=set_random_prompt)
 
         # Number of inference steps
         num_inference_steps_one = st.slider("Number of Inference Steps", 
@@ -76,15 +113,8 @@ with tab1:
             negative_prompt_one = st.text_area("Enter your negative prompt:", 
                             "cartoon, graphic, text, painting, crayon, graphite, abstract glitch, blurry")
 
-        
-
-
-
-
-
-    if st.button('Generate Image'):
-        with st.spinner(f'Generating Image with {num_inference_steps_one} iterations'):
-            with right_column:
+        if st.button('Generate Image'):
+            with st.spinner(f'Generating Image with {num_inference_steps_one} iterations'):
                 start_time = time.time()
                 # ===============
                 # Example input data
@@ -94,7 +124,8 @@ with tab1:
                         "num_inference_steps": num_inference_steps_one,
                         "seed": seed_one,
                         "negative_prompt": negative_prompt_one
-                    }
+                    },
+                    "endpoint": "huggingface-pytorch-inference-neuronx-2023-11-14-21-22-10-388"
                 }
 
                 # Make API request
@@ -105,20 +136,93 @@ with tab1:
                     result_one = response_one.json()
                     # st.success(f"Prediction result: {result}")
                     image_one = display_image(decode_base64_image(result_one["generated_images"][0]))
-                    st.image(image_one, 
+                    cont.image(image_one, 
                         caption=f"{prompt_one}") 
                     end_time = time.time()
                     total_time = round(end_time - start_time, 2)
-                    st.text(f"Prompt: {prompt_one}")
-                    st.text(f"Number of Iterations: {num_inference_steps_one}")
-                    st.text(f"Random Seed: {seed_one}")
-                    st.text(f'Total time taken: {total_time} seconds')
+                    cont.text(f"Prompt: {prompt_one}")
+                    cont.text(f"Number of Iterations: {num_inference_steps_one}")
+                    cont.text(f"Random Seed: {seed_one}")
+                    cont.text(f'Total time taken: {total_time} seconds')
                     # Calculate and display the time per iteration in milliseconds
                     time_per_iteration_ms = (total_time / num_inference_steps_one)
-                    st.text(f'Time per iteration: {time_per_iteration_ms:.2f} seconds')
+                    cont.text(f'Time per iteration: {time_per_iteration_ms:.2f} seconds')
                 else:
                     st.error(f"Error: {response_one.text}")
 
+# with tab2:
+   
+# st.title('Llama 2 7B Text Generation with AWS Inferentia 2')
+
+# params = {
+#     "do_sample" : True,
+#     "top_p": 0.6,
+#     "temperature": 0.9,
+#     "top_k": 50,
+#     "max_new_tokens": 512,
+#     "repetition_penalty": 1.03,
+# }
+    
+# if "messages" not in st.session_state:
+#     st.session_state.messages = [
+#         {"role": "system", "content": "You are a helpful Travel Planning Assistant. You respond with only 1-2 sentences."},
+#         {'role': 'user', 'content': 'Where can I travel in the fall for cloudy, rainy, and beautiful views?'},
+#     ]
+
+# for message in st.session_state.messages:
+#     with st.chat_message(message["role"]):
+#         st.markdown(message["content"])
+
+# with st.chat_message("assistant"):
+#     message_placeholder = st.empty()
+#     full_response = ""
+#     prompt_input_one = {
+#             "prompt": st.session_state.messages,
+#             "parameters": params,
+#             "endpoint": "huggingface-pytorch-inference-neuronx-2023-11-28-16-09-51-708"
+#         }
+    
+#     response_one = requests.post(api_url, json=prompt_input_one)
+
+#     if response_one.status_code == 200:
+#         result_one = response_one.json()
+#         # st.success(f"Prediction result: {result}")
+#         full_response += result_one["generation"]
+#     else:
+#         st.error(f"Error: {response_one.text}")
+
+#     message_placeholder.markdown(full_response)
+# st.session_state.messages.append({"role": "assistant", "content": full_response})
+
+# if prompt := st.chat_input("What is up?"):
+#     st.session_state.messages.append({"role": "user", "content": prompt})
+#     print(st.session_state.messages)
+#     with st.chat_message("user"):
+#         st.markdown(prompt)
+    
+#     with st.chat_message("assistant"):
+#         message_placeholder = st.empty()
+#         new_response = ""
+#         prompt_input_one = {
+#                 "prompt": st.session_state.messages,
+#                 "parameters": params,
+#                 "endpoint": "huggingface-pytorch-inference-neuronx-2023-11-28-16-09-51-708"
+#             }
+        
+#         response_one = requests.post(api_url, json=prompt_input_one)
+
+#         if response_one.status_code == 200:
+#             result_one = response_one.json()
+#             # st.success(f"Prediction result: {result}")
+#             new_response += result_one["generation"]
+#         else:
+#             st.error(f"Error: {response_one.text}")
+
+#         message_placeholder.markdown(new_response)
+#     st.session_state.messages.append({"role": "assistant", "content": new_response})
+
+
+
 
 with tab2:        
     # ===========