Spaces:
Running on Zero
Running on Zero
| import spaces | |
| 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 | |
| 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) | |
| 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() | |
| img.save("/tmp/last_overlay.png") | |
| return img |