yml/generate-editor.py

import os, sys, requests, re, time
import httpx

# =========================================================
# رفع قطعی مشکل Time Out رانر گیت‌هاب بر اساس اسپیس قدیمی
original_client_init = httpx.Client.__init__
def patched_client_init(self, *args, **kwargs):
    kwargs['timeout'] = httpx.Timeout(300.0)
    original_client_init(self, *args, **kwargs)
httpx.Client.__init__ = patched_client_init

original_async_init = httpx.AsyncClient.__init__
def patched_async_init(self, *args, **kwargs):
    kwargs['timeout'] = httpx.Timeout(300.0)
    original_async_init(self, *args, **kwargs)
httpx.AsyncClient.__init__ = patched_async_init
# =========================================================

from gradio_client import Client, handle_file

prompt = os.environ.get('PROMPT', '')
image_url_raw = os.environ.get('IMAGE_URL', '')

# تجزیه و تفکیک لینک‌ها از متغیر مشترک گیت‌هاب
urls = image_url_raw.split(',')
image_url = urls[0] if len(urls) > 0 else ""
image_url2 = urls[1] if len(urls) > 1 else ""
image_url3 = urls[2] if len(urls) > 2 else ""

run_id = os.environ.get('RUN_ID')
space_url = os.environ.get('SPACE_URL')
github_run_id = os.environ.get('GITHUB_RUN_ID')

# هماهنگی کامل با سیستم هوشمند تلاش مجدد و بررسی خطاهای شبکه
def report_failure(error_msg):
    try:
        requests.post(
            f"{space_url}/api/webhook/fail",
            json={
                "run_id": run_id,
                "error": error_msg,
                "event_type": "generate-editor",
                "client_payload": {
                    "prompt": prompt,
                    "image_url": image_url,
                    "image_url2": image_url2,
                    "image_url3": image_url3,
                    "run_id": run_id,
                    "space_url": space_url
                },
                "github_run_id": github_run_id
            },
            timeout=15
        )
    except Exception as e:
        print(f"Failed to report failure: {e}")

def check_existing_edited_image(space_url, run_id):
    """بررسی اینکه آیا تصویر ویرایش شده در تلاش قبلی آپلود شده است یا خیر"""
    for ext in ['png', 'webp', 'jpg']:
        url = f"{space_url}/static/images/{run_id}.{ext}"
        try:
            # استفاده از GET سبک با stream=True برای سازگاری کامل با تمامی سرورها
            res = requests.get(url, stream=True, timeout=10)
            if res.status_code == 200:
                return url, ext
        except:
            pass
    return None, None

print('1. Checking if edited image already exists on Docker Space...')
existing_url, existing_ext = None, None
if run_id:
    existing_url, existing_ext = check_existing_edited_image(space_url, run_id)

filename = 'output.png'
ext = 'png'

# اگر تصویر ویرایش شده از قبل موجود بود، نیازی به فرآیند ادیت مجدد نیست
if existing_url:
    print(f'   -> Edited image already exists: {existing_url}')
    print('   -> Skipping Omni Editor pipeline. Proceeding directly to CodeFormer upscaling...')
    ext = existing_ext
    filename = f'output.{ext}'
    try:
        img_req = requests.get(existing_url, timeout=60)
        with open(filename, 'wb') as f:
            f.write(img_req.content)
        print(f'   -> Downloaded existing edited image as {filename}')
    except Exception as dl_err:
        print(f'   -> Failed to download existing image: {dl_err}. Re-running Omni Editor...')
        existing_url = None

if not existing_url:
    print('1. Downloading input images from Docker Space...')
    try:
        req = requests.get(image_url, timeout=30)
        with open('input.png', 'wb') as f:
            f.write(req.content)
        print('   -> Input image 1 ready.')
    except Exception as e:
        err_str = f"Download failed for image 1: {e}"
        print(err_str)
        report_failure(err_str)
        sys.exit(1)

    # دانلود اختیاری تصویر دوم
    has_image2 = False
    if image_url2 and image_url2.strip():
        try:
            req2 = requests.get(image_url2, timeout=30)
            with open('input2.png', 'wb') as f:
                f.write(req2.content)
            has_image2 = True
            print('   -> Input image 2 ready.')
        except Exception as e:
            print(f"   -> Image 2 optional download warning (skipped): {e}")

    # دانلود اختیاری تصویر سوم
    has_image3 = False
    if image_url3 and image_url3.strip():
        try:
            req3 = requests.get(image_url3, timeout=30)
            with open('input3.png', 'wb') as f:
                f.write(req3.content)
            has_image3 = True
            print('   -> Input image 3 ready.')
        except Exception as e:
            print(f"   -> Image 3 optional download warning (skipped): {e}")

    print('2. Connecting to Omni Editor and Processing (Timeout extended to 300s)...')
    result_str = ''
    success = False

    # مکانیزم تلاش مجدد در صورت شلوغی سرور
    for attempt in range(3):
        try:
            print(f'   -> Attempt {attempt + 1} of 3...')
            client = Client('selfit-camera/omni-image-editor')
            
            # اگر فقط تصویر اول آپلود شده باشد، مثل سابق رفتار کن
            if not (has_image2 or has_image3):
                print('   -> Processing with single-image pipeline (exact same as before)...')
                result = client.predict(
                    handle_file('input.png'), 
                    prompt,
                    api_name='/edit_image_interface'
                )
            # اگر چند تصویر آپلود شده باشند، از هوک چند تصویری استفاده کن
            else:
                print('   -> Processing with multi-image pipeline (fn_index=10)...')
                result = client.predict(
                    handle_file('input.png'),
                    handle_file('input2.png') if has_image2 else None,
                    handle_file('input3.png') if has_image3 else None,
                    prompt,
                    'Auto', # نسبت تصویر پیشفرض
                    fn_index=10
                )
            
            result_str = str(result)
            if 'src=' in result_str or 'http' in result_str:
                success = True
                break
            else:
                print(f'   -> Missing valid response format: {result_str[:100]}...')
        except Exception as client_err:
            print(f'   -> Attempt {attempt + 1} failed: {client_err}')
            time.sleep(5)

    if not success:
        err_str = f"All attempts failed. Raw result: {result_str}"
        print(f'CRITICAL ERROR: {err_str}')
        report_failure(err_str)
        sys.exit(1)

    print('3. Parsing result string...')
    urls = re.findall(r'src=[\'\"]([^\'\"]+)[\'\"]', result_str)
    final_remote_url = None

    for url in urls:
        if url.startswith('http'):
            final_remote_url = url
            break
        elif url.startswith('/'):
            # برای تبدیل آدرس‌های نسبی محلی گرادیو به آدرس کامل و قابل دانلود
            final_remote_url = f"https://selfit-camera-omni-image-editor.hf.space{url}"
            break

    if not final_remote_url:
        direct_urls = re.findall(r'(https?://[^\s\'\"]+\.(?:png|jpg|jpeg|webp|gif))', result_str)
        if direct_urls:
            final_remote_url = direct_urls[0]

    if not final_remote_url:
        err_str = f"Could not find valid URL. Raw result: {result_str}"
        print(f'ERROR: {err_str}')
        report_failure(err_str)
        sys.exit(1)
        
    print(f'   -> Found Final Image URL: {final_remote_url}')

    print('4. Downloading final image from remote host...')
    try:
        img_req = requests.get(final_remote_url, timeout=60)
        
        content_type = img_req.headers.get('Content-Type', '')
        ext = 'png'
        if 'image/webp' in content_type:
            ext = 'webp'
        elif 'image/jpeg' in content_type or 'image/jpg' in content_type:
            ext = 'jpg'
            
        filename = f'output.{ext}'
        with open(filename, 'wb') as f:
            f.write(img_req.content)
        print(f'   -> Saved locally as {filename}')
    except Exception as dl_err:
        err_str = f"Failed to download processed image: {dl_err}"
        print(err_str)
        report_failure(err_str)
        sys.exit(1)

    print(f'4.2. Uploading edited image to Docker Space: {space_url}')
    try:
        with open(filename, 'rb') as f:
            res = requests.post(
                f'{space_url}/api/webhook/upload', 
                data={'run_id': run_id, 'github_run_id': github_run_id, 'ext': ext}, 
                files={'file': f},
                timeout=60
            )
        if res.status_code == 200:
            print('   -> Edited image uploaded successfully.')
        else:
            print(f'   -> Edited image upload failed: {res.status_code} - {res.text}')
    except Exception as up_err:
        print(f'   -> Edited image upload exception: {up_err}')

# =========================================================
# بخش بهبودیافته: تلاش چندجانبه برای ارتباط موفق با API اسپیس CodeFormer
print('4.5. Connecting to CodeFormer to restore and upscale (Factor=4, Fidelity=1.0)...')
cf_success = False
try:
    cf_client = Client('sczhou/CodeFormer')
    upscale_result = None
    
    # ساختار مانیتورینگ متدهای متوالی برای جلوگیری از بروز خطای api_name
    methods_to_try = [
        # روش ۱: استفاده از ایندکس صفر به همراه ساختار ۶ پارامتری (ویرایش جدید اسپیس)
        {"args": [handle_file(filename), True, True, True, 4, 1.0], "kwargs": {"fn_index": 0}},
        # روش ۲: استفاده از ایندکس صفر به همراه ساختار ۵ پارامتری (نسخه‌های پایدار قدیمی)
        {"args": [handle_file(filename), True, True, 4, 1.0], "kwargs": {"fn_index": 0}},
        # روش ۳: تلاش با ایندکس یک به همراه ساختار ۶ پارامتری
        {"args": [handle_file(filename), True, True, True, 4, 1.0], "kwargs": {"fn_index": 1}},
        # روش ۴: تلاش با ایندکس یک به همراه ساختار ۵ پارامتری
        {"args": [handle_file(filename), True, True, 4, 1.0], "kwargs": {"fn_index": 1}},
        # روش ۵: فراخوانی اسمی بدون اسلش (predict) با ساختار ۶ پارامتری
        {"args": [handle_file(filename), True, True, True, 4, 1.0], "kwargs": {"api_name": "predict"}},
        # روش ۶: فراخوانی اسمی بدون اسلش (predict) با ساختار ۵ پارامتری
        {"args": [handle_file(filename), True, True, 4, 1.0], "kwargs": {"api_name": "predict"}},
        # روش ۷: فراخوانی اسمی با اسلش (/predict) با ساختار ۶ پارامتری
        {"args": [handle_file(filename), True, True, True, 4, 1.0], "kwargs": {"api_name": "/predict"}},
        # روش ۸: فراخوانی اسمی با اسلش (/predict) با ساختار ۵ پارامتری
        {"args": [handle_file(filename), True, True, 4, 1.0], "kwargs": {"api_name": "/predict"}},
    ]
    
    for idx, method in enumerate(methods_to_try):
        try:
            print(f'   -> Trying execution path {idx + 1} of {len(methods_to_try)}...')
            upscale_result = cf_client.predict(*method["args"], **method["kwargs"])
            if upscale_result:
                print(f'   -> Success! CodeFormer handled by path {idx + 1}.')
                break
        except Exception as method_err:
            print(f'   -> Path {idx + 1} bypassed: {method_err}')
            
    upscale_path = None
    if isinstance(upscale_result, list) and len(upscale_result) > 0:
        upscale_path = upscale_result[0]
    elif isinstance(upscale_result, tuple) and len(upscale_result) > 0:
        upscale_path = upscale_result[0]
    elif isinstance(upscale_result, dict):
        upscale_path = upscale_result.get('image') or upscale_result.get('path')
    else:
        upscale_path = upscale_result

    if upscale_path and os.path.exists(str(upscale_path)):
        print(f'   -> CodeFormer Success! Output saved at {upscale_path}')
        upscale_ext = 'png'
        if str(upscale_path).endswith('.webp'):
            upscale_ext = 'webp'
        elif str(upscale_path).endswith('.jpg') or str(upscale_path).endswith('.jpeg'):
            upscale_ext = 'jpg'
            
        print(f'   -> Uploading quality-enhanced image as: {run_id}_upscaled')
        with open(upscale_path, 'rb') as f_up:
            res_up = requests.post(
                f'{space_url}/api/webhook/upload', 
                data={'run_id': f"{run_id}_upscaled", 'github_run_id': github_run_id, 'ext': upscale_ext}, 
                files={'file': f_up},
                timeout=60
            )
        if res_up.status_code == 200:
            print('   -> Enhanced image uploaded successfully.')
            cf_success = True
        else:
            print(f'   -> Enhanced image upload failed: {res_up.status_code} - {res_up.text}')
    else:
        print('   -> Warning: CodeFormer did not return a valid file path.')
except Exception as cf_err:
    print(f'   -> Warning: CodeFormer process bypassed or failed: {cf_err}')
# =========================================================

if not cf_success:
    err_str = "CodeFormer upscaling failed or was bypassed due to ZeroGPU/Gradio connection limits."
    print(f'CRITICAL WARNING: {err_str}')
    report_failure(err_str)
    sys.exit(1) # بازگرداندن وضعیت خطا به رانر گیت‌هاب تا فرآیند تلاش مجدد را تا ۲۰ بار تکرار کند

print('6. SUCCESS! Process finished successfully.')