"""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.')