Create setup.py

setup.py

@@ -0,0 +1,420 @@
+# AI Image Creator: Enhanced UI and UX
+# Part 1: Core Setup, Model Classes and API Configuration
+
+import gradio as gr
+import logging
+import sys
+import random
+import time
+from huggingface_hub import InferenceClient
+from PIL import Image
+import io
+import base64
+
+os.getenv("HF_API_KEY")
+
+# Set up logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
+    handlers=[logging.StreamHandler(sys.stdout)]
+)
+logger = logging.getLogger("ai_image_creator")
+
+# =============== MODEL CLIENTS SETUP ===============
+def setup_client(api_key, provider=None):
+    """Initialize and return API client"""
+    try:
+        if provider:
+            client = InferenceClient(provider=provider, api_key=api_key)
+            logger.info(f"{provider} client initialized successfully")
+        else:
+            client = InferenceClient(api_key=api_key)
+            logger.info("Hugging Face client initialized successfully")
+        return client
+    except Exception as e:
+        logger.error(f"Error initializing client: {str(e)}")
+        return None
+
+# Initialize clients
+try:
+    # Replace with your actual HF API key
+    hf_api_key = os.getenv("HF_API_KEY")
+    hf_client = setup_client(hf_api_key)
+    logger.info("Hugging Face client created successfully")
+    
+    # Set up Llama client if API key is provided
+    llama_api_key = os.getenv("HF_API_KEY")  # Replace with actual key if available
+    try:
+        llama_client = setup_client(llama_api_key, "sambanova")
+        use_llama = True
+        logger.info("Llama client created successfully")
+    except Exception as e:
+        logger.warning(f"Llama client not available: {str(e)}. Will use fallback enhancement.")
+        llama_client = None
+        use_llama = False
+except Exception as e:
+    logger.error(f"Failed to create Hugging Face client: {str(e)}")
+    hf_client = None
+    llama_client = None
+    use_llama = False
+
+# =============== DATA MODELS ===============
+
+# Image Models with friendly names, descriptions and icons
+IMAGE_MODELS = {
+    "stabilityai/stable-diffusion-xl-base-1.0": {
+        "display_name": "SDXL 1.0", 
+        "description": "Best overall quality, slower generation",
+        "icon": "⭐",
+        "speed": "slow",
+        "quality": "excellent"
+    },
+    "runwayml/stable-diffusion-v1-5": {
+        "display_name": "SD 1.5", 
+        "description": "Good for general purpose, faster generation",
+        "icon": "🚀",
+        "speed": "fast",
+        "quality": "good"
+    },
+    "stabilityai/stable-diffusion-2-1": {
+        "display_name": "SD 2.1", 
+        "description": "Improved details, balanced speed and quality",
+        "icon": "✨",
+        "speed": "medium",
+        "quality": "very good"
+    },
+    "prompthero/openjourney": {
+        "display_name": "OpenJourney", 
+        "description": "Midjourney-like stylized results",
+        "icon": "🎨",
+        "speed": "medium",
+        "quality": "stylized"
+    },
+    "dreamlike-art/dreamlike-diffusion-1.0": {
+        "display_name": "Dreamlike", 
+        "description": "Artistic style with dreamy aesthetics",
+        "icon": "💫",
+        "speed": "medium",
+        "quality": "artistic"
+    }
+}
+
+# Creation types with icons and detailed descriptions
+CREATION_TYPES = {
+    "Realistic Photo": {
+        "description": "Create a photorealistic image with natural details and lighting",
+        "icon": "📷",
+        "prompt_hint": "Try to include details about lighting, time of day, and environment"
+    },
+    "Digital Art": {
+        "description": "Create colorful digital artwork with clean lines and vibrant colors",
+        "icon": "🖌️",
+        "prompt_hint": "Consider specifying color palette and mood for better results"
+    },
+    "Fantasy Illustration": {
+        "description": "Create magical and fantastical scenes with otherworldly elements",
+        "icon": "🧙",
+        "prompt_hint": "Describe magical elements, creatures, and environments in detail"
+    },
+    "Concept Art": {
+        "description": "Create professional concept art for characters, environments or objects",
+        "icon": "🎮",
+        "prompt_hint": "Include details about perspective, purpose, and design influences"
+    },
+    "Anime/Manga": {
+        "description": "Create Japanese anime or manga style illustration",
+        "icon": "🍙",
+        "prompt_hint": "Specify anime aesthetics like shading style and character features"
+    },
+    "Oil Painting": {
+        "description": "Create an image with oil painting textures and artistic brushstrokes",
+        "icon": "🖼️",
+        "prompt_hint": "Consider describing texture, brushwork style, and canvas feel"
+    },
+    "Watercolor": {
+        "description": "Create a soft watercolor illustration with subtle color blending",
+        "icon": "💧",
+        "prompt_hint": "Mention color blending, paper texture, and watercolor-specific effects"
+    },
+    "Sketch": {
+        "description": "Create a detailed sketch or drawing with line art focus",
+        "icon": "✏️",
+        "prompt_hint": "Describe line weight, hatching style, and sketch medium (pencil, charcoal, etc.)"
+    },
+    "3D Rendering": {
+        "description": "Create an image that looks like a 3D rendered scene with realistic lighting",
+        "icon": "💻",
+        "prompt_hint": "Include details about lighting setup, materials, and camera perspective"
+    },
+    "Pixel Art": {
+        "description": "Create retro-style pixel art with limited color palette",
+        "icon": "👾",
+        "prompt_hint": "Specify resolution, color limitations, and pixel art style (e.g., 16-bit, 8-bit)"
+    }
+}
+
+# Art styles with icons and detailed descriptions
+ART_STYLES = {
+    "Photorealistic": {
+        "description": "Detailed realistic style that resembles a photograph with accurate lighting and textures",
+        "icon": "📸",
+        "examples": "Works by Chuck Close, Richard Estes, or modern 3D renderings"
+    },
+    "Impressionist": {
+        "description": "Soft brushstrokes that capture light and atmosphere over precise details, like Monet",
+        "icon": "🌈",
+        "examples": "Claude Monet, Pierre-Auguste Renoir, Camille Pissarro"
+    },
+    "Surrealist": {
+        "description": "Dreamlike quality with impossible or irrational scenes, like Salvador Dali",
+        "icon": "🌀",
+        "examples": "Salvador Dali, René Magritte, Frida Kahlo"
+    },
+    "Pop Art": {
+        "description": "Bold colors, sharp lines and popular culture references, like Andy Warhol",
+        "icon": "🎭",
+        "examples": "Andy Warhol, Roy Lichtenstein, Keith Haring"
+    },
+    "Minimalist": {
+        "description": "Simplified forms, limited color palette, and clean composition with minimal elements",
+        "icon": "⬜",
+        "examples": "Piet Mondrian, Kazimir Malevich, Agnes Martin"
+    },
+    "Abstract": {
+        "description": "Non-representational style using shapes, colors, and forms to express ideas",
+        "icon": "🔶",
+        "examples": "Wassily Kandinsky, Jackson Pollock, Mark Rothko"
+    },
+    "Cubist": {
+        "description": "Geometric shapes and multiple perspectives shown simultaneously, like Picasso",
+        "icon": "📐",
+        "examples": "Pablo Picasso, Georges Braque, Juan Gris"
+    },
+    "Art Nouveau": {
+        "description": "Ornate, flowing lines inspired by natural forms with decorative elegance",
+        "icon": "🌿",
+        "examples": "Alphonse Mucha, Gustav Klimt, Antoni Gaudí"
+    },
+    "Gothic": {
+        "description": "Dark, medieval-inspired aesthetic with dramatic lighting and architectural elements",
+        "icon": "🏰",
+        "examples": "Zdzisław Beksiński, H.R. Giger, medieval architecture"
+    },
+    "Cyberpunk": {
+        "description": "Futuristic dystopian style with neon colors, technology, and urban decay",
+        "icon": "🤖",
+        "examples": "Blade Runner, Ghost in the Shell, Akira"
+    },
+    "Steampunk": {
+        "description": "Victorian-era aesthetic combined with steam-powered technology and brass elements",
+        "icon": "⚙️",
+        "examples": "Works by James Ng, Keith Thompson, retrofuturistic Jules Verne adaptations"
+    },
+    "Retro/Vintage": {
+        "description": "Nostalgic style reminiscent of past decades with period-appropriate elements",
+        "icon": "📺",
+        "examples": "1950s advertisements, vintage travel posters, pulp magazine covers"
+    },
+    "Art Deco": {
+        "description": "Geometric patterns, bold colors, and luxurious materials in a symmetrical style",
+        "icon": "🏢",
+        "examples": "Works from the 1920s-30s, Chrysler Building, Tamara de Lempicka paintings"
+    },
+    "Baroque": {
+        "description": "Dramatic, ornate style with rich details, contrast, and dynamic composition",
+        "icon": "👑",
+        "examples": "Caravaggio, Rembrandt, Peter Paul Rubens"
+    },
+    "Ukiyo-e": {
+        "description": "Traditional Japanese woodblock print style with flat areas of color and strong outlines",
+        "icon": "🌊",
+        "examples": "Hokusai's Great Wave, Hiroshige's landscapes, traditional Japanese prints"
+    },
+    "Comic Book": {
+        "description": "Bold outlines, bright colors, and action-oriented composition like classic comics",
+        "icon": "💥",
+        "examples": "Jack Kirby, Steve Ditko, modern Marvel/DC art styles"
+    },
+    "Psychedelic": {
+        "description": "Vibrant, swirling colors with abstract patterns inspired by 1960s art",
+        "icon": "🌈",
+        "examples": "1960s concert posters, Peter Max, Alex Grey"
+    },
+    "Vaporwave": {
+        "description": "Glitch aesthetics with pastel colors, 80s/90s nostalgia and digital elements",
+        "icon": "📼",
+        "examples": "Retrowave aesthetics, 80s digital graphics, glitch art"
+    },
+    "Studio Ghibli": {
+        "description": "Whimsical, detailed animation style inspired by Japanese animated films",
+        "icon": "🐉",
+        "examples": "Spirited Away, My Neighbor Totoro, Howl's Moving Castle"
+    },
+    "Hyperrealism": {
+        "description": "Extremely detailed realism that exceeds photograph-like precision",
+        "icon": "🔍",
+        "examples": "Works by Roberto Bernardi, Denis Peterson, Gottfried Helnwein"
+    }
+}
+
+# Moods with icons and descriptions
+MOODS = {
+    "Happy": {
+        "description": "Bright, cheerful atmosphere with warm colors",
+        "icon": "😊",
+        "color_palette": "Warm and vibrant colors: yellows, bright oranges, light blues"
+    },
+    "Sad": {
+        "description": "Melancholic atmosphere with muted colors",
+        "icon": "😢",
+        "color_palette": "Muted blues, grays, desaturated colors, cool tones"
+    },
+    "Mysterious": {
+        "description": "Enigmatic atmosphere with shadows and hidden elements",
+        "icon": "🔮",
+        "color_palette": "Deep purples, dark blues, hints of teal, selective lighting"
+    },
+    "Peaceful": {
+        "description": "Serene, calm atmosphere with gentle lighting",
+        "icon": "🕊️",
+        "color_palette": "Soft blues, gentle greens, pastel colors, balanced light"
+    },
+    "Tense": {
+        "description": "Suspenseful atmosphere with dramatic lighting",
+        "icon": "😰",
+        "color_palette": "High contrast, dark shadows, selective reds, strong highlights"
+    },
+    "Whimsical": {
+        "description": "Playful, imaginative atmosphere with fanciful elements",
+        "icon": "🦄",
+        "color_palette": "Pastels, candy colors, unexpected color combinations"
+    },
+    "Dark": {
+        "description": "Gloomy atmosphere with deep shadows and low lighting",
+        "icon": "🌑",
+        "color_palette": "Dark blues, blacks, deep greens, minimal highlights"
+    },
+    "Energetic": {
+        "description": "Dynamic, vibrant atmosphere with strong colors and movement",
+        "icon": "⚡",
+        "color_palette": "Saturated primary colors, bold contrasts, vibrant hues"
+    },
+    "Romantic": {
+        "description": "Soft, dreamy atmosphere with warm, gentle lighting",
+        "icon": "❤️",
+        "color_palette": "Soft pinks, gentle reds, golden highlights, warm tones"
+    },
+    "Epic": {
+        "description": "Grand, impressive atmosphere with dramatic scale and lighting",
+        "icon": "🏔️",
+        "color_palette": "Bold colors, dramatic contrast, atmospheric lighting, expansive scale"
+    }
+}
+
+# Example prompts with rich metadata
+EXAMPLE_PROMPTS = [
+    {
+        "text": "A serene lake at sunset with mountains in the background and a small wooden boat floating nearby",
+        "thumbnail_desc": "Peaceful lake scene at sunset",
+        "creation_type": "Realistic Photo",
+        "art_style": "Photorealistic", 
+        "mood": "Peaceful",
+        "tags": ["nature", "landscape", "water", "sunset"]
+    },
+    {
+        "text": "A futuristic cityscape with flying cars, neon lights, and tall skyscrapers under a night sky with two moons",
+        "thumbnail_desc": "Futuristic city with flying cars",
+        "creation_type": "Concept Art",
+        "art_style": "Cyberpunk",
+        "mood": "Mysterious",
+        "tags": ["scifi", "future", "urban", "night"]
+    },
+    {
+        "text": "A close-up portrait of an elderly craftsman with weathered hands working on an intricate wooden carving in his workshop",
+        "thumbnail_desc": "Elderly craftsman working with wood",
+        "creation_type": "Oil Painting",
+        "art_style": "Hyperrealism",
+        "mood": "Peaceful",
+        "tags": ["portrait", "craftsmanship", "elderly", "detail"]
+    },
+    {
+        "text": "A magical forest with glowing mushrooms, fairy lights, and a small cottage with smoke coming from the chimney",
+        "thumbnail_desc": "Magical forest with glowing elements",
+        "creation_type": "Fantasy Illustration",
+        "art_style": "Studio Ghibli",
+        "mood": "Whimsical",
+        "tags": ["fantasy", "forest", "magic", "cottage"]
+    },
+    {
+        "text": "A cybernetic samurai with glowing blue circuits standing in a rainy Tokyo street at night",
+        "thumbnail_desc": "Cybernetic samurai in rainy Tokyo",
+        "creation_type": "Digital Art",
+        "art_style": "Cyberpunk",
+        "mood": "Dark",
+        "tags": ["character", "cyberpunk", "samurai", "rain"]
+    },
+    {
+        "text": "A cute cat with dragon wings and tiny horns sleeping on a pile of gold coins",
+        "thumbnail_desc": "Cat with dragon features on gold",
+        "creation_type": "Fantasy Illustration",
+        "art_style": "Comic Book",
+        "mood": "Whimsical",
+        "tags": ["creature", "fantasy", "cute", "treasure"]
+    },
+    {
+        "text": "An ancient temple covered in vines and moss, partially sunken into a crystal-clear cenote in the jungle",
+        "thumbnail_desc": "Ancient temple in jungle cenote",
+        "creation_type": "Concept Art",
+        "art_style": "Photorealistic",
+        "mood": "Mysterious",
+        "tags": ["architecture", "ruins", "jungle", "water"]
+    },
+    {
+        "text": "A cozy coffee shop interior with rain falling outside the windows, soft lighting, and a few people reading books",
+        "thumbnail_desc": "Cozy rainy day in coffee shop",
+        "creation_type": "Digital Art",
+        "art_style": "Impressionist",
+        "mood": "Peaceful",
+        "tags": ["interior", "rainy", "cozy", "urban"]
+    }
+]
+
+# CSS for enhanced UI styling - Will be included in part 2
+
+# =============== HELPER FUNCTIONS ===============
+
+# Helper function to format dropdown choices with icons
+def format_dropdown_choices(options_dict):
+    return [f"{info.get('icon', '•')} {key}" for key in options_dict.keys()]
+
+# Helper function to extract the key from formatted choice
+def extract_key(formatted_choice):
+    # Skip the icon and space at the beginning
+    parts = formatted_choice.split(' ', 1)
+    if len(parts) > 1:
+        return parts[1]
+    return formatted_choice
+
+# Function to load example prompt
+def load_example(example_index):
+    if example_index < 0 or example_index >= len(EXAMPLE_PROMPTS):
+        return "", "", "", ""
+    
+    example = EXAMPLE_PROMPTS[example_index]
+    creation = f"{CREATION_TYPES[example['creation_type']]['icon']} {example['creation_type']}"
+    art = f"{ART_STYLES[example['art_style']]['icon']} {example['art_style']}"
+    mood = f"{MOODS[example['mood']]['icon']} {example['mood']}"
+    
+    return example["text"], creation, art, mood
+
+# Get model key from formatted display name
+def get_model_key(formatted_name):
+    # Extract display name without the icon
+    if ' ' in formatted_name:
+        display_name = formatted_name.split(' ', 1)[1]
+        # Find the corresponding key
+        for key, info in IMAGE_MODELS.items():
+            if info['display_name'] == display_name:
+                return key
+    return list(IMAGE_MODELS.keys())[0]  # Default to first model if not found