| """Isolated Qwen Image 2.1 research worker; the 2511 handler stays unchanged.""" |
| import io |
| import json |
| import os |
| import time |
| import traceback |
| import uuid |
| from pathlib import Path |
|
|
| os.environ['HF_HUB_CACHE'] = '/runpod-volume/huggingface-cache/hub' |
| os.environ['HF_HUB_OFFLINE'] = '1' |
| os.environ['TRANSFORMERS_OFFLINE'] = '1' |
|
|
| import requests |
| from PIL import Image |
|
|
| _pipe = None |
| _bucket = None |
|
|
|
|
| def sampling(job): |
| allowed = {'job_id', 'image_url', 'image_urls', 'prompt', 'negative_prompt', |
| 'seed', 'size', 'num_inference_steps', 'true_cfg_scale'} |
| if set(job) - allowed: |
| raise ValueError('Unsupported job fields: ' + ', '.join(sorted(set(job) - allowed))) |
| width, height = map(int, job.get('size', '1024*1536').split('*')) |
| steps = int(job.get('num_inference_steps', 40)) |
| cfg = float(job.get('true_cfg_scale', 1.0)) |
| if not (64 <= width <= 4096 and 64 <= height <= 4096 and |
| width % 16 == height % 16 == 0 and 1 <= steps <= 60 and 1 <= cfg <= 5): |
| raise ValueError('Invalid resolution or sampling parameters') |
| if not job.get('prompt', '').strip(): |
| raise ValueError('An explicit editing prompt is required') |
| return {'prompt': job['prompt'], 'negative_prompt': job.get('negative_prompt'), |
| 'width': width, 'height': height, 'num_inference_steps': steps, |
| 'true_cfg_scale': cfg, 'use_kv_cache': True, |
| 'output_resolution': 1024} |
|
|
|
|
| def cached_model_path(): |
| folder = Path(os.environ['HF_HUB_CACHE']) / ( |
| 'models--' + os.environ['QWEN21_MODEL_REPO'].replace('/', '--')) / ( |
| 'snapshots/' + os.environ['QWEN21_MODEL_REVISION']) |
| if not (folder / 'model_index.json').is_file(): |
| raise FileNotFoundError('Pinned model is missing from RunPod model cache: ' + str(folder)) |
| return str(folder) |
|
|
|
|
| def pipeline(): |
| global _pipe |
| if _pipe is None: |
| import torch |
| from diffusers import QwenImage21Pipeline |
| _pipe = QwenImage21Pipeline.from_pretrained( |
| cached_model_path(), local_files_only=True, |
| torch_dtype=torch.bfloat16) |
| _pipe.enable_model_cpu_offload() |
| return _pipe |
|
|
|
|
| def upload(photo): |
| global _bucket |
| if _bucket is None: |
| from google.cloud import storage |
| from google.oauth2 import service_account |
| raw = os.environ['JSON_PATH_KEY'] |
| try: |
| credentials = service_account.Credentials.from_service_account_info(json.loads(raw)) |
| except json.JSONDecodeError: |
| credentials = service_account.Credentials.from_service_account_file(raw) |
| _bucket = storage.Client(credentials=credentials).bucket(os.environ.get('GCS_BUCKET_NAME', 'shoupapp')) |
| buf = io.BytesIO() |
| photo.save(buf, 'PNG') |
| blob = _bucket.blob('research/qwen21/' + uuid.uuid4().hex + '.png') |
| blob.upload_from_string(buf.getvalue(), content_type='image/png') |
| return blob.public_url |
|
|
|
|
| def handler(event): |
| data = event['input'] |
| if data.get('task') != 'edit' or not 1 <= len(data.get('jobs', [])) <= 4: |
| raise ValueError('Research worker accepts one to four explicit edit jobs') |
| results = [] |
| for job in data['jobs']: |
| try: |
| args = sampling(job) |
| urls = job.get('image_urls') or [job.get('image_url')] |
| if not 1 <= len(urls) <= 10 or not all(urls): |
| raise ValueError('One to ten source image URLs are required') |
| images = [] |
| for url in urls: |
| response = requests.get(url, timeout=60) |
| response.raise_for_status() |
| images.append(Image.open(io.BytesIO(response.content)).convert('RGBA')) |
| import torch |
| pipe = pipeline() |
| torch.cuda.synchronize() |
| torch.cuda.reset_peak_memory_stats() |
| started = time.monotonic() |
| photo = pipe(image=images, generator=torch.Generator('cuda').manual_seed( |
| int(job.get('seed', 20261040))), **args).images[0] |
| torch.cuda.synchronize() |
| inference_ms = round((time.monotonic() - started) * 1000) |
| results.append({'job_id': job.get('job_id'), 'ok': True, |
| 'output_url': upload(photo.convert('RGB')), |
| 'native_alpha_range': photo.getchannel('A').getextrema() if photo.mode == 'RGBA' else None, |
| 'inference_ms': inference_ms, |
| 'denoising_steps': len(pipe.scheduler.timesteps), |
| 'gpu': torch.cuda.get_device_name(), |
| 'peak_reserved_gib': round(torch.cuda.max_memory_reserved() / 1024**3, 2), |
| 'model_revision': os.environ['QWEN21_MODEL_REVISION']}) |
| except Exception as exc: |
| detail = traceback.format_exc() |
| for key in ('HF_TOKEN', 'JSON_PATH_KEY'): |
| if os.environ.get(key): |
| detail = detail.replace(os.environ[key], '[REDACTED]') |
| print('Research job failed:', type(exc).__name__, detail, flush=True) |
| results.append({'job_id': job.get('job_id'), 'ok': False, |
| 'error': type(exc).__name__ + ': edit failed; inspect private worker logs'}) |
| return {'results': results} |
|
|
|
|
| if __name__ == '__main__': |
| import runpod |
| runpod.serverless.start({'handler': handler}) |
|
|
