Create image_generator.py

agents/image_generator.py

@@ -0,0 +1,47 @@
+"""
+Image Generator Agent
+Generates keyframe images using FLUX.1-dev.
+"""
+import os
+import torch
+from diffusers import FluxPipeline
+from PIL import Image
+
+class ImageGeneratorAgent:
+    def __init__(self, gpu_id: int = 5):
+        self.gpu_id = gpu_id
+        self.device = f"cuda:{gpu_id}"
+        self.pipe = None
+        self._load_model()
+
+    def _load_model(self):
+        try:
+            model_id = "black-forest-labs/FLUX.1-dev"
+            self.pipe = FluxPipeline.from_pretrained(
+                model_id,
+                torch_dtype=torch.float8_e4m3fn,
+                variant="fp8",
+                cache_dir="/workspace/.cache/huggingface"
+            ).to(self.device)
+            self.pipe.safety_checker = None
+            print(f"[IMAGE_GENERATOR] FLUX loaded on {self.device}")
+        except Exception as e:
+            print(f"[IMAGE_GENERATOR] Failed to load FLUX: {e}")
+
+    async def generate(self, prompt: str, job_id: str) -> str:
+        if not self.pipe:
+            raise RuntimeError("FLUX pipeline not loaded")
+
+        output_path = f"/workspace/outputs/{job_id}_keyframe.png"
+
+        image = self.pipe(
+            prompt=prompt,
+            height=1024,
+            width=1024,
+            num_inference_steps=30,
+            guidance_scale=3.5,
+            max_sequence_length=512
+        ).images[0]
+
+        image.save(output_path)
+        return output_path