Spaces:
Running
Running
| """Test NVIDIA NIM FLUX.1 Kontext-dev for consistent frame variant generation.""" | |
| import asyncio | |
| import base64 | |
| import os | |
| import time | |
| from pathlib import Path | |
| import httpx | |
| from dotenv import load_dotenv | |
| load_dotenv() | |
| NVIDIA_API_KEY = os.getenv("NVIDIA_API_KEY") | |
| KONTEXT_URL = "https://ai.api.nvidia.com/v1/genai/black-forest-labs/flux.1-kontext-dev" | |
| # Try to find an existing generated image | |
| WORKDIR = Path("workdir/projects/2/images") | |
| async def generate_test_keyframe() -> bytes: | |
| """Generate a fresh manhwa-style keyframe via HF FLUX.1-schnell.""" | |
| hf_token = os.getenv("HF_TOKEN") | |
| prompt = ( | |
| "manhwa style, korean webtoon art, sharp lineart, cel-shaded, " | |
| "close-up shot, young adult male, jet black messy hair, deep crimson eyes, " | |
| "black high-collar coat, dramatic lighting, vibrant colors, detailed background" | |
| ) | |
| async with httpx.AsyncClient(timeout=30) as client: | |
| resp = await client.post( | |
| "https://huggingface.co/proxy/api-inference.huggingface.co/models/black-forest-labs/FLUX.1-schnell", | |
| headers={"Authorization": f"Bearer {hf_token}"}, | |
| json={"inputs": prompt}, | |
| ) | |
| resp.raise_for_status() | |
| return resp.content | |
| async def kontext_edit(image_b64: str, prompt: str) -> bytes: | |
| """Send an image to Kontext for editing.""" | |
| async with httpx.AsyncClient(timeout=60) as client: | |
| resp = await client.post( | |
| KONTEXT_URL, | |
| headers={ | |
| "Authorization": f"Bearer {NVIDIA_API_KEY}", | |
| "Content-Type": "application/json", | |
| "Accept": "application/json", | |
| }, | |
| json={ | |
| "image": image_b64, | |
| "prompt": prompt, | |
| "steps": 25, | |
| "cfg_scale": 3.5, | |
| }, | |
| ) | |
| resp.raise_for_status() | |
| data = resp.json() | |
| # Response format: {"artifacts": [{"base64": "..."}]} or {"data": [{"b64_json": "..."}]} | |
| if "artifacts" in data: | |
| return base64.b64decode(data["artifacts"][0]["base64"]) | |
| elif "data" in data: | |
| return base64.b64decode(data["data"][0]["b64_json"]) | |
| else: | |
| # Try direct b64 field | |
| for key in ("b64_json", "base64", "image"): | |
| if key in data: | |
| return base64.b64decode(data[key]) | |
| raise ValueError(f"Unknown response format: {list(data.keys())}") | |
| async def main(): | |
| if not NVIDIA_API_KEY: | |
| print("ERROR: NVIDIA_API_KEY not set in .env") | |
| return | |
| output_dir = Path("test_kontext_output") | |
| output_dir.mkdir(exist_ok=True) | |
| # Step 1: Get a keyframe image | |
| print("Step 1: Getting keyframe image...") | |
| keyframe_path = None | |
| # Try existing images first | |
| if WORKDIR.exists(): | |
| images = sorted(WORKDIR.glob("*.png")) | |
| if images: | |
| keyframe_path = images[0] | |
| print(f" Using existing image: {keyframe_path}") | |
| if keyframe_path: | |
| with open(keyframe_path, "rb") as f: | |
| keyframe_bytes = f.read() | |
| else: | |
| print(" Generating fresh keyframe via HF FLUX.1-schnell...") | |
| keyframe_bytes = await generate_test_keyframe() | |
| keyframe_path = output_dir / "00_keyframe.png" | |
| with open(keyframe_path, "wb") as f: | |
| f.write(keyframe_bytes) | |
| print(f" Saved keyframe: {keyframe_path}") | |
| # Save original | |
| orig_path = output_dir / "00_original.png" | |
| with open(orig_path, "wb") as f: | |
| f.write(keyframe_bytes) | |
| keyframe_b64 = base64.b64encode(keyframe_bytes).decode() | |
| # Step 2: Generate 3 variants | |
| edit_prompts = [ | |
| ("01_smile", "same scene, character now has a slight smile"), | |
| ("02_hand_raise", "same scene, character raises their right hand"), | |
| ("03_wind", "same scene, wind blowing through hair, clothes rippling"), | |
| ] | |
| for filename, prompt in edit_prompts: | |
| print(f"\nStep 2: Kontext edit - '{prompt}'...") | |
| t0 = time.time() | |
| try: | |
| variant_bytes = await kontext_edit(keyframe_b64, prompt) | |
| elapsed = time.time() - t0 | |
| out_path = output_dir / f"{filename}.png" | |
| with open(out_path, "wb") as f: | |
| f.write(variant_bytes) | |
| print(f" OK in {elapsed:.1f}s -> {out_path}") | |
| except Exception as e: | |
| elapsed = time.time() - t0 | |
| print(f" FAILED in {elapsed:.1f}s: {e}") | |
| print(f"\nDone! Check {output_dir}/ for results.") | |
| print("Visual check: Are faces/proportions/colors consistent enough for RIFE interpolation?") | |
| if __name__ == "__main__": | |
| asyncio.run(main()) | |