Expanded generation options. Multiple outfits support.

app.py

@@ -6,7 +6,7 @@ import requests
 import random
 from flask import Flask, render_template, request, redirect, url_for, flash, session
 from werkzeug.utils import secure_filename
-from models import db, Character
+from models import db, Character, Settings
 
 app = Flask(__name__)
 app.config['SQLALCHEMY_DATABASE_URI'] = 'sqlite:///database.db'
@@ -17,11 +17,21 @@ app.config['CHARACTERS_DIR'] = 'characters'
 app.config['COMFYUI_URL'] = 'http://127.0.0.1:8188'
 app.config['ILLUSTRIOUS_MODELS_DIR'] = '/mnt/alexander/AITools/Image Models/Stable-diffusion/Illustrious/'
 app.config['NOOB_MODELS_DIR'] = '/mnt/alexander/AITools/Image Models/Stable-diffusion/Noob/'
+app.config['LORA_DIR'] = '/mnt/alexander/AITools/Image Models/lora/Illustrious/Looks/'
 
 db.init_app(app)
 
 ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'gif', 'webp'}
 
+def get_available_loras():
+    loras = []
+    if os.path.exists(app.config['LORA_DIR']):
+        for f in os.listdir(app.config['LORA_DIR']):
+            if f.endswith('.safetensors'):
+                # Using the format seen in character JSONs
+                loras.append(f"Illustrious/Looks/{f}")
+    return sorted(loras)
+
 def get_available_checkpoints():
     checkpoints = []
     
@@ -42,7 +52,7 @@ def get_available_checkpoints():
 def allowed_file(filename):
     return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS
 
-def build_prompt(data, selected_fields=None, default_fields=None):
+def build_prompt(data, selected_fields=None, default_fields=None, active_outfit='default'):
     def is_selected(section, key):
         # Priority: 
         # 1. Manual selection from form (if list is not empty)
@@ -55,7 +65,17 @@ def build_prompt(data, selected_fields=None, default_fields=None):
         return True
 
     identity = data.get('identity', {})
-    wardrobe = data.get('wardrobe', {})
+    
+    # Get wardrobe - handle both new nested format and legacy flat format
+    wardrobe_data = data.get('wardrobe', {})
+    if 'default' in wardrobe_data and isinstance(wardrobe_data.get('default'), dict):
+        # New nested format - get active outfit
+        wardrobe = wardrobe_data.get(active_outfit or 'default', wardrobe_data.get('default', {}))
+    else:
+        # Legacy flat format
+        wardrobe = wardrobe_data
+    
+    defaults = data.get('defaults', {})
     
     # Pre-calculate Hand/Glove priority
     hand_val = ""
@@ -71,16 +91,22 @@ def build_prompt(data, selected_fields=None, default_fields=None):
     if char_tag and is_selected('special', 'name'):
         parts.append(char_tag)
     
-    for key in ['base_specs', 'hair', 'eyes', 'expression', 'distinguishing_marks']:
+    for key in ['base_specs', 'hair', 'eyes', 'extra']:
         val = identity.get(key)
         if val and is_selected('identity', key):
             parts.append(val)
     
+    # Add defaults (expression, pose, scene)
+    for key in ['expression', 'pose', 'scene']:
+        val = defaults.get(key)
+        if val and is_selected('defaults', key):
+            parts.append(val)
+    
     # Add hand priority value to main prompt
     if hand_val:
         parts.append(hand_val)
     
-    for key in ['outer_layer', 'inner_layer', 'lower_body', 'footwear', 'accessories']:
+    for key in ['top', 'headwear', 'legwear', 'footwear', 'accessories']:
         val = wardrobe.get(key)
         if val and is_selected('wardrobe', key):
             parts.append(val)
@@ -101,7 +127,7 @@ def build_prompt(data, selected_fields=None, default_fields=None):
     face_parts = []
     if char_tag and is_selected('special', 'name'): face_parts.append(char_tag)
     if identity.get('eyes') and is_selected('identity', 'eyes'): face_parts.append(identity.get('eyes'))
-    if identity.get('expression') and is_selected('identity', 'expression'): face_parts.append(identity.get('expression'))
+    if defaults.get('expression') and is_selected('defaults', 'expression'): face_parts.append(defaults.get('expression'))
     
     # 3. Hand Prompt: Hand value (Gloves or Hands)
     hand_parts = [hand_val] if hand_val else []
@@ -160,6 +186,7 @@ def sync_characters():
                         character.data = data
                         character.name = name
                         character.slug = slug
+                        character.filename = filename
                         
                         # Check if cover image still exists
                         if character.image_path:
@@ -174,6 +201,7 @@ def sync_characters():
                         new_char = Character(
                             character_id=char_id,
                             slug=slug,
+                            filename=filename,
                             name=name,
                             data=data
                         )
@@ -189,6 +217,66 @@ def sync_characters():
             
     db.session.commit()
 
+def call_llm(prompt, system_prompt="You are a creative assistant."):
+    settings = Settings.query.first()
+    if not settings or not settings.openrouter_api_key:
+        raise ValueError("OpenRouter API Key not configured. Please configure it in Settings.")
+    
+    headers = {
+        "Authorization": f"Bearer {settings.openrouter_api_key}",
+        "Content-Type": "application/json"
+    }
+    data = {
+        "model": settings.openrouter_model or 'google/gemini-2.0-flash-001',
+        "messages": [
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": prompt}
+        ]
+    }
+    
+    try:
+        response = requests.post("https://openrouter.ai/api/v1/chat/completions", headers=headers, json=data)
+        response.raise_for_status()
+        result = response.json()
+        return result['choices'][0]['message']['content']
+    except requests.exceptions.RequestException as e:
+        raise RuntimeError(f"LLM API request failed: {str(e)}") from e
+    except (KeyError, IndexError) as e:
+        raise RuntimeError(f"Unexpected LLM response format: {str(e)}") from e
+
+@app.route('/get_openrouter_models', methods=['POST'])
+def get_openrouter_models():
+    api_key = request.form.get('api_key')
+    if not api_key:
+        return {'error': 'API key is required'}, 400
+    
+    headers = {"Authorization": f"Bearer {api_key}"}
+    try:
+        response = requests.get("https://openrouter.ai/api/v1/models", headers=headers)
+        response.raise_for_status()
+        models = response.json().get('data', [])
+        # Return simplified list of models
+        return {'models': [{'id': m['id'], 'name': m.get('name', m['id'])} for m in models]}
+    except Exception as e:
+        return {'error': str(e)}, 500
+
+@app.route('/settings', methods=['GET', 'POST'])
+def settings():
+    settings = Settings.query.first()
+    if not settings:
+        settings = Settings()
+        db.session.add(settings)
+        db.session.commit()
+        
+    if request.method == 'POST':
+        settings.openrouter_api_key = request.form.get('api_key')
+        settings.openrouter_model = request.form.get('model')
+        db.session.commit()
+        flash('Settings updated successfully!')
+        return redirect(url_for('settings'))
+        
+    return render_template('settings.html', settings=settings)
+
 @app.route('/')
 def index():
     characters = Character.query.order_by(Character.name).all()
@@ -278,6 +366,323 @@ def detail(slug):
     
     return render_template('detail.html', character=character, preferences=preferences, preview_image=preview_image)
 
+@app.route('/create', methods=['GET', 'POST'])
+def create_character():
+    if request.method == 'POST':
+        name = request.form.get('name')
+        slug = request.form.get('filename')
+        prompt = request.form.get('prompt')
+        
+        # Validate slug
+        safe_slug = re.sub(r'[^a-zA-Z0-9_]', '', slug)
+        if not safe_slug:
+            flash("Invalid filename.")
+            return redirect(request.url)
+            
+        # Check if exists
+        if os.path.exists(os.path.join(app.config['CHARACTERS_DIR'], f"{safe_slug}.json")):
+            flash("Character with this filename already exists.")
+            return redirect(request.url)
+
+        # Generate JSON with LLM
+        system_prompt = """You are a JSON generator. output ONLY valid JSON matching this exact structure. Do not wrap in markdown blocks.
+        Structure:
+        {
+          "character_id": "WILL_BE_REPLACED",
+          "character_name": "WILL_BE_REPLACED",
+          "identity": {
+            "base_specs": "string (e.g. 1girl, build, skin)",
+            "hair": "string",
+            "eyes": "string",
+            "hands": "string",
+            "arms": "string",
+            "torso": "string",
+            "pelvis": "string",
+            "legs": "string",
+            "feet": "string",
+            "extra": "string"
+          },
+          "defaults": {
+            "expression": "",
+            "pose": "",
+            "scene": ""
+          },
+          "wardrobe": {
+            "headwear": "string",
+            "top": "string",
+            "legwear": "string",
+            "footwear": "string",
+            "hands": "string",
+            "accessories": "string"
+          },
+          "styles": {
+            "aesthetic": "string",
+            "primary_color": "string",
+            "secondary_color": "string",
+            "tertiary_color": "string"
+          },
+          "lora": {
+            "lora_name": "",
+            "lora_weight": 1.0,
+            "lora_triggers": ""
+          },
+          "tags": ["string", "string"]
+        }
+        Fill the fields based on the user's description. Use Danbooru-style tags for the values (e.g. 'long hair', 'blue eyes'). Keep values concise. Leave defaults fields empty."""
+        
+        try:
+            llm_response = call_llm(f"Create a character profile for '{name}' based on this description: {prompt}", system_prompt)
+            
+            # Clean response (remove markdown if present)
+            clean_json = llm_response.replace('```json', '').replace('```', '').strip()
+            char_data = json.loads(clean_json)
+            
+            # Enforce IDs
+            char_data['character_id'] = safe_slug
+            char_data['character_name'] = name
+            
+            # Save file
+            file_path = os.path.join(app.config['CHARACTERS_DIR'], f"{safe_slug}.json")
+            with open(file_path, 'w') as f:
+                json.dump(char_data, f, indent=2)
+                
+            # Add to DB
+            new_char = Character(
+                character_id=safe_slug,
+                slug=safe_slug,
+                filename=f"{safe_slug}.json",
+                name=name,
+                data=char_data
+            )
+            db.session.add(new_char)
+            db.session.commit()
+            
+            flash('Character created successfully!')
+            return redirect(url_for('detail', slug=safe_slug))
+                
+        except Exception as e:
+            print(f"LLM/Save error: {e}")
+            flash(f"Failed to create character: {e}")
+            return redirect(request.url)
+
+    return render_template('create.html')
+
+@app.route('/character/<path:slug>/edit', methods=['GET', 'POST'])
+def edit_character(slug):
+    character = Character.query.filter_by(slug=slug).first_or_404()
+    loras = get_available_loras()
+    
+    if request.method == 'POST':
+        try:
+            # 1. Update basic fields
+            character.name = request.form.get('character_name')
+            
+            # 2. Rebuild the data dictionary
+            new_data = character.data.copy()
+            new_data['character_name'] = character.name
+            
+            # Update nested sections (non-wardrobe)
+            for section in ['identity', 'defaults', 'styles', 'lora']:
+                if section in new_data:
+                    for key in new_data[section]:
+                        form_key = f"{section}_{key}"
+                        if form_key in request.form:
+                            val = request.form.get(form_key)
+                            # Handle numeric weight
+                            if key == 'lora_weight':
+                                try: val = float(val)
+                                except: val = 1.0
+                            new_data[section][key] = val
+            
+            # Handle wardrobe - support both nested and flat formats
+            wardrobe = new_data.get('wardrobe', {})
+            if 'default' in wardrobe and isinstance(wardrobe.get('default'), dict):
+                # New nested format - update each outfit
+                for outfit_name in wardrobe.keys():
+                    for key in wardrobe[outfit_name].keys():
+                        form_key = f"wardrobe_{outfit_name}_{key}"
+                        if form_key in request.form:
+                            wardrobe[outfit_name][key] = request.form.get(form_key)
+                new_data['wardrobe'] = wardrobe
+            else:
+                # Legacy flat format
+                if 'wardrobe' in new_data:
+                    for key in new_data['wardrobe'].keys():
+                        form_key = f"wardrobe_{key}"
+                        if form_key in request.form:
+                            new_data['wardrobe'][key] = request.form.get(form_key)
+            
+            # Update Tags (comma separated string to list)
+            tags_raw = request.form.get('tags', '')
+            new_data['tags'] = [t.strip() for f in tags_raw.split(',') for t in [f.strip()] if t]
+            
+            character.data = new_data
+            flag_modified(character, "data")
+            
+            # 3. Write back to JSON file
+            # Use the filename we stored during sync, or fallback to a sanitized ID
+            char_file = character.filename or f"{re.sub(r'[^a-zA-Z0-9_]', '', character.character_id)}.json"
+            file_path = os.path.join(app.config['CHARACTERS_DIR'], char_file)
+            
+            with open(file_path, 'w') as f:
+                json.dump(new_data, f, indent=2)
+            
+            db.session.commit()
+            flash('Character profile updated successfully!')
+            return redirect(url_for('detail', slug=slug))
+            
+        except Exception as e:
+            print(f"Edit error: {e}")
+            flash(f"Error saving changes: {str(e)}")
+            
+    return render_template('edit.html', character=character, loras=loras)
+
+@app.route('/character/<path:slug>/outfit/switch', methods=['POST'])
+def switch_outfit(slug):
+    """Switch the active outfit for a character."""
+    character = Character.query.filter_by(slug=slug).first_or_404()
+    outfit_name = request.form.get('outfit', 'default')
+    
+    # Validate outfit exists
+    available_outfits = character.get_available_outfits()
+    if outfit_name in available_outfits:
+        character.active_outfit = outfit_name
+        db.session.commit()
+        flash(f'Switched to "{outfit_name}" outfit.')
+    else:
+        flash(f'Outfit "{outfit_name}" not found.', 'error')
+    
+    return redirect(url_for('detail', slug=slug))
+
+@app.route('/character/<path:slug>/outfit/add', methods=['POST'])
+def add_outfit(slug):
+    """Add a new outfit to a character."""
+    character = Character.query.filter_by(slug=slug).first_or_404()
+    outfit_name = request.form.get('outfit_name', '').strip()
+    
+    if not outfit_name:
+        flash('Outfit name cannot be empty.', 'error')
+        return redirect(url_for('edit_character', slug=slug))
+    
+    # Sanitize outfit name for use as key
+    safe_name = re.sub(r'[^a-zA-Z0-9_]', '_', outfit_name.lower())
+    
+    # Get wardrobe data
+    wardrobe = character.data.get('wardrobe', {})
+    
+    # Ensure wardrobe is in new nested format
+    if 'default' not in wardrobe or not isinstance(wardrobe.get('default'), dict):
+        # Convert legacy format
+        wardrobe = {'default': wardrobe}
+    
+    # Check if outfit already exists
+    if safe_name in wardrobe:
+        flash(f'Outfit "{safe_name}" already exists.', 'error')
+        return redirect(url_for('edit_character', slug=slug))
+    
+    # Create new outfit (copy from default as template)
+    default_outfit = wardrobe.get('default', {
+        'headwear': '', 'top': '', 'legwear': '', 
+        'footwear': '', 'hands': '', 'accessories': ''
+    })
+    wardrobe[safe_name] = default_outfit.copy()
+    
+    # Update character data
+    character.data['wardrobe'] = wardrobe
+    flag_modified(character, 'data')
+    
+    # Save to JSON file
+    char_file = character.filename or f"{re.sub(r'[^a-zA-Z0-9_]', '', character.character_id)}.json"
+    file_path = os.path.join(app.config['CHARACTERS_DIR'], char_file)
+    with open(file_path, 'w') as f:
+        json.dump(character.data, f, indent=2)
+    
+    db.session.commit()
+    flash(f'Added new outfit "{safe_name}".')
+    
+    return redirect(url_for('edit_character', slug=slug))
+
+@app.route('/character/<path:slug>/outfit/delete', methods=['POST'])
+def delete_outfit(slug):
+    """Delete an outfit from a character."""
+    character = Character.query.filter_by(slug=slug).first_or_404()
+    outfit_name = request.form.get('outfit', '')
+    
+    wardrobe = character.data.get('wardrobe', {})
+    
+    # Cannot delete default
+    if outfit_name == 'default':
+        flash('Cannot delete the default outfit.', 'error')
+        return redirect(url_for('edit_character', slug=slug))
+    
+    if outfit_name not in wardrobe:
+        flash(f'Outfit "{outfit_name}" not found.', 'error')
+        return redirect(url_for('edit_character', slug=slug))
+    
+    # Delete outfit
+    del wardrobe[outfit_name]
+    character.data['wardrobe'] = wardrobe
+    flag_modified(character, 'data')
+    
+    # Switch active outfit if deleted was active
+    if character.active_outfit == outfit_name:
+        character.active_outfit = 'default'
+    
+    # Save to JSON file
+    char_file = character.filename or f"{re.sub(r'[^a-zA-Z0-9_]', '', character.character_id)}.json"
+    file_path = os.path.join(app.config['CHARACTERS_DIR'], char_file)
+    with open(file_path, 'w') as f:
+        json.dump(character.data, f, indent=2)
+    
+    db.session.commit()
+    flash(f'Deleted outfit "{outfit_name}".')
+    
+    return redirect(url_for('edit_character', slug=slug))
+
+@app.route('/character/<path:slug>/outfit/rename', methods=['POST'])
+def rename_outfit(slug):
+    """Rename an outfit."""
+    character = Character.query.filter_by(slug=slug).first_or_404()
+    old_name = request.form.get('old_name', '')
+    new_name = request.form.get('new_name', '').strip()
+    
+    if not new_name:
+        flash('New name cannot be empty.', 'error')
+        return redirect(url_for('edit_character', slug=slug))
+    
+    # Sanitize new name
+    safe_name = re.sub(r'[^a-zA-Z0-9_]', '_', new_name.lower())
+    
+    wardrobe = character.data.get('wardrobe', {})
+    
+    if old_name not in wardrobe:
+        flash(f'Outfit "{old_name}" not found.', 'error')
+        return redirect(url_for('edit_character', slug=slug))
+    
+    if safe_name in wardrobe and safe_name != old_name:
+        flash(f'Outfit "{safe_name}" already exists.', 'error')
+        return redirect(url_for('edit_character', slug=slug))
+    
+    # Rename (copy to new key, delete old)
+    wardrobe[safe_name] = wardrobe.pop(old_name)
+    character.data['wardrobe'] = wardrobe
+    flag_modified(character, 'data')
+    
+    # Update active outfit if renamed was active
+    if character.active_outfit == old_name:
+        character.active_outfit = safe_name
+    
+    # Save to JSON file
+    char_file = character.filename or f"{re.sub(r'[^a-zA-Z0-9_]', '', character.character_id)}.json"
+    file_path = os.path.join(app.config['CHARACTERS_DIR'], char_file)
+    with open(file_path, 'w') as f:
+        json.dump(character.data, f, indent=2)
+    
+    db.session.commit()
+    flash(f'Renamed outfit "{old_name}" to "{safe_name}".')
+    
+    return redirect(url_for('edit_character', slug=slug))
+
 @app.route('/character/<path:slug>/upload', methods=['POST'])
 def upload_image(slug):
     character = Character.query.filter_by(slug=slug).first_or_404()
@@ -416,8 +821,8 @@ def _queue_generation(character, action='preview', selected_fields=None, client_
     with open('comfy_workflow.json', 'r') as f:
         workflow = json.load(f)
     
-    # 2. Build prompts
-    prompts = build_prompt(character.data, selected_fields, character.default_fields)
+    # 2. Build prompts with active outfit
+    prompts = build_prompt(character.data, selected_fields, character.default_fields, character.active_outfit)
     
     # 3. Prepare workflow
     workflow = _prepare_workflow(workflow, character, prompts)
@@ -540,5 +945,18 @@ if __name__ == '__main__':
     with app.app_context():
         os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)
         db.create_all()
+        
+        # Migration: Add active_outfit column if it doesn't exist
+        try:
+            from sqlalchemy import text
+            db.session.execute(text('ALTER TABLE character ADD COLUMN active_outfit VARCHAR(100) DEFAULT \'default\''))
+            db.session.commit()
+            print("Added active_outfit column to character table")
+        except Exception as e:
+            if 'duplicate column name' in str(e).lower() or 'already exists' in str(e).lower():
+                print("active_outfit column already exists")
+            else:
+                print(f"Migration note: {e}")
+        
         sync_characters()
     app.run(debug=True, port=5000)