""" šØ AI Image Editor Pro - Streamlit Version ============================================= A private, self-hosted AI image editing tool using open-source models. Runs on Hugging Face Spaces with Streamlit SDK. Now with advanced Gemini-style instruction understanding! """ import os import gc import re import torch import numpy as np import streamlit as st from PIL import Image from typing import Tuple, Optional, Dict, List from io import BytesIO # ============================================================================ # PAGE CONFIG (must be first Streamlit command) # ============================================================================ st.set_page_config( page_title="šØ AI Image Editor Pro", page_icon="šØ", layout="wide", initial_sidebar_state="expanded" ) # ============================================================================ # CONFIGURATION # ============================================================================ DEVICE = "cuda" if torch.cuda.is_available() else "cpu" DTYPE = torch.float16 if DEVICE == "cuda" else torch.float32 INPAINT_MODEL = "runwayml/stable-diffusion-inpainting" CLIPSEG_MODEL = "CIDAS/clipseg-rd64-refined" # ============================================================================ # ADVANCED INSTRUCTION PARSER - GEMINI-STYLE # ============================================================================ class GeminiStyleParser: """ Advanced natural language parser that understands complex editing instructions like Google Gemini. Handles various phrasings, synonyms, and compound commands. """ # Comprehensive action patterns with synonyms REMOVE_KEYWORDS = [ "remove", "delete", "erase", "get rid of", "take out", "eliminate", "clear", "wipe", "clean up", "take away", "disappear", "vanish", "make disappear", "get away", "rid of", "cut out", "crop out", "hide", "discard", "throw away", "dispose", "extract", "pull out", "subtract", "minus", "without", "lose", "drop", "ditch", "nix", "scratch", "strike", "zap", "nuke", "kill", "destroy", "obliterate" ] REPLACE_KEYWORDS = [ "replace", "swap", "switch", "substitute", "exchange", "trade", "put", "place", "add", "insert", "set", "change to", "turn into", "transform to", "convert to", "make it", "make this", "transform into", "morph into", "become", "evolve into", "shift to" ] CHANGE_KEYWORDS = [ "change", "modify", "alter", "adjust", "edit", "transform", "convert", "turn", "make", "update", "recolor", "repaint", "tint", "color", "paint", "dye", "shade", "hue", "tone", "brighten", "darken", "lighten", "saturate", "desaturate" ] ADD_KEYWORDS = [ "add", "insert", "put", "place", "include", "attach", "append", "introduce", "bring", "create", "generate", "draw", "paint", "render", "give", "apply", "overlay" ] ENHANCE_KEYWORDS = [ "enhance", "improve", "beautify", "upgrade", "refine", "polish", "perfect", "optimize", "boost", "amplify", "sharpen", "clarify", "fix", "repair", "restore" ] # Prepositions and connectors PREPOSITIONS = [ "with", "to", "into", "as", "by", "for", "from", "using", "via", "through", "in place of", "instead of" ] # Color mappings for better understanding COLORS = { "red": "vibrant red colored", "blue": "deep blue colored", "green": "lush green colored", "yellow": "bright yellow colored", "orange": "warm orange colored", "purple": "rich purple colored", "pink": "soft pink colored", "black": "pure black colored", "white": "clean white colored", "gold": "shimmering golden colored", "silver": "metallic silver colored", "brown": "natural brown colored", "gray": "neutral gray colored", "grey": "neutral grey colored", "cyan": "cyan turquoise colored", "magenta": "vivid magenta colored", "teal": "elegant teal colored", "navy": "deep navy blue colored", "maroon": "rich maroon colored", "olive": "earthy olive colored", "coral": "beautiful coral colored", "beige": "soft beige colored", "tan": "warm tan colored", "cream": "creamy off-white colored", "mint": "fresh mint green colored", "lavender": "delicate lavender colored", "rose": "romantic rose colored", "burgundy": "deep burgundy colored", "bronze": "warm bronze colored" } # Object synonyms for better detection OBJECT_SYNONYMS = { "person": ["person", "human", "man", "woman", "people", "guy", "girl", "boy", "lady", "gentleman", "individual", "figure", "someone", "somebody", "pedestrian"], "sky": ["sky", "clouds", "heaven", "atmosphere", "air above", "skyline"], "car": ["car", "vehicle", "automobile", "auto", "ride", "wheels", "sedan", "suv", "truck", "van"], "background": ["background", "backdrop", "behind", "scenery", "setting", "surroundings", "environment"], "text": ["text", "words", "letters", "writing", "inscription", "watermark", "logo", "signature", "label", "caption"], "grass": ["grass", "lawn", "turf", "field", "meadow", "greenery"], "tree": ["tree", "plant", "vegetation", "foliage", "bush", "shrub"], "water": ["water", "ocean", "sea", "lake", "river", "pond", "pool", "stream"], "building": ["building", "house", "structure", "architecture", "construction", "edifice"], "animal": ["animal", "pet", "creature", "dog", "cat", "bird"], "face": ["face", "facial", "head", "portrait", "visage"], "hair": ["hair", "hairstyle", "locks", "mane", "tresses"], "clothes": ["clothes", "clothing", "outfit", "dress", "shirt", "pants", "garment", "attire", "wear"], "wall": ["wall", "walls", "surface", "partition"], "floor": ["floor", "ground", "flooring", "surface below"], "window": ["window", "glass", "pane", "windowpane"], "door": ["door", "doorway", "entrance", "entry", "gate"] } # Scene/style transformations STYLE_TRANSFORMS = { "sunset": "beautiful golden sunset sky with orange and pink clouds, dramatic lighting", "sunrise": "stunning sunrise with warm golden light, peaceful morning atmosphere", "night": "dark nighttime scene with stars, moonlit atmosphere", "day": "bright daylight, clear blue sky, natural sunlight", "winter": "snowy winter scene, frost covered, cold atmosphere", "summer": "bright summer day, warm sunny atmosphere", "autumn": "fall colors, orange and brown leaves, autumn atmosphere", "spring": "fresh spring scene, blooming flowers, new growth", "rain": "rainy weather, wet surfaces, overcast sky", "snow": "heavy snowfall, white snow covered, winter wonderland", "foggy": "misty foggy atmosphere, soft diffused light", "stormy": "dramatic stormy sky, dark clouds, lightning", "vintage": "vintage retro aesthetic, warm sepia tones, nostalgic feel", "cyberpunk": "neon cyberpunk aesthetic, futuristic, glowing lights", "fantasy": "magical fantasy scene, ethereal atmosphere, dreamlike", "realistic": "photorealistic, natural, lifelike quality", "cartoon": "cartoon animated style, colorful, illustrated", "anime": "anime style, japanese animation aesthetic", "watercolor": "watercolor painting style, soft brushstrokes", "oil painting": "oil painting style, rich textures, artistic", "sketch": "pencil sketch style, hand-drawn look", "cinematic": "cinematic movie quality, dramatic lighting, film-like", "hdr": "high dynamic range, vivid colors, enhanced contrast", "dreamy": "soft dreamy atmosphere, ethereal glow, romantic", "dramatic": "dramatic lighting, high contrast, intense mood", "peaceful": "calm peaceful atmosphere, serene, tranquil", "scary": "dark scary atmosphere, horror aesthetic, ominous", "happy": "bright cheerful atmosphere, joyful, vibrant colors", "sad": "melancholic atmosphere, muted colors, somber mood" } def __init__(self): self.last_confidence = 0.0 self.interpretation = "" def normalize_text(self, text: str) -> str: """Normalize input text for better parsing.""" text = text.lower().strip() # Remove extra whitespace text = re.sub(r'\s+', ' ', text) # Remove common punctuation that doesn't affect meaning text = re.sub(r'[.,!?;:]+$', '', text) # Handle contractions text = text.replace("don't", "do not") text = text.replace("can't", "cannot") text = text.replace("won't", "will not") text = text.replace("i'd", "i would") text = text.replace("i'm", "i am") text = text.replace("it's", "it is") return text def extract_target_object(self, text: str) -> str: """Extract the target object from the instruction.""" # Remove common filler words filler_words = ["the", "a", "an", "this", "that", "those", "these", "my", "your", "please", "kindly", "can you", "could you", "would you", "i want to", "i'd like to", "i would like to"] result = text for filler in filler_words: result = re.sub(r'\b' + filler + r'\b', '', result, flags=re.IGNORECASE) return result.strip() def find_best_synonym(self, target: str) -> str: """Find the best matching object for CLIPSeg detection.""" target_lower = target.lower() # Check if target matches any known synonym for main_object, synonyms in self.OBJECT_SYNONYMS.items(): for synonym in synonyms: if synonym in target_lower or target_lower in synonym: return main_object return target def enhance_prompt(self, prompt: str) -> str: """Enhance the replacement prompt for better results.""" prompt_lower = prompt.lower() # Check for style transformations for style_key, style_value in self.STYLE_TRANSFORMS.items(): if style_key in prompt_lower: return f"{style_value}, high quality, detailed, professional" # Check for colors and enhance for color_key, color_value in self.COLORS.items(): if color_key in prompt_lower: prompt = prompt.replace(color_key, color_value) # Add quality modifiers if not present quality_terms = ["high quality", "detailed", "professional", "beautiful", "stunning"] has_quality = any(term in prompt_lower for term in quality_terms) if not has_quality: prompt = f"{prompt}, high quality, detailed, professional photography" return prompt def detect_action_type(self, text: str) -> str: """Detect the type of editing action requested.""" text_lower = text.lower() for keyword in self.REMOVE_KEYWORDS: if keyword in text_lower: return "remove" for keyword in self.ADD_KEYWORDS: if keyword in text_lower: return "add" for keyword in self.REPLACE_KEYWORDS: if keyword in text_lower: return "replace" for keyword in self.CHANGE_KEYWORDS: if keyword in text_lower: return "change" for keyword in self.ENHANCE_KEYWORDS: if keyword in text_lower: return "enhance" return "general" def parse(self, instruction: str) -> Tuple[str, str, float]: """ Parse the instruction and return (target, replacement_prompt, confidence). This is the main parsing method that handles all types of instructions. """ original = instruction normalized = self.normalize_text(instruction) action_type = self.detect_action_type(normalized) target = "" replacement = "" confidence = 0.5 # ===== REMOVE ACTION ===== if action_type == "remove": for keyword in self.REMOVE_KEYWORDS: if keyword in normalized: target = normalized.split(keyword, 1)[-1].strip() break target = self.extract_target_object(target) target = self.find_best_synonym(target) replacement = "clean empty background, seamless natural texture, nothing there, blank space" confidence = 0.85 self.interpretation = f"šļø Remove: Detecting and removing '{target}'" # ===== ADD ACTION ===== elif action_type == "add": for keyword in self.ADD_KEYWORDS: if keyword in normalized: parts = normalized.split(keyword, 1) if len(parts) > 1: target = "main subject area" replacement = parts[1].strip() break replacement = self.extract_target_object(replacement) replacement = self.enhance_prompt(replacement) confidence = 0.75 self.interpretation = f"ā Add: Adding '{replacement}' to the image" # ===== REPLACE ACTION ===== elif action_type == "replace": # Try to find "X with Y" or "X to Y" patterns preposition_found = False for prep in self.PREPOSITIONS: if f" {prep} " in normalized: parts = normalized.split(f" {prep} ", 1) # Extract target from first part first_part = parts[0] for keyword in self.REPLACE_KEYWORDS + self.CHANGE_KEYWORDS: first_part = first_part.replace(keyword, "") target = self.extract_target_object(first_part) target = self.find_best_synonym(target) # Extract replacement from second part replacement = self.extract_target_object(parts[1]) replacement = self.enhance_prompt(replacement) preposition_found = True confidence = 0.9 break if not preposition_found: # Fallback: try to extract target and use generic replacement for keyword in self.REPLACE_KEYWORDS: if keyword in normalized: target = normalized.split(keyword, 1)[-1].strip() target = self.extract_target_object(target) target = self.find_best_synonym(target) replacement = "something different, new object, alternative" confidence = 0.6 break self.interpretation = f"š Replace: Replacing '{target}' with '{replacement[:50]}...'" # ===== CHANGE ACTION ===== elif action_type == "change": # Look for patterns like "change X to Y" or "make X Y" preposition_found = False for prep in ["to", "into", "as"]: if f" {prep} " in normalized: parts = normalized.split(f" {prep} ", 1) # Extract target from first part first_part = parts[0] for keyword in self.CHANGE_KEYWORDS: first_part = first_part.replace(keyword, "") target = self.extract_target_object(first_part) target = self.find_best_synonym(target) # Extract new state from second part new_state = self.extract_target_object(parts[1]) # Combine target with new state for replacement replacement = f"{target} that is {new_state}, {self.enhance_prompt(new_state)}" preposition_found = True confidence = 0.85 break if not preposition_found: # Check for color changes like "make it red" for color in self.COLORS.keys(): if color in normalized: target = "main subject" replacement = f"{self.COLORS[color]}, high quality, detailed" confidence = 0.8 preposition_found = True break if not preposition_found: target = "main subject" replacement = self.enhance_prompt(normalized) confidence = 0.6 self.interpretation = f"āļø Change: Modifying '{target}' ā '{replacement[:50]}...'" # ===== ENHANCE ACTION ===== elif action_type == "enhance": target = "main subject" replacement = "enhanced improved professional high quality detailed stunning beautiful" confidence = 0.7 self.interpretation = f"⨠Enhance: Improving overall image quality" # ===== GENERAL/UNKNOWN ACTION ===== else: # Try to intelligently guess from the instruction # Check if it's just a noun/object (user wants to remove it) words = normalized.split() if len(words) <= 3: target = self.find_best_synonym(normalized) replacement = "clean empty background, seamless natural texture" confidence = 0.5 self.interpretation = f"š¤ Guessing: You might want to remove '{target}'?" else: # Treat as a creative prompt target = "main subject area" replacement = self.enhance_prompt(normalized) confidence = 0.5 self.interpretation = f"šØ Creative: Applying '{replacement[:50]}...'" # Final cleanup target = target.strip() if target else "main subject" replacement = replacement.strip() if replacement else "improved version" # Store confidence self.last_confidence = confidence return target, replacement, confidence # Create global parser instance gemini_parser = GeminiStyleParser() def parse_instruction(instruction: str) -> Tuple[str, str]: """ Enhanced parsing function that uses the GeminiStyleParser. Maintains backward compatibility with existing code. """ target, replacement, _ = gemini_parser.parse(instruction) return target, replacement # ============================================================================ # MODEL CACHING # ============================================================================ @st.cache_resource def load_inpaint_pipeline(): """Load and cache the inpainting pipeline.""" from diffusers import StableDiffusionInpaintPipeline pipe = StableDiffusionInpaintPipeline.from_pretrained( INPAINT_MODEL, torch_dtype=DTYPE, safety_checker=None, requires_safety_checker=False ) pipe = pipe.to(DEVICE) if DEVICE == "cuda": pipe.enable_attention_slicing() try: pipe.enable_xformers_memory_efficient_attention() except Exception: pass else: pipe.enable_attention_slicing(1) return pipe @st.cache_resource def load_clipseg(): """Load and cache CLIPSeg for automatic mask generation.""" from transformers import CLIPSegProcessor, CLIPSegForImageSegmentation processor = CLIPSegProcessor.from_pretrained(CLIPSEG_MODEL) model = CLIPSegForImageSegmentation.from_pretrained(CLIPSEG_MODEL) model = model.to(DEVICE) model.eval() return processor, model # ============================================================================ # MASK GENERATION (Enhanced) # ============================================================================ def generate_mask_clipseg( image: Image.Image, target_text: str, threshold: float = 0.3, expand_pixels: int = 10 ) -> Optional[Image.Image]: """Generate a segmentation mask using CLIPSeg with enhanced detection.""" try: processor, model = load_clipseg() # Try multiple variations of the target text for better detection target_variations = [ target_text, f"a {target_text}", f"the {target_text}", f"{target_text} in photo", f"photo of {target_text}" ] best_mask = None best_score = 0 for variation in target_variations: inputs = processor( text=[variation], images=[image], padding=True, return_tensors="pt" ) inputs = {k: v.to(DEVICE) for k, v in inputs.items()} with torch.no_grad(): outputs = model(**inputs) preds = outputs.logits pred = torch.sigmoid(preds[0]).cpu().numpy() score = pred.max() if score > best_score: best_score = score best_mask = pred if best_mask is None: return None # Resize to original image size pred_pil = Image.fromarray((best_mask * 255).astype(np.uint8)) pred_resized = pred_pil.resize(image.size, Image.BILINEAR) pred_array = np.array(pred_resized) # Apply threshold mask = (pred_array > (threshold * 255)).astype(np.uint8) * 255 # Expand mask if expand_pixels > 0: from PIL import ImageFilter mask_image = Image.fromarray(mask, mode="L") mask_image = mask_image.filter( ImageFilter.MaxFilter(size=expand_pixels * 2 + 1) ) mask_image = mask_image.filter( ImageFilter.GaussianBlur(radius=3) ) return mask_image return Image.fromarray(mask, mode="L") except Exception as e: st.error(f"Mask generation error: {str(e)}") return None def process_manual_mask(mask_image: Image.Image, target_size: Tuple[int, int]) -> Image.Image: """Process a manually uploaded mask.""" mask = mask_image.convert("L") mask = mask.resize(target_size, Image.LANCZOS) mask_array = np.array(mask) mask_array = ((mask_array > 127) * 255).astype(np.uint8) return Image.fromarray(mask_array, mode="L") # ============================================================================ # IMAGE INPAINTING (Enhanced) # ============================================================================ def inpaint_image( image: Image.Image, mask: Image.Image, prompt: str, negative_prompt: str = "blurry, bad quality, distorted, ugly, deformed, low resolution, pixelated, jpeg artifacts, watermark, text, logo", num_inference_steps: int = 30, guidance_scale: float = 7.5 ) -> Optional[Image.Image]: """Inpaint the masked region of an image with enhanced prompts.""" try: pipe = load_inpaint_pipeline() # Resize for SD (512x512) original_size = image.size target_size = (512, 512) image_resized = image.resize(target_size, Image.LANCZOS) mask_resized = mask.resize(target_size, Image.NEAREST) if image_resized.mode != "RGB": image_resized = image_resized.convert("RGB") # Adjust steps for CPU if DEVICE == "cpu": num_inference_steps = min(num_inference_steps, 20) # Enhanced prompt engineering enhanced_prompt = f"{prompt}, masterpiece, best quality, highly detailed, sharp focus, professional" with torch.inference_mode(): result = pipe( prompt=enhanced_prompt, negative_prompt=negative_prompt, image=image_resized, mask_image=mask_resized, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale ).images[0] result = result.resize(original_size, Image.LANCZOS) if DEVICE == "cpu": gc.collect() return result except Exception as e: st.error(f"Inpainting error: {str(e)}") return None # ============================================================================ # CUSTOM CSS FOR PRO LOOK # ============================================================================ def inject_custom_css(): """Inject custom CSS for a more professional look.""" st.markdown(""" """, unsafe_allow_html=True) # ============================================================================ # MAIN APP # ============================================================================ def main(): inject_custom_css() st.markdown("""