Spaces:
Running
Running
| """ | |
| Gradio Web Interface for Stencil Image Generator | |
| This module provides a web-based UI for the Stencil Generator using Gradio. | |
| Run this file to launch the interactive web interface. | |
| """ | |
| import gradio as gr | |
| from Stencil import StencilGenerator | |
| from StencilCV import StencilCV | |
| import torch | |
| from typing import Optional | |
| import numpy as np | |
| import os | |
| MAX_IMAGES = 4 | |
| class StencilApp: | |
| """Wrapper class for the Gradio application.""" | |
| def __init__(self): | |
| """Initialize the Stencil Generator.""" | |
| self.generator = None | |
| self.current_model_type = None | |
| self.original_images = [] # Store original images for toggling | |
| self.outlined_status = [] # Track which images have outline applied | |
| def load_model(self, model_type: str = "Standard SD 2.1"): | |
| """ | |
| Lazy load the model when first needed or reload if model type changed. | |
| Args: | |
| model_type: Type of model to load ("Standard SD 2.1", "Checkpoint-500", "Checkpoint-1000") | |
| """ | |
| # Reload if model type changed or first load | |
| if self.generator is None or self.current_model_type != model_type: | |
| print(f"Initializing Stencil Generator with {model_type}...") | |
| # Determine checkpoint path based on model type | |
| # Can be local path or HuggingFace Hub model ID | |
| checkpoint_path = None | |
| if model_type == "Checkpoint-500": | |
| # Try local path first, fallback to HuggingFace Hub | |
| checkpoint_path = "./Fine-tuning/checkpoint-500" | |
| if not os.path.exists(checkpoint_path): | |
| checkpoint_path = "mrpink925/stencilai-checkpoint-500" | |
| elif model_type == "Checkpoint-1000": | |
| # Try local path first, fallback to HuggingFace Hub | |
| checkpoint_path = "./Fine-tuning/checkpoint-1000" | |
| if not os.path.exists(checkpoint_path): | |
| checkpoint_path = "mrpink925/stencilai-checkpoint-1000" | |
| self.generator = StencilGenerator( | |
| model_id="Manojb/stable-diffusion-2-1-base", | |
| checkpoint_path=checkpoint_path, | |
| use_fp16=torch.cuda.is_available() | |
| ) | |
| self.current_model_type = model_type | |
| return self.generator | |
| def generate_stencil( | |
| self, | |
| prompt: str, | |
| model_type: str, | |
| negative_prompt: Optional[str], | |
| num_images: int, | |
| num_inference_steps: int, | |
| guidance_scale: float, | |
| width: int, | |
| height: int, | |
| seed: int, | |
| use_seed: bool, | |
| add_stencil_suffix: bool, | |
| clean_background: bool | |
| ): | |
| """ | |
| Generate stencil images based on user inputs. | |
| This is the main function called by the Gradio interface. | |
| """ | |
| if not prompt or prompt.strip() == "": | |
| return [], "Please enter a prompt!" | |
| try: | |
| # Load model (will reload if model type changed) | |
| generator = self.load_model(model_type) | |
| # Generate the image(s) | |
| images = generator.generate( | |
| prompt=prompt, | |
| negative_prompt=negative_prompt if negative_prompt else None, | |
| num_images=num_images, | |
| num_inference_steps=num_inference_steps, | |
| guidance_scale=guidance_scale, | |
| width=width, | |
| height=height, | |
| seed=seed if use_seed else None, | |
| add_stencil_suffix=add_stencil_suffix, | |
| clean_background=clean_background | |
| ) | |
| # Ensure images is a list | |
| if not isinstance(images, list): | |
| images = [images] | |
| # Store original images and reset outlined status | |
| self.original_images = [img.copy() for img in images] | |
| self.outlined_status = [False] * len(images) | |
| return images, f"Generation successful! Created {len(images)} image(s)." | |
| except Exception as e: | |
| return [], f"Error: {str(e)}" | |
| def apply_outline(self, gallery_data, selected_index): | |
| """ | |
| Toggle outline processing on a selected image using StencilCV. | |
| If the image has outline applied, revert to original. Otherwise, apply outline. | |
| Args: | |
| gallery_data: Gallery data from Gradio (list of images or tuples) | |
| selected_index: Index of the selected image (from gr.Gallery select event) | |
| Returns: | |
| Updated gallery and status message | |
| """ | |
| # print(f"DEBUG: apply_outline called") | |
| # print(f"DEBUG: gallery_data type: {type(gallery_data)}") | |
| # print(f"DEBUG: gallery_data length: {len(gallery_data) if gallery_data else 0}") | |
| # print(f"DEBUG: selected_index: {selected_index}") | |
| if not gallery_data: | |
| return gallery_data, "No images to process!" | |
| # If there's only 1 image and no selection, default to index 0 | |
| if selected_index is None: | |
| if len(self.original_images) == 1: | |
| selected_index = 0 | |
| else: | |
| return gallery_data, "Please select an image first by clicking on it!" | |
| if selected_index >= len(self.original_images): | |
| return gallery_data, "Error: Image index out of range!" | |
| try: | |
| # Create a copy of the gallery data | |
| updated_gallery = list(gallery_data) | |
| # Check if this image already has outline applied | |
| if self.outlined_status[selected_index]: | |
| # Revert to original | |
| # print(f"DEBUG: Reverting image {selected_index} to original") | |
| updated_gallery[selected_index] = self.original_images[selected_index].copy() | |
| self.outlined_status[selected_index] = False | |
| return updated_gallery, f"Reverted image {selected_index + 1} to original." | |
| else: | |
| # Apply outline | |
| # print(f"DEBUG: Applying outline to image {selected_index}") | |
| # Initialize StencilCV processor | |
| processor = StencilCV() | |
| # Get the original image (not the gallery one, to ensure consistency) | |
| original_img = self.original_images[selected_index] | |
| # print(f"DEBUG: Applying edge_stencil...") | |
| # Apply outline to the original image | |
| outlined = processor.edge_stencil(original_img) | |
| # print(f"DEBUG: Outline applied successfully!") | |
| # Update gallery with outlined version | |
| updated_gallery[selected_index] = outlined | |
| self.outlined_status[selected_index] = True | |
| return updated_gallery, f"Applied outline to image {selected_index + 1}. Click again to revert." | |
| except Exception as e: | |
| import traceback | |
| print("DEBUG: Exception occurred:") | |
| traceback.print_exc() | |
| return gallery_data, f"Error applying outline: {str(e)}" | |
| def create_interface(): | |
| """Create and configure the Gradio interface.""" | |
| app = StencilApp() | |
| # Define the interface | |
| with gr.Blocks(title="Stencil Image Generator", theme=gr.themes.Soft()) as interface: | |
| gr.Markdown( | |
| """ | |
| # 🎨 Stencil Image Generator | |
| Generate black and white stencil-style images using AI. Perfect for creating | |
| cutting templates, vector art, and silhouette designs. | |
| """ | |
| ) | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| # Input controls | |
| prompt = gr.Textbox( | |
| label="Prompt", | |
| placeholder="e.g., a cat sitting, a tree with spreading branches...", | |
| lines=3 | |
| ) | |
| model_selector = gr.Radio( | |
| choices=["Standard SD 2.1", "Checkpoint-500", "Checkpoint-1000"], | |
| value="Checkpoint-1000", | |
| label="Model Type", | |
| info="Choose between standard model or fine-tuned checkpoints (trained on sketch-style images)" | |
| ) | |
| num_images = gr.Slider( | |
| minimum=1, | |
| maximum=MAX_IMAGES, | |
| value=2, | |
| step=1, | |
| label="Number of Images", | |
| info="Generate multiple variations to choose from" | |
| ) | |
| with gr.Accordion("Advanced Settings", open=False): | |
| negative_prompt = gr.Textbox( | |
| label="Negative Prompt (optional)", | |
| placeholder="Things to avoid in the generation...", | |
| lines=2 | |
| ) | |
| add_stencil_suffix = gr.Checkbox( | |
| label="Add stencil styling suffix(recommended)", | |
| value=True, | |
| info="Automatically adds stencil-specific styling to your prompt (prompt decorations)" | |
| ) | |
| clean_background = gr.Checkbox( | |
| label="Clean white background (recommended)", | |
| value=True, | |
| info="Post-process to ensure pure white background and remove artifacts" | |
| ) | |
| num_inference_steps = gr.Slider( | |
| minimum=10, | |
| maximum=50, | |
| value=25, | |
| step=5, | |
| label="Inference Steps", | |
| info="Higher = better quality but slower" | |
| ) | |
| guidance_scale = gr.Slider( | |
| minimum=1, | |
| maximum=15, | |
| value=7.5, | |
| step=0.5, | |
| label="Guidance Scale", | |
| info="How closely to follow the prompt (7-8 recommended)" | |
| ) | |
| with gr.Row(): | |
| width = gr.Slider( | |
| minimum=256, | |
| maximum=1024, | |
| value=512, | |
| step=64, | |
| label="Width" | |
| ) | |
| height = gr.Slider( | |
| minimum=256, | |
| maximum=1024, | |
| value=512, | |
| step=64, | |
| label="Height" | |
| ) | |
| with gr.Row(): | |
| use_seed = gr.Checkbox( | |
| label="Use fixed seed", | |
| value=False, | |
| info="Enable for reproducible results" | |
| ) | |
| seed = gr.Number( | |
| label="Seed", | |
| value=42, | |
| precision=0 | |
| ) | |
| generate_btn = gr.Button("Generate Stencil", variant="primary", size="lg") | |
| # Example prompts | |
| gr.Examples( | |
| examples=[ | |
| ["a cat sitting"], | |
| ["a tree with spreading branches"], | |
| ["a bicycle"], | |
| ["a coffee cup"], | |
| ["a bird in flight"], | |
| ["a deer with antlers"], | |
| ["a mountain landscape"], | |
| ["a lighthouse by the sea"] | |
| ], | |
| inputs=prompt, | |
| label="Example Prompts" | |
| ) | |
| with gr.Column(scale=1): | |
| # Output - Gallery for multiple images | |
| output_gallery = gr.Gallery( | |
| label="Generated Stencils (click to select)", | |
| show_label=True, | |
| columns=2, | |
| rows=2, | |
| height="auto", | |
| object_fit="contain" | |
| ) | |
| status_text = gr.Textbox( | |
| label="Status", | |
| interactive=False, | |
| lines=1 | |
| ) | |
| # Hidden state to track selected image | |
| selected_image_index = gr.State(value=None) | |
| # Post-processing section | |
| with gr.Accordion("Post-Processing Options", open=False): | |
| gr.Markdown( | |
| """ | |
| **Outline Generation**: Click an image above to select it, then click the button below | |
| to toggle outline processing. This creates a line-art effect and works best on images | |
| with clear subjects. Click the button again to revert to the original. | |
| You can toggle outline on/off for each image independently to compare styles. | |
| """ | |
| ) | |
| apply_outline_btn = gr.Button("Toggle Outline on Selected Image", variant="secondary") | |
| gr.Markdown( | |
| """ | |
| ### Tips for Best Results: | |
| - Keep prompts simple and descriptive | |
| - **Standard SD 2.1**: Best for general stencils with detailed prompt engineering | |
| - **Checkpoint models**: Fine-tuned for sketch-style stencils (automatically adds "sketch of" prefix) | |
| - Generate multiple images to see variations | |
| - Use negative prompts to avoid unwanted features (works best with Standard SD 2.1) | |
| - Try the outline option after generation for different styles | |
| - Higher inference steps = better quality (but slower) | |
| """ | |
| ) | |
| # Connect the generate button | |
| generate_btn.click( | |
| fn=app.generate_stencil, | |
| inputs=[ | |
| prompt, | |
| model_selector, | |
| negative_prompt, | |
| num_images, | |
| num_inference_steps, | |
| guidance_scale, | |
| width, | |
| height, | |
| seed, | |
| use_seed, | |
| add_stencil_suffix, | |
| clean_background | |
| ], | |
| outputs=[output_gallery, status_text] | |
| ) | |
| # Track when user selects an image in the gallery | |
| def update_selection(evt: gr.SelectData): | |
| print(f"DEBUG: Gallery selection event - index: {evt.index}") | |
| return evt.index | |
| output_gallery.select( | |
| fn=update_selection, | |
| outputs=selected_image_index | |
| ) | |
| # Connect the outline button | |
| apply_outline_btn.click( | |
| fn=app.apply_outline, | |
| inputs=[output_gallery, selected_image_index], | |
| outputs=[output_gallery, status_text] | |
| ) | |
| gr.Markdown( | |
| """ | |
| --- | |
| Built with [Stable Diffusion](https://stability.ai/) and [Gradio](https://gradio.app/) | |
| """ | |
| ) | |
| return interface | |
| def launch( | |
| share: bool = False, | |
| server_name: str = "0.0.0.0", | |
| server_port: int = 7860, | |
| **kwargs | |
| ): | |
| """ | |
| Launch the Gradio interface. | |
| Args: | |
| share: Whether to create a public shareable link | |
| server_name: Server host (0.0.0.0 for public access) | |
| server_port: Port to run the server on | |
| **kwargs: Additional arguments passed to gradio.launch() | |
| """ | |
| interface = create_interface() | |
| interface.launch( | |
| share=share, | |
| server_name=server_name, | |
| server_port=server_port, | |
| **kwargs | |
| ) | |
| if __name__ == "__main__": | |
| # Launch with default settings | |
| # Set share=True to create a public link | |
| launch(share=True, pwa=True) | |