import gradio as gr
import torch
import argparse
import os
import datetime
from diffusers import FluxPipeline
from lib_layerdiffuse.pipeline_flux_img2img import FluxImg2ImgPipeline
from lib_layerdiffuse.vae import TransparentVAE, pad_rgb
import numpy as np
from torchvision import transforms
from safetensors.torch import load_file
from PIL import Image, ImageDraw, ImageFont
import spaces

HF_TOKEN = os.getenv("HF_TOKEN")
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

def seed_everything(seed: int) -> torch.Generator:
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    np.random.seed(seed)
    generator = torch.Generator()
    generator.manual_seed(seed)
    return generator

t2i_pipe = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    torch_dtype=torch.bfloat16,
    use_auth_token=HF_TOKEN
).to(device)

trans_vae = TransparentVAE(t2i_pipe.vae, t2i_pipe.vae.dtype)
trans_vae.load_state_dict(torch.load("./models/TransparentVAE.pth"), strict=False)
trans_vae.to(device)

@spaces.GPU(duration=120)
def t2i_gen(
    prompt: str,
    negative_prompt: str = None,
    seed: int = 1111,
    width: int = 1024,
    height: int = 1024,
    guidance_scale: float = 3.5,
    num_inference_steps: int = 50,
):
    t2i_pipe.load_lora_weights("RedAIGC/Flux-version-LayerDiffuse", weight_name="layerlora.safetensors")
    latents = t2i_pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
        height=height,
        width=width,
        num_inference_steps=num_inference_steps,
        output_type="latent",
        generator=seed_everything(seed),
        guidance_scale=guidance_scale,
    ).images

    latents = t2i_pipe._unpack_latents(latents, height, width, t2i_pipe.vae_scale_factor)
    latents = (latents / t2i_pipe.vae.config.scaling_factor) + t2i_pipe.vae.config.shift_factor

    with torch.no_grad():
        original_x, x = trans_vae.decode(latents)

    x = x.clamp(0, 1)
    x = x.permute(0, 2, 3, 1)
    img = Image.fromarray((x*255).float().cpu().numpy().astype(np.uint8)[0])
    torch.cuda.empty_cache()

    return img