inference_sd3.py

import torch
from pipelines.pipeline_sd3_multidiffusion import SD3MultiDiffusionPipeline
from pathlib import Path

#Path + name 
output_path = Path("./outputs")
output_path.mkdir(parents=True, exist_ok=True)
img_name = "res_sd3"


# 1. Define the model ID (Stable Diffusion 3 Medium or 3.5)
model_id = "stabilityai/stable-diffusion-3-medium-diffusers"

print(f"Loading model: {model_id}...")

# 2. Load the custom MultiDiffusion pipeline
pipe = SD3MultiDiffusionPipeline.from_pretrained(
    model_id,
    torch_dtype=torch.float16
)

# Move pipeline to GPU
pipe = pipe.to("cuda")

# Optional: Enable memory optimizations if you have limited VRAM
pipe.enable_model_cpu_offload(0)

# 3. Define your prompt and MultiDiffusion parameters
prompt = (
    "dense tropical rainforest canopy viewed from above, lush green leaves "
    "of various sizes overlapping, dappled sunlight filtering through, "
    "rich saturated greens, seamless botanical pattern, 8k, photorealistic, "
    "no sky, no ground, full coverage"
)
negative_prompt = "sky, dirt, path, blurry, low quality, artifacts, seam, repetition, grid"

# Tile dimensions (the resolution SD3 processes per block)
# 1024x1024 is the native optimal size for SD3
tile_height = 1024
tile_width = 1024

# Target final dimensions (the total size of the output panorama)
# Width is increased to create an ultra-wide shot
target_height = 2048
target_width = 2048 

generator = torch.Generator(device="cuda").manual_seed(42)

print(f"Generating image with target resolution {target_width}x{target_height}...")

# 4. Run Inference
image = pipe(
    prompt=prompt,
    negative_prompt=negative_prompt,
    height=tile_height,         # Block height
    width=tile_width,           # Block width
    total_height=target_height, # Output height
    total_width=target_width,   # Output width
    num_inference_steps=28,
    guidance_scale=7.0,
    generator=generator,
    overlap=128,                # Overlap between blocks to prevent seams
    view_batch_size=1           # Reduce this number if you run out of VRAM (e.g., to 1 or 2)
).images[0]

# 5. Save the output
output_filename = output_path / f"{img_name}.png"
image.save(output_filename)
print(f"Success! Image saved as {output_filename}")