Update app.py

app.py

@@ -53,15 +53,15 @@ STYLES = {
         'prompt': 'fix the jpeg compression',
         'lora_url': 'https://huggingface.co/fofr/flux-kontext-dev-jpeg-compression-fix-lora/resolve/main/flux-kontext-dev-jpeg-compression-fix-lora.safetensors'
     },
-        'gfx': {
+    'gfx': {
         'prompt': 'render this image into a gfx image, fill in the background perfect and edit the image to look like awesome graphic gfx image',
         'lora_url': 'https://huggingface.co/jerrrycans/gfx/resolve/main/flux-kontext-gfx-lora.safetensors'
     },
-        'watermarkremover': {
+    'watermarkremover': {
         'prompt': 'remove all the watermarks from this image, all watermarks that are over this image',
         'lora_url': 'https://huggingface.co/jerrrycans/watermark20000/resolve/main/flux-kontext-watermark20000-lora.safetensors'
     },
-        'gfxultra': {
+    'gfxultra': {
         'prompt': 'render this image into a gfx image, fill in the background perfect and edit the image to look like awesome graphic gfx image',
         'lora_url': 'https://huggingface.co/jerrrycans/gfx20000/resolve/main/flux-kontext-gfx20000-lora.safetensors'
     },
@@ -86,11 +86,9 @@ STYLES = {
 def upload_base64_image(base64_data):
     """Upload base64 image to jerrrycans-file.hf.space"""
     try:
-        # Convert base64 to bytes
         header, data = base64_data.split(',', 1)
         image_data = base64.b64decode(data)
         
-        # Upload to jerrrycans
         files = {'file': ('generated_image.png', io.BytesIO(image_data), 'image/png')}
         headers = {
             'Origin': 'https://jerrrycans-file.hf.space',
@@ -108,19 +106,20 @@ def upload_base64_image(base64_data):
         print(f"Upload error: {e}")
     return None
 
-def get_style_config(style, custom_lora_url=None):
-    """Get style configuration"""
+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
         return {
-            'prompt': style_config['prompt'],
+            'prompt': custom_prompt or style_config['prompt'],
             'lora_url': custom_lora_url or style_config['lora_url']
         }
     else:
-        # Custom prompt
+        # No style found: treat 'style' as the prompt (legacy behavior)
         return {
-            'prompt': style,
+            'prompt': custom_prompt or style,
             'lora_url': custom_lora_url
         }
 
@@ -133,7 +132,7 @@ def process_stream_response(response):
         if chunk:
             buffer += chunk
             lines = buffer.split('\n')
-            buffer = lines.pop()  # Keep incomplete line in buffer
+            buffer = lines.pop()
             
             for line in lines:
                 if line.startswith('data: ') and len(line) > 6:
@@ -144,7 +143,6 @@ def process_stream_response(response):
                             
                         data = json.loads(data_content)
                         
-                        # Check for progress or completed data with images
                         if (data.get('json', {}).get('type') in ['progress', 'completed'] and 
                             data.get('json', {}).get('data', {}).get('images')):
                             
@@ -161,32 +159,27 @@ def process_stream_response(response):
 @app.route('/api/transform', methods=['POST'])
 def transform_image():
     try:
-        # Get JSON data from request
         data = request.get_json()
         
         if not data:
             return jsonify({'success': False, 'error': 'No JSON data provided'}), 400
         
-        # Extract required fields
         image_url = data.get('image_url')
         style = data.get('style')
+        custom_prompt = data.get('custom_prompt') # New field
         custom_lora_url = data.get('lora_url')
         
-        # Validate required fields
         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
         
-        # Validate image URL
         if not image_url.startswith(('http://', 'https://')):
             return jsonify({'success': False, 'error': 'Invalid image URL'}), 400
         
-        # Get style configuration
-        style_config = get_style_config(style, custom_lora_url)
+        style_config = get_style_config(style, custom_prompt, custom_lora_url)
         
-        # Prepare generation parameters
         generate_params = {
             'json': {
                 'imageUrl': image_url,
@@ -197,20 +190,19 @@ def transform_image():
         if style_config['lora_url']:
             generate_params['json']['loraUrl'] = style_config['lora_url']
         
-        # Make request to FAL API
         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
         
-        # Process streaming response
         uploaded_images = process_stream_response(response)
         
         return jsonify({
             'success': True,
             'originalImage': image_url,
             'style': style,
+            'promptUsed': style_config['prompt'],
             'generatedImages': uploaded_images,
             'count': len(uploaded_images)
         })
@@ -222,7 +214,8 @@ def transform_image():
 def index():
     return jsonify({
         'message': 'Image Style Transfer API',
-        'usage': 'POST /api/transform with JSON body: {"image_url": "...", "style": "...", "lora_url": "..."}',
+        '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.',
         'styles': list(STYLES.keys())
     })