updated_1

app.py

@@ -1,70 +1,112 @@
 import streamlit as st
-import torch
+import os
 import random
 import subprocess
-import os
+import io
+import numpy as np
+
 from PIL import Image
+import torch
 from diffusers import StableDiffusionPipeline, UNet2DConditionModel
-# ... any other imports you need ...
+from torchvision import transforms
+
+# If you're using Zero123++:
+from diffusers import DiffusionPipeline, EulerAncestralDiscreteScheduler
 
 # ------------------------------------------------------------------------------
-# 1. Load your models ONCE in global scope (so they don't reload on every run).
+# 0. GLOBAL CONFIG & UTILS
 # ------------------------------------------------------------------------------
 
-@st.cache_resource
-def load_sd_pipeline(base_model_path: str, fine_tuned_path: str):
-    # Safety checker dummy function for demonstration:
-    def dummy_safety_checker(images, clip_input):
-        return images, False
+# Provide your base SD model path & fine-tuned UNet path here.
+# (In a HF Space, you might store them in a local folder or load from HF repos.)
+BASE_MODEL_PATH = "runwayml/stable-diffusion-v1-5"
+FINE_TUNED_PATH = "my_finetuned_unet"   # e.g., local folder or HF Hub ID
+
+# If you want to use Zero123++ from a local clone:
+ZERO123_MODEL_ID = "sudo-ai/zero123plus-v1.2"
+
+# Example safety checker dummy, as used in your snippet:
+def dummy_safety_checker(images, clip_input):
+    return images, False
+
+# Make sure to remove or comment out any "!pip install ..." lines and rely
+# on your requirements.txt in the environment.
 
+# ------------------------------------------------------------------------------
+# 1. LOAD MODELS & PIPELINES
+# ------------------------------------------------------------------------------
+
+@st.cache_resource
+def load_sd_pipeline():
+    """Load the base stable diffusion pipeline with fine-tuned UNet attached."""
     pipe = StableDiffusionPipeline.from_pretrained(
-        base_model_path,
+        BASE_MODEL_PATH,
         torch_dtype=torch.float16
     )
     pipe.to("cuda")
-
-    # Load the fine-tuned UNet
+    
+    # Load and replace UNet
     unet = UNet2DConditionModel.from_pretrained(
-        fine_tuned_path,
+        FINE_TUNED_PATH,
         subfolder="unet",
         torch_dtype=torch.float16
-    ).to('cuda')
-
+    ).to("cuda")
+    
     pipe.unet = unet
     pipe.safety_checker = dummy_safety_checker
-
     return pipe
 
-# Similarly, if you want to load Zero123++ or other pipelines:
 @st.cache_resource
 def load_zero123_pipeline():
-    from diffusers import DiffusionPipeline, EulerAncestralDiscreteScheduler
-
+    """Load Zero123++ pipeline (v1.2) with EulerAncestralDiscreteScheduler."""
     pipeline = DiffusionPipeline.from_pretrained(
-        "sudo-ai/zero123plus-v1.2",
+        ZERO123_MODEL_ID,
         custom_pipeline="sudo-ai/zero123plus-pipeline",
         torch_dtype=torch.float16
     )
     pipeline.scheduler = EulerAncestralDiscreteScheduler.from_config(
-        pipeline.scheduler.config, timestep_spacing='trailing'
+        pipeline.scheduler.config, 
+        timestep_spacing='trailing'
     )
     pipeline.to("cuda")
     return pipeline
 
+# ------------------------------------------------------------------------------
+# 2. HELPER FUNCTIONS
+# ------------------------------------------------------------------------------
 
-# Example placeholders for the SyncDreamer command or internal functions:
-def run_syncdreamer(input_path: str, output_dir: str = "syncdreamer_output"):
-    """Runs SyncDreamer on input_path and places results into output_dir."""
-    st.info("Running SyncDreamer... (placeholder)")
-    # This is where your actual command would go:
-    # subprocess.run([...], check=True)
+def generate_funko_image(pipe, prompt: str, steps: int = 50):
+    """Generate a Funko image using the loaded Stable Diffusion pipeline."""
+    with torch.autocast("cuda"):
+        image = pipe(prompt, num_inference_steps=steps).images[0]
+    return image
+
+def run_syncdreamer(input_path: str, output_dir: str):
+    """
+    Placeholder for the SyncDreamer command-line call.
+    You would adapt this to run your real command. For example:
+
+      syncdreamer_cmd = [
+          "python", "generate.py",
+          "--ckpt", "ckpt/syncdreamer-pretrain.ckpt",
+          "--input", input_path,
+          "--output", output_dir,
+          "--sample_num", "4",
+          "--cfg_scale", "2.0",
+          ...
+      ]
+      subprocess.run(syncdreamer_cmd, check=True)
+    """
+    st.info("Running SyncDreamer... (this is a placeholder call)")
     os.makedirs(output_dir, exist_ok=True)
-    # (In a real scenario, you'd handle .jpg to .png conversion, etc.)
-    st.success(f"SyncDreamer completed. Results in: {output_dir}")
-
+    # In real usage, call the above commented command via subprocess
+    st.success(f"SyncDreamer completed. Output in: {output_dir}")
 
-# Helper function for Zero123++ pipeline
 def make_square_min_dim(image: Image.Image, min_side: int = 320) -> Image.Image:
+    """
+    Resize 'image' so that neither dimension is < min_side,
+    then pad to a square with white background.
+    """
     w, h = image.size
     scale = max(min_side / w, min_side / h, 1.0)
     new_w, new_h = int(w * scale), int(h * scale)
@@ -77,34 +119,77 @@ def make_square_min_dim(image: Image.Image, min_side: int = 320) -> Image.Image:
     new_img.paste(image, (offset_x, offset_y))
     return new_img
 
+def run_zero123(pipeline, input_image: Image.Image, steps: int = 50):
+    """Generate a 640x960 grid from Zero123++ pipeline."""
+    cond = make_square_min_dim(input_image, min_side=320)
+    with torch.autocast("cuda"):
+        result_grid = pipeline(cond, num_inference_steps=steps).images[0]
+    return result_grid
+
+def crop_zero123_grid(grid_img: Image.Image):
+    """
+    Zero123++ default output for 6-views is 640x960 (2 columns, 3 rows).
+    Crop into six 320x320 sub-images.
+    """
+    coords = [
+        (0,   0,   320, 320),
+        (320, 0,   640, 320),
+        (0,   320, 320, 640),
+        (320, 320, 640, 640),
+        (0,   640, 320, 960),
+        (320, 640, 640, 960),
+    ]
+    sub_images = []
+    for x1, y1, x2, y2 in coords:
+        sub_img = grid_img.crop((x1, y1, x2, y2))
+        sub_images.append(sub_img)
+    return sub_images
+
+# Example background compositing if desired:
+def create_mask(image, bg_color=(255,255,255), threshold=30):
+    arr = np.array(image)
+    diff = np.abs(arr - np.array(bg_color))
+    diff = diff.max(axis=2)
+    mask = (diff > threshold) * 255
+    return Image.fromarray(mask.astype(np.uint8), mode="L")
+
+def composite_foreground_background(fg, bg, bg_color=(255,255,255), threshold=30):
+    fg = fg.convert("RGBA")
+    bg = bg.convert("RGBA").resize(fg.size)
+    mask = create_mask(fg.convert("RGB"), bg_color=bg_color, threshold=threshold)
+    result = Image.composite(fg, bg, mask)
+    return result
+
+def get_bg_color(image):
+    corner_pixel = image.getpixel((0, 0))
+    # Heuristic: if corner pixel is near-white, treat as white background
+    if sum(corner_pixel) / 3 > 240:
+        return (255, 255, 255)
+    else:
+        return (200, 200, 200)
 
 # ------------------------------------------------------------------------------
-# 2. Streamlit application.
+# 3. STREAMLIT UI
 # ------------------------------------------------------------------------------
-
 def main():
-    st.title("Funko Generator Demo")
+    st.title("Funko Generator (SD + SyncDreamer + Zero123)")
 
-    # Let’s load pipelines in the background:
-    base_model_path = "runwayml/stable-diffusion-v1-5"
-    fine_tuned_path = "/content/drive/MyDrive/CC_Project/checkpoint-3000"  # adapt if needed
-    sd_pipe = load_sd_pipeline(base_model_path, fine_tuned_path)
+    # Load pipelines once
+    sd_pipe = load_sd_pipeline()
+    zero123_pipe = load_zero123_pipeline()
 
-    zero123_pipe = load_zero123_pipeline()  # For multi-view generation
-
-    # Session state to hold:
+    # Session state to store images
     if "latest_image" not in st.session_state:
         st.session_state["latest_image"] = None
     if "original_prompt" not in st.session_state:
         st.session_state["original_prompt"] = ""
 
-    # --------------------------------------------------------------------------
-    # A) Prompt input & initial generation
-    # --------------------------------------------------------------------------
-    st.subheader("1. Enter your Funko prompt")
+    # ---------------------------
+    # A) Prompt Input
+    # ---------------------------
+    st.subheader("1. Enter your initial Funko prompt")
 
-    # Show examples in the UI
-    with st.expander("Examples of valid prompts"):
+    with st.expander("Prompt Examples"):
         st.write("""
         - A standing plain human Funko in a blue shirt and blue pants with round black eyes with glasses with a belt.
         - A sitting angry animal Funko with squint black eyes.
@@ -112,217 +197,204 @@ def main():
         - ...
         """)
 
-    user_prompt = st.text_area("Type your Funko prompt here:", 
-                               value="A standing plain human Funko in a blue shirt and blue pants with round black eyes with glasses.")
-    generate_button = st.button("Generate Initial Funko")
+    user_prompt = st.text_area(
+        "Type your Funko prompt here:",
+        value="A standing plain human Funko in a blue shirt and blue pants with round black eyes with glasses."
+    )
+    generate_initial = st.button("Generate Initial Funko")
 
-    if generate_button:
+    if generate_initial:
         st.session_state["original_prompt"] = user_prompt
-        with st.spinner("Generating image..."):
-            with torch.autocast("cuda"):
-                image = sd_pipe(user_prompt, num_inference_steps=50).images[0]
-            st.session_state["latest_image"] = image
-        
+        with st.spinner("Generating initial Funko..."):
+            out_img = generate_funko_image(sd_pipe, user_prompt, steps=50)
+            st.session_state["latest_image"] = out_img
         st.success("Image generated!")
 
     if st.session_state["latest_image"] is not None:
-        st.image(st.session_state["latest_image"], caption="Latest Generated Image", use_column_width=True)
+        st.image(st.session_state["latest_image"], caption="Latest Funko Image", use_column_width=True)
 
-    # --------------------------------------------------------------------------
-    # B) Change the Funko (attributes)
-    # --------------------------------------------------------------------------
-    st.subheader("2. Modify Funko Attributes")
-    st.write("Select new attributes below. If you choose 'none', that attribute will be ignored/omitted in the prompt.")
+    # ---------------------------
+    # B) Modify Funko Attributes
+    # ---------------------------
+    st.subheader("2. Modify the Funko (attributes)")
+
+    st.write("Pick new attributes. If you choose 'none', we won't override that attribute.")
 
-    # Possible attributes (from your code) — including 'none'
     characters = ['none', 'animal', 'human', 'robot']
-    eyes_shape = ['none', 'anime', 'black', 'closed', 'round', 'square', 'squint']
-    eyes_color = ['none', 'black', 'blue', 'brown', 'green', 'grey', 'orange', 'pink', 'purple', 'red', 'white', 'yellow']
-    eyewear = ['none', 'eyepatch', 'glasses', 'goggles', 'helmet', 'mask', 'sunglasses']
-    hair_color = ['none', 'black', 'blonde', 'blue', 'brown', 'green', 'grey', 'orange', 'pink', 'purple', 'red', 'white', 'yellow']
-    emotion = ['none', 'angry', 'happy', 'plain', 'sad']
-    shirt_color = ['none', 'black', 'blue', 'brown', 'green', 'grey', 'orange', 'pink', 'purple', 'red', 'white', 'yellow']
-    pants_color = ['none', 'black', 'blue', 'brown', 'green', 'grey', 'orange', 'pink', 'purple', 'red', 'white', 'yellow']
+    eyes_shapes = ['none', 'anime', 'black', 'closed', 'round', 'square', 'squint']
+    eyes_colors = ['none', 'black', 'blue', 'brown', 'green', 'grey', 'orange', 'pink', 'purple', 'red', 'white', 'yellow']
+    eyewears = ['none', 'eyepatch', 'glasses', 'goggles', 'helmet', 'mask', 'sunglasses']
+    hair_colors = ['none', 'black', 'blonde', 'blue', 'brown', 'green', 'grey', 'orange', 'pink', 'purple', 'red', 'white', 'yellow']
+    emotions = ['none', 'angry', 'happy', 'plain', 'sad']
+    shirt_colors = ['none', 'black', 'blue', 'brown', 'green', 'grey', 'orange', 'pink', 'purple', 'red', 'white', 'yellow']
+    pants_colors = ['none', 'black', 'blue', 'brown', 'green', 'grey', 'orange', 'pink', 'purple', 'red', 'white', 'yellow']
     accessories = ['none', 'bag', 'ball', 'belt', 'bird', 'book', 'cape', 'guitar', 'hat', 'helmet', 'sword', 'wand', 'wings']
-    pose = ['none', 'sitting', 'standing']
-
-    # Create selection widgets:
-    chosen_char = st.selectbox("Character:", characters)
-    chosen_eyes_shape = st.selectbox("Eyes Shape:", eyes_shape)
-    chosen_eyes_color = st.selectbox("Eyes Color:", eyes_color)
-    chosen_eyewear = st.selectbox("Eyewear:", eyewear)
-    chosen_hair_color = st.selectbox("Hair Color:", hair_color)
-    chosen_emotion = st.selectbox("Emotion:", emotion)
-    chosen_shirt_color = st.selectbox("Shirt Color:", shirt_color)
-    chosen_pants_color = st.selectbox("Pants Color:", pants_color)
-    chosen_accessories = st.selectbox("Accessories:", accessories)
-    chosen_pose = st.selectbox("Pose:", pose)
-
-    # Now we form a modified prompt. For demonstration,
-    # let's do something simple: we take the original prompt, parse it, and
-    # replace only the attributes that are not 'none'.
-    def modify_prompt(base_prompt: str):
-        # A simple example: we can build a new prompt from scratch, ignoring the old text.
-        # In reality, you might parse the old text or do something more sophisticated.
-        new_prompt_segments = []
+    poses = ['none', 'sitting', 'standing']
+
+    chosen_char = st.selectbox("Character", characters)
+    chosen_eyes_shape = st.selectbox("Eyes Shape", eyes_shapes)
+    chosen_eyes_color = st.selectbox("Eyes Color", eyes_colors)
+    chosen_eyewear = st.selectbox("Eyewear", eyewears)
+    chosen_hair_color = st.selectbox("Hair Color", hair_colors)
+    chosen_emotion = st.selectbox("Emotion", emotions)
+    chosen_shirt_color = st.selectbox("Shirt Color", shirt_colors)
+    chosen_pants_color = st.selectbox("Pants Color", pants_colors)
+    chosen_accessory = st.selectbox("Accessories", accessories)
+    chosen_pose = st.selectbox("Pose", poses)
+
+    def build_modified_prompt():
+        # Simple new prompt builder
+        # If 'none', we do not override the attribute (use fallback or skip).
+        tokens = []
 
         # Pose
         if chosen_pose != 'none':
-            new_prompt_segments.append(f"A {chosen_pose}")
+            tokens.append(f"A {chosen_pose}")
         else:
-            new_prompt_segments.append("A standing")  # default fallback
+            tokens.append("A standing")
 
         # Emotion + Character
         if chosen_emotion != 'none':
-            new_prompt_segments.append(chosen_emotion)
+            tokens.append(chosen_emotion)
         else:
-            new_prompt_segments.append("plain")  # fallback
-
+            tokens.append("plain")
         if chosen_char != 'none':
-            new_prompt_segments.append(chosen_char + " Funko")
+            tokens.append(chosen_char + " Funko")
         else:
-            new_prompt_segments.append("human Funko")
+            tokens.append("human Funko")
 
         # Shirt color
         if chosen_shirt_color != 'none':
-            new_prompt_segments.append(f"in a {chosen_shirt_color} shirt")
+            tokens.append(f"in a {chosen_shirt_color} shirt")
         else:
-            new_prompt_segments.append("in a blue shirt")
+            tokens.append("in a blue shirt")
 
         # Pants color
         if chosen_pants_color != 'none':
-            new_prompt_segments.append(f"and {chosen_pants_color} pants")
+            tokens.append(f"and {chosen_pants_color} pants")
         else:
-            new_prompt_segments.append("and blue pants")
+            tokens.append("and blue pants")
 
         # Eyes
-        eye_text = []
+        eye_desc = []
         if chosen_eyes_shape != 'none':
-            eye_text.append(f"{chosen_eyes_shape}")
+            eye_desc.append(chosen_eyes_shape)
         else:
-            eye_text.append("round")
+            eye_desc.append("round")
         if chosen_eyes_color != 'none':
-            eye_text.append(f"{chosen_eyes_color}")
+            eye_desc.append(chosen_eyes_color)
         else:
-            eye_text.append("black")
-        eye_text.append("eyes")
-        new_prompt_segments.append("with " + " ".join(eye_text))
+            eye_desc.append("black")
+        eye_desc.append("eyes")
 
-        # Eyewear
-        if chosen_eyewear != 'none':
-            new_prompt_segments.append(f"with {chosen_eyewear}")
+        tokens.append("with " + " ".join(eye_desc))
 
-        # Hair
+        if chosen_eyewear != 'none':
+            tokens.append(f"with {chosen_eyewear}")
         if chosen_hair_color != 'none':
-            new_prompt_segments.append(f"with {chosen_hair_color} hair")
-
-        # Accessories
-        if chosen_accessories != 'none':
-            new_prompt_segments.append(f"with a {chosen_accessories}")
+            tokens.append(f"with {chosen_hair_color} hair")
+        if chosen_accessory != 'none':
+            tokens.append(f"with a {chosen_accessory}")
 
-        return " ".join(new_prompt_segments) + "."
+        return " ".join(tokens) + "."
 
     if st.button("Generate Modified Funko"):
-        if not st.session_state["original_prompt"]:
-            st.warning("Please generate an initial Funko (step 1) before modifying it.")
+        if st.session_state["original_prompt"] == "":
+            st.warning("Please generate an initial Funko first.")
         else:
-            new_prompt = modify_prompt(st.session_state["original_prompt"])
-            st.write(f"**New Prompt**: {new_prompt}")
-
+            new_prompt = build_modified_prompt()
+            st.write("**New Prompt**:", new_prompt)
             with st.spinner("Generating modified image..."):
-                with torch.autocast("cuda"):
-                    image = sd_pipe(new_prompt, num_inference_steps=50).images[0]
-                st.session_state["latest_image"] = image
-
-            st.image(st.session_state["latest_image"], caption="Modified Image", use_column_width=True)
+                out_img = generate_funko_image(sd_pipe, new_prompt, steps=50)
+                st.session_state["latest_image"] = out_img
+            st.image(st.session_state["latest_image"], caption="Modified Funko", use_column_width=True)
 
-    # --------------------------------------------------------------------------
-    # C) Animate the Funko with SyncDreamer
-    # --------------------------------------------------------------------------
-    st.subheader("3. Animate the Funko (SyncDreamer)")
-    st.write("Click the button to run SyncDreamer on the last generated image. (Demo)")
+    # ---------------------------
+    # C) Animate with SyncDreamer
+    # ---------------------------
+    st.subheader("3. Animate the Funko with SyncDreamer")
 
-    if st.button("Animate with SyncDreamer"):
+    st.write("Click to run SyncDreamer on the last generated image (placeholder).")
+    if st.button("Animate Funko"):
         if st.session_state["latest_image"] is None:
-            st.warning("No image found. Please generate a Funko first.")
+            st.warning("No image to animate. Generate a Funko first.")
         else:
-            # Save latest image locally so SyncDreamer can process it
+            # Save the current image
             input_path = "latest_funko.png"
             st.session_state["latest_image"].save(input_path)
-            run_syncdreamer(input_path, output_dir="syncdreamer_output")
+            output_dir = "syncdreamer_output"
+            run_syncdreamer(input_path, output_dir=output_dir)
 
-            # Optionally display a placeholder or actual frames/GIF
-            # ...
-            st.success("SyncDreamer animation completed (placeholder).")
+            st.success("SyncDreamer run complete (demo). Check output directory for results.")
 
-    # --------------------------------------------------------------------------
-    # D) Multi-View 3D Funko (Zero123++)
-    # --------------------------------------------------------------------------
-    st.subheader("4. Generate Multi-View 3D Funko (Zero123++)")
+    # ---------------------------
+    # D) Multi-View with Zero123++
+    # ---------------------------
+    st.subheader("4. Generate Multi-View Funko (Zero123++)")
 
     if st.button("Generate Multi-View 3D"):
         if st.session_state["latest_image"] is None:
-            st.warning("No image found. Please generate a Funko first.")
+            st.warning("No image to process. Generate a Funko first.")
         else:
-            # Save the last image as input for Zero123
-            input_path = "funko_for_zero123.png"
-            st.session_state["latest_image"].save(input_path)
-
-            # Make sure image is at least 320x320 and square
-            original_img = Image.open(input_path).convert("RGB")
-            cond = make_square_min_dim(original_img, min_side=320)
-
-            # Inference
-            st.info("Running Zero123++ pipeline... Please wait.")
-            with torch.autocast("cuda"):
-                result_grid = zero123_pipe(cond, num_inference_steps=50).images[0]
-
-            result_grid.save("zero123_grid.png")
-            st.image(result_grid, caption="Zero123++ Multi-View Grid (640x960)")
-
-            # Optionally crop and display sub-views
-            # Here we crop 6 sub-images of 320x320 from the 640x960 grid:
-            coords = [
-                (0,   0,   320, 320),
-                (320, 0,   640, 320),
-                (0,   320, 320, 640),
-                (320, 320, 640, 640),
-                (0,   640, 320, 960),
-                (320, 640, 640, 960),
-            ]
-            st.write("### Generated Views:")
-            for i, (x1, y1, x2, y2) in enumerate(coords):
-                sub_img = result_grid.crop((x1, y1, x2, y2))
-                sub_path = f"zero123_view_{i}.png"
-                sub_img.save(sub_path)
-                st.image(sub_path, width=256)
-
-    # --------------------------------------------------------------------------
-    # E) Integrate a New Background
-    # --------------------------------------------------------------------------
-    st.subheader("5. Apply a New Background to Each View")
-
-    st.write("Upload a background image, then apply it to each previously generated view.")
-    bg_file = st.file_uploader("Upload Background Image", type=["png", "jpg", "jpeg"])
+            # Save for Zero123
+            zero123_input_path = "funko_for_zero123.png"
+            st.session_state["latest_image"].save(zero123_input_path)
+
+            with st.spinner("Running Zero123++..."):
+                full_image = run_zero123(zero123_pipe, st.session_state["latest_image"], steps=50)
+            
+            # Display the 640x960 grid
+            st.image(full_image, caption="Zero123++ Grid (640x960)", use_column_width=True)
+
+            # Crop sub-images
+            sub_images = crop_zero123_grid(full_image)
+            st.write("Six sub-views:")
+            for i, s_img in enumerate(sub_images):
+                st.image(s_img, width=256, caption=f"View {i+1}")
+
+    # ---------------------------
+    # E) Background Compositing
+    # ---------------------------
+    st.subheader("5. Apply Background to Each View")
+
+    bg_file = st.file_uploader("Upload a background image (PNG/JPG)", type=["png","jpg","jpeg"])
     if bg_file is not None:
-        st.image(bg_file, caption="Selected Background", width=200)
+        st.image(bg_file, caption="Your Background", width=200)
 
-    if st.button("Apply Background to Multi-View"):
+    if st.button("Composite Background onto Views"):
         if bg_file is None:
             st.warning("No background uploaded.")
         else:
-            # Save background to disk:
-            bg_path = "background.png"
-            with open(bg_path, "wb") as f:
-                f.write(bg_file.read())
-
-            # In a real implementation, you would do the compositing described
-            # in your original code with threshold-based masking, etc.
-            # For demonstration, let's just say "Applied!"
-            st.success("Background compositing placeholder done. Check your images in the output folder.")
-
-    st.write("End of the Demo. Adjust code as needed for your pipeline paths and logic.")
-
+            # We assume you already did "Generate Multi-View 3D" so we have "Zero123++ Grid"
+            # In a real scenario, you might store sub-images in session_state after generation
+            # For this example, let's assume we re-run the pipeline or re-crop a stored grid.
+            if st.session_state["latest_image"] is None:
+                st.warning("No Funko image found. Generate or do multi-view first.")
+            else:
+                # We'll read the background
+                bg = Image.open(bg_file).convert("RGBA")
+
+                # Suppose we have a stored "zero123_grid.png" from the step above
+                # This is a simplistic approach. You might track them in session state.
+                if not os.path.exists("zero123_grid.png"):
+                    st.warning("No zero123_grid.png found. Please run Zero123++ step first.")
+                else:
+                    grid_img = Image.open("zero123_grid.png").convert("RGB")
+                    sub_images = crop_zero123_grid(grid_img)
+
+                    # Composite each sub-image
+                    st.write("Applying background to each sub-view...")
+                    for i, fg_img in enumerate(sub_images):
+                        # Detect background color from Funko sub-view
+                        bg_color = get_bg_color(fg_img)
+                        comp = composite_foreground_background(fg_img, bg, bg_color=bg_color, threshold=30)
+                        st.image(comp, width=256, caption=f"Composite View {i+1}")
+
+    st.write("---")
+    st.write("End of the demo. Adapt paths and code to your environment as needed.")
 
+# ------------------------------------------------------------------------------
+# 4. ENTRY POINT
+# ------------------------------------------------------------------------------
 if __name__ == "__main__":
     main()