311 lines
12 KiB
Python
311 lines
12 KiB
Python
import os
|
|
import base64
|
|
import uuid
|
|
import glob
|
|
import json
|
|
from datetime import datetime
|
|
from io import BytesIO
|
|
from send2trash import send2trash
|
|
from flask import Flask, render_template, request, jsonify, url_for
|
|
from google import genai
|
|
from google.genai import types
|
|
from PIL import Image, PngImagePlugin
|
|
|
|
app = Flask(__name__)
|
|
app.config['SEND_FILE_MAX_AGE_DEFAULT'] = 0
|
|
|
|
# Ensure generated directory exists inside Flask static folder
|
|
GENERATED_DIR = os.path.join(app.static_folder, 'generated')
|
|
os.makedirs(GENERATED_DIR, exist_ok=True)
|
|
|
|
# Ensure uploads directory exists
|
|
UPLOADS_DIR = os.path.join(app.static_folder, 'uploads')
|
|
os.makedirs(UPLOADS_DIR, exist_ok=True)
|
|
|
|
@app.route('/')
|
|
def index():
|
|
return render_template('index.html')
|
|
|
|
@app.route('/generate', methods=['POST'])
|
|
def generate_image():
|
|
multipart = request.content_type and 'multipart/form-data' in request.content_type
|
|
|
|
if multipart:
|
|
form = request.form
|
|
prompt = form.get('prompt')
|
|
aspect_ratio = form.get('aspect_ratio')
|
|
resolution = form.get('resolution', '2K')
|
|
api_key = form.get('api_key') or os.environ.get('GOOGLE_API_KEY')
|
|
reference_files = request.files.getlist('reference_images')
|
|
reference_paths_json = form.get('reference_image_paths')
|
|
else:
|
|
data = request.get_json() or {}
|
|
prompt = data.get('prompt')
|
|
aspect_ratio = data.get('aspect_ratio')
|
|
resolution = data.get('resolution', '2K')
|
|
api_key = data.get('api_key') or os.environ.get('GOOGLE_API_KEY')
|
|
reference_files = []
|
|
reference_paths_json = data.get('reference_image_paths')
|
|
|
|
if not prompt:
|
|
return jsonify({'error': 'Prompt is required'}), 400
|
|
|
|
if not api_key:
|
|
return jsonify({'error': 'API Key is required.'}), 401
|
|
|
|
try:
|
|
client = genai.Client(api_key=api_key)
|
|
|
|
image_config_args = {
|
|
"image_size": resolution
|
|
}
|
|
|
|
if aspect_ratio and aspect_ratio != 'Auto':
|
|
image_config_args["aspect_ratio"] = aspect_ratio
|
|
|
|
# Process reference paths and files
|
|
final_reference_paths = []
|
|
contents = [prompt]
|
|
|
|
# Parse reference paths from frontend
|
|
frontend_paths = []
|
|
if reference_paths_json:
|
|
try:
|
|
frontend_paths = json.loads(reference_paths_json)
|
|
except json.JSONDecodeError:
|
|
pass
|
|
|
|
# If no paths provided but we have files (legacy or simple upload), treat all as new uploads
|
|
# But we need to handle the mix.
|
|
# Strategy: Iterate frontend_paths. If it looks like a path/URL, keep it.
|
|
# If it doesn't (or is null), consume from reference_files.
|
|
|
|
file_index = 0
|
|
|
|
# If frontend_paths is empty but we have files, just use the files
|
|
if not frontend_paths and reference_files:
|
|
for _ in reference_files:
|
|
frontend_paths.append(None) # Placeholder for each file
|
|
|
|
for path in frontend_paths:
|
|
if path and (path.startswith('/') or path.startswith('http')):
|
|
# Existing path/URL
|
|
final_reference_paths.append(path)
|
|
# We also need to add the image content to the prompt
|
|
# We need to fetch it or read it if it's local (server-side local)
|
|
# If it's a URL we generated, it's in static/generated or static/uploads
|
|
# path might be "http://localhost:8888/static/generated/..." or "/static/generated/..."
|
|
|
|
# Extract relative path to open file
|
|
# Assuming path contains '/static/'
|
|
try:
|
|
if '/static/' in path:
|
|
rel_path = path.split('/static/')[1]
|
|
abs_path = os.path.join(app.static_folder, rel_path)
|
|
if os.path.exists(abs_path):
|
|
img = Image.open(abs_path)
|
|
contents.append(img)
|
|
else:
|
|
print(f"Warning: Reference file not found at {abs_path}")
|
|
else:
|
|
print(f"Warning: Could not resolve local path for {path}")
|
|
except Exception as e:
|
|
print(f"Error loading reference from path {path}: {e}")
|
|
|
|
elif file_index < len(reference_files):
|
|
# New upload
|
|
file = reference_files[file_index]
|
|
file_index += 1
|
|
|
|
try:
|
|
# Save to uploads
|
|
ext = os.path.splitext(file.filename)[1]
|
|
if not ext:
|
|
ext = '.png'
|
|
filename = f"{uuid.uuid4()}{ext}"
|
|
filepath = os.path.join(UPLOADS_DIR, filename)
|
|
|
|
# We need to read the file for Gemini AND save it
|
|
# file.stream is a stream.
|
|
file.stream.seek(0)
|
|
file_bytes = file.read()
|
|
|
|
with open(filepath, 'wb') as f:
|
|
f.write(file_bytes)
|
|
|
|
# Add to contents
|
|
image = Image.open(BytesIO(file_bytes))
|
|
contents.append(image)
|
|
|
|
# Add to final paths
|
|
# URL for the uploaded file
|
|
rel_path = os.path.join('uploads', filename)
|
|
file_url = url_for('static', filename=rel_path)
|
|
final_reference_paths.append(file_url)
|
|
|
|
except Exception as e:
|
|
print(f"Error processing uploaded file: {e}")
|
|
continue
|
|
|
|
model_name = "gemini-3-pro-image-preview"
|
|
response = client.models.generate_content(
|
|
model=model_name,
|
|
contents=contents,
|
|
config=types.GenerateContentConfig(
|
|
response_modalities=['IMAGE'],
|
|
image_config=types.ImageConfig(**image_config_args),
|
|
)
|
|
)
|
|
|
|
for part in response.parts:
|
|
if part.inline_data:
|
|
image_bytes = part.inline_data.data
|
|
|
|
image = Image.open(BytesIO(image_bytes))
|
|
png_info = PngImagePlugin.PngInfo()
|
|
|
|
date_str = datetime.now().strftime("%Y%m%d")
|
|
|
|
# Find existing files to determine next ID
|
|
search_pattern = os.path.join(GENERATED_DIR, f"{model_name}_{date_str}_*.png")
|
|
existing_files = glob.glob(search_pattern)
|
|
max_id = 0
|
|
for f in existing_files:
|
|
try:
|
|
basename = os.path.basename(f)
|
|
name_without_ext = os.path.splitext(basename)[0]
|
|
id_part = name_without_ext.split('_')[-1]
|
|
id_num = int(id_part)
|
|
if id_num > max_id:
|
|
max_id = id_num
|
|
except ValueError:
|
|
continue
|
|
|
|
next_id = max_id + 1
|
|
filename = f"{model_name}_{date_str}_{next_id}.png"
|
|
filepath = os.path.join(GENERATED_DIR, filename)
|
|
rel_path = os.path.join('generated', filename)
|
|
image_url = url_for('static', filename=rel_path)
|
|
|
|
metadata = {
|
|
'prompt': prompt,
|
|
'aspect_ratio': aspect_ratio or 'Auto',
|
|
'resolution': resolution,
|
|
'reference_images': final_reference_paths,
|
|
}
|
|
|
|
png_info.add_text('sdvn_meta', json.dumps(metadata))
|
|
|
|
buffer = BytesIO()
|
|
image.save(buffer, format='PNG', pnginfo=png_info)
|
|
final_bytes = buffer.getvalue()
|
|
|
|
# Save image to file
|
|
with open(filepath, 'wb') as f:
|
|
f.write(final_bytes)
|
|
|
|
image_data = base64.b64encode(final_bytes).decode('utf-8')
|
|
return jsonify({
|
|
'image': image_url,
|
|
'image_data': image_data,
|
|
'metadata': metadata,
|
|
})
|
|
|
|
return jsonify({'error': 'No image generated'}), 500
|
|
|
|
except Exception as e:
|
|
return jsonify({'error': str(e)}), 500
|
|
|
|
@app.route('/delete_image', methods=['POST'])
|
|
def delete_image():
|
|
data = request.get_json()
|
|
filename = data.get('filename')
|
|
|
|
if not filename:
|
|
return jsonify({'error': 'Filename is required'}), 400
|
|
|
|
# Security check: ensure filename is just a basename, no paths
|
|
filename = os.path.basename(filename)
|
|
filepath = os.path.join(GENERATED_DIR, filename)
|
|
|
|
if os.path.exists(filepath):
|
|
try:
|
|
send2trash(filepath)
|
|
return jsonify({'success': True})
|
|
except Exception as e:
|
|
return jsonify({'error': str(e)}), 500
|
|
else:
|
|
return jsonify({'error': 'File not found'}), 404
|
|
|
|
@app.route('/gallery')
|
|
def get_gallery():
|
|
# List all png files in generated dir, sorted by modification time (newest first)
|
|
files = glob.glob(os.path.join(GENERATED_DIR, '*.png'))
|
|
files.sort(key=os.path.getmtime, reverse=True)
|
|
|
|
image_urls = [url_for('static', filename=f'generated/{os.path.basename(f)}') for f in files]
|
|
response = jsonify({'images': image_urls})
|
|
response.headers["Cache-Control"] = "no-cache, no-store, must-revalidate"
|
|
return response
|
|
|
|
@app.route('/prompts')
|
|
def get_prompts():
|
|
category = request.args.get('category')
|
|
|
|
try:
|
|
# Read prompts.json file
|
|
prompts_path = os.path.join(os.path.dirname(__file__), 'prompts.json')
|
|
with open(prompts_path, 'r', encoding='utf-8') as f:
|
|
prompts = json.load(f)
|
|
|
|
# Filter by category if specified
|
|
if category:
|
|
prompts = [p for p in prompts if p.get('category') == category]
|
|
|
|
response = jsonify({'prompts': prompts})
|
|
response.headers["Cache-Control"] = "no-cache, no-store, must-revalidate"
|
|
return response
|
|
except Exception as e:
|
|
return jsonify({'error': str(e)}), 500
|
|
|
|
@app.route('/refine_prompt', methods=['POST'])
|
|
def refine_prompt():
|
|
data = request.get_json()
|
|
current_prompt = data.get('current_prompt')
|
|
instruction = data.get('instruction')
|
|
api_key = data.get('api_key') or os.environ.get('GOOGLE_API_KEY')
|
|
|
|
if not api_key:
|
|
return jsonify({'error': 'API Key is required.'}), 401
|
|
|
|
if not instruction:
|
|
return jsonify({'error': 'Instruction is required'}), 400
|
|
|
|
try:
|
|
client = genai.Client(api_key=api_key)
|
|
|
|
system_instruction = "You are an expert prompt engineer for image generation AI. Rewrite the prompt to incorporate the user's instruction while maintaining the original intent and improving quality. Return ONLY the new prompt text, no explanations."
|
|
|
|
prompt_content = f"Current prompt: {current_prompt}\nUser instruction: {instruction}\nNew prompt:"
|
|
|
|
print(f"Refining prompt with instruction: {instruction}")
|
|
|
|
response = client.models.generate_content(
|
|
model="gemini-2.5-flash",
|
|
contents=[prompt_content],
|
|
config=types.GenerateContentConfig(
|
|
system_instruction=system_instruction,
|
|
temperature=0.7,
|
|
)
|
|
)
|
|
|
|
if response.text:
|
|
return jsonify({'refined_prompt': response.text.strip()})
|
|
else:
|
|
return jsonify({'error': 'No response from AI'}), 500
|
|
|
|
except Exception as e:
|
|
return jsonify({'error': str(e)}), 500
|
|
|
|
if __name__ == '__main__':
|
|
app.run(debug=True, port=8888)
|