Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import torch | |
| import os | |
| import numpy as np | |
| from lib_layerdiffuse.pipeline_flux_img2img import FluxImg2ImgPipeline | |
| from lib_layerdiffuse.vae import TransparentVAE, pad_rgb | |
| from torchvision import transforms | |
| from PIL import Image | |
| 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(device=device) | |
| generator.manual_seed(seed) | |
| return generator | |
| # Initialize the pipeline | |
| i2i_pipe = FluxImg2ImgPipeline.from_pretrained( | |
| "black-forest-labs/FLUX.1-dev", | |
| torch_dtype=torch.bfloat16, | |
| use_auth_token=HF_TOKEN | |
| ).to(device) | |
| # Load the LoRA weights | |
| i2i_pipe.load_lora_weights("RedAIGC/Flux-version-LayerDiffuse", weight_name="layerlora.safetensors") | |
| # Initialize the transparent VAE | |
| trans_vae = TransparentVAE(i2i_pipe.vae, i2i_pipe.vae.dtype) | |
| trans_vae.load_state_dict(torch.load("./models/TransparentVAE.pth"), strict=False) | |
| trans_vae.to(device) | |
| def i2i_gen( | |
| input_image, | |
| prompt: str, | |
| seed: int = 1111, | |
| guidance_scale: float = 7.0, | |
| num_inference_steps: int = 50, | |
| strength: float = 0.8, | |
| ): | |
| if input_image is None: | |
| return None | |
| # Process the input image | |
| original_image = (transforms.ToTensor()(input_image)).unsqueeze(0) | |
| # Get dimensions from the input image | |
| height, width = original_image.shape[2], original_image.shape[3] | |
| # Prepare the image for processing | |
| padding_feed = [x for x in original_image.movedim(1, -1).float().cpu().numpy()] | |
| list_of_np_rgb_padded = [pad_rgb(x) for x in padding_feed] | |
| rgb_padded_bchw_01 = torch.from_numpy(np.stack(list_of_np_rgb_padded, axis=0)).float().movedim(-1, 1).to(device) | |
| original_image_feed = original_image.clone() | |
| original_image_feed[:, :3, :, :] = original_image_feed[:, :3, :, :] * 2.0 - 1.0 | |
| original_image_rgb = original_image_feed[:, :3, :, :] * original_image_feed[:, 3:, :, :] | |
| original_image_feed = original_image_feed.to(device) | |
| original_image_rgb = original_image_rgb.to(device) | |
| # Generate the initial latent | |
| initial_latent = trans_vae.encode(original_image_feed, original_image_rgb, rgb_padded_bchw_01, use_offset=True) | |
| # Generate the image | |
| latents = i2i_pipe( | |
| latents=initial_latent, | |
| image=original_image, | |
| prompt=prompt, | |
| height=height, | |
| width=width, | |
| num_inference_steps=num_inference_steps, | |
| output_type="latent", | |
| generator=seed_everything(seed), | |
| guidance_scale=guidance_scale, | |
| strength=strength, | |
| ).images | |
| # Process the latents | |
| latents = i2i_pipe._unpack_latents(latents, height, width, i2i_pipe.vae_scale_factor) | |
| latents = (latents / i2i_pipe.vae.config.scaling_factor) + i2i_pipe.vae.config.shift_factor | |
| # Decode the latents | |
| with torch.no_grad(): | |
| original_x, x = trans_vae.decode(latents) | |
| # Convert to image | |
| x = x.clamp(0, 1) | |
| x = x.permute(0, 2, 3, 1) | |
| img = Image.fromarray((x*255).float().cpu().numpy().astype(np.uint8)[0]) | |
| # Clean up | |
| torch.cuda.empty_cache() | |
| return img |