VTSTech-SlopGen

VTSTech-SlopGen.py

# -*- coding: utf-8 -*-
import os, sys, warnings, logging, argparse, time, multiprocessing

# --- 1. ENVIRONMENT & PERSISTENT CACHE SETUP ---
# These MUST be set before torch is imported to catch low-level kernels
os.environ["CUDA_VISIBLE_DEVICES"] = "-1" # Force CPU
os.environ["TORCHINDUCTOR_FX_GRAPH_CACHE"] = "1"
os.environ["TORCHINDUCTOR_CACHE_DIR"] = os.path.abspath("./models/pt_inductor_cache")
os.environ["TRITON_CACHE_DIR"] = os.path.abspath("./models/pt_triton_cache")
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"

warnings.filterwarnings("ignore")

import torch
from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler

# Silence loggers
for logger in ["transformers", "diffusers", "huggingface_hub", "optimum"]:
    logging.getLogger(logger).setLevel(logging.CRITICAL)

def expand_prompt(user_prompt):
    enhancements = "highly detailed, masterpiece, sharp focus, professional lighting, 8k resolution, cinematic"
    return f"{user_prompt}, {enhancements}"

def main():
    parser = argparse.ArgumentParser(description="VTSTech SlopGen R5")
    parser.add_argument("--model", type=str, default="segmind/tiny-sd")
    parser.add_argument("--steps", type=int, default=0)
    parser.add_argument("--threads", type=int, default=multiprocessing.cpu_count())
    parser.add_argument("--prompt", type=str, default="Portrait of a cyberpunk woman wearing goggles")
    parser.add_argument("--expand", action="store_true")
    parser.add_argument("--vino", action="store_true")
    parser.add_argument("--onnx", action="store_true")
    parser.add_argument("--gguf", action="store_true")
    parser.add_argument("--save-cache", action="store_true", help="Save PT compile artifacts after generation")
    args = parser.parse_args()
    
    torch.set_num_threads(args.threads)
    inference_steps = args.steps if args.steps > 0 else (4 if "turbo" in args.model.lower() else 15)
    final_prompt = expand_prompt(args.prompt) if args.expand else args.prompt
    safe_name = args.model.replace("/", "--")
    
    # Unified Cache Path for PyTorch Mega-Cache
    pt_cache_file = os.path.abspath(f"./models/{safe_name}_pt_cache.bin")
    os.makedirs("./models", exist_ok=True)

    print(f"--- 🚀 VTSTech SlopGen R5 ---")
    print(f"💻 Model: {args.model}")
    print(f"⏱️  Steps: {inference_steps} | Threads: {args.threads}")

    pipe = None

    # --- BACKEND 1: GGUF ---
    if args.gguf:
        from stable_diffusion_cpp import StableDiffusion
        gguf_dir = f"./models/{safe_name}_gguf"
        gguf_files = [f for f in os.listdir(gguf_dir) if f.endswith(".gguf")] if os.path.exists(gguf_dir) else []
        if not gguf_files:
            print(f"❌ Error: Place a .gguf file in {gguf_dir}")
            return
        pipe = StableDiffusion(model_path=os.path.join(gguf_dir, gguf_files[0]), n_threads=args.threads)
        print(f"📦 Loaded GGUF. Generating...")
        start_time = time.time()
        image = pipe.txt_to_img(prompt=final_prompt, steps=inference_steps)[0]

    # --- BACKEND 2: ONNX ---
    elif args.onnx:
        from optimum.onnxruntime import ORTStableDiffusionPipeline
        onnx_path = f"./models/{safe_name}_onnx"
        if os.path.exists(os.path.join(onnx_path, "model_index.json")):
            pipe = ORTStableDiffusionPipeline.from_pretrained(onnx_path)
        else:
            print(f"⚙️ Exporting to ONNX...")
            pipe = ORTStableDiffusionPipeline.from_pretrained(args.model, export=True)
            pipe.save_pretrained(onnx_path)

    # --- BACKEND 3: OpenVINO ---
    elif args.vino:
        from optimum.intel import OVStableDiffusionPipeline
        vino_path = f"./models/{safe_name}_vino"
        if os.path.exists(os.path.join(vino_path, "model_index.json")):
            pipe = OVStableDiffusionPipeline.from_pretrained(vino_path, device="CPU")
        else:
            print(f"⚙️ Exporting to OpenVINO...")
            pipe = OVStableDiffusionPipeline.from_pretrained(args.model, export=True, device="CPU", task="stable-diffusion")
            pipe.save_pretrained(vino_path)
        pipe.reshape(batch_size=1, height=512, width=512, num_images_per_prompt=1)
        pipe.compile()

    # --- BACKEND 4: PYTORCH (DEFAULT) ---
    else:
        print(f"📦 Mode: PyTorch + Mega-Cache")
        pipe = StableDiffusionPipeline.from_pretrained(args.model, torch_dtype=torch.float32)
        pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
        
        # Optimize
        pipe.unet = torch.compile(pipe.unet)

        # Try to load existing binary cache
        if os.path.exists(pt_cache_file):
            try:
                with open(pt_cache_file, "rb") as f:
                    torch.compiler.load_cache_artifacts(f.read())
                print("🔥 Warm start: Loaded persistent binary cache.")
            except Exception as e:
                print(f"⚠️ Cache Load Failed: {e}")

    # --- SHARED GENERATION (Non-GGUF) ---
    if not args.gguf:
        print("🎨 Generating...")
        start_time = time.time()
        image = pipe(final_prompt, num_inference_steps=inference_steps).images[0]

    # --- POST-GENERATION ACTIONS ---
    duration = time.time() - start_time
    
    # Save Mega-Cache (Only for PyTorch mode if requested)
    if args.save_cache and not (args.vino or args.onnx or args.gguf):
        print(f"💾 Saving compiler artifacts to {pt_cache_file}...")
        artifacts = torch.compiler.save_cache_artifacts()
        if artifacts:
            artifact_bytes, _ = artifacts
            with open(pt_cache_file, "wb") as f:
                f.write(artifact_bytes)

    # Save Image
    os.makedirs("outputs", exist_ok=True)
    out_path = f"outputs/slop_{int(time.time())}.png"
    image.save(out_path)
    print(f"✅ Success! Saved to {out_path} ({duration:.1f}s)")

if __name__ == "__main__":
    main()