Configure video pipeline and text encoder GPUs with memory guidance
This commit is contained in:
+21
-6
@@ -3,10 +3,16 @@ import json
|
||||
import os
|
||||
import sys
|
||||
|
||||
# Keep stdout exclusively for bounded protocol replies, discard library output.
|
||||
protocol=os.fdopen(os.dup(sys.stdout.fileno()),'w',buffering=1)
|
||||
os.dup2(os.open(os.devnull,os.O_WRONLY),sys.stdout.fileno())
|
||||
def reply(value):protocol.write(json.dumps(value)+'\n')
|
||||
|
||||
class EncoderOnDevice:
|
||||
"""Encode on a separate GPU; move all returned conditioning to the diffusion GPU."""
|
||||
def __init__(self,encoder,source,target):self.encoder=encoder;self.source=source;self.target=target
|
||||
def __call__(self,*args,**kwargs):
|
||||
import torch
|
||||
with torch.cuda.device(self.source):outputs=self.encoder(*args,**kwargs)
|
||||
return [type(output)(*(value.to(self.target) if value is not None else None for value in output)) for output in outputs]
|
||||
|
||||
def main():
|
||||
import torch
|
||||
from ltx_pipelines.distilled import DistilledPipeline
|
||||
@@ -17,6 +23,11 @@ def main():
|
||||
config=json.loads(sys.stdin.readline())
|
||||
paths=ModelPaths.from_split(**config['paths'])
|
||||
pipeline=DistilledPipeline(model_paths=paths,spatial_upsampler_path=config['spatial_upsampler'],loras=(),device=torch.device('cuda:0'),offload_mode=OffloadMode.DISK)
|
||||
encoder_device=torch.device(config.get('text_encoder_device','cuda:0'))
|
||||
if encoder_device!=pipeline.device:
|
||||
from ltx_pipelines.utils.blocks import PromptEncoder
|
||||
encoder=PromptEncoder(paths,torch.bfloat16,encoder_device,offload_mode=OffloadMode.DISK)
|
||||
pipeline.prompt_encoder=EncoderOnDevice(encoder,encoder_device,pipeline.device)
|
||||
# Validate all split pack metadata without materializing BF16 weights in RAM.
|
||||
from ltx_core.text_encoders.gemma.gemma_assets import GemmaAssets
|
||||
GemmaAssets.load(config['paths']['text_encoder_path'])
|
||||
@@ -24,7 +35,7 @@ def main():
|
||||
def __init__(self,target,label):self.target=target;self.label=label
|
||||
def __getattr__(self,name):return getattr(self.target,name)
|
||||
def __call__(self,*args,**kwargs):
|
||||
reply({'state':'progress','phase':self.label})
|
||||
reply({'state':'progress','phase':self.label+(' · '+config.get('device_names',{}).get('text_encoder' if self.label.startswith('Textencoder') else 'video',''))})
|
||||
return self.target(*args,**kwargs)
|
||||
for attr,label in [('prompt_encoder','Textencoder · Prompt verarbeiten'),('stage','Videogewichte laden und Diffusion berechnen'),('upsampler','Video hochskalieren'),('video_decoder','Videobilder dekodieren'),('audio_decoder','Audiospur dekodieren')]:setattr(pipeline,attr,Phase(getattr(pipeline,attr),label))
|
||||
reply({'state':'ready'})
|
||||
@@ -37,5 +48,9 @@ def main():
|
||||
reply({'state':'complete'})
|
||||
except torch.OutOfMemoryError:reply({'state':'failed','error':'GPU-Speicher erschöpft (CUDA OOM). Auflösung oder Bildanzahl reduzieren.'});return
|
||||
except Exception:reply({'state':'failed','error':'LTX-Generierung fehlgeschlagen. Modellformat, Laufzeit und Speicher prüfen.'});return
|
||||
try:main()
|
||||
except Exception:reply({'state':'failed','error':'LTX-Profil konnte nicht vorbereitet werden. Komponenten oder Laufzeit inkompatibel.'})
|
||||
if __name__=='__main__':
|
||||
# Keep stdout exclusively for bounded protocol replies, discard library output.
|
||||
protocol=os.fdopen(os.dup(sys.stdout.fileno()),'w',buffering=1)
|
||||
os.dup2(os.open(os.devnull,os.O_WRONLY),sys.stdout.fileno())
|
||||
try:main()
|
||||
except Exception:reply({'state':'failed','error':'LTX-Profil konnte nicht vorbereitet werden. Komponenten oder Laufzeit inkompatibel.'})
|
||||
|
||||
Reference in New Issue
Block a user