1from pathlib import Path
2from typing import Dict, List, Optional, Tuple
3
4import torch
5from PIL import Image
6from ultraflux.pipeline_flux import FluxPipeline
7from ultraflux.transformer_flux_visionyarn import FluxTransformer2DModel
8from ultraflux.autoencoder_kl import AutoencoderKL
9import os
10
11local_vae = AutoencoderKL.from_pretrained("Owen777/UltraFlux-v1",subfolder="vae", torch_dtype=torch.bfloat16)
12transformer = FluxTransformer2DModel.from_pretrained("Owen777/UltraFlux-v1",subfolder="transformer",torch_dtype=torch.bfloat16)
13# transformer = FluxTransformer2DModel.from_pretrained("Owen777/UltraFlux-v1-1-Transformer",torch_dtype=torch.bfloat16) # NOTE:uncomment this line to use UltraFlux-v1.1
14pipe = FluxPipeline.from_pretrained("Owen777/UltraFlux-v1", vae=local_vae, torch_dtype=torch.bfloat16, transformer=transformer)
15from diffusers import FlowMatchEulerDiscreteScheduler
16pipe.scheduler.config.use_dynamic_shifting = False
17pipe.scheduler.config.time_shift = 4
18pipe = pipe.to("cuda")
19
20os.makedirs("results", exist_ok=True)
21prompts = [
22 "A vast rocky landscape dominated by towering, weathered stone formations, bathed in the ethereal glow of a vibrant night sky filled with a sea of stars, the Milky Way stretching across the heavens, captured from a low angle to emphasize the immense scale of the rocks against the expansive cosmos above. The scene is illuminated by soft, cool moonlight, casting long, dramatic shadows on the textured rock surfaces. The color palette is rich with deep blues, purples, and silvery whites, creating a serene, otherworldly atmosphere.",
23 "A breathtaking scene of snow-capped mountains encircling a serene lake, their towering peaks perfectly mirrored in the still water under a twilight sky adorned with soft, colorful clouds. The gentle mist rises from the lake's surface, adding a mystical touch to the tranquil landscape, with the golden hues of the setting sun casting a warm glow over the scene. The composition captures the vastness of the mountains, the calm water, and the ethereal atmosphere, with a focus on soft, natural lighting and a cool color palette that enhances the peaceful mood.",
24 "A serene snow-covered countryside unfolds in soft morning light, where a small group of woolly sheep graze in the crisp foreground, their breath visible in the cold air, while in the gently blurred midground, a charming, snow-dusted church with a stone steeple rises among rustic timber houses nestled among tall pine trees, all bathed in a pale blue winter palette under an overcast sky, captured with a 50mm lens for a shallow depth of field and a cinematic, filmic warmth.",
25 "A confident woman with short, dark hair stands in a lush forest, her black dress flowing slightly in the breeze. She holds an owl with large, outstretched wings, its feathers sharp and detailed against the dappled sunlight. The light filters softly through the tall trees above, casting intricate shadows on her skin and highlighting her intricate tattoo along her arm. The forest around her feels serene and mystical, with the air thick with the scents of nature. The color palette blends deep greens, earthy browns, and the rich contrast of black and white.",
26]
27
28from pathlib import Path
29
30for idx, prompt in enumerate(prompts, start=1):
31 out_path = Path("results") / f"ultra_flux_{idx:02d}.jpeg"
32 image = pipe(
33 prompt,
34 height=4096,
35 width=4096,
36 guidance_scale=4,
37 num_inference_steps=50,
38 max_sequence_length=512,
39 generator=torch.Generator("cpu").manual_seed(0)
40 ).images[0]
41
42 image.save(f"results/ultra_flux_{idx:02d}.jpeg")
