import gradio as gr
import spaces
import random
import numpy as np
#from extras.expansion import FooocusExpansion
from feifeilib.feifeimodload import feifeimodload
from feifeilib.feifeiprompt import feifeiprompt
from feifeilib.feifeisharpened import feifeisharpened

pipe = feifeimodload()
MAX_SEED = np.iinfo(np.int32).max

@spaces.GPU()
def feifeitexttoimg(
        prompt,
        quality_select=False,
        sharpened_select=False,
        styles_Radio=["(None)"],
        FooocusExpansion_select=False,
        seed=random.randint(0, MAX_SEED),
        randomize_seed=False,
        width=896,
        height=1152,
        num_inference_steps=4,
        guidance_scale=3.5,
        num_strength=0.35,
        num_feifei=0.85,
        progress=gr.Progress(track_tqdm=True),
):
    prompt,generator = feifeiprompt(randomize_seed,seed,prompt,quality_select,styles_Radio,FooocusExpansion_select)

    pipe.set_adapters(
        ["feifei"],
        adapter_weights=[num_feifei],
    )
    
    pipe.fuse_lora(
        adapter_name=["feifei"],
        lora_scale=1.0,
    )
    
    image = pipe(
        prompt="flux,",
        prompt_2=prompt,
        width=width,
        height=height,
        num_inference_steps=num_inference_steps,
        generator=generator,
        guidance_scale=guidance_scale,
        output_type="pil",
    ).images[0]

    if sharpened_select:
        feifeisharpened(image,num_strength)
        
    return image, prompt