app.py

from flask import Flask, request, jsonify
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.',
        'lora_url': 'https://huggingface.co/Owen777/Kontext-Style-Loras/resolve/main/Pixel_lora_weights.safetensors'
    },
    'snoopy': {
        'prompt': 'Turn this image into the Snoopy style.',
        'lora_url': 'https://huggingface.co/Owen777/Kontext-Style-Loras/resolve/main/Snoopy_lora_weights.safetensors'
    },
    'jojo': {
        'prompt': 'Turn this image into the JoJo style.',
        'lora_url': 'https://huggingface.co/Owen777/Kontext-Style-Loras/resolve/main/Jojo_lora_weights.safetensors'
    },
    'clay': {
        'prompt': 'Turn this image into the Clay style.',
        'lora_url': 'https://huggingface.co/Owen777/Kontext-Style-Loras/resolve/main/Clay_Toy_lora_weights.safetensors'
    },
    'ghibli': {
        'prompt': 'Turn this image into the Ghibli style.',
        'lora_url': 'https://huggingface.co/Owen777/Kontext-Style-Loras/resolve/main/Ghibli_lora_weights.safetensors'
    },
    'american-cartoon': {
        'prompt': 'Turn this image into the American Cartoon style.',
        'lora_url': 'https://huggingface.co/Owen777/Kontext-Style-Loras/resolve/main/American_Cartoon_lora_weights.safetensors'
    },
    'lego': {
        'prompt': 'convert to lego style',
        'lora_url': 'https://huggingface.co/Owen777/Kontext-Style-Loras/resolve/main/LEGO_lora_weights.safetensors'
    },
    'broccoli-hair': {
        'prompt': 'Change hair to a broccoli haircut',
        'lora_url': 'https://huggingface.co/fal/Broccoli-Hair-Kontext-Dev-LoRA/resolve/main/broccoli-hair-kontext-dev-lora.safetensors'
    },
    'plushie': {
        'prompt': 'Convert to plushie style',
        'lora_url': 'https://huggingface.co/fal/Plushie-Kontext-Dev-LoRA/resolve/main/plushie-kontext-dev-lora.safetensors'
    },
    'wojak': {
        'prompt': 'Convert to wojak style drawing',
        'lora_url': 'https://huggingface.co/fal/Wojak-Kontext-Dev-LoRA/resolve/main/wojak-kontext-dev-lora.safetensors'
    },
    'upscalecompression': {
        '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': {
        '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'
    },
    'fluffy': {
        'prompt': 'make this object fluffy',
        'lora_url': None
    },
    'glass': {
        'prompt': 'make the character/object look like it was made out of glass, black background',
        'lora_url': None
    },
    'simpsons': {
        'prompt': 'convert to Simpsons cartoon style',
        'lora_url': None
    },
    'anime': {
        'prompt': 'convert to anime art style with large eyes and stylized features',
        'lora_url': None
    }
}

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',
            '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)
        
        if response.status_code == 200:
            result = response.json()
            return f"https://jerrrycans-file.hf.space{result['url']}"
    except Exception as e:
        print(f"Upload error: {e}")
    return None

def get_style_config(style, custom_lora_url=None):
    """Get style configuration"""
    style_config = STYLES.get(style.lower())
    
    if style_config:
        return {
            'prompt': style_config['prompt'],
            'lora_url': custom_lora_url or style_config['lora_url']
        }
    else:
        # Custom prompt
        return {
            'prompt': style,
            'lora_url': custom_lora_url
        }

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()  # Keep incomplete line in buffer
            
            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)
                        
                        # 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')):
                            
                            for image in data['json']['data']['images']:
                                if image.get('url', '').startswith('data:image/'):
                                    uploaded_url = upload_base64_image(image['url'])
                                    if uploaded_url:
                                        uploaded_images.append(uploaded_url)
                    except:
                        continue
    
    return uploaded_images

@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_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)
        
        # Prepare generation parameters
        generate_params = {
            'json': {
                'imageUrl': image_url,
                'prompt': style_config['prompt']
            }
        }
        
        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,
            'generatedImages': uploaded_images,
            'count': len(uploaded_images)
        })
        
    except Exception as e:
        return jsonify({'success': False, 'error': str(e)}), 500

@app.route('/')
def index():
    return jsonify({
        'message': 'Image Style Transfer API',
        'usage': 'POST /api/transform with JSON body: {"image_url": "...", "style": "...", "lora_url": "..."}',
        'styles': list(STYLES.keys())
    })

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=7860, debug=True)