import gradio as gr
import spaces
import torch
import random
from diffusers import DiffusionPipeline, AutoPipelineForText2Image, FluxPipeline
import os

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

image_model = os.getenv("MODEL_NAME", "Heartsync/NSFW-Uncensored")
hf_token = os.getenv("HF_TOKEN", None)


@spaces.GPU(duration=120)
def generate(prompt, negative_prompt, model=image_model):
    print("Generating image...")
    pipe = None
    if negative_prompt == "" or negative_prompt == None:
        negative_prompt = "ugly, deformed, disfigured, poor quality, low resolution"
    if model == 'enhanceaiteam/Flux-uncensored-v2':
        # Load the base model
        pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16, token=hf_token).to('cuda')
        
        # Load the uncensored LoRA weights
        pipe.load_lora_weights('enhanceaiteam/Flux-uncensored-v2', weight_name='lora.safetensors')

    else:
        pipe = DiffusionPipeline.from_pretrained(
            model,
            torch_dtype=torch.float16
        )
        pipe.to(device)
    print(f"Using model: {image_model}")            
    return pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
        guidance_scale=8.0,
        num_inference_steps=50,
        width=1024,
        height=1024,
    ).images


gr.Interface(
    fn=generate,
    inputs=[
        gr.Text(label="Prompt"),
        gr.Text("", label="Negative Prompt"),
        gr.Dropdown(
            ["Heartsync/NSFW-Uncensored", "UnfilteredAI/NSFW-gen-v2", 'enhanceaiteam/Flux-uncensored-v2'], label="Image model", info="Select the image model:"
        ),
    ],
    outputs=gr.Gallery(),
).launch(show_api=True)