Multi-Model implementation of ImageGen via single prompt using Multi GPUs

mutli-gen.py

import torch
from diffusers import SanaPipeline,StableDiffusion3Pipeline, FluxPipeline,Transformer2DModel, PixArtSigmaPipeline, AuraFlowPipeline, Kandinsky3Pipeline, LuminaText2ImgPipeline, HunyuanDiTPipeline, LuminaText2ImgPipeline
from app.sana_pipeline import SanaPipeline
from torchvision.utils import save_image
from models import SwittiPipeline
from onediffusion.diffusion.pipelines.onediffusion import OneDiffusionPipeline

# FLUX, SD, Aura, Hunay, Sana, Pix Sigma, 

master_prompt = "A serene and emotive scene depicting a college  girl weeping under a large, lush tree, with her loyal dog sitting close by, offering comfort. In the background, a small camp is situated , illuminated by the gentle glow of a campfire around which several people are gathered, sitting on benches and engaging in quiet conversation. The setting is in a forest clearing, during twilight, with the sky painted in soft shades of pink and blue, creating a tranquil yet poignant atmosphere,"

pipe = StableDiffusion3Pipeline.from_pretrained("stabilityai/stable-diffusion-3.5-large", torch_dtype=torch.bfloat16)
pipe = pipe.to("cuda")

image = pipe(
   master_prompt,
    num_inference_steps=28,
    guidance_scale=3.5,
).images[0]
image.save("capybara.png")

# 1. FLUX
pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16)
#pipe.enable_model_cpu_offload() #save some VRAM by offloading the model to CPU. Remove this if you have enough GPU power

seed = 42
image = pipe(
    master_prompt,
    output_type="pil",
    num_inference_steps=30, #use a lower number if you are using [schinel]
    generator=torch.Generator("cpu").manual_seed(seed)
).images[0]
image.save("flux-dev.png")

#========

# 2. PixArt 
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
weight_dtype = torch.float16

pipe = PixArtSigmaPipeline.from_pretrained(
    "PixArt-alpha/PixArt-Sigma-XL-2-1024-MS", 
    torch_dtype=weight_dtype,
    use_safetensors=True,
)
pipe.to(device)

# Enable memory optimizations.
# pipe.enable_model_cpu_offload()

image = pipe(master_prompt).images[0]
image.save("./catcus.png")

#======

# 3. auraflow

pipe = AuraFlowPipeline.from_pretrained("fal/AuraFlow", torch_dtype=torch.float16)
pipe = pipe.to("cuda")
image = pipe(master_prompt).images[0]
image.save("aura_flow.png")

#======

#4. kandisky

pipeline = Kandinsky3Pipeline.from_pretrained("kandinsky-community/kandinsky-3", variant="fp16", torch_dtype=torch.float16)
pipeline.enable_model_cpu_offload()

image = pipeline(master_prompt).images[0]
image.save("kandi.png")

#=====

# 5. Hunay

pipe = HunyuanDiTPipeline.from_pretrained(
    "Tencent-Hunyuan/HunyuanDiT-Diffusers", torch_dtype=torch.float16
)
pipe.to("cuda")

# You may also use English prompt as HunyuanDiT supports both English and Chinese

image = pipe(master_prompt).images[0]
image.save("hunay.png")

#=====

# 6. lumina

pipeline = LuminaText2ImgPipeline.from_pretrained("Alpha-VLLM/Lumina-Next-SFT-diffusers", torch_dtype=torch.bfloat16).to("cuda")

# or you can download the model using code directly
# pipeline = LuminaText2ImgPipeline.from_pretrained("Alpha-VLLM/Lumina-Next-SFT-diffusers", torch_dtype=torch.bfloat16).to("cuda")

image = pipeline(prompt=master_prompt).images[0]

image.save("lumina.png")

#======

# 7. nviida sana : https://github.com/NVlabs/Sana

pipe = SanaPipeline.from_pretrained(
"Efficient-Large-Model/Sana_1600M_1024px_BF16_diffusers", torch_dtype=torch.float32)
pipe.to("cuda")
pipe.text_encoder.to(torch.bfloat16)
pipe.transformer = pipe.transformer.to(torch.bfloat16)

image = pipe(prompt=master_prompt)[0]
image[0].save("sana.png")

#======

# 8. switti : https://github.com/yandex-research/switti

import torch
from models import SwittiPipeline

device = 'cuda:0'
model_path = "yresearch/Switti"

pipe = SwittiPipeline.from_pretrained(model_path, device=device, torch_dtype=torch.bfloat16)

image = pipe(master_prompt,
              cfg=6.0,
              top_k=400,
              top_p=0.95,
              more_smooth=True,
              return_pil=True,
              smooth_start_si=2,
              turn_on_cfg_start_si=0,
              turn_off_cfg_start_si=8,
              last_scale_temp=0.1,
              seed=59,
             )
image.save("switti.png")

# =====

# 9. one diffusion : https://github.com/lehduong/OneDiffusion

device = torch.device('cuda:0')

pipeline = OneDiffusionPipeline.from_pretrained("lehduong/OneDiffusion").to(device=device, dtype=torch.bfloat16)

NEGATIVE_PROMPT = "monochrome, greyscale, low-res, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry, artist name, poorly drawn, bad anatomy, wrong anatomy, extra limb, missing limb, floating limbs, disconnected limbs, mutation, mutated, ugly, disgusting, blurry, amputation"

output = pipeline(
    prompt=f"[[text2image]] {master_prompt}", 
    negative_prompt=NEGATIVE_PROMPT, 
    num_inference_steps=50,
    guidance_scale=4,
    height=1024, 
    width=1024,
)
output.images[0].save('text2image_output.jpg')

parallel.py

import torch
from concurrent.futures import ThreadPoolExecutor
from diffusers import (
    FluxPipeline,
    PixArtSigmaPipeline,
    AuraFlowPipeline,
    Kandinsky3Pipeline,
    HunyuanDiTPipeline,
    LuminaText2ImgPipeline,
    SanaPipeline
)
from app.sana_pipeline import SanaPipeline
from models import SwittiPipeline
from onediffusion.diffusion.pipelines.onediffusion import OneDiffusionPipeline
from torchvision.utils import save_image

# Master prompt for all pipelines
master_prompt = "A serene and emotive scene depicting a college girl weeping under a large, lush tree, with her loyal dog sitting close by, offering comfort. In the background, a small camp is situated, illuminated by the gentle glow of a campfire around which several people are gathered, sitting on benches and engaging in quiet conversation. The setting is in a forest clearing, during twilight, with the sky painted in soft shades of pink and blue, creating a tranquil yet poignant atmosphere."

# 1. FLUX
def run_flux_pipeline(device, seed=42):
    # Load pipeline and specify dtype
    pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16)
    pipe.to(device)  # Assign to GPU
    # Pipe execution with manual seed
    generator = torch.Generator(device=device).manual_seed(seed)
    image = pipe(master_prompt, output_type="pil", num_inference_steps=30, generator=generator).images[0]
    image.save("flux-dev.png")

# 2. PixArt Sigma
def run_pixart_pipeline(device):
    pipe = PixArtSigmaPipeline.from_pretrained(
        "PixArt-alpha/PixArt-Sigma-XL-2-1024-MS", 
        torch_dtype=torch.float16,
        use_safetensors=True,  # Use SafeTensors for security
    )
    pipe.to(device)
    # Enable memory optimizations.
    # pipe.enable_model_cpu_offload()
    image = pipe(master_prompt).images[0]
    image.save("pixart.png")

# 3. AuraFlow
def run_auraflow_pipeline(device):
    pipe = AuraFlowPipeline.from_pretrained("fal/AuraFlow", torch_dtype=torch.float16)
    pipe.to(device)
    image = pipe(master_prompt).images[0]
    image.save("aura_flow.png")

# 4. Kandinsky 3
def run_kandinsky_pipeline(device):
    pipeline = Kandinsky3Pipeline.from_pretrained("kandinsky-community/kandinsky-3", variant="fp16", torch_dtype=torch.float16)
    # Optional: Enable CPU offloading for VRAM optimization
    pipeline.enable_model_cpu_offload()
    pipeline.to(device)
    image = pipeline(master_prompt).images[0]
    image.save("kandi.png")

# 5. Hunyuan DiT
def run_hunyuan_pipeline(device):
    pipe = HunyuanDiTPipeline.from_pretrained("Tencent-Hunyuan/HunyuanDiT-Diffusers", torch_dtype=torch.float16)
    pipe.to(device)
    # You may also use English prompt as HunyuanDiT supports both English and Chinese
    image = pipe(master_prompt).images[0]
    image.save("hunay.png")

# 6. Lumina
def run_lumina_pipeline(device):
    # Load the Lumina model
    pipeline = LuminaText2ImgPipeline.from_pretrained("Alpha-VLLM/Lumina-Next-SFT-diffusers", torch_dtype=torch.bfloat16)
    pipeline.to(device)
    image = pipeline(prompt=master_prompt).images[0]
    image.save("lumina.png")

# 7. Nvidia Sana
def run_sana_pipeline(device):
    pipe = SanaPipeline.from_pretrained(
    "Efficient-Large-Model/Sana_1600M_1024px_BF16_diffusers", torch_dtype=torch.float32)
    pipe.to("cuda")
    pipe.text_encoder.to(torch.bfloat16)
    pipe.transformer = pipe.transformer.to(torch.bfloat16)

    image = pipe(prompt=master_prompt)[0]
    image[0].save("sana.png")


# 8. Switti
def run_switti_pipeline(device):
    # Load Switti pipeline
    model_path = "yresearch/Switti"
    pipe = SwittiPipeline.from_pretrained(model_path, device=device, torch_dtype=torch.bfloat16)
    image = pipe(
        master_prompt,
        cfg=6.0,
        top_k=400,
        top_p=0.95,
        more_smooth=True,
        return_pil=True,
        smooth_start_si=2,
        turn_on_cfg_start_si=0,
        turn_off_cfg_start_si=8,
        last_scale_temp=0.1,
        seed=59,
    )
    image.save("switti.png")

# 9. OneDiffusion
def run_onediffusion_pipeline(device):
    pipe = OneDiffusionPipeline.from_pretrained("lehduong/OneDiffusion").to(device=device, dtype=torch.bfloat16)
    NEGATIVE_PROMPT = "monochrome, greyscale, low-res, bad anatomy, bad hands, text, error, missing fingers, extra digit, cropped"
    output = pipe(
        prompt=f"[[text2image]] {master_prompt}",
        negative_prompt=NEGATIVE_PROMPT,
        num_inference_steps=50,
        guidance_scale=4,
        height=1024,
        width=1024,
    )
    output.images[0].save("onediffusion.png")

# Task-to-GPU Mapping
pipelines = [
    (run_flux_pipeline, "cuda:0"),
    (run_pixart_pipeline, "cuda:1"),
    (run_auraflow_pipeline, "cuda:2"),
    (run_kandinsky_pipeline, "cuda:3"),
    (run_hunyuan_pipeline, "cuda:4"),
    (run_lumina_pipeline, "cuda:5"),
    (run_sana_pipeline, "cuda:6"),
    (run_switti_pipeline, "cuda:7"),
    (run_onediffusion_pipeline, "cuda:0"),  # Reuse cuda:0 for OneDiffusion
]

# Parallel Execution
def main():
    with ThreadPoolExecutor(max_workers=8) as executor:
        futures = [executor.submit(func, device) for func, device in pipelines]
        for future in futures:
            future.result()  # Ensure all tasks complete

if __name__ == "__main__":
    main()

pre_download.py

import os
from huggingface_hub import snapshot_download

# Base directory where models will be stored
BASE_DIR = "./pre_downloaded_models"
os.makedirs(BASE_DIR, exist_ok=True)

# List of models and their Hugging Face repository paths
MODEL_REPOS = {
    "flux": "black-forest-labs/FLUX.1-dev",
    "pixart_sigma": "PixArt-alpha/PixArt-Sigma-XL-2-1024-MS",
    "auraflow": "fal/AuraFlow",
    "kandinsky3": "kandinsky-community/kandinsky-3",
    "hunyuan_dit": "Tencent-Hunyuan/HunyuanDiT-Diffusers",
    "lumina": "Alpha-VLLM/Lumina-Next-SFT-diffusers",
    "sana_checkpoint": "Efficient-Large-Model/Sana_1600M_1024px",  # Hugging Face repo
    "switti": "yresearch/Switti",
    "one_diffusion": "lehduong/OneDiffusion"
}

# Function to download models
def download_model(repo_name, output_dir):
    print(f"Downloading {repo_name}...")
    model_path = snapshot_download(repo_id=repo_name, local_dir=output_dir)
    print(f"Downloaded: {model_path}")
    return model_path

# Pre-download all models
def main():
    model_paths = {}
    for model_name, repo in MODEL_REPOS.items():
        model_dir = os.path.join(BASE_DIR, model_name)
        os.makedirs(model_dir, exist_ok=True)
        # Download the model
        model_path = download_model(repo, model_dir)
        model_paths[model_name] = model_path
    
    print("\nModel download complete! Paths to pre-downloaded models:")
    for name, path in model_paths.items():
        print(f"{name}: {path}")

if __name__ == "__main__":
    main()