File size: 5,386 Bytes
3173362 33644f7 3173362 f4ade94 3173362 33644f7 3173362 f4ade94 3173362 f4ade94 33644f7 3173362 f4ade94 3173362 df262ae 3173362 f4ade94 3173362 33644f7 3173362 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 | """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})
|