from gradio_client import Client

class FluxImageGenerator:
    def __init__(
        self,
        space_name="K00B404/FluxCapacitor2",
        api_key=None,
        use_dev=False,
        default_sampler="DPM++ 2M Karras",
        default_steps=35,
        default_cfg=7,
        default_strength=0.7,
        default_negative_prompt=(
            "(deformed, distorted, disfigured), poorly drawn, bad anatomy, "
            "wrong anatomy, extra limb, missing limb, floating limbs, "
            "(mutated hands and fingers), disconnected limbs, mutation, mutated, "
            "ugly, disgusting, blurry, amputation, misspellings, typos"
        ),
        api_name="/query"
    ):
        self.client = Client(space_name)
        self.api_key = api_key
        self.use_dev = use_dev
        self.default_sampler = default_sampler
        self.default_steps = default_steps
        self.default_cfg = default_cfg
        self.default_strength = default_strength
        self.default_negative_prompt = default_negative_prompt
        self.api_name = api_name

    def generate_image(
        self,
        prompt: str,
        steps: int = None,
        cfg_scale: float = None,
        sampler: str = None,
        seed: int = -1,
        strength: float = None,
        enhance_prompt_option: bool = False,
        use_mistral_nemo: bool = False,
    ):
        prompt = prompt.strip()
        steps = steps or self.default_steps
        cfg_scale = cfg_scale or self.default_cfg
        sampler = sampler or self.default_sampler
        strength = strength or self.default_strength

        return self.client.predict(
            prompt=prompt,
            is_negative=self.default_negative_prompt,
            steps=steps,
            cfg_scale=cfg_scale,
            sampler=sampler,
            seed=seed,
            strength=strength,
            huggingface_api_key=self.api_key,
            use_dev=self.use_dev,
            enhance_prompt_style=prompt,
            enhance_prompt_option=enhance_prompt_option,
            nemo_enhance_prompt_style=prompt,
            use_mistral_nemo=use_mistral_nemo,
            api_name=self.api_name
        )