Layer-streaming inference for large diffusion models — Under 4 GB VRAM, no quantization

Some Benchmarks(For RTX 3050 Having <4GB Vram)

Supported Model List

Image Generation Models

Image Edit Models

Video Generation Models

TODOs

Quick Start

Local Setup

git clone https://github.com/Jit-Roy/weellm.git
cd weellm
pip install -r requirements.txt

Kaggle / Colab Setup

!git clone https://github.com/Jit-Roy/weellm.git
%cd weellm
!pip install -r requirements.txt

Kaggle / Colab Setup (via pip, no git clone needed)

!pip install -q git+https://github.com/Jit-Roy/weellm

Running the CLI

# Run directly from a Hugging Face repository ID!
python main.py \
 --model "black-forest-labs/FLUX.1-dev" \
 --prompt "A majestic lion at golden hour" \
 --negative_prompt "blurry, distorted, low quality" \
 --height 1024 \
 --width 1024 \
 --steps 20 \
 --guidance_scale 3.5 \
 --seed 42 \
 --dtype bfloat16 \
 --vram_budget 4 \
 --ram_budget 4 \
 --output "flux_lion.png" \
 --verbose
# Run Text-to-Video from a Hugging Face repository ID!
python main.py \
 --model "Lightricks/LTX-Video" \
 --prompt "A drone flying over a snowy mountain peak at sunrise." \
 --negative_prompt "worst quality, inconsistent motion, blurry" \
 --height 512 \
 --width 704 \
 --steps 40 \
 --guidance_scale 3.0 \
 --seed 12345 \
 --dtype bfloat16 \
 --vram_budget 4 \
 --ram_budget 4 \
 --output "output_video.mp4" \
 --verbose

⚠️ Cloud vs. Local Execution Warning

Warning

This streaming architecture is specifically optimized for local execution on modern PCs equipped with NVMe SSDs and modern GPUs. You will experience slowdowns on standard cloud instances due to restricted disk I/O speeds.

⚠️ Autoregressive & Vision-Language Model Limitations

Python API

Text to Image

from weellm import WeePipeline
import torch
pipe = WeePipeline.from_pretrained(
 "Tongyi-MAI/Z-Image-Turbo", 
 device="cuda", 
 torch_dtype=torch.bfloat16,
 vram_budget=4, 
 ram_budget=4
)
image = pipe.generate(
 prompt="A serene Japanese zen garden at sunrise",
 negative_prompt="blurry, low quality",
 height=512,
 width=512,
 num_inference_steps=4,
 guidance_scale=1.5,
 seed=42,
)
image.save("output.png")

Text + Image to Image (Image-to-Image)

from weellm import WeePipeline
from PIL import Image
import torch
pipe = WeePipeline.from_pretrained(
 "black-forest-labs/FLUX.1-dev", 
 device="cuda", 
 torch_dtype=torch.bfloat16,
 vram_budget=4, 
 ram_budget=4
)
init_image = Image.open("input.jpg").convert("RGB")
image = pipe.generate(
 prompt="A futuristic cyberpunk city",
 negative_prompt="blurry, distorted",
 image=init_image,
 height=1024,
 width=1024,
 num_inference_steps=20,
 guidance_scale=3.5,
 seed=42,
)
image.save("output_i2i.png")

Text to Video

from weellm import WeePipeline
from weellm.utils import export_to_video
import torch
pipe = WeePipeline.from_pretrained(
 "Lightricks/LTX-Video", 
 device="cuda", 
 torch_dtype=torch.bfloat16,
 vram_budget=4, 
 ram_budget=4
)
video_frames = pipe.generate(
 prompt="A drone flying over a snowy mountain peak at sunrise.",
 negative_prompt="worst quality, inconsistent motion",
 height=512,
 width=704,
 num_inference_steps=40,
 guidance_scale=3.0,
 seed=42,
)
export_to_video(video_frames, "output_video.mp4", fps=24)

Text + Image to Video

from weellm import WeePipeline
from weellm.utils import export_to_video
from PIL import Image
import torch
pipe = WeePipeline.from_pretrained(
 "Lightricks/LTX-Video", 
 device="cuda", 
 torch_dtype=torch.bfloat16,
 vram_budget=4, 
 ram_budget=4
)
start_image = Image.open("start_frame.jpg").convert("RGB")
video_frames = pipe.generate(
 prompt="The camera pans slowly across the room.",
 negative_prompt="worst quality, inconsistent motion",
 image=start_image,
 height=512,
 width=704,
 num_inference_steps=40,
 guidance_scale=3.0,
 seed=42,
)
export_to_video(video_frames, "output_i2v.mp4", fps=24)

Using GGUF Models

from weellm import WeePipeline
import torch
pipe = WeePipeline.from_pretrained(
 "black-forest-labs/FLUX.2-klein-4B", 
 transformer_path="unsloth/FLUX.2-klein-4B-GGUF/flux-2-klein-4b-Q4_K_M.gguf",
 text_encoder_path="unsloth/Qwen3-4B-GGUF/Qwen3-4B-Q5_K_M.gguf",
 # If using SD/SDXL models, use `unet_path` instead of `transformer_path`
 # unet_path="path/to/juggernaut-xl-v9-Q8_0.gguf",
 device="cuda", 
 torch_dtype=torch.bfloat16,
 vram_budget=4, 
 ram_budget=4,
 prefetch=False # Recommended when using GGUF to avoid GPU contention
)
image = pipe.generate(
 prompt="A majestic lion at golden hour",
 height=1024,
 width=1024,
 num_inference_steps=4,
)
image.save("flux_klein_gguf.png")

Star History