import logging import time import traceback from pathlib import Path from typing import Optional, Tuple, Dict, Any, List from PIL import Image import numpy as np import cv2 import gradio as gr import spaces from scene_weaver_core import SceneWeaverCore from css_styles import CSSStyles from scene_templates import SceneTemplateManager from inpainting_templates import InpaintingTemplateManager logger = logging.getLogger(__name__) logger.setLevel(logging.INFO) logging.basicConfig( level=logging.INFO, format='%(asctime)s [%(name)s] %(levelname)s: %(message)s', datefmt='%H:%M:%S' ) class UIManager: """ Gradio UI Manager with support for background generation and inpainting. Provides a professional interface with mode switching, template selection, and advanced parameter controls. Attributes: sceneweaver: SceneWeaverCore instance template_manager: Scene template manager inpainting_template_manager: Inpainting template manager """ def __init__(self): self.sceneweaver = SceneWeaverCore() self.template_manager = SceneTemplateManager() self.inpainting_template_manager = InpaintingTemplateManager() self.generation_history = [] self.inpainting_history = [] self._preview_sensitivity = 0.5 self._current_mode = "background" # "background" or "inpainting" def apply_template(self, display_name: str, current_negative: str) -> Tuple[str, str, float]: """ Apply a scene template to the prompt fields. Args: display_name: The display name from dropdown (e.g., "🏢 Modern Office") current_negative: Current negative prompt value Returns: Tuple of (prompt, negative_prompt, guidance_scale) """ if not display_name: return "", current_negative, 7.5 # Convert display name to template key template_key = self.template_manager.get_template_key_from_display(display_name) if not template_key: return "", current_negative, 7.5 template = self.template_manager.get_template(template_key) if template: prompt = template.prompt negative = self.template_manager.get_negative_prompt_for_template( template_key, current_negative ) guidance = template.guidance_scale return prompt, negative, guidance return "", current_negative, 7.5 def quick_preview( self, uploaded_image: Optional[Image.Image], sensitivity: float = 0.5 ) -> Optional[Image.Image]: """ Generate quick foreground preview using lightweight traditional methods. Args: uploaded_image: Uploaded PIL Image sensitivity: Detection sensitivity (0.0 - 1.0) Returns: Preview image with colored overlay or None """ if uploaded_image is None: return None try: logger.info(f"Generating quick preview (sensitivity={sensitivity:.2f})") img_array = np.array(uploaded_image.convert('RGB')) height, width = img_array.shape[:2] max_preview_size = 512 if max(width, height) > max_preview_size: scale = max_preview_size / max(width, height) new_w = int(width * scale) new_h = int(height * scale) img_array = cv2.resize(img_array, (new_w, new_h), interpolation=cv2.INTER_AREA) height, width = new_h, new_w gray = cv2.cvtColor(img_array, cv2.COLOR_RGB2GRAY) blurred = cv2.GaussianBlur(gray, (5, 5), 0) low_threshold = int(30 + (1 - sensitivity) * 50) high_threshold = int(100 + (1 - sensitivity) * 100) edges = cv2.Canny(blurred, low_threshold, high_threshold) kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (7, 7)) dilated = cv2.dilate(edges, kernel, iterations=2) contours, _ = cv2.findContours(dilated, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) mask = np.zeros((height, width), dtype=np.uint8) if contours: sorted_contours = sorted(contours, key=cv2.contourArea, reverse=True) min_area = (width * height) * 0.01 * (1 - sensitivity) for contour in sorted_contours: if cv2.contourArea(contour) > min_area: cv2.fillPoly(mask, [contour], 255) kernel_close = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (11, 11)) mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel_close) overlay = img_array.copy().astype(np.float32) fg_mask = mask > 127 overlay[fg_mask] = overlay[fg_mask] * 0.5 + np.array([0, 255, 0]) * 0.5 bg_mask = mask <= 127 overlay[bg_mask] = overlay[bg_mask] * 0.5 + np.array([255, 0, 0]) * 0.5 overlay = np.clip(overlay, 0, 255).astype(np.uint8) original_size = uploaded_image.size preview_image = Image.fromarray(overlay) if preview_image.size != original_size: preview_image = preview_image.resize(original_size, Image.LANCZOS) logger.info("Quick preview generated successfully") return preview_image except Exception as e: logger.error(f"Quick preview failed: {e}") return None def _save_result(self, combined_image: Image.Image, prompt: str): """Save result with memory-conscious history management""" if not combined_image: return output_dir = Path("outputs") output_dir.mkdir(exist_ok=True) combined_image.save(output_dir / "latest_combined.png") self.generation_history.append({ "prompt": prompt, "timestamp": time.time() }) max_history = self.sceneweaver.max_history if len(self.generation_history) > max_history: self.generation_history = self.generation_history[-max_history:] @spaces.GPU(duration=240) def generate_handler( self, uploaded_image: Optional[Image.Image], prompt: str, combination_mode: str, focus_mode: str, negative_prompt: str, steps: int, guidance: float, progress=gr.Progress() ): """Enhanced generation handler with memory management and ZeroGPU support""" if uploaded_image is None: return None, None, None, "Please upload an image to get started!", gr.update(visible=False) if not prompt.strip(): return None, None, None, "Please describe the background scene you'd like!", gr.update(visible=False) try: if not self.sceneweaver.is_initialized: progress(0.05, desc="Loading AI models (first time may take 2-3 minutes)...") def init_progress(msg, pct): if pct < 30: desc = "Loading image analysis models..." elif pct < 60: desc = "Loading Stable Diffusion XL..." elif pct < 90: desc = "Applying memory optimizations..." else: desc = "Almost ready..." progress(0.05 + (pct/100) * 0.2, desc=desc) self.sceneweaver.load_models(progress_callback=init_progress) def gen_progress(msg, pct): if pct < 20: desc = "Analyzing your image..." elif pct < 50: desc = "Generating background scene..." elif pct < 80: desc = "Blending foreground and background..." elif pct < 95: desc = "Applying final touches..." else: desc = "Complete!" progress(0.25 + (pct/100) * 0.75, desc=desc) result = self.sceneweaver.generate_and_combine( original_image=uploaded_image, prompt=prompt, combination_mode=combination_mode, focus_mode=focus_mode, negative_prompt=negative_prompt, num_inference_steps=int(steps), guidance_scale=float(guidance), progress_callback=gen_progress ) if result["success"]: combined = result["combined_image"] generated = result["generated_scene"] original = result["original_image"] self._save_result(combined, prompt) status_msg = "Image created successfully!" return combined, generated, original, status_msg, gr.update(visible=True) else: error_msg = result.get("error", "Something went wrong") return None, None, None, f"Error: {error_msg}", gr.update(visible=False) except Exception as e: error_traceback = traceback.format_exc() logger.error(f"Generation handler error: {str(e)}") logger.error(f"Traceback:\n{error_traceback}") return None, None, None, f"Error: {str(e)}", gr.update(visible=False) def create_interface(self): """Create professional user interface""" self._css = CSSStyles.get_main_css() # Check Gradio version for API compatibility self._gradio_version = gr.__version__ self._gradio_major = int(self._gradio_version.split('.')[0]) # Compatible with Gradio 4.44.0+ # Use minimal constructor arguments for maximum compatibility with gr.Blocks() as interface: # Inject CSS (compatible with all Gradio versions) gr.HTML(f"") # Header gr.HTML("""

🎨 SceneWeaver

AI-powered background generation and inpainting with professional edge processing

""") # Main Tabs for Mode Selection with gr.Tabs(elem_id="main-mode-tabs") as main_tabs: # Background Generation Tab with gr.Tab("Background Generation", elem_id="bg-gen-tab"): with gr.Row(): # Left Column - Input controls with gr.Column(scale=1, min_width=350, elem_classes=["feature-card"]): gr.HTML("""

📸 Upload & Generate

""") uploaded_image = gr.Image( label="Upload Your Image", type="pil", height=280, elem_classes=["input-field"] ) # Scene Template Selector (without Accordion to fix dropdown positioning in Gradio 5.x) template_dropdown = gr.Dropdown( label="Scene Templates", choices=[""] + self.template_manager.get_template_choices_sorted(), value="", info="24 curated scenes sorted A-Z (optional)", elem_classes=["template-dropdown"] ) prompt_input = gr.Textbox( label="Background Scene Description", placeholder="Select a template above or describe your own scene...", lines=3, elem_classes=["input-field"] ) combination_mode = gr.Dropdown( label="Composition Mode", choices=["center", "left_half", "right_half", "full"], value="center", info="center=Smart Center | left_half=Left Half | right_half=Right Half | full=Full Image", elem_classes=["input-field"] ) focus_mode = gr.Dropdown( label="Focus Mode", choices=["person", "scene"], value="person", info="person=Tight Crop | scene=Include Surrounding Objects", elem_classes=["input-field"] ) with gr.Accordion("Advanced Options", open=False): negative_prompt = gr.Textbox( label="Negative Prompt", value="blurry, low quality, distorted, people, characters", lines=2, elem_classes=["input-field"] ) steps_slider = gr.Slider( label="Quality Steps", minimum=15, maximum=50, value=25, step=5, elem_classes=["input-field"] ) guidance_slider = gr.Slider( label="Guidance Scale", minimum=5.0, maximum=15.0, value=7.5, step=0.5, elem_classes=["input-field"] ) generate_btn = gr.Button( "Generate Background", variant="primary", size="lg", elem_classes=["primary-button"] ) # Right Column - Results display with gr.Column(scale=2, elem_classes=["feature-card"], elem_id="results-gallery-centered"): gr.HTML("""

🎭 Results Gallery

""") # Loading notice gr.HTML("""
⏱️ First-time users: Initial model loading takes 1-2 minutes. Subsequent generations are much faster (~30s).
""") # Quick start guide gr.HTML("""
💡 Quick Start Guide

Step 1: Upload any image with a clear subject

Step 2: Describe or Choose your desired background scene

Step 3: Choose composition mode (center works best)

Step 4: Click Generate and wait for the magic!

Tip: For dark clothing, ensure good lighting in original photo.

""") with gr.Tabs(): with gr.TabItem("Final Result"): combined_output = gr.Image( label="Your Generated Image", elem_classes=["result-gallery"], show_label=False ) with gr.TabItem("Background"): generated_output = gr.Image( label="Generated Background", elem_classes=["result-gallery"], show_label=False ) with gr.TabItem("Original"): original_output = gr.Image( label="Processed Original", elem_classes=["result-gallery"], show_label=False ) status_output = gr.Textbox( label="Status", value="Ready to create! Upload an image and describe your vision.", interactive=False, elem_classes=["status-panel", "status-ready"] ) with gr.Row(): download_btn = gr.DownloadButton( "Download Result", value=None, visible=False, elem_classes=["secondary-button"] ) clear_btn = gr.Button( "Clear All", elem_classes=["secondary-button"] ) memory_btn = gr.Button( "Clean Memory", elem_classes=["secondary-button"] ) # Event handlers for Background Generation Tab # Template selection handler template_dropdown.change( fn=self.apply_template, inputs=[template_dropdown, negative_prompt], outputs=[prompt_input, negative_prompt, guidance_slider] ) generate_btn.click( fn=self.generate_handler, inputs=[ uploaded_image, prompt_input, combination_mode, focus_mode, negative_prompt, steps_slider, guidance_slider ], outputs=[ combined_output, generated_output, original_output, status_output, download_btn ] ) clear_btn.click( fn=lambda: (None, None, None, "Ready to create!", gr.update(visible=False)), outputs=[combined_output, generated_output, original_output, status_output, download_btn] ) memory_btn.click( fn=lambda: self.sceneweaver._ultra_memory_cleanup() or "Memory cleaned!", outputs=[status_output] ) combined_output.change( fn=lambda img: gr.update(value="outputs/latest_combined.png", visible=True) if (img is not None) else gr.update(visible=False), inputs=[combined_output], outputs=[download_btn] ) # End of Background Generation Tab # Inpainting Tab self.create_inpainting_tab() # Footer with tech credits (outside tabs) gr.HTML("""

Powered By

Stable Diffusion XL OpenCLIP BiRefNet rembg PyTorch Gradio

SceneWeaver © 2025  |  Built with SDXL and OpenCLIP

""") return interface def launch(self, share: bool = True, debug: bool = False): """Launch the UI interface""" interface = self.create_interface() # Launch kwargs compatible with Gradio 4.44.0+ # Keep minimal for maximum compatibility launch_kwargs = { "share": share, "debug": debug, "show_error": True, "quiet": False } return interface.launch(**launch_kwargs) # INPAINTING UI METHODS def apply_inpainting_template( self, display_name: str, current_prompt: str ) -> Tuple[str, float, int, str]: """ Apply an inpainting template to the UI fields. Parameters ---------- display_name : str Template display name from dropdown current_prompt : str Current prompt content Returns ------- tuple (prompt, conditioning_scale, feather_radius, conditioning_type) """ if not display_name: return current_prompt, 0.7, 8, "canny" template_key = self.inpainting_template_manager.get_template_key_from_display(display_name) if not template_key: return current_prompt, 0.7, 8, "canny" template = self.inpainting_template_manager.get_template(template_key) if template: params = self.inpainting_template_manager.get_parameters_for_template(template_key) return ( current_prompt, params.get('controlnet_conditioning_scale', 0.7), params.get('feather_radius', 8), params.get('preferred_conditioning', 'canny') ) return current_prompt, 0.7, 8, "canny" def extract_mask_from_editor(self, editor_output: Dict[str, Any]) -> Optional[Image.Image]: """ Extract mask from Gradio ImageEditor output. Handles different Gradio versions' output formats. Parameters ---------- editor_output : dict Output from gr.ImageEditor component Returns ------- PIL.Image or None Extracted mask as grayscale image """ if editor_output is None: return None try: # Gradio 5.x format if isinstance(editor_output, dict): # Check for 'layers' key (Gradio 5.x ImageEditor) if 'layers' in editor_output and editor_output['layers']: # Get the first layer as mask layer = editor_output['layers'][0] if isinstance(layer, np.ndarray): mask_array = layer elif isinstance(layer, Image.Image): mask_array = np.array(layer) else: return None # Check for 'composite' key elif 'composite' in editor_output: composite = editor_output['composite'] if isinstance(composite, np.ndarray): mask_array = composite elif isinstance(composite, Image.Image): mask_array = np.array(composite) else: return None else: return None elif isinstance(editor_output, np.ndarray): mask_array = editor_output elif isinstance(editor_output, Image.Image): mask_array = np.array(editor_output) else: logger.warning(f"Unexpected editor output type: {type(editor_output)}") return None # Convert to grayscale mask if len(mask_array.shape) == 3: if mask_array.shape[2] == 4: # RGBA format - extract white brush strokes from RGB channels # White brush strokes have high RGB values AND high alpha rgb_part = mask_array[:, :, :3] alpha_part = mask_array[:, :, 3] # Convert RGB to grayscale to detect white areas gray = cv2.cvtColor(rgb_part, cv2.COLOR_RGB2GRAY) # Combine: white areas (high gray value) with opacity (high alpha) # This captures white brush strokes mask_gray = np.minimum(gray, alpha_part) elif mask_array.shape[2] == 3: # RGB - convert to grayscale (white areas become white in mask) mask_gray = cv2.cvtColor(mask_array, cv2.COLOR_RGB2GRAY) else: mask_gray = mask_array[:, :, 0] else: # Already grayscale mask_gray = mask_array return Image.fromarray(mask_gray.astype(np.uint8), mode='L') except Exception as e: logger.error(f"Failed to extract mask from editor: {e}") return None @spaces.GPU(duration=420) def inpainting_handler( self, image: Optional[Image.Image], mask_editor: Dict[str, Any], prompt: str, template_dropdown: str, conditioning_type: str, conditioning_scale: float, feather_radius: int, guidance_scale: float, num_steps: int, progress: gr.Progress = gr.Progress() ) -> Tuple[Optional[Image.Image], Optional[Image.Image], Optional[Image.Image], str]: """ Handle inpainting generation request. Parameters ---------- image : PIL.Image Original image to inpaint mask_editor : dict Mask editor output prompt : str Text description of desired content template_dropdown : str Selected template (optional) conditioning_type : str ControlNet conditioning type conditioning_scale : float ControlNet influence strength feather_radius : int Mask feathering radius guidance_scale : float Guidance scale for generation num_steps : int Number of inference steps progress : gr.Progress Progress callback Returns ------- tuple (result_image, control_image, status_message) """ if image is None: return None, None, "⚠️ Please upload an image first" # Extract mask mask = self.extract_mask_from_editor(mask_editor) if mask is None: return None, None, "⚠️ Please draw a mask on the image" # Validate mask mask_array = np.array(mask) coverage = np.count_nonzero(mask_array > 127) / mask_array.size if coverage < 0.01: return None, None, "⚠️ Mask too small - please select a larger area" if coverage > 0.95: return None, None, "⚠️ Mask too large - consider using background generation instead" def progress_callback(msg: str, pct: int): progress(pct / 100, desc=msg) try: start_time = time.time() # Get template key if selected template_key = None if template_dropdown: template_key = self.inpainting_template_manager.get_template_key_from_display( template_dropdown ) # Execute inpainting through SceneWeaverCore facade result = self.sceneweaver.execute_inpainting( image=image, mask=mask, prompt=prompt, preview_only=False, template_key=template_key, conditioning_type=conditioning_type, controlnet_conditioning_scale=conditioning_scale, feather_radius=feather_radius, guidance_scale=guidance_scale, num_inference_steps=num_steps, progress_callback=progress_callback ) elapsed = time.time() - start_time if result.get('success'): # Store in history self.inpainting_history.append({ 'result': result.get('combined_image'), 'prompt': prompt, 'time': elapsed }) if len(self.inpainting_history) > 3: self.inpainting_history.pop(0) quality_score = result.get('quality_score', 0) # Clean, simple status message status = f"✅ Inpainting complete in {elapsed:.1f}s" if quality_score > 0: status += f" | Quality: {quality_score:.0f}/100" return ( result.get('combined_image'), result.get('control_image'), status ) else: error_msg = result.get('error', 'Unknown error') return None, None, f"❌ Inpainting failed: {error_msg}" except Exception as e: logger.error(f"Inpainting handler error: {e}") logger.error(traceback.format_exc()) return None, None, f"❌ Error: {str(e)}" def create_inpainting_tab(self) -> gr.Tab: """ Create the inpainting tab UI. Returns ------- gr.Tab Configured inpainting tab component """ with gr.Tab("Inpainting", elem_id="inpainting-tab") as tab: gr.HTML("""

ControlNet Inpainting BETA

Draw a mask to select the area you want to regenerate

⚠️ Beta Feature - Continuously Optimizing
Results may vary depending on complexity. Use templates and detailed prompts for best results. Advanced features (like Add Accessories) may require multiple attempts.

""") with gr.Row(): # Left column - Input with gr.Column(scale=1): # Image upload inpaint_image = gr.Image( label="Upload Image", type="pil", height=300 ) # Mask editor mask_editor = gr.ImageEditor( label="Draw Mask (white = area to inpaint)", type="pil", height=300, brush=gr.Brush(colors=["#FFFFFF"], default_size=20), eraser=gr.Eraser(default_size=20), layers=True, sources=["upload"], image_mode="RGBA" ) # Template selection with gr.Accordion("Inpainting Templates", open=False): inpaint_template = gr.Dropdown( choices=[""] + self.inpainting_template_manager.get_template_choices_sorted(), value="", label="Select Template", elem_classes=["template-dropdown"] ) template_tips = gr.Markdown("") # Prompt inpaint_prompt = gr.Textbox( label="Prompt", placeholder="Describe what you want to generate in the masked area...", lines=2 ) # Right column - Settings and Output with gr.Column(scale=1): # Settings with gr.Accordion("Generation Settings", open=True): conditioning_type = gr.Radio( choices=["canny", "depth"], value="canny", label="ControlNet Mode" ) conditioning_scale = gr.Slider( minimum=0.05, maximum=1.0, value=0.7, step=0.05, label="ControlNet Strength" ) feather_radius = gr.Slider( minimum=0, maximum=20, value=8, step=1, label="Feather Radius (px)" ) with gr.Accordion("Advanced Settings", open=False): inpaint_guidance = gr.Slider( minimum=5.0, maximum=15.0, value=7.5, step=0.5, label="Guidance Scale" ) inpaint_steps = gr.Slider( minimum=15, maximum=50, value=25, step=5, label="Inference Steps" ) # Generate button inpaint_btn = gr.Button( "Generate Inpainting", variant="primary", elem_classes=["primary-button"] ) # Processing time reminder gr.Markdown( """

⏳ Please be patient! Inpainting typically takes 5-7 minutes depending on GPU availability and image complexity. Please don't refresh the page while processing.

🔄 Want to make more changes? After each generation, please re-upload your image and draw a new mask if you want to apply additional edits.

""" ) # Status inpaint_status = gr.Textbox( label="Status", value="Ready for inpainting", interactive=False ) # Output row with gr.Row(): with gr.Column(scale=1): inpaint_result = gr.Image( label="Result", type="pil", height=400 ) with gr.Column(scale=1): # Control image (structure guidance visualization) inpaint_control = gr.Image( label="Control Image (Structure Guidance)", type="pil", height=400 ) # Event handlers inpaint_template.change( fn=self.apply_inpainting_template, inputs=[inpaint_template, inpaint_prompt], outputs=[inpaint_prompt, conditioning_scale, feather_radius, conditioning_type] ) inpaint_template.change( fn=lambda x: self._get_template_tips(x), inputs=[inpaint_template], outputs=[template_tips] ) # Copy uploaded image to mask editor inpaint_image.change( fn=lambda x: x, inputs=[inpaint_image], outputs=[mask_editor] ) inpaint_btn.click( fn=self.inpainting_handler, inputs=[ inpaint_image, mask_editor, inpaint_prompt, inpaint_template, conditioning_type, conditioning_scale, feather_radius, inpaint_guidance, inpaint_steps ], outputs=[ inpaint_result, inpaint_control, inpaint_status ] ) return tab def _get_template_tips(self, display_name: str) -> str: """Get usage tips for selected template.""" if not display_name: return "" template_key = self.inpainting_template_manager.get_template_key_from_display(display_name) if not template_key: return "" tips = self.inpainting_template_manager.get_usage_tips(template_key) if tips: return "**Tips:**\n" + "\n".join(f"- {tip}" for tip in tips) return ""