Qwen-image code loop

gistfile1.txt

from diffusers import DiffusionPipeline, BitsAndBytesConfig, AutoModel
import torch
import time

model_name = "Qwen/Qwen-Image"

torch_dtype = torch.bfloat16
device = "cuda"

pipe = DiffusionPipeline.from_pretrained(model_name,
                                         torch_dtype=torch_dtype)
pipe = pipe.to(device)

positive_magic = {
    "en": "\nUltra HD, 4K, cinematic composition.", # for english prompt,
    "zh": "超清，4K，电影级构图" # for chinese prompt,
}

negative_prompt = " "

aspect_ratios = {
    "1:1": (1328, 1328),
    "16:9": (1664, 928),
    "9:16": (928, 1664),
    "4:3": (1472, 1140),
    "3:4": (1140, 1472),
    "big 1:1":(1536,1536),
    "big 16:10":(1920,1200),
}

i=0

while(True):
    prompt = input("New prompt:")

    print("\nChoose an aspect ratio:")
    aspect_ratio_list = list(aspect_ratios.keys())

    while True:
        for idx, ratio in enumerate(aspect_ratio_list, 1):
            resolution = aspect_ratios[ratio]
            print(f"{idx}. {ratio} ({resolution[0]}x{resolution[1]})")
        try:
            choice = int(input("Enter your choice (1-{}): ".format(len(aspect_ratio_list))))
            if 1 <= choice <= len(aspect_ratio_list):
                selected_ratio = aspect_ratio_list[choice - 1]
                width, height = aspect_ratios[selected_ratio]
                print(f"Selected: {selected_ratio} ({width}x{height})")
                break
            else:
                print(f"Please enter a number between 1 and {len(aspect_ratio_list)}")
        except ValueError:
            print("Please enter a valid number")

    start_time = time.perf_counter()

    image = pipe(
        #prompt=prompt + positive_magic["en"],
        prompt,
        negative_prompt=negative_prompt,
        width=width,
        height=height,
        num_inference_steps=30,
        true_cfg_scale=4.0,
        #generator=torch.Generator(device="cuda").manual_seed(42)
    ).images[0] # type: ignore
    total_time = time.perf_counter() - start_time
    print(f"took {total_time:.2f}s")

    image.save(f"example_{i}.png")
    with open(f"example_{i}.txt", "w", encoding="utf-8") as f:
        f.write(prompt)
    i += 1

If you get errors for any missing python packages at that point just install them

If you rent a GPU on cloud, just skip the venv part (first two lines). On linux use "source venv\bin\activate" instead but if you use linux you probably already know what you're doing anyway.