agents/image_generator.py

"""
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