Upload 2 files

utils/i2i.py

@@ -0,0 +1,100 @@
+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)
+
+@spaces.GPU(duration=120)
+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 

utils/t2i.py

@@ -0,0 +1,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