Spaces:
Running
Running
| import random | |
| import string | |
| import requests | |
| from flask import Flask, request, Response, stream_with_context, send_file | |
| from io import BytesIO | |
| app = Flask(__name__) | |
| # --- Utility functions --- | |
| def random_user_agent(): | |
| agents = [ | |
| 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/124.0.0.0 Safari/537.36', | |
| 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/605.1.15 (KHTML, like Gecko) Version/17.0 Safari/605.1.15', | |
| 'Mozilla/5.0 (Linux; Android 13; SM-G991B) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/124.0.0.0 Mobile Safari/537.36', | |
| 'Mozilla/5.0 (iPhone; CPU iPhone OS 17_0 like Mac OS X) AppleWebKit/605.1.15 (KHTML, like Gecko) Version/17.0 Mobile/15E148 Safari/604.1', | |
| 'Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/124.0.0.0 Safari/537.36' | |
| ] | |
| return random.choice(agents) | |
| def random_ip(): | |
| return '.'.join(str(random.randint(0, 255)) for _ in range(4)) | |
| def random_session_hash(): | |
| return ''.join(random.choices(string.ascii_lowercase + string.digits, k=12)) | |
| # --- SSE Streaming Helper --- | |
| def sse_stream(generator): | |
| return Response(stream_with_context(generator), mimetype='text/event-stream') | |
| # --- UNO-FLUX Endpoint --- | |
| def generate_flux(): | |
| data = request.json or {} | |
| prompt = data.get('prompt', '') | |
| width = int(data.get('width', 1280)) | |
| height = int(data.get('height', 720)) | |
| guidance = data.get('guidance', 4) | |
| num_steps = data.get('num_steps', 25) | |
| seed = data.get('seed', -1) | |
| image_prompt1 = data.get('image_prompt1') | |
| image_prompt2 = data.get('image_prompt2') | |
| image_prompt3 = data.get('image_prompt3') | |
| image_prompt4 = data.get('image_prompt4') | |
| aspect = data.get('aspect', '16:9') | |
| if not prompt: | |
| return sse_stream(lambda: iter([f"data: {{\"type\":\"error\",\"message\":\"Prompt is required\"}}\n\n"])) | |
| # Aspect ratio logic | |
| if aspect == '9:16': | |
| width, height = 720, 1280 | |
| elif aspect == '1:1': | |
| width, height = 1024, 1024 | |
| session_hash = random_session_hash() | |
| join_payload = { | |
| "data": [ | |
| prompt, | |
| width, | |
| height, | |
| guidance, | |
| num_steps, | |
| seed, | |
| image_prompt1, | |
| image_prompt2, | |
| image_prompt3, | |
| image_prompt4 | |
| ], | |
| "event_data": None, | |
| "fn_index": 0, | |
| "trigger_id": 25, | |
| "session_hash": session_hash | |
| } | |
| def event_gen(): | |
| # Join queue | |
| try: | |
| join_resp = requests.post( | |
| 'https://bytedance-research-uno-flux.hf.space/gradio_api/queue/join?__theme=system', | |
| json=join_payload, | |
| headers={ | |
| 'Content-Type': 'application/json', | |
| 'Referer': 'https://bytedance-research-uno-flux.hf.space/', | |
| 'Origin': 'https://bytedance-research-uno-flux.hf.space', | |
| 'User-Agent': random_user_agent(), | |
| 'X-Forwarded-For': random_ip() | |
| }, | |
| timeout=30 | |
| ) | |
| if join_resp.status_code != 200: | |
| yield f"data: {{\"type\":\"error\",\"message\":\"Failed to join UNO-FLUX queue\",\"status\":{join_resp.status_code}}}\n\n" | |
| return | |
| except Exception as e: | |
| yield f"data: {{\"type\":\"error\",\"message\":\"Failed to join UNO-FLUX queue\",\"details\":\"{str(e)}\"}}\n\n" | |
| return | |
| # Poll for results | |
| max_attempts = 90 | |
| for attempt in range(max_attempts): | |
| try: | |
| poll_resp = requests.get( | |
| f'https://bytedance-research-uno-flux.hf.space/gradio_api/queue/data?session_hash={session_hash}', | |
| headers={ | |
| 'Accept': 'text/event-stream', | |
| 'Referer': 'https://bytedance-research-uno-flux.hf.space/', | |
| 'Origin': 'https://bytedance-research-uno-flux.hf.space', | |
| 'User-Agent': random_user_agent(), | |
| 'X-Forwarded-For': random_ip() | |
| }, | |
| timeout=30 | |
| ) | |
| if poll_resp.status_code != 200: | |
| yield f"data: {{\"type\":\"error\",\"message\":\"Polling failed\",\"status\":{poll_resp.status_code}}}\n\n" | |
| return | |
| lines = poll_resp.text.split('\n') | |
| event_data = '' | |
| for line in lines: | |
| if line.startswith('data: '): | |
| event_data += line[6:] | |
| elif line == '': | |
| if event_data: | |
| try: | |
| json_event = eval(event_data, {}, {}) if event_data.strip().startswith('{') else None | |
| except Exception: | |
| json_event = None | |
| if json_event: | |
| if json_event.get('msg') == 'estimation': | |
| yield f"data: {{\"type\":\"estimation\",\"queueSize\":{json_event.get('queue_size')},\"eta\":{json_event.get('rank_eta')}}}\n\n" | |
| elif json_event.get('msg') == 'process_starts': | |
| yield f"data: {{\"type\":\"processing\"}}\n\n" | |
| elif json_event.get('msg') == 'process_generating': | |
| yield f"data: {{\"type\":\"generating\",\"output\":{json_event.get('output')}}}\n\n" | |
| elif json_event.get('msg') == 'process_failed': | |
| yield f"data: {{\"type\":\"error\",\"message\":\"Generation failed on server\"}}\n\n" | |
| return | |
| elif json_event.get('msg') == 'process_completed' and json_event.get('output', {}).get('data', [{}])[0].get('url'): | |
| image_url = json_event['output']['data'][0]['url'] | |
| proxies = f"/api/proxy-image?url={image_url}" | |
| yield f"data: {{\"type\":\"success\",\"imageUrl\":\"{image_url}\",\"proxies\":\"{proxies}\"}}\n\n" | |
| return | |
| elif json_event.get('msg') == 'process_completed' and json_event.get('output', {}).get('error'): | |
| yield f"data: {{\"type\":\"error\",\"message\":\"{json_event['output']['error']}\"}}\n\n" | |
| return | |
| event_data = '' | |
| elif line.startswith(':'): | |
| continue | |
| elif line != '': | |
| event_data += line | |
| except Exception as e: | |
| yield f"data: {{\"type\":\"error\",\"message\":\"Polling error\",\"details\":\"{str(e)}\"}}\n\n" | |
| return | |
| yield f"data: {{\"type\":\"error\",\"message\":\"Generation timed out\"}}\n\n" | |
| return sse_stream(event_gen()) | |
| # --- Heartsync Endpoint --- | |
| def generate_heartsync(): | |
| data = request.json or {} | |
| prompt = data.get('prompt', '') | |
| if not prompt: | |
| return sse_stream(lambda: iter([f"data: {{\"type\":\"error\",\"message\":\"Prompt is required\"}}\n\n"])) | |
| session_hash = random_session_hash() | |
| join_payload = { | |
| "data": [ | |
| prompt, | |
| "text, talk bubble, low quality, watermark, signature", | |
| 0, | |
| True, | |
| 1024, | |
| 1024, | |
| 7, | |
| 28 | |
| ], | |
| "event_data": None, | |
| "fn_index": 2, | |
| "trigger_id": 14, | |
| "session_hash": session_hash | |
| } | |
| def event_gen(): | |
| try: | |
| join_resp = requests.post( | |
| 'https://heartsync-nsfw-uncensored.hf.space/gradio_api/queue/join', | |
| json=join_payload, | |
| headers={ | |
| 'Content-Type': 'application/json', | |
| 'Referer': 'https://heartsync-nsfw-uncensored.hf.space/', | |
| 'User-Agent': random_user_agent(), | |
| 'X-Forwarded-For': random_ip() | |
| }, | |
| timeout=30 | |
| ) | |
| if join_resp.status_code != 200: | |
| yield f"data: {{\"type\":\"error\",\"message\":\"Failed to join Heartsync queue\",\"status\":{join_resp.status_code}}}\n\n" | |
| return | |
| except Exception as e: | |
| yield f"data: {{\"type\":\"error\",\"message\":\"Failed to join Heartsync queue\",\"details\":\"{str(e)}\"}}\n\n" | |
| return | |
| max_attempts = 60 | |
| for attempt in range(max_attempts): | |
| try: | |
| poll_resp = requests.get( | |
| f'https://heartsync-nsfw-uncensored.hf.space/gradio_api/queue/data?session_hash={session_hash}', | |
| headers={ | |
| 'Accept': 'text/event-stream', | |
| 'Referer': 'https://heartsync-nsfw-uncensored.hf.space/', | |
| 'User-Agent': random_user_agent(), | |
| 'X-Forwarded-For': random_ip() | |
| }, | |
| timeout=30 | |
| ) | |
| if poll_resp.status_code != 200: | |
| yield f"data: {{\"type\":\"error\",\"message\":\"Polling failed\",\"status\":{poll_resp.status_code}}}\n\n" | |
| return | |
| lines = poll_resp.text.split('\n') | |
| event_data = '' | |
| for line in lines: | |
| if line.startswith('data: '): | |
| event_data += line[6:] | |
| elif line == '': | |
| if event_data: | |
| try: | |
| json_event = eval(event_data, {}, {}) if event_data.strip().startswith('{') else None | |
| except Exception: | |
| json_event = None | |
| if json_event: | |
| if json_event.get('msg') == 'estimation': | |
| yield f"data: {{\"type\":\"estimation\",\"queueSize\":{json_event.get('queue_size')},\"eta\":{json_event.get('rank_eta')}}}\n\n" | |
| elif json_event.get('msg') == 'process_starts': | |
| yield f"data: {{\"type\":\"processing\"}}\n\n" | |
| elif json_event.get('msg') == 'process_generating': | |
| yield f"data: {{\"type\":\"generating\",\"output\":{json_event.get('output')}}}\n\n" | |
| elif json_event.get('msg') == 'process_failed': | |
| yield f"data: {{\"type\":\"error\",\"message\":\"Generation failed on server\"}}\n\n" | |
| return | |
| elif json_event.get('msg') == 'process_completed' and json_event.get('output', {}).get('data', [{}])[0].get('url'): | |
| image_url = json_event['output']['data'][0]['url'] | |
| proxies = f"/api/proxy-image?url={image_url}" | |
| yield f"data: {{\"type\":\"success\",\"originalUrl\":\"{image_url}\"}}\n\n" | |
| return | |
| elif json_event.get('msg') == 'process_completed' and json_event.get('output', {}).get('error'): | |
| yield f"data: {{\"type\":\"error\",\"message\":\"{json_event['output']['error']}\"}}\n\n" | |
| return | |
| event_data = '' | |
| elif line.startswith(':'): | |
| continue | |
| elif line != '': | |
| event_data += line | |
| except Exception as e: | |
| yield f"data: {{\"type\":\"error\",\"message\":\"Polling error\",\"details\":\"{str(e)}\"}}\n\n" | |
| return | |
| yield f"data: {{\"type\":\"error\",\"message\":\"Generation timed out\"}}\n\n" | |
| return sse_stream(event_gen()) | |
| # --- Proxy Image Endpoint --- | |
| def proxy_image(): | |
| url = request.args.get('url') | |
| if not url: | |
| return 'Missing url parameter', 400 | |
| try: | |
| resp = requests.get(url, stream=True, timeout=30) | |
| if resp.status_code != 200: | |
| return 'Failed to fetch image', resp.status_code | |
| return Response(resp.raw, content_type=resp.headers.get('content-type', 'image/png')) | |
| except Exception as e: | |
| return f'Proxy error: {str(e)}', 500 | |
| if __name__ == '__main__': | |
| app.run(host='0.0.0.0', port=7860, debug=True) |