Create utils/t2i.py

utils/t2i.py

@@ -0,0 +1,67 @@
+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=75)
+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,
+        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])
+
+    return img