Update llm_diffusion_serving_app

.gitignore

@@ -30,6 +30,7 @@ test/model_store/
 test/ts_console.log
 test/config.properties
 
+model-store-local/
 
 .vscode
 .scratch/

examples/usecases/llm_diffusion_serving_app/Readme.md

@@ -1,5 +1,5 @@
 
-## Multi-Image Generation Streamlit App by chaining LLaMA & Stable Diffusion using TorchServe, torch.compile & OpenVINO
+## Multi-Image Generation Streamlit App: Chaining LLaMA & Stable Diffusion using TorchServe, torch.compile & OpenVINO
 
 This Multi-Image Generation Streamlit app is designed to generate multiple images based on a provided text prompt. Instead of using Stable Diffusion directly, this app chains LLaMA and Stable Diffusion to enhance the image generation process. Here’s how it works:
 - The app takes a user prompt and uses [Meta-LLaMA-3.2](https://huggingface.co/meta-llama) to create multiple interesting and relevant prompts.
@@ -115,4 +115,18 @@ Collecting usage statistics. To deactivate, set browser.gatherUsageStats to fals
   Local URL: http://localhost:8084
   Network URL: http://123.11.0.2:8084
   External URL: http://123.123.12.34:8084
-```
+```
+
+## App Screenshots
+
+### Server App Screenshots
+
+| Screenshot 1 | Screenshot 2 | Screenshot 3 |
+| --- | --- | --- |
+| <img src="./docker/img/server-app-screen-1.png" width="400"> | <img src="./docker/img/server-app-screen-2.png" width="400"> | <img src="./docker/img/server-app-screen-3.png" width="400"> |
+
+### Client App Screenshots
+
+| Screenshot 1 | Screenshot 2 | Screenshot 3 |
+| --- | --- | --- |
+| <img src="./docker/img/client-app-screen-1.png" width="400"> | <img src="./docker/img/client-app-screen-2.png" width="400"> | <img src="./docker/img/client-app-screen-3.png" width="400"> |

examples/usecases/llm_diffusion_serving_app/docker/client_app.py

@@ -10,6 +10,9 @@
 import numpy as np
 from PIL import Image
 
+# App title
+st.set_page_config(page_title="Multi-image Gen App")
+
 logger = logging.getLogger(__name__)
 
 MODEL_NAME_LLM = os.environ["MODEL_NAME_LLM"]
@@ -25,9 +28,6 @@
 st.session_state.llm_prompts = st.session_state.get("llm_prompts", None)
 st.session_state.llm_time = st.session_state.get("llm_time", 0)
 
-# App title
-st.set_page_config(page_title="Multi-image Gen App")
-
 with st.sidebar:
     st.title("Image Generation with LLaMA, SDXL and OpenVINO")
 
@@ -244,15 +244,15 @@ def generate_llm_model_response(prompt_template_with_user_input, user_prompt):
     ### Multi-Image Generation App with TorchServe and OpenVINO
     Welcome to the Multi-Image Generation Client App. This app allows you to generate multiple images 
     from a single text prompt. Simply input your prompt, and the app will enhance it using a LLM (LLaMA) and 
-    generate images in parallel using the Stable Diffusion model.
+    generate images in parallel using the Stable Diffusion with latent-consistency/lcm-sdxl model.
     See [GitHub](https://github.com/pytorch/serve/tree/master/examples/usecases/llm_diffusion_serving_app) for details.
     """,
         unsafe_allow_html=True,
     )
     st.image("./img/workflow-2.png")
 
     st.markdown(
-        """<div style='background-color: #f0f0f0; font-size: 14px; padding: 10px; 
+        """<div style='background-color: #232628; font-size: 14px; padding: 10px; 
                 border: 1px solid #ddd; border-radius: 5px;'>
             NOTE: Initial image generation may take longer due to model warm-up. Subsequent generations will be faster !
             </div>""",
@@ -309,7 +309,8 @@ def display_prompts():
     elif len(st.session_state.llm_prompts) < num_images:
         prompt_container.warning(
             f"""Insufficient prompts. Regenerate prompts ! 
-                                 Num Images Requested: {num_images}, Prompts Generated: {len(st.session_state.llm_prompts)}""",
+                Num Images Requested: {num_images}, Prompts Generated: {len(st.session_state.llm_prompts)}
+                {f"Consider increasing the max_new_tokens parameter !" if num_images > 4 else ""}""",
             icon="⚠️",
         )
     else:

examples/usecases/llm_diffusion_serving_app/docker/server_app.py

@@ -9,6 +9,9 @@
 import streamlit as st
 import logging
 
+# App title
+st.set_page_config(page_title="Multi-Image Gen App Control Center")
+
 logger = logging.getLogger(__name__)
 
 MODEL_NAME_LLM = os.environ["MODEL_NAME_LLM"]
@@ -19,9 +22,9 @@
 MODEL_NAME_SD = MODEL_NAME_SD.replace("/", "---")
 MODEL_SD = MODEL_NAME_SD.split("---")[1]
 
-# Initialize Session State variables
+# Init Session State variables
 st.session_state.started = st.session_state.get("started", False)
-st.session_state.stopped = st.session_state.get("stopped", True)
+st.session_state.stopped = st.session_state.get("stopped", True)    
 st.session_state.registered = st.session_state.get(
     "registered",
     {
@@ -30,21 +33,36 @@
     },
 )
 
-# App title
-st.set_page_config(page_title="Multi-Image Gen App Control Center")
+def is_server_running():
+    """Check if the TorchServe server is running."""
+    try:
+        res = requests.get("http://localhost:8080/ping")
+        return res.status_code == 200
+    except requests.exceptions.ConnectionError:
+        return False
 
+def init_model_registrations():
+    for model_name in [MODEL_NAME_LLM, MODEL_NAME_SD]:
+        url = f"http://localhost:8081/models/{model_name}"
+        try:
+            res = requests.get(url)
+            if res.status_code == 200:
+                st.session_state.registered[model_name] = True
+            else:
+                st.session_state.registered[model_name] = False
+        except requests.exceptions.RequestException as e:
+            logger.info(f"Error checking model registration: {e}")
+            st.session_state.registered[model_name] = False
 
-def start_torchserve_server():
-    """Starts the TorchServe server if it's not already running."""
+# Update Session State variables
+if is_server_running():
+    st.session_state.started = True
+    st.session_state.stopped = False    
+    init_model_registrations()
 
-    def is_server_running():
-        """Check if the TorchServe server is running."""
-        try:
-            res = requests.get("http://localhost:8080/ping")
-            return res.status_code == 200
-        except requests.exceptions.ConnectionError:
-            return False
 
+def start_torchserve_server():
+    """Starts the TorchServe server if it's not already running."""
     def launch_server():
         """Launch the TorchServe server with the specified configurations."""
         subprocess.run(
@@ -82,10 +100,7 @@ def launch_server():
     # Update session state variables if server started successfully
     st.session_state.started = True
     st.session_state.stopped = False
-    st.session_state.registered = {
-        MODEL_NAME_LLM: False,
-        MODEL_NAME_SD: False,
-    }
+    init_model_registrations()
 
 
 def stop_server():
@@ -301,13 +316,13 @@ def get_sw_versions():
     ### Multi-Image Generation App Control Center
     Manage the Multi-Image Generation App workflow with this administrative interface. 
     Use this app to Start/stop TorchServe, load/register models, scale up/down workers, 
-    and perform other tasks to ensure smooth operation.
+    and review TorchServe Server and Model info.
     See [GitHub](https://github.com/pytorch/serve/tree/master/examples/usecases/llm_diffusion_serving_app) for details.
     """,
         unsafe_allow_html=True,
     )
     st.markdown(
-        """<div style='background-color: #f0f0f0; font-size: 14px; padding: 10px; 
+        """<div style='background-color: #232628; font-size: 14px; padding: 10px; 
                 border: 1px solid #ddd; border-radius: 5px;'>
         <b>NOTE</b>: After Starting TorchServe and Registering models, proceed to Client App running at port 8085.
         </div>""",