import gradio as gr
import spaces
import random
import numpy as np

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)"],
    nsfw_select=False,
    nsfw_slider=0.45,
    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,
    progress=gr.Progress(track_tqdm=True),
):
    prompt, generator = feifeiprompt(
        randomize_seed,
        seed,
        prompt,
        quality_select,
        styles_Radio,
    )

    #if nsfw_select:
        #pipe.set_adapters(
        #    ["sldr_flux_nsfw_v2"],
        #    adapter_weights=[nsfw_slider],
        #)
        #pipe.fuse_lora(
        #    adapter_name=["sldr_flux_nsfw_v2"],
        #    lora_scale=1.0,
        #)

    image = pipe(
        prompt="flux, 8k, ",
        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:
        image = feifeisharpened(image, num_strength)
        
        
    return image, prompt