Update app.py

app.py

@@ -1,137 +1,164 @@
-import os
-import base64
-import json
-from flask import Flask, render_template, request, jsonify
-from flask_cors import CORS
-import requests
-from werkzeug.utils import secure_filename
-import tempfile
-
-app = Flask(__name__)
-CORS(app)
-
-# Configuration
-app.config['MAX_CONTENT_LENGTH'] = 25 * 1024 * 1024  # 25MB max file size
-ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'webp'}
-
-# Get API key from environment variable
-OPENAI_API_KEY = os.environ.get('OPENAI_API_KEY')
-API_BASE_URL = os.environ.get('API_BASE_URL', 'https://api.openai.com/v1')
-
-def allowed_file(filename):
-    return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS
-
-@app.route('/')
-def index():
-    return render_template('index.html')
-
-@app.route('/generate', methods=['POST'])
-def generate_image():
-    try:
-        data = request.json
-        prompt = data.get('prompt', '')
-        
-        if not prompt:
-            return jsonify({'error': 'Prompt is required'}), 400
-        
-        headers = {
-            'Content-Type': 'application/json',
-            'Authorization': f'Bearer {OPENAI_API_KEY}'
-        }
-        
-        payload = {
-            'model': 'dall-e-3',
-            'prompt': prompt,
-            'n': 1,
-            'size': data.get('size', '1024x1024'),
-            'quality': data.get('quality', 'standard'),
-            'style': data.get('style', 'vivid')
-        }
-        
-        response = requests.post(
-            f'{API_BASE_URL}/images/generations',
-            headers=headers,
-            json=payload
-        )
-        
-        if response.status_code == 200:
-            result = response.json()
-            # For dall-e-3, the response contains URLs
-            image_url = result['data'][0]['url']
-            
-            # Download the image and convert to base64
-            image_response = requests.get(image_url)
-            if image_response.status_code == 200:
-                image_base64 = base64.b64encode(image_response.content).decode('utf-8')
-                return jsonify({
-                    'success': True,
-                    'image': f'data:image/png;base64,{image_base64}',
-                    'usage': result.get('usage', {})
-                })
-            else:
-                return jsonify({'error': 'Failed to download generated image'}), 500
-        else:
-            return jsonify({'error': response.json()}), response.status_code
-            
-    except Exception as e:
-        return jsonify({'error': str(e)}), 500
-
-@app.route('/edit', methods=['POST'])
-def edit_image():
-    try:
-        # Check if files were uploaded
-        if 'images' not in request.files:
-            return jsonify({'error': 'No images provided'}), 400
-        
-        files = request.files.getlist('images')
-        prompt = request.form.get('prompt', '')
-        
-        if not prompt:
-            return jsonify({'error': 'Prompt is required'}), 400
-        
-        if not files or len(files) == 0:
-            return jsonify({'error': 'At least one image is required'}), 400
-        
-        # Prepare multipart form data
-        headers = {
-            'Authorization': f'Bearer {OPENAI_API_KEY}'
-        }
-        
-        # Build files list for multipart upload
-        files_list = [('model', (None, 'dall-e-2'))]
-        files_list.append(('prompt', (None, prompt)))
-        files_list.append(('n', (None, '1')))
-        files_list.append(('size', (None, request.form.get('size', '1024x1024'))))
-        
-        # Add images
-        for idx, file in enumerate(files):
-            if file and allowed_file(file.filename):
-                # Read file content
-                file_content = file.read()
-                files_list.append(('image[]', (file.filename, file_content, file.content_type)))
-        
-        response = requests.post(
-            f'{API_BASE_URL}/images/edits',
-            headers=headers,
-            files=files_list
-        )
-        
-        if response.status_code == 200:
-            result = response.json()
-            # The response contains base64 encoded image
-            image_base64 = result['data'][0]['b64_json']
-            return jsonify({
-                'success': True,
-                'image': f'data:image/png;base64,{image_base64}',
-                'usage': result.get('usage', {})
-            })
-        else:
-            return jsonify({'error': response.json()}), response.status_code
-            
-    except Exception as e:
-        return jsonify({'error': str(e)}), 500
-
-if __name__ == '__main__':
-    if not OPENAI_API_KEY:
-        print("Warning: OPENAI_API_KEY environment variable not set!")
+import os
+import base64
+import json
+from flask import Flask, render_template, request, jsonify
+from flask_cors import CORS
+import requests
+from werkzeug.utils import secure_filename
+import tempfile
+
+app = Flask(__name__)
+CORS(app)
+
+# Configuration
+app.config['MAX_CONTENT_LENGTH'] = 25 * 1024 * 1024  # 25MB max file size
+ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'webp'}
+
+# Get API key from environment variable
+OPENAI_API_KEY = os.environ.get('OPENAI_API_KEY')
+API_BASE_URL = os.environ.get('API_BASE_URL', 'https://api.openai.com/v1')
+
+def allowed_file(filename):
+    return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS
+
+@app.route('/')
+def index():
+    return render_template('index.html')
+
+@app.route('/generate', methods=['POST'])
+def generate_image():
+    try:
+        data = request.json
+        prompt = data.get('prompt', '')
+        
+        if not prompt:
+            return jsonify({'error': 'Prompt is required'}), 400
+        
+        headers = {
+            'Content-Type': 'application/json',
+            'Authorization': f'Bearer {OPENAI_API_KEY}'
+        }
+        
+        payload = {
+            'model': 'dall-e-3',
+            'prompt': prompt,
+            'n': 1,
+            'size': data.get('size', '1024x1024'),
+            'quality': data.get('quality', 'standard'),
+            'style': data.get('style', 'vivid')
+        }
+        
+        response = requests.post(
+            f'{API_BASE_URL}/images/generations',
+            headers=headers,
+            json=payload
+        )
+        
+        if response.status_code == 200:
+            result = response.json()
+            # For dall-e-3, the response contains URLs
+            image_url = result['data'][0]['url']
+            
+            # Download the image and convert to base64
+            image_response = requests.get(image_url)
+            if image_response.status_code == 200:
+                image_base64 = base64.b64encode(image_response.content).decode('utf-8')
+                return jsonify({
+                    'success': True,
+                    'image': f'data:image/png;base64,{image_base64}',
+                    'usage': result.get('usage', {})
+                })
+            else:
+                return jsonify({'error': 'Failed to download generated image'}), 500
+        else:
+            return jsonify({'error': response.json()}), response.status_code
+            
+    except Exception as e:
+        return jsonify({'error': str(e)}), 500
+
+@app.route('/edit', methods=['POST'])
+def edit_image():
+    try:
+        # Check if file was uploaded
+        if 'image' not in request.files:
+            return jsonify({'error': 'No image provided'}), 400
+        
+        image_file = request.files['image']
+        prompt = request.form.get('prompt', '')
+        
+        if not prompt:
+            return jsonify({'error': 'Prompt is required'}), 400
+        
+        if not image_file or image_file.filename == '':
+            return jsonify({'error': 'Image is required'}), 400
+        
+        if not allowed_file(image_file.filename):
+            return jsonify({'error': 'Invalid image format'}), 400
+        
+        # Prepare multipart form data
+        headers = {
+            'Authorization': f'Bearer {OPENAI_API_KEY}'
+        }
+        
+        # Read image content
+        image_content = image_file.read()
+        
+        # Build files list for multipart upload
+        files = {
+            'model': (None, 'dall-e-2'),
+            'prompt': (None, prompt),
+            'n': (None, '1'),
+            'size': (None, request.form.get('size', '1024x1024')),
+            'image': (image_file.filename, image_content, image_file.content_type)
+        }
+        
+        # Add mask if provided
+        if 'mask' in request.files:
+            mask_file = request.files['mask']
+            if mask_file and mask_file.filename != '':
+                mask_content = mask_file.read()
+                files['mask'] = (mask_file.filename, mask_content, mask_file.content_type)
+        
+        response = requests.post(
+            f'{API_BASE_URL}/images/edits',
+            headers=headers,
+            files=files
+        )
+        
+        if response.status_code == 200:
+            result = response.json()
+            # For dall-e-2, the response format depends on response_format parameter
+            # Default is URL format
+            if 'b64_json' in result['data'][0]:
+                image_data = f'data:image/png;base64,{result["data"][0]["b64_json"]}'
+            else:
+                # Download image from URL
+                image_url = result['data'][0]['url']
+                image_response = requests.get(image_url)
+                if image_response.status_code == 200:
+                    image_base64 = base64.b64encode(image_response.content).decode('utf-8')
+                    image_data = f'data:image/png;base64,{image_base64}'
+                else:
+                    return jsonify({'error': 'Failed to download generated image'}), 500
+            
+            return jsonify({
+                'success': True,
+                'image': image_data,
+                'usage': result.get('usage', {})
+            })
+        else:
+            error_msg = 'Request failed'
+            try:
+                error_data = response.json()
+                error_msg = error_data.get('error', {}).get('message', error_msg)
+            except:
+                error_msg = response.text
+            return jsonify({'error': error_msg}), response.status_code
+            
+    except Exception as e:
+        return jsonify({'error': str(e)}), 500
+
+if __name__ == '__main__':
+    if not OPENAI_API_KEY:
+        print("Warning: OPENAI_API_KEY environment variable not set!")
     app.run(host='0.0.0.0', port=7860, debug=True)