import os 
from gradio_client import Client

token=os.getenv('HF_WRITE_TOKEN')

class Capacitor:
    def __init__(self, client="K00B404/FluxCapacitor" ):
        self.token = token
        self.client = Client(client, hf_token=self.token)
        

    def generate(self, prompt, neg=None, steps=35,cfg_scale=7,sampler="DPM++ 2M Karras",seed=-1,strength=0.7,use_dev=False,enhance_prompt_style="Hello!!",enhance_prompt_option=False,nemo_enhance_prompt_style="Hello!!",use_mistral_nemo=False,):
        img_path,seed,used_prompt = client.predict(
        		prompt=prompt,
        		is_negative="(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" if not neg else neg,
        		steps=steps,
        		cfg_scale=cfg_scale,
        		sampler=sampler,
        		seed=seed,
        		strength=v,
        		huggingface_api_key=self.token,
        		use_dev=use_dev,
        		enhance_prompt_style=enhance_prompt_style,
        		enhance_prompt_option=enhance_prompt_option,
        		nemo_enhance_prompt_style=nemo_enhance_prompt_style,
        		use_mistral_nemo=use_mistral_nemo,
        		api_name="/query"
        )
        print(img_path,seed,used_prompt )
        return img_path,seed,used_prompt