Add new features and fixes

- add new features like :
1. choose between multiple models
2. add optimized model for CPU using LCM method
3. add raw prompt input
4. add progress bar for generation

- fix some issues :
1. CPU/GPU compatibility and add force CPU mode for testing
2. fix About section and other documentation issues

README.md

@@ -23,11 +23,19 @@ AI-powered meme generator using Stable Diffusion and LoRA fine-tuning.
 
 ## 🌟 Features
 
+- **Multiple Model Support**: Switch between fine-tuned LoRA and base models
+  - Pepe Fine-tuned (LoRA) - Custom trained model
+  - Base SD 1.5 - Standard Stable Diffusion
+  - Dreamlike Photoreal 2.0 - Photorealistic style
+  - Openjourney v4 - Artistic Midjourney-style
+- **Raw Prompt Mode**: Use exact prompts without automatic enhancements
 - Generate **custom Pepe memes** from text prompts  
 - Multiple **style presets** (happy, sad, smug, angry, etc.)  
-- **Add meme text overlays** and download results  
+- **Add meme text overlays** with automatic "MJ" signature
+- **Real-time progress tracking** for each generation step
 - Adjustable generation parameters (CFG, steps, seed, etc.)  
-- Batch generation and meme gallery system  
+- Batch generation and meme gallery system
+- **GPU & CPU compatible** with automatic optimization  
 
 ---
 
@@ -54,22 +62,6 @@ pip install -r requirements.txt
 streamlit run src/app.py
 ```
 
----
-
-## 🚀 Deployment on Hugging Face Spaces
-
-This app is optimized for deployment on Hugging Face Spaces with the following fixes:
-
-- **CPU Compatibility**: Uses `torch.float32` on CPU deployments to avoid dtype errors
-- **Memory Optimization**: Automatically enables attention and VAE slicing
-- **Error Handling**: Proper exception handling for optional dependencies like xformers
-- **Docker Support**: Updated Dockerfile with Python 3.11 and necessary system packages
-
-### Deployment Fixes Applied:
-- Fixed mixed dtype errors when running on CPU-only environments
-- Removed autocast context that can cause tensor type mismatches
-- Added proper device detection and dtype selection
-- Enhanced error handling for optional GPU optimizations
 
 ---
 

src/app.py

@@ -41,12 +41,33 @@ def init_session_state():
         st.session_state.generated_images = []
     if 'generation_count' not in st.session_state:
         st.session_state.generation_count = 0
+    if 'current_model' not in st.session_state:
+        st.session_state.current_model = None
 
 
 @st.cache_resource
-def load_generator():
-    """Load and cache the generator"""
-    return PepeGenerator()
+def load_generator(model_name: str = "Pepe Fine-tuned (LoRA)"):
+    """Load and cache the generator based on selected model"""
+    config = ModelConfig()
+    model_config = config.AVAILABLE_MODELS[model_name]
+    
+    # Update config with selected model settings
+    config.BASE_MODEL = model_config['base']
+    config.LORA_PATH = model_config.get('lora')
+    config.USE_LORA = model_config.get('use_lora', False)
+    config.TRIGGER_WORD = model_config.get('trigger_word', 'pepe the frog')
+    
+    # LCM settings
+    config.USE_LCM = model_config.get('use_lcm', False)
+    config.LCM_LORA_PATH = model_config.get('lcm_lora')
+    
+    # Log which model is being loaded
+    import logging
+    logger = logging.getLogger(__name__)
+    logger.info(f"Loading model: {model_name}")
+    logger.info(f"Base: {config.BASE_MODEL}, LoRA: {config.USE_LORA}, LCM: {config.USE_LCM}")
+    
+    return PepeGenerator(config)
 
 
 def get_example_prompts():
@@ -64,14 +85,41 @@ def main():
     """Main application"""
     init_session_state()
     
+    # Sidebar (needs to be first to define selected_model)
+    st.sidebar.header("⚙️ Settings")
+    
+    # Model selection
+    st.sidebar.subheader("🤖 Model Selection")
+    config = ModelConfig()
+    available_models = list(config.AVAILABLE_MODELS.keys())
+    selected_model = st.sidebar.selectbox(
+        "Choose Model",
+        available_models,
+        index=0,
+        help="Select which model to use for generation"
+    )
+    
+    # Detect model change and auto-clear cache
+    if st.session_state.current_model is not None and st.session_state.current_model != selected_model:
+        st.cache_resource.clear()
+        st.sidebar.success(f"✅ Switched to: {selected_model}")
+    
+    # Update current model in session state
+    st.session_state.current_model = selected_model
+    
+    # Show LCM mode indicator if enabled
+    model_config = config.AVAILABLE_MODELS[selected_model]
+    if model_config.get('use_lcm', False):
+        st.sidebar.success("⚡ LCM Mode: 8x Faster! (6-8 steps optimal)")
+    
     # Header
     st.title("🐸 Pepe the Frog Meme Generator")
     st.markdown("Create custom Pepe memes using AI! Powered by Stable Diffusion.")
     
-    # Sidebar
-    st.sidebar.header("⚙️ Settings")
+    st.sidebar.divider()
     
     # Style selection
+    st.sidebar.subheader("🎨 Style & Prompt")
     style_options = {
         "Default": "default",
         "😊 Happy": "happy",
@@ -88,10 +136,28 @@ def main():
     )
     style = style_options[selected_style]
     
-    # Advanced settings
+    # Raw prompt mode
+    use_raw_prompt = st.sidebar.checkbox(
+        "Raw Prompt Mode",
+        help="Use your exact prompt without trigger words or style modifiers"
+    )
+    
+    # Advanced settings - adjust defaults based on LCM mode
+    is_lcm_mode = model_config.get('use_lcm', False)
+    
     with st.sidebar.expander("🔧 Advanced Settings"):
-        steps = st.slider("Steps", 20, 100, 50, 5)
-        guidance = st.slider("Guidance Scale", 1.0, 20.0, 7.5, 0.5)
+        if is_lcm_mode:
+            # LCM needs fewer steps and lower guidance
+            steps = st.slider("Steps", 4, 12, 6, 1, 
+                            help="⚡ LCM Mode: 4-8 steps optimal. Recommended: 6")
+            guidance = st.slider("Guidance Scale", 1.0, 2.5, 1.5, 0.1,
+                               help="⚡ LCM Mode: Lower guidance (1.0-2.0). Recommended: 1.5")
+        else:
+            # Normal mode settings
+            steps = st.slider("Steps", 15, 50, 25, 5, 
+                            help="Fewer steps = faster generation. 20-25 recommended for CPU")
+            guidance = st.slider("Guidance Scale", 1.0, 20.0, 7.5, 0.5)
+        
         use_seed = st.checkbox("Fixed Seed")
         seed = st.number_input("Seed", 0, 999999, 42) if use_seed else None
     
@@ -133,7 +199,7 @@ def main():
         if st.session_state.generated_images:
             placeholder.image(
                 st.session_state.generated_images[-1],
-                use_container_width=True
+                width='stretch'
             )
         else:
             placeholder.info("Your meme will appear here...")
@@ -141,40 +207,58 @@ def main():
     # Generate
     if generate and prompt:
         try:
-            generator = load_generator()
+            generator = load_generator(selected_model)
+            processor = ImageProcessor()
+            
+            # Overall progress for multiple images
+            overall_progress = st.progress(0)
+            overall_status = st.empty()
             
-            progress = st.progress(0)
-            status = st.empty()
+            # Progress for current image generation steps
+            step_progress = st.progress(0)
+            step_status = st.empty()
             
             for i in range(num_vars):
-                status.text(f"Generating {i+1}/{num_vars}...")
-                progress.progress((i + 1) / num_vars)
+                overall_status.text(f"Generating image {i+1}/{num_vars}...")
                 
-                # Generate
+                # Define callback for step-by-step progress
+                def progress_callback(current_step: int, total_steps: int):
+                    step_progress.progress(current_step / total_steps)
+                    step_status.text(f"Step {current_step}/{total_steps}")
+                
+                # Generate with progress callback
                 image = generator.generate(
                     prompt=prompt,
                     style=style,
                     num_inference_steps=steps,
                     guidance_scale=guidance,
-                    seed=seed
+                    seed=seed,
+                    callback=progress_callback,
+                    raw_prompt=use_raw_prompt
                 )
                 
                 # Add text if requested
                 if add_text and (top_text or bottom_text):
-                    processor = ImageProcessor()
                     image = processor.add_meme_text(image, top_text, bottom_text)
                 
+                # Always add MJ signature
+                image = processor.add_signature(image, signature="MJaheen", font_size=10, opacity=200)
+                
                 st.session_state.generated_images.append(image)
                 st.session_state.generation_count += 1
+                
+                # Update overall progress
+                overall_progress.progress((i + 1) / num_vars)
             
-            progress.empty()
-            status.empty()
-            
-            st.success("✅ Meme generated!")
+            # Clear progress indicators
+            overall_progress.empty()
+            overall_status.empty()
+            step_progress.empty()
+            step_status.empty()
             
             # Show result
             if num_vars == 1:
-                placeholder.image(image, use_container_width=True)
+                placeholder.image(image, width='stretch')
                 
                 # Download
                 buf = io.BytesIO()
@@ -190,7 +274,7 @@ def main():
                 cols = st.columns(min(num_vars, 2))
                 for idx, img in enumerate(st.session_state.generated_images[-num_vars:]):
                     with cols[idx % 2]:
-                        st.image(img, use_container_width=True)
+                        st.image(img, width='stretch')
         
         except Exception as e:
             st.error(f"Error: {str(e)}")
@@ -205,7 +289,7 @@ def main():
             cols = st.columns(4)
             for idx, img in enumerate(reversed(st.session_state.generated_images[-8:])):
                 with cols[idx % 4]:
-                    st.image(img, use_container_width=True)
+                    st.image(img, width='stretch')
     
     # Footer
     st.divider()
@@ -219,6 +303,30 @@ def main():
             st.session_state.generated_images = []
             st.session_state.generation_count = 0
             st.rerun()
+    
+    # Personal Information
+    st.divider()
+    st.markdown("### 👨‍💻 About the Engineer")
+    info_col1, info_col2 = st.columns(2)
+    
+    with info_col1:
+        st.markdown("""
+        **Contact Information:**
+        - 📧 Email: [[email protected]](mailto:[email protected])
+        - 🔗 LinkedIn: [Mohamed Jaheen](https://www.linkedin.com/in/mohamedjaheen/)
+        """)
+    
+    with info_col2:
+        st.markdown("""
+        **About this App:**
+        - supported by worldquant university
+        - Built with Streamlit & Stable Diffusion
+        - Fine-tuned Pepe model available
+        - Open source and customizable
+        - MIT licences
+        """)
+    
+    st.caption("© 2025 - AI Meme Generator (Pepe the Frog) | Made with ❤️ using Python and MJ")
 
 
 if __name__ == "__main__":

src/model/config.py

@@ -8,12 +8,79 @@ from typing import Optional
 class ModelConfig:
     """Model configuration parameters"""
     
-    # Model paths
-    BASE_MODEL: str ="runwayml/stable-diffusion-v1-5"
+    # Available models
+    AVAILABLE_MODELS: dict = None
+    
+    def __post_init__(self):
+        if self.AVAILABLE_MODELS is None:
+            self.AVAILABLE_MODELS = {
+                "Pepe Fine-tuned (LoRA)": {
+                    "base": "runwayml/stable-diffusion-v1-5",
+                    "lora": "MJaheen/Pepe_The_Frog_model_v1_lora",
+                    "trigger_word": "pepe_style_frog",
+                    "use_lora": True,
+                    "use_lcm": False
+                },
+                "Pepe + LCM (FAST)": {
+                    "base": "runwayml/stable-diffusion-v1-5",
+                    "lora": "MJaheen/Pepe_The_Frog_model_v1_lora",
+                    "lcm_lora": "latent-consistency/lcm-lora-sdv1-5",
+                    "trigger_word": "pepe_style_frog",
+                    "use_lora": True,
+                    "use_lcm": True
+                },
+                "Base SD 1.5": {
+                    "base": "runwayml/stable-diffusion-v1-5",
+                    "lora": None,
+                    "trigger_word": "pepe the frog",
+                    "use_lora": False,
+                    "use_lcm": False
+                },
+                "Dreamlike Photoreal 2.0": {
+                    "base": "dreamlike-art/dreamlike-photoreal-2.0",
+                    "lora": None,
+                    "trigger_word": "pepe the frog",
+                    "use_lora": False,
+                    "use_lcm": False
+                },
+                "Openjourney v4": {
+                    "base": "prompthero/openjourney-v4",
+                    "lora": None,
+                    "trigger_word": "pepe the frog",
+                    "use_lora": False,
+                    "use_lcm": False
+                },
+                "Tiny SD (Fast CPU)": {
+                    "base": "segmind/tiny-sd",
+                    "lora": None,
+                    "trigger_word": "pepe the frog",
+                    "use_lora": False,
+                    "use_lcm": False
+                },
+                "Small SD (Balanced CPU)": {
+                    "base": "segmind/small-sd",
+                    "lora": None,
+                    "trigger_word": "pepe the frog",
+                    "use_lora": False,
+                    "use_lcm": False
+                }
+            }
+    
+    # Default model selection
+    SELECTED_MODEL: str = "Pepe Fine-tuned (LoRA)"
+    
+    # Model paths (will be set based on selection)
+    BASE_MODEL: str = "runwayml/stable-diffusion-v1-5"
     LORA_PATH: str = "MJaheen/Pepe_The_Frog_model_v1_lora"
+    USE_LORA: bool = True
+    TRIGGER_WORD: str = "pepe_style_frog"
+    
+    # LCM settings
+    USE_LCM: bool = False
+    LCM_LORA_PATH: Optional[str] = None
     
     # Default generation parameters
-    DEFAULT_STEPS: int = 50
+    DEFAULT_STEPS: int = 25  # Reduced for faster CPU inference (was 50)
     DEFAULT_GUIDANCE: float = 7.5
     DEFAULT_WIDTH: int = 512
     DEFAULT_HEIGHT: int = 512
@@ -27,6 +94,7 @@ class ModelConfig:
     # Performance
     ENABLE_ATTENTION_SLICING: bool = True
     ENABLE_VAE_SLICING: bool = True
+    FORCE_CPU: bool = True  # Set to True to force CPU, False to use GPU if available
     
     # Available styles
     AVAILABLE_STYLES: tuple = (

src/model/generator.py

@@ -1,11 +1,12 @@
 """Pepe Meme Generator - Core generation logic"""
 
-from typing import Optional, List
+from typing import Optional, List, Callable
 import torch
-from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
+from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler, LCMScheduler
 import streamlit as st
 from PIL import Image
 import logging
+import os
 
 from .config import ModelConfig
 
@@ -14,40 +15,118 @@ logger = logging.getLogger(__name__)
 
 class PepeGenerator:
     """Main generator class for creating Pepe memes"""
-    
+
     def __init__(self, config: Optional[ModelConfig] = None):
         """Initialize the generator"""
         self.config = config or ModelConfig()
-        self.device = self._get_device()
-        self.pipe = self._load_model()
+        self.device = self._get_device(self.config.FORCE_CPU)
+        self.pipe = self._load_model(
+            self.config.BASE_MODEL,
+            self.config.USE_LORA,
+            self.config.LORA_PATH,
+            self.config.FORCE_CPU,
+            self.config.USE_LCM,
+            self.config.LCM_LORA_PATH
+        )
         logger.info(f"PepeGenerator initialized on {self.device}")
-    
+
     @staticmethod
     @st.cache_resource
-    def _load_model() -> StableDiffusionPipeline:
-        """Load and cache the Stable Diffusion model"""
-        logger.info("Loading Stable Diffusion model...")
-        
+    def _load_model(base_model: str, use_lora: bool, lora_path: Optional[str], 
+                    force_cpu: bool = False, use_lcm: bool = False, 
+                    lcm_lora_path: Optional[str] = None) -> StableDiffusionPipeline:
+        """Load and cache the Stable Diffusion model with LoRA and LCM support"""
+        logger.info("="*60)
+        logger.info("LOADING NEW MODEL PIPELINE")
+        logger.info(f"Base Model: {base_model}")
+        logger.info(f"LoRA Enabled: {use_lora}")
+        if use_lora and lora_path:
+            logger.info(f"LoRA Path: {lora_path}")
+        logger.info(f"LCM Enabled: {use_lcm}")
+        if use_lcm and lcm_lora_path:
+            logger.info(f"LCM-LoRA Path: {lcm_lora_path}")
+        logger.info(f"Force CPU: {force_cpu}")
+        logger.info("="*60)
+
         # Determine appropriate dtype based on device
-        device = "cuda" if torch.cuda.is_available() else "cpu"
-        torch_dtype = torch.float16 if device == "cuda" else torch.float32
+        if force_cpu:
+            device = "cpu"
+            logger.info("🔧 FORCED CPU MODE - GPU disabled for testing")
+        else:
+            device = "cuda" if torch.cuda.is_available() else "cpu"
         
+        torch_dtype = torch.float16 if (device == "cuda" and not force_cpu) else torch.float32
+        logger.info(f"Using device: {device}, dtype: {torch_dtype}")
+
         pipe = StableDiffusionPipeline.from_pretrained(
-            ModelConfig.BASE_MODEL,
+            base_model,
             torch_dtype=torch_dtype,
             safety_checker=None,  # Disabled for meme generation - users must comply with SD license
         )
+
+        # Load LoRA weights if configured
+        if use_lora and lora_path:
+            logger.info(f"Loading LoRA weights from: {lora_path}")
+            try:
+                # Check if it's a local path or Hugging Face model ID
+                # Explicitly name it "pepe" to avoid "default_0" naming
+                if os.path.exists(lora_path):
+                    # Local path
+                    pipe.load_lora_weights(lora_path, adapter_name="pepe")
+                    logger.info("LoRA weights loaded successfully from local path")
+                elif "/" in lora_path:
+                    # Hugging Face model ID (format: username/model_name)
+                    pipe.load_lora_weights(lora_path, adapter_name="pepe")
+                    logger.info(f"✅ LoRA weights loaded successfully from Hugging Face: {lora_path}")
+                else:
+                    logger.warning(f"Invalid LoRA path format: {lora_path}")
+                
+                # If not using LCM, set Pepe LoRA as the active adapter
+                if not use_lcm:
+                    pipe.set_adapters(["pepe"])
+                    logger.info("✅ Pepe LoRA active")
+            except Exception as e:
+                logger.error(f"Failed to load LoRA weights: {e}")
+                logger.info("Continuing without LoRA weights...")
         
-        # Optimize scheduler
-        pipe.scheduler = DPMSolverMultistepScheduler.from_config(
-            pipe.scheduler.config
-        )
-        
+        # Load LCM-LoRA on top if configured (this enables fast inference!)
+        if use_lcm and lcm_lora_path:
+            logger.info(f"Loading LCM-LoRA from: {lcm_lora_path}")
+            try:
+                # Load LCM-LoRA as a separate adapter
+                pipe.load_lora_weights(lcm_lora_path, adapter_name="lcm")
+                logger.info("✅ LCM-LoRA loaded successfully")
+                
+                # If we have both Pepe LoRA and LCM-LoRA, fuse them
+                if use_lora:
+                    logger.info("Fusing Pepe LoRA + LCM-LoRA adapters...")
+                    # Use the correct adapter names: "pepe" and "lcm"
+                    pipe.set_adapters(["pepe", "lcm"], adapter_weights=[1.0, 1.0])
+                    logger.info("✅ Both LoRAs fused successfully (pepe + lcm)")
+                else:
+                    # Only LCM, set it as active
+                    pipe.set_adapters(["lcm"])
+                    logger.info("✅ LCM-LoRA active (solo mode)")
+            except Exception as e:
+                logger.error(f"Failed to load LCM-LoRA: {e}")
+                logger.info("Continuing without LCM...")
+                use_lcm = False
+
+        # Set appropriate scheduler based on LCM mode
+        if use_lcm:
+            pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
+            logger.info("⚡ Using LCM Scheduler (few-step mode)")
+        else:
+            pipe.scheduler = DPMSolverMultistepScheduler.from_config(
+                pipe.scheduler.config
+            )
+            logger.info("🔧 Using DPM Solver Scheduler (standard mode)")
+
         # Enable memory optimizations
         pipe.enable_attention_slicing()
         pipe.enable_vae_slicing()
-        
-        if device == "cuda":
+
+        if device == "cuda" and not force_cpu:
             pipe = pipe.to("cuda")
             try:
                 pipe.enable_xformers_memory_efficient_attention()
@@ -56,16 +135,21 @@ class PepeGenerator:
             except Exception as e:
                 logger.warning(f"Could not enable xformers: {e}")
         else:
-            logger.info("Running on CPU - memory optimizations applied")
-        
+            if force_cpu:
+                logger.info("Running on CPU - FORCED for testing")
+            else:
+                logger.info("Running on CPU - memory optimizations applied")
+
         logger.info("Model loaded successfully")
         return pipe
-    
+
     @staticmethod
-    def _get_device() -> str:
+    def _get_device(force_cpu: bool = False) -> str:
         """Determine the best available device"""
+        if force_cpu:
+            return "cpu"
         return "cuda" if torch.cuda.is_available() else "cpu"
-    
+
     def generate(
         self,
         prompt: str,
@@ -76,23 +160,42 @@ class PepeGenerator:
         seed: Optional[int] = None,
         width: int = 512,
         height: int = 512,
+        callback: Optional[Callable[[int, int], None]] = None,
+        raw_prompt: bool = False,
     ) -> Image.Image:
-        """Generate a single Pepe meme image"""
-        
-        # Apply style preset
-        enhanced_prompt = self._apply_style_preset(prompt, style)
+        """Generate a single Pepe meme image
         
+        Args:
+            callback: Optional callback function (current_step, total_steps)
+            raw_prompt: If True, use prompt as-is without modifications
+        """
+
+        # Apply style preset or use raw prompt
+        if raw_prompt:
+            enhanced_prompt = prompt
+        else:
+            enhanced_prompt = self._apply_style_preset(prompt, style)
+
         # Set default negative prompt
         if negative_prompt is None:
             negative_prompt = self.config.DEFAULT_NEGATIVE_PROMPT
-        
+
         # Set seed for reproducibility
         generator = None
         if seed is not None:
             generator = torch.Generator(device=self.device).manual_seed(seed)
-        
+
         logger.info(f"Generating: {enhanced_prompt[:50]}...")
-        
+        logger.debug(f"Full prompt: {enhanced_prompt}")
+        logger.debug(f"Model config - Base: {self.config.BASE_MODEL}, LoRA: {self.config.USE_LORA}")
+
+        # Create callback wrapper if provided (using new API)
+        callback_on_step_end_fn = None
+        if callback:
+            def callback_on_step_end_fn(pipe, step, timestep, callback_kwargs):
+                callback(step + 1, num_inference_steps)
+                return callback_kwargs
+
         # Generate image (removed autocast for CPU compatibility)
         output = self.pipe(
             prompt=enhanced_prompt,
@@ -102,32 +205,32 @@ class PepeGenerator:
             generator=generator,
             width=width,
             height=height,
+            callback_on_step_end=callback_on_step_end_fn,
         )
-        
+
         return output.images[0]
-    
+
     def generate_batch(
         self,
         prompt: str,
         num_images: int = 4,
         **kwargs
     ) -> List[Image.Image]:
-        """Generate multiple variations"""
+        """Generate multiple variations with callback support"""
         images = []
         for i in range(num_images):
             if 'seed' not in kwargs:
                 kwargs['seed'] = torch.randint(0, 100000, (1,)).item()
-            
+
             image = self.generate(prompt, **kwargs)
             images.append(image)
-            
+
             if 'seed' in kwargs:
                 del kwargs['seed']
-        
+
         return images
-    
-    @staticmethod
-    def _apply_style_preset(prompt: str, style: str) -> str:
+
+    def _apply_style_preset(self, prompt: str, style: str) -> str:
         """Apply style-specific prompt enhancements"""
         style_modifiers = {
             "happy": "cheerful, smiling, joyful",
@@ -138,11 +241,14 @@ class PepeGenerator:
             "surprised": "shocked, amazed, wide eyes",
         }
         
-        base = f"pepe the frog, {prompt}"
+        # Use trigger word from config
+        trigger_word = self.config.TRIGGER_WORD
         
+        base = f"{trigger_word}, {prompt}"
+
         if style in style_modifiers:
             base = f"{base}, {style_modifiers[style]}"
-        
+
         base = f"{base}, high quality, detailed, meme art"
-        
+
         return base

src/utils/image_processor.py

@@ -72,6 +72,87 @@ class ImageProcessor:
         # Draw main text
         draw.text(position, text, font=font, fill="white", anchor="mm")
     
+    @staticmethod
+    def add_signature(
+        image: Image.Image,
+        signature: str = "MJ",
+        position: str = "bottom-right",
+        font_size: int = 20,
+        opacity: int = 180,
+    ) -> Image.Image:
+        """Add a small signature/watermark to the image
+        
+        Args:
+            image: Input image
+            signature: Text to add as signature
+            position: Position of signature (bottom-right, bottom-left, top-right, top-left)
+            font_size: Size of the signature font
+            opacity: Opacity of the signature (0-255)
+        """
+        img = image.copy()
+        
+        # Create a transparent overlay
+        overlay = Image.new('RGBA', img.size, (255, 255, 255, 0))
+        draw = ImageDraw.Draw(overlay)
+        
+        # Load font
+        try:
+            font = ImageFont.truetype("arial.ttf", font_size)
+        except:
+            try:
+                font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", font_size)
+            except:
+                font = ImageFont.load_default()
+                logger.warning("Using default font for signature")
+        
+        # Calculate text size and position
+        bbox = draw.textbbox((0, 0), signature, font=font)
+        text_width = bbox[2] - bbox[0]
+        text_height = bbox[3] - bbox[1]
+        
+        padding = 10
+        
+        if position == "bottom-right":
+            x = img.width - text_width - padding
+            y = img.height - text_height - padding
+        elif position == "bottom-left":
+            x = padding
+            y = img.height - text_height - padding
+        elif position == "top-right":
+            x = img.width - text_width - padding
+            y = padding
+        elif position == "top-left":
+            x = padding
+            y = padding
+        else:
+            x = img.width - text_width - padding
+            y = img.height - text_height - padding
+        
+        # Draw signature with semi-transparent background
+        bg_padding = 5
+        draw.rectangle(
+            [x - bg_padding, y - bg_padding, 
+             x + text_width + bg_padding, y + text_height + bg_padding],
+            fill=(0, 0, 0, opacity // 2)
+        )
+        
+        # Draw text
+        draw.text((x, y), signature, font=font, fill=(255, 255, 255, opacity))
+        
+        # Convert to RGB if needed and composite
+        if img.mode != 'RGBA':
+            img = img.convert('RGBA')
+        
+        img = Image.alpha_composite(img, overlay)
+        
+        # Convert back to RGB
+        if img.mode == 'RGBA':
+            rgb_img = Image.new('RGB', img.size, (255, 255, 255))
+            rgb_img.paste(img, mask=img.split()[3])
+            return rgb_img
+        
+        return img
+    
     @staticmethod
     def enhance_image(
         image: Image.Image,