import streamlit as st import os import random import subprocess import io import numpy as np from PIL import Image import torch from diffusers import StableDiffusionPipeline, UNet2DConditionModel from torchvision import transforms # If you're using Zero123++: from diffusers import DiffusionPipeline, EulerAncestralDiscreteScheduler # ------------------------------------------------------------------------------ # 0. GLOBAL CONFIG & UTILS # ------------------------------------------------------------------------------ # 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, torch_dtype=torch.float16 ) pipe.to("cuda") # Load and replace UNet unet = UNet2DConditionModel.from_pretrained( FINE_TUNED_PATH, subfolder="unet", torch_dtype=torch.float16 ).to("cuda") pipe.unet = unet pipe.safety_checker = dummy_safety_checker return pipe @st.cache_resource def load_zero123_pipeline(): """Load Zero123++ pipeline (v1.2) with EulerAncestralDiscreteScheduler.""" pipeline = DiffusionPipeline.from_pretrained( 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.to("cuda") return pipeline # ------------------------------------------------------------------------------ # 2. HELPER FUNCTIONS # ------------------------------------------------------------------------------ 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 real usage, call the above commented command via subprocess st.success(f"SyncDreamer completed. Output in: {output_dir}") 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) image = image.resize((new_w, new_h), Image.LANCZOS) side = max(new_w, new_h) new_img = Image.new(mode="RGB", size=(side, side), color=(255, 255, 255)) offset_x = (side - new_w) // 2 offset_y = (side - new_h) // 2 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) # ------------------------------------------------------------------------------ # 3. STREAMLIT UI # ------------------------------------------------------------------------------ def main(): st.title("Funko Generator (SD + SyncDreamer + Zero123)") # Load pipelines once sd_pipe = load_sd_pipeline() zero123_pipe = load_zero123_pipeline() # 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 # --------------------------- st.subheader("1. Enter your initial Funko prompt") 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. - A standing happy robot Funko in a brown shirt and grey pants with squint black eyes with cane and monocle. - ... """) 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_initial: st.session_state["original_prompt"] = user_prompt 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 Funko Image", use_column_width=True) # --------------------------- # 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.") characters = ['none', 'animal', 'human', 'robot'] 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'] 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': tokens.append(f"A {chosen_pose}") else: tokens.append("A standing") # Emotion + Character if chosen_emotion != 'none': tokens.append(chosen_emotion) else: tokens.append("plain") if chosen_char != 'none': tokens.append(chosen_char + " Funko") else: tokens.append("human Funko") # Shirt color if chosen_shirt_color != 'none': tokens.append(f"in a {chosen_shirt_color} shirt") else: tokens.append("in a blue shirt") # Pants color if chosen_pants_color != 'none': tokens.append(f"and {chosen_pants_color} pants") else: tokens.append("and blue pants") # Eyes eye_desc = [] if chosen_eyes_shape != 'none': eye_desc.append(chosen_eyes_shape) else: eye_desc.append("round") if chosen_eyes_color != 'none': eye_desc.append(chosen_eyes_color) else: eye_desc.append("black") eye_desc.append("eyes") tokens.append("with " + " ".join(eye_desc)) if chosen_eyewear != 'none': tokens.append(f"with {chosen_eyewear}") if chosen_hair_color != 'none': tokens.append(f"with {chosen_hair_color} hair") if chosen_accessory != 'none': tokens.append(f"with a {chosen_accessory}") return " ".join(tokens) + "." if st.button("Generate Modified Funko"): if st.session_state["original_prompt"] == "": st.warning("Please generate an initial Funko first.") else: new_prompt = build_modified_prompt() st.write("**New Prompt**:", new_prompt) with st.spinner("Generating modified image..."): 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 with SyncDreamer # --------------------------- st.subheader("3. Animate the Funko 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 to animate. Generate a Funko first.") else: # Save the current image input_path = "latest_funko.png" st.session_state["latest_image"].save(input_path) output_dir = "syncdreamer_output" run_syncdreamer(input_path, output_dir=output_dir) st.success("SyncDreamer run complete (demo). Check output directory for results.") # --------------------------- # 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 to process. Generate a Funko first.") else: # 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="Your Background", width=200) if st.button("Composite Background onto Views"): if bg_file is None: st.warning("No background uploaded.") else: # 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()