generate-image-by-sana-sprint-and-display-in-terminal

play-with-sana.py

from diffusers import SanaSprintPipeline
import torch
import sys
import os
import json
from datetime import datetime

# Set environment variable to avoid tokenizers fork warning
os.environ["TOKENIZERS_PARALLELISM"] = "false"

print("Loading model...")
pipeline = SanaSprintPipeline.from_pretrained(
    #"Efficient-Large-Model/Sana_Sprint_1.6B_1024px_diffusers",
    "Efficient-Large-Model/Sana_Sprint_0.6B_1024px_diffusers",
    torch_dtype=torch.bfloat16
)

# Use DC-AE-Lite for faster speed.
# from diffusers import AutoencoderDC
# vae = AutoencoderDC.from_pretrained("mit-han-lab/dc-ae-lite-f32c32-sana-1.1-diffusers")
# pipeline.vae = vae

pipeline.to("mps")
print("Model loaded successfully!")

image_counter = 1

# Open meta.jsonl file for logging
meta_file = open("meta.jsonl", "a", encoding="utf-8")

try:
    while True:
        try:
            prompt = input("\nEnter prompt (or 'quit'/'exit' to stop): ").strip()

            if not prompt or prompt.lower() in ['quit', 'exit']:
                print("Goodbye!")
                break

            print(f"Generating 2 images for: {prompt}")

            # Generate first image
            print("Generating image 1/2...")
            image1 = pipeline(prompt=prompt, num_inference_steps=2).images[0]

            # Generate second image
            print("Generating image 2/2...")
            image2 = pipeline(prompt=prompt, num_inference_steps=2).images[0]

            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            filename1 = f"generated_{image_counter:03d}a_{timestamp}.png"
            filename2 = f"generated_{image_counter:03d}b_{timestamp}.png"

            image1.save(filename1)
            image2.save(filename2)

            print(f"Images saved as: {filename1}, {filename2}")

            # Log both images to meta.jsonl
            meta_entry1 = {
                "filename": filename1,
                "prompt": prompt,
                "timestamp": timestamp,
                "counter": f"{image_counter}a"
            }
            meta_entry2 = {
                "filename": filename2,
                "prompt": prompt,
                "timestamp": timestamp,
                "counter": f"{image_counter}b"
            }
            meta_file.write(json.dumps(meta_entry1, ensure_ascii=False) + "\n")
            meta_file.write(json.dumps(meta_entry2, ensure_ascii=False) + "\n")
            meta_file.flush()

            # Display images side by side using magick + viu pipeline
            try:
                import subprocess
                magick_proc = subprocess.Popen(
                    ["magick", filename1, filename2, "+append", "png:-"],
                    stdout=subprocess.PIPE
                )
                viu_proc = subprocess.Popen(
                    ["viu", "-"],
                    stdin=magick_proc.stdout
                )
                magick_proc.stdout.close()
                viu_proc.wait()
            except subprocess.CalledProcessError:
                print("Error: Failed to display images with magick+viu pipeline")
            except FileNotFoundError:
                print("Error: magick or viu not found. Install with: brew install imagemagick viu")

            image_counter += 1

        except EOFError:
            print("\nReceived EOF (Ctrl+D). Goodbye!")
            break
        except KeyboardInterrupt:
            print("\nInterrupted. Goodbye!")
            break
        except Exception as e:
            print(f"Error generating image: {e}")
            continue

finally:
    meta_file.close()