last try, perplexity

Dockerfile

@@ -3,6 +3,9 @@ FROM nvidia/cuda:12.5.1-cudnn-devel-ubi9
 
 # Set non-interactive mode
 ENV DEBIAN_FRONTEND=noninteractive
+# Add these environment variables early
+ENV TF_FORCE_GPU_ALLOW_GROWTH=true
+ENV MAX_MEMORY=0.95
 
 # Update and install necessary system packages with microdnf
 RUN microdnf update -y && \
@@ -33,8 +36,12 @@ RUN curl -L https://www.python.org/ftp/python/3.12.9/Python-3.12.9.tgz -o Python
 # Upgrade pip to the latest version
 RUN pip install --upgrade pip
 
-# Install PyTorch with CUDA 12.5 support
-RUN pip install torch torchvision --extra-index-url https://download.pytorch.org/whl/cu125
+# Install xformers dependencies
+RUN microdnf install -y ninja-build && \
+    pip install -U xformers --index-url https://download.pytorch.org/whl/cu125
+
+# Modify PyTorch installation
+RUN pip install torch==2.3.0 torchvision==0.18.0 --index-url https://download.pytorch.org/whl/cu125
 
 # Set the working directory
 WORKDIR /app
@@ -45,6 +52,9 @@ COPY requirements.txt .
 # Install Python dependencies
 RUN pip install -r requirements.txt
 
+# Add memory management tools
+RUN pip install nvidia-ml-py3 pynvml
+
 # Copy the rest of your repository into the container
 COPY . .
 

Image-Morpher/main.py

@@ -10,6 +10,7 @@ from model import DiffMorpherPipeline
 import time
 import logging
 import gc
+from memory_manager import check_vram, limit_precision
 
 os.environ["HF_HOME"] = "/app/hf_cache"
 os.environ["DIFFUSERS_CACHE"] = "/app/hf_cache"
@@ -117,7 +118,13 @@ pipeline = DiffMorpherPipeline.from_pretrained(args.model_path, torch_dtype=torc
 pipeline.enable_vae_slicing()
 pipeline.enable_attention_slicing()
 
-pipeline.to("cuda")
+pipeline.enable_xformers_memory_efficient_attention()
+
+# Replace manual cache clearing with automatic management
+
+from accelerate import init_empty_weights
+with init_empty_weights():
+    pipeline.to("cuda", dtype=torch.float32)
 
 # Add these AFTER device movement
 torch.backends.cudnn.benchmark = True # finds efficient convolution algo by running short benchmark, minimal speed-up.
@@ -135,6 +142,12 @@ if args.use_lcm:
     # set CFG (range allowed by legacy code: 0 to 1, 1 performs best)
     args.guidance_scale = 1
 
+pipeline = limit_precision(pipeline)
+
+# Add memory checks during execution
+check_vram(0.85)  # Check before critical operations
+
+
 # Run the pipeline inference using existing parameters
 images = pipeline(
     img_path_0=args.image_path_0,

memory_manager.py

@@ -0,0 +1,19 @@
+import torch
+import pynvml
+import gc
+
+def check_vram(threshold=0.9):
+    pynvml.nvmlInit()
+    handle = pynvml.nvmlDeviceGetHandleByIndex(0)
+    info = pynvml.nvmlDeviceGetMemoryInfo(handle)
+    used_percent = info.used / info.total
+    if used_percent > threshold:
+        torch.cuda.empty_cache()
+        gc.collect()
+    return used_percent
+
+def limit_precision(model):
+    for param in model.parameters():
+        if param.dtype == torch.float32:
+            param.data = param.data.half()
+    return model

requirements.txt

@@ -10,10 +10,13 @@ Pillow==10.1.0
 safetensors==0.4.0
 tqdm==4.65.0
 transformers==4.34.1
-# torch
-# torchvision
+torch
+torchvision
 lpips
 # peft
 tensorflow==2.18.0
 tensorflow_hub==0.16.1
 opencv_python
+xformers
+nvidia-ml-py3
+pynvml