Create inference_v42.py

inference_v42.py

@@ -0,0 +1,54 @@
+# TinyFlux Inference - Colab Cell
+# 
+# Install and run:
+#   !pip install git+https://github.com/AbstractPhil/tinyflux.git
+#   # Then run this cell
+
+#@title TinyFlux Inference
+prompt = "a big lion sitting on a couch in a living room, walls, table, ceiling" #@param {type:"string"}
+negative_prompt = "glass, cup" #@param {type:"string"}
+seed = 42 #@param {type:"integer"}
+steps = 50 #@param {type:"slider", min:4, max:50, step:1}
+guidance = 3 #@param {type:"slider", min:1.0, max:15.0, step:0.5}
+
+import torch
+from tinyflux.model.zoo import ModelZoo
+from tinyflux.trainer.sampling import Sampler
+from PIL import Image
+from IPython.display import display
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+dtype = torch.bfloat16 if device == "cuda" else torch.float32
+
+# Load models
+print("Loading models...")
+zoo = ModelZoo(device=device, dtype=dtype)
+zoo.load_vae()
+zoo.load_clip()
+zoo.load_t5()
+model = zoo.load_tinyflux(
+    source="AbstractPhil/tiny-flux-deep",
+    load_ema=False,
+    ema_path="step_407682.safetensors",
+)
+print("✓ Models loaded")
+
+# Create sampler
+sampler = Sampler(
+    zoo=zoo,
+    model=model,
+    num_steps=steps,
+    guidance_scale=guidance,
+    shift=3.0,
+    device=device,
+    dtype=dtype,
+)
+
+# Generate
+print(f"\nGenerating: {prompt}")
+images = sampler.generate([prompt], seed=seed, negative_prompt=negative_prompt)
+
+# Display
+img = (images[0].permute(1, 2, 0).cpu().float().numpy() * 255).astype("uint8")
+display(Image.fromarray(img))
+print(f"Seed: {seed}, Steps: {steps}, CFG: {guidance}")