Create inference.py

inference.py

@@ -0,0 +1,372 @@
+# ============================================================================
+# TinyFlux Inference Cell - Euler Discrete Flow Matching
+# ============================================================================
+# Run the model cell before this one (defines TinyFlux, TinyFluxConfig)
+# Loads from: AbstractPhil/tiny-flux or local checkpoint
+# ============================================================================
+
+import torch
+from huggingface_hub import hf_hub_download
+from safetensors.torch import load_file
+from transformers import T5EncoderModel, T5Tokenizer, CLIPTextModel, CLIPTokenizer
+from diffusers import AutoencoderKL
+from PIL import Image
+import numpy as np
+import os
+
+# ============================================================================
+# CONFIG
+# ============================================================================
+DEVICE = "cuda"
+DTYPE = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
+
+# Model loading
+HF_REPO = "AbstractPhil/tiny-flux"
+LOAD_FROM = "hub"  # "hub", "hub:step_1000", "local:/path/to/weights.safetensors"
+
+# Generation settings
+NUM_STEPS = 20          # Euler steps (20-50 typical)
+GUIDANCE_SCALE = 3.5    # CFG scale (1.0 = no guidance, 3-7 typical)
+HEIGHT = 512            # Output height
+WIDTH = 512             # Output width
+SEED = None             # None for random
+
+# ============================================================================
+# LOAD TEXT ENCODERS
+# ============================================================================
+print("Loading text encoders...")
+t5_tok = T5Tokenizer.from_pretrained("google/flan-t5-base")
+t5_enc = T5EncoderModel.from_pretrained("google/flan-t5-base", torch_dtype=DTYPE).to(DEVICE).eval()
+
+clip_tok = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14")
+clip_enc = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14", torch_dtype=DTYPE).to(DEVICE).eval()
+
+# ============================================================================
+# LOAD VAE
+# ============================================================================
+print("Loading Flux VAE...")
+vae = AutoencoderKL.from_pretrained(
+    "black-forest-labs/FLUX.1-schnell",
+    subfolder="vae",
+    torch_dtype=DTYPE
+).to(DEVICE).eval()
+
+# ============================================================================
+# LOAD TINYFLUX MODEL
+# ============================================================================
+print(f"Loading TinyFlux from: {LOAD_FROM}")
+
+config = TinyFluxConfig()
+model = TinyFlux(config).to(DEVICE).to(DTYPE)
+
+if LOAD_FROM == "hub":
+    # Load best model from hub
+    weights_path = hf_hub_download(repo_id=HF_REPO, filename="model.safetensors")
+    weights = load_file(weights_path)
+    model.load_state_dict(weights)
+    print(f"✓ Loaded from {HF_REPO}/model.safetensors")
+elif LOAD_FROM.startswith("hub:"):
+    # Load specific checkpoint from hub
+    ckpt_name = LOAD_FROM[4:]
+    if not ckpt_name.endswith(".safetensors"):
+        ckpt_name = f"checkpoints/{ckpt_name}.safetensors"
+    weights_path = hf_hub_download(repo_id=HF_REPO, filename=ckpt_name)
+    weights = load_file(weights_path)
+    model.load_state_dict(weights)
+    print(f"✓ Loaded from {HF_REPO}/{ckpt_name}")
+elif LOAD_FROM.startswith("local:"):
+    # Load local file
+    weights_path = LOAD_FROM[6:]
+    weights = load_file(weights_path)
+    model.load_state_dict(weights)
+    print(f"✓ Loaded from {weights_path}")
+else:
+    raise ValueError(f"Unknown LOAD_FROM: {LOAD_FROM}")
+
+model.eval()
+print(f"Model params: {sum(p.numel() for p in model.parameters()):,}")
+
+# ============================================================================
+# ENCODING FUNCTIONS
+# ============================================================================
+@torch.no_grad()
+def encode_prompt(prompt: str, max_length: int = 128):
+    """Encode prompt with flan-t5-base and CLIP-L."""
+    # T5 encoding (sequence)
+    t5_in = t5_tok(
+        prompt, 
+        max_length=max_length, 
+        padding="max_length", 
+        truncation=True, 
+        return_tensors="pt"
+    ).to(DEVICE)
+    t5_out = t5_enc(
+        input_ids=t5_in.input_ids, 
+        attention_mask=t5_in.attention_mask
+    ).last_hidden_state  # (1, L, 768)
+    
+    # CLIP encoding (pooled)
+    clip_in = clip_tok(
+        prompt, 
+        max_length=77, 
+        padding="max_length", 
+        truncation=True, 
+        return_tensors="pt"
+    ).to(DEVICE)
+    clip_out = clip_enc(
+        input_ids=clip_in.input_ids, 
+        attention_mask=clip_in.attention_mask
+    )
+    clip_pooled = clip_out.pooler_output  # (1, 768)
+    
+    return t5_out, clip_pooled
+
+# ============================================================================
+# EULER DISCRETE FLOW MATCHING SAMPLER
+# ============================================================================
+@torch.no_grad()
+def euler_sample(
+    model,
+    prompt: str,
+    negative_prompt: str = "",
+    num_steps: int = 20,
+    guidance_scale: float = 3.5,
+    height: int = 512,
+    width: int = 512,
+    seed: int = None,
+):
+    """
+    Euler discrete sampler for flow matching.
+    
+    Flow matching formulation:
+        x_t = (1 - t) * x_0 + t * x_1
+        where x_0 = noise, x_1 = data
+        velocity v = x_1 - x_0 = data - noise
+        
+    Sampling (t: 0 -> 1, noise -> data):
+        x_{t+dt} = x_t + v_pred * dt
+    
+    With Flux shift for improved sampling distribution.
+    """
+    # Set seed
+    if seed is not None:
+        torch.manual_seed(seed)
+        generator = torch.Generator(device=DEVICE).manual_seed(seed)
+    else:
+        generator = None
+    
+    # Latent dimensions (VAE downscales by 8)
+    H_lat = height // 8  # 64 for 512
+    W_lat = width // 8   # 64 for 512
+    C_lat = 16           # Flux VAE channels
+    
+    # Encode prompts
+    t5_cond, clip_cond = encode_prompt(prompt)
+    if guidance_scale > 1.0 and negative_prompt is not None:
+        t5_uncond, clip_uncond = encode_prompt(negative_prompt)
+    else:
+        t5_uncond, clip_uncond = None, None
+    
+    # Start from pure noise (t=0 in flow matching convention)
+    # Shape: (1, H*W, C)
+    x = torch.randn(1, H_lat * W_lat, C_lat, device=DEVICE, dtype=DTYPE, generator=generator)
+    
+    # Create image position IDs for RoPE
+    img_ids = TinyFlux.create_img_ids(1, H_lat, W_lat, DEVICE)
+    
+    # Timesteps: 0 -> 1 (noise -> data)
+    # We use uniform spacing, model handles flux shift internally for training
+    # For inference, linear timesteps work well
+    timesteps = torch.linspace(0, 1, num_steps + 1, device=DEVICE, dtype=DTYPE)
+    
+    print(f"Sampling with {num_steps} Euler steps...")
+    
+    for i in range(num_steps):
+        t_curr = timesteps[i]
+        t_next = timesteps[i + 1]
+        dt = t_next - t_curr
+        
+        t_batch = t_curr.unsqueeze(0)  # (1,)
+        
+        # Guidance embedding (used during training with random values 1-5)
+        guidance_embed = torch.tensor([guidance_scale], device=DEVICE, dtype=DTYPE)
+        
+        # Conditional prediction
+        v_cond = model(
+            hidden_states=x,
+            encoder_hidden_states=t5_cond,
+            pooled_projections=clip_cond,
+            timestep=t_batch,
+            img_ids=img_ids,
+            guidance=guidance_embed,
+        )
+        
+        # Classifier-free guidance
+        if guidance_scale > 1.0 and t5_uncond is not None:
+            v_uncond = model(
+                hidden_states=x,
+                encoder_hidden_states=t5_uncond,
+                pooled_projections=clip_uncond,
+                timestep=t_batch,
+                img_ids=img_ids,
+                guidance=guidance_embed,
+            )
+            v = v_uncond + guidance_scale * (v_cond - v_uncond)
+        else:
+            v = v_cond
+        
+        # Euler step: x_{t+dt} = x_t + v * dt
+        x = x + v * dt
+        
+        if (i + 1) % 5 == 0 or i == num_steps - 1:
+            print(f"  Step {i+1}/{num_steps}, t={t_next.item():.3f}")
+    
+    # Reshape to image format: (1, H*W, C) -> (1, C, H, W)
+    latents = x.reshape(1, H_lat, W_lat, C_lat).permute(0, 3, 1, 2)
+    
+    return latents
+
+# ============================================================================
+# DECODE LATENTS TO IMAGE
+# ============================================================================
+@torch.no_grad()
+def decode_latents(latents):
+    """Decode VAE latents to PIL Image."""
+    # Flux VAE scaling
+    latents = latents / vae.config.scaling_factor
+    
+    # Decode
+    image = vae.decode(latents.float()).sample
+    
+    # Normalize to [0, 1]
+    image = (image / 2 + 0.5).clamp(0, 1)
+    
+    # To PIL
+    image = image[0].permute(1, 2, 0).cpu().numpy()
+    image = (image * 255).astype(np.uint8)
+    
+    return Image.fromarray(image)
+
+# ============================================================================
+# MAIN GENERATION FUNCTION
+# ============================================================================
+def generate(
+    prompt: str,
+    negative_prompt: str = "",
+    num_steps: int = NUM_STEPS,
+    guidance_scale: float = GUIDANCE_SCALE,
+    height: int = HEIGHT,
+    width: int = WIDTH,
+    seed: int = SEED,
+    save_path: str = None,
+):
+    """
+    Generate an image from a text prompt.
+    
+    Args:
+        prompt: Text description of desired image
+        negative_prompt: What to avoid (empty string for none)
+        num_steps: Number of Euler steps (20-50)
+        guidance_scale: CFG scale (1.0=none, 3-7 typical)
+        height: Output height in pixels (must be divisible by 8)
+        width: Output width in pixels (must be divisible by 8)
+        seed: Random seed (None for random)
+        save_path: Path to save image (None to skip saving)
+    
+    Returns:
+        PIL.Image
+    """
+    print(f"\nGenerating: '{prompt}'")
+    print(f"Settings: {num_steps} steps, cfg={guidance_scale}, {width}x{height}, seed={seed}")
+    
+    # Sample latents
+    latents = euler_sample(
+        model=model,
+        prompt=prompt,
+        negative_prompt=negative_prompt,
+        num_steps=num_steps,
+        guidance_scale=guidance_scale,
+        height=height,
+        width=width,
+        seed=seed,
+    )
+    
+    # Decode to image
+    print("Decoding latents...")
+    image = decode_latents(latents)
+    
+    # Save if requested
+    if save_path:
+        image.save(save_path)
+        print(f"✓ Saved to {save_path}")
+    
+    print("✓ Done!")
+    return image
+
+# ============================================================================
+# BATCH GENERATION
+# ============================================================================
+def generate_batch(
+    prompts: list,
+    negative_prompt: str = "",
+    num_steps: int = NUM_STEPS,
+    guidance_scale: float = GUIDANCE_SCALE,
+    height: int = HEIGHT,
+    width: int = WIDTH,
+    seed: int = SEED,
+    output_dir: str = "./outputs",
+):
+    """Generate multiple images."""
+    os.makedirs(output_dir, exist_ok=True)
+    images = []
+    
+    for i, prompt in enumerate(prompts):
+        # Increment seed for variety if seed is set
+        img_seed = seed + i if seed is not None else None
+        
+        image = generate(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            num_steps=num_steps,
+            guidance_scale=guidance_scale,
+            height=height,
+            width=width,
+            seed=img_seed,
+            save_path=os.path.join(output_dir, f"{i:03d}.png"),
+        )
+        images.append(image)
+    
+    return images
+
+# ============================================================================
+# QUICK TEST
+# ============================================================================
+if __name__ == "__main__" or True:  # Always run in Colab
+    print("\n" + "="*60)
+    print("TinyFlux Inference Ready!")
+    print("="*60)
+    print(f"""
+Usage:
+    # Single image
+    image = generate("a photo of a cat")
+    image.show()
+    
+    # With options
+    image = generate(
+        prompt="a beautiful sunset over mountains",
+        negative_prompt="blurry, low quality",
+        num_steps=30,
+        guidance_scale=4.0,
+        height=512,
+        width=512,
+        seed=42,
+        save_path="output.png"
+    )
+    
+    # Batch generation
+    images = generate_batch([
+        "a red sports car",
+        "a blue ocean wave", 
+        "a green forest path",
+    ], output_dir="./my_outputs")
+""")