chatQwenne / flux_client_wrapper.py
K00B404's picture
Rename flux_client_wrapper .py to flux_client_wrapper.py
c52c5e3 verified
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
)