Update app.py

app.py

@@ -3,11 +3,9 @@ import requests
 import base64
 import io
 import json
-from urllib.parse import unquote
 
 app = Flask(__name__)
 
-# Style configurations
 STYLES = {
     'pixel': {
         'prompt': 'Turn this image into the Pixel style.',
@@ -84,21 +82,16 @@ STYLES = {
 }
 
 def upload_base64_image(base64_data):
-    """Upload base64 image to jerrrycans-file.hf.space"""
     try:
         header, data = base64_data.split(',', 1)
         image_data = base64.b64decode(data)
-        
         files = {'file': ('generated_image.png', io.BytesIO(image_data), 'image/png')}
         headers = {
             'Origin': 'https://jerrrycans-file.hf.space',
             'Referer': 'https://jerrrycans-file.hf.space/',
             'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36'
         }
-        
-        response = requests.post('https://jerrrycans-file.hf.space/upload', 
-                               files=files, headers=headers)
-        
+        response = requests.post('https://jerrrycans-file.hf.space/upload', files=files, headers=headers)
         if response.status_code == 200:
             result = response.json()
             return f"https://jerrrycans-file.hf.space{result['url']}"
@@ -106,46 +99,35 @@ def upload_base64_image(base64_data):
         print(f"Upload error: {e}")
     return None
 
-def get_style_config(style, custom_prompt=None, custom_lora_url=None):
-    """Get style configuration, allowing custom prompt to override style prompt."""
-    style_config = STYLES.get(style.lower())
-    
-    if style_config:
-        # Style found: use its LoRA, but allow prompt to be overridden
+def get_style_config(style=None, custom_prompt=None, custom_lora_url=None):
+    if style and style.lower() in STYLES:
+        sc = STYLES[style.lower()]
         return {
-            'prompt': custom_prompt or style_config['prompt'],
-            'lora_url': custom_lora_url or style_config['lora_url']
+            'prompt': custom_prompt or sc['prompt'],
+            'lora_url': custom_lora_url or sc['lora_url']
         }
+    elif custom_prompt:
+        return {'prompt': custom_prompt, 'lora_url': custom_lora_url}
     else:
-        # No style found: treat 'style' as the prompt (legacy behavior)
-        return {
-            'prompt': custom_prompt or style,
-            'lora_url': custom_lora_url
-        }
+        return None
 
 def process_stream_response(response):
-    """Process streaming response from FAL API"""
     uploaded_images = []
     buffer = ''
-    
     for chunk in response.iter_content(chunk_size=1024, decode_unicode=True):
         if chunk:
             buffer += chunk
             lines = buffer.split('\n')
             buffer = lines.pop()
-            
             for line in lines:
                 if line.startswith('data: ') and len(line) > 6:
                     try:
                         data_content = line[6:].strip()
                         if not data_content:
                             continue
-                            
                         data = json.loads(data_content)
-                        
                         if (data.get('json', {}).get('type') in ['progress', 'completed'] and 
                             data.get('json', {}).get('data', {}).get('images')):
-                            
                             for image in data['json']['data']['images']:
                                 if image.get('url', '').startswith('data:image/'):
                                     uploaded_url = upload_base64_image(image['url'])
@@ -153,51 +135,38 @@ def process_stream_response(response):
                                         uploaded_images.append(uploaded_url)
                     except:
                         continue
-    
     return uploaded_images
 
 @app.route('/api/transform', methods=['POST'])
 def transform_image():
     try:
         data = request.get_json()
-        
         if not data:
             return jsonify({'success': False, 'error': 'No JSON data provided'}), 400
-        
         image_url = data.get('image_url')
         style = data.get('style')
-        custom_prompt = data.get('custom_prompt') # New field
+        custom_prompt = data.get('custom_prompt')
         custom_lora_url = data.get('lora_url')
-        
-        if not image_url:
-            return jsonify({'success': False, 'error': 'image_url is required'}), 400
-        
-        if not style:
-            return jsonify({'success': False, 'error': 'style is required'}), 400
-        
-        if not image_url.startswith(('http://', 'https://')):
-            return jsonify({'success': False, 'error': 'Invalid image URL'}), 400
-        
+        if not image_url or not image_url.startswith(('http://', 'https://')):
+            return jsonify({'success': False, 'error': 'Valid image_url is required'}), 400
+        if not (style or custom_prompt):
+            return jsonify({'success': False, 'error': 'Either style or custom_prompt is required'}), 400
         style_config = get_style_config(style, custom_prompt, custom_lora_url)
-        
+        if not style_config or not style_config.get('prompt'):
+            return jsonify({'success': False, 'error': 'Valid style or custom_prompt required'}), 400
         generate_params = {
             'json': {
                 'imageUrl': image_url,
                 'prompt': style_config['prompt']
             }
         }
-        
-        if style_config['lora_url']:
+        if style_config.get('lora_url'):
             generate_params['json']['loraUrl'] = style_config['lora_url']
-        
         generate_url = f"https://fal-kontext-demo.vercel.app/api/trpc/generateImageStream?input={requests.utils.quote(json.dumps(generate_params))}"
-        
         response = requests.get(generate_url, stream=True)
         if response.status_code != 200:
             return jsonify({'success': False, 'error': 'Failed to generate styled image'}), 500
-        
         uploaded_images = process_stream_response(response)
-        
         return jsonify({
             'success': True,
             'originalImage': image_url,
@@ -206,7 +175,6 @@ def transform_image():
             'generatedImages': uploaded_images,
             'count': len(uploaded_images)
         })
-        
     except Exception as e:
         return jsonify({'success': False, 'error': str(e)}), 500
 
@@ -215,7 +183,7 @@ def index():
     return jsonify({
         'message': 'Image Style Transfer API',
         'usage': 'POST /api/transform with JSON body: {"image_url": "...", "style": "...", "custom_prompt": "(optional)", "lora_url": "(optional)"}',
-        'note': 'If "custom_prompt" is provided, it will override the default prompt for the selected "style", while still using the style\'s LoRA.',
+        'note': 'Either "style" or "custom_prompt" must be provided. If "custom_prompt" is provided, it overrides the default prompt for a known "style"; if only "custom_prompt" is given, style can be omitted.',
         'styles': list(STYLES.keys())
     })