58 lines
3.8 KiB
Python
58 lines
3.8 KiB
Python
"""One-shot official Qwen PE inference or pinned snapshot download.
|
|
|
|
Runs in Deck's isolated image Python environment. Never retains model weights
|
|
after a request; no prompt or image is written to an application log.
|
|
"""
|
|
import json
|
|
import os
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
def download(data):
|
|
from huggingface_hub import snapshot_download
|
|
snapshot_download(repo_id=data['repo'],revision=data['revision'],local_dir=data['directory'],token=os.environ.get('HF_TOKEN') or None,allow_patterns=['*.safetensors','*.json','*.jinja','system_prompt.txt','LICENSE'],max_workers=2)
|
|
directory=Path(data['directory'])
|
|
required=['model.safetensors.index.json','system_prompt.txt','config.json','tokenizer.json']
|
|
if any(not (directory/name).is_file() for name in required) or len(list(directory.glob('model-*.safetensors')))<4:
|
|
raise RuntimeError('Prompt-Aufwerter-Dateien unvollständig.')
|
|
|
|
def rewrite(data):
|
|
import torch
|
|
from PIL import Image
|
|
from transformers import AutoModelForCausalLM,AutoModelForImageTextToText,AutoProcessor,AutoTokenizer
|
|
directory=data['directory'];task=data['task'];images=data['images']
|
|
system=Path(directory,'system_prompt.txt').read_text().strip()
|
|
max_memory=None
|
|
if data['device']!='cpu':
|
|
if torch.cuda.device_count()!=len(data['gpu_limits']):raise ValueError('Gewählte GPU-Geräte stehen der PE-Laufzeit nicht zur Verfügung.')
|
|
max_memory={i:f'{limit}GiB' for i,limit in enumerate(data['gpu_limits'])}
|
|
max_memory['cpu']='24GiB'
|
|
if task=='t2i':
|
|
tokenizer=AutoTokenizer.from_pretrained(directory,local_files_only=True,trust_remote_code=False)
|
|
model=AutoModelForCausalLM.from_pretrained(directory,dtype=torch.bfloat16,device_map='auto' if data['device']!='cpu' else 'cpu',max_memory=max_memory,local_files_only=True,trust_remote_code=False).eval()
|
|
prompt=tokenizer.apply_chat_template([{'role':'system','content':system},{'role':'user','content':data['prompt']}],tokenize=False,add_generation_prompt=True,enable_thinking=True)
|
|
inputs=tokenizer(prompt,return_tensors='pt').to(model.device)
|
|
elif task=='i2i':
|
|
processor=AutoProcessor.from_pretrained(directory,local_files_only=True,trust_remote_code=False)
|
|
model=AutoModelForImageTextToText.from_pretrained(directory,dtype=torch.bfloat16,device_map='auto' if data['device']!='cpu' else 'cpu',max_memory=max_memory,local_files_only=True,trust_remote_code=False).eval()
|
|
content=[{'type':'image','image':Image.open(path).convert('RGB')} for path in images]+[{'type':'text','text':data['prompt']}]
|
|
inputs=processor.apply_chat_template([{'role':'system','content':[{'type':'text','text':system}]},{'role':'user','content':content}],add_generation_prompt=True,tokenize=True,return_dict=True,return_tensors='pt',enable_thinking=True).to(model.device)
|
|
tokenizer=processor.tokenizer
|
|
else:raise ValueError('Unbekannte Aufwerter-Aufgabe.')
|
|
with torch.inference_mode():
|
|
output=model.generate(**inputs,max_new_tokens=8192,do_sample=True,temperature=1.0,top_p=.95,top_k=20)
|
|
generated=tokenizer.decode(output[0,inputs['input_ids'].shape[1]:],skip_special_tokens=True)
|
|
answer=generated.partition('</think>')[2] or generated
|
|
start=answer.find('{');end=answer.rfind('}')
|
|
if start<0 or end<=start:raise ValueError('Aufwerter lieferte kein JSON-Ergebnis.')
|
|
result=json.loads(answer[start:end+1]);rewritten=result.get('rewritten_prompt')
|
|
if not isinstance(rewritten,str) or not 1<=len(rewritten.strip())<=4000:raise ValueError('Überarbeiteter Prompt ist leer oder länger als 4000 Zeichen.')
|
|
ratio=result.get('wh_ratio','')
|
|
return {'prompt':rewritten.strip(),'wh_ratio':ratio if isinstance(ratio,str) and len(ratio)<=16 else ''}
|
|
|
|
if __name__=='__main__':
|
|
command=sys.argv[1];data=json.load(sys.stdin)
|
|
if command=='download':download(data)
|
|
elif command=='rewrite':print(json.dumps(rewrite(data),ensure_ascii=False),flush=True)
|
|
else:raise ValueError('Unbekannter Aufwerter-Befehl.')
|