Spaces:
Runtime error
Runtime error
File size: 2,080 Bytes
6e99f64 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 |
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,
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 |