from diffusers import DiffusionPipeline

pipe_id = "stabilityai/stable-diffusion-xl-base-1.0"
pipe = DiffusionPipeline.from_pretrained(pipe_id, torch_dtype=torch.float16).to("cuda")

prompt = "a blue hair gril"
image = pipe(prompt, num_inference_steps=45, guidance_scale=7.5, height=1024, width=1024).images[0]

image.save("output.jpg")

from diffusers import DiffusionPipeline

pipe_id = "stabilityai/stable-diffusion-xl-base-1.0"
pipe = DiffusionPipeline.from_pretrained(pipe_id, torch_dtype=torch.float16).to("cuda")
pipe.load_lora_weights("sd-gbf-lora", weight_name="default_0")
pipeline.set_adapters("default_0", 1.0) 

prompt = "a blue hair gril"
lora_scale = 0.9
image = pipe(prompt, num_inference_steps=45, guidance_scale=7.5, cross_attention_kwargs={"scale": lora_scale}, height=1024, width=1024).images[0]

image.save("output.jpg")

import torch
from PIL import Image
from diffusers import StableDiffusionXLImg2ImgPipeline

pipe_id = "stabilityai/stable-diffusion-xl-base-1.0"
pipe = StableDiffusionXLImg2ImgPipeline.from_pretrained(pipe_id, torch_dtype=torch.float16).to("cuda")
pipe.load_lora_weights("sd-gbf-lora")

# 加载输入图像
input_image_path = "examples/lubi.jpg"  
input_image = Image.open(input_image_path).convert("RGB")

prompt = "gbfhero"  
negative_prompt = "low quality, bad quality"  

with torch.no_grad():
    output_image = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
        guidance_scale=7.5,
        cross_attention_kwargs={"scale": 0.9}, 
        height=1024, width=1024,
        image=input_image,
        strength=0.5  # 控制生成图像与输入图像的相似程度，范围为0到1
    ).images[0]

output_image.save(f"outputs/1.jpg")

import os
import cv2
import torch
import numpy as np
from PIL import Image
from diffusers import StableDiffusionXLControlNetImg2ImgPipeline, ControlNetModel

output_folder = "outputs"
os.makedirs(output_folder, exist_ok=True)

# 加载模型和Controlnet
pipe_id = "stabilityai/stable-diffusion-xl-base-1.0"
controlnet_id = "diffusers/controlnet-canny-sdxl-1.0"
controlnet = ControlNetModel.from_pretrained(controlnet_id, torch_dtype=torch.float16)
pipe = StableDiffusionXLControlNetImg2ImgPipeline.from_pretrained(pipe_id, controlnet=controlnet, torch_dtype=torch.float16).to("cuda")
pipe.load_lora_weights("sd-gbf-lora3")

# 加载图片
input_image_path = "examples/leishen.jpeg"  
input_image = Image.open(input_image_path).convert("RGB")
np_image = np.array(input_image)
# 生成 edges
np_image = cv2.Canny(np_image, 100, 200)
np_image = np_image[:, :, None]
np_image = np.concatenate([np_image, np_image, np_image], axis=2)
canny_image = Image.fromarray(np_image)
canny_image.save(f'{output_folder}/tmp_edge.png')

prompt = "gbfhero, clean background"
negative_prompt = "low quality, bad quality"
lora_scale = 0.9
image = pipe(prompt, 
    negative_prompt=negative_prompt,
    guidance_scale=7.5,
    cross_attention_kwargs={"scale": lora_scale}, 
    controlnet_conditioning_scale=0.5,
    image=input_image,
    strength=0.9,
    control_image=canny_image,
    height=1024, width=1024).images[0]

image.save(f"{output_folder}/5.jpg")

from sd_embed.embedding_funcs import get_weighted_text_embeddings_sdxl

prompt = "..."
negative_prompt = "..."
seed = 481167465

pipeline = StableDiffusionXLPipeline.from_single_file(
    model_path,
    torch_dtype=torch.float16,
    use_safetensors=True,
    variant="fp16",
    scheduler = scheduler
).to("cuda")

( 
  prompt_embeds,
  prompt_neg_embeds,
  pooled_prompt_embeds,
  negative_pooled_prompt_embeds
) = get_weighted_text_embeddings_sdxl(
    pipeline,
    prompt=prompt,
    neg_prompt=negative_prompt
)

# 生成图像
image = pipeline(
    prompt_embeds=prompt_embeds,
    negative_prompt_embeds=prompt_neg_embeds,
    pooled_prompt_embeds=pooled_prompt_embeds,
    negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
    num_inference_steps=28,
    guidance_scale=5.0,
    generator=torch.Generator(device=pipeline.device).manual_seed(seed),
    height=1216,
    width=832
).images[0]

import torch
from diffusers import AutoencoderKL
from PIL import Image
import numpy as np

vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="vae", torch_dtype=torch.bfloat16)
device = "cuda" if torch.cuda.is_available() else "cpu"
vae = vae.to(device)

image_path = "0.png"
image = Image.open(image_path).convert("RGB")

# Resize to a size divisible by 16 (e.g., 1024x1024 for Flux)
image = image.resize((1024, 1024))

# Preprocess image to tensor in [-1, 1]
img_array = np.array(image).astype(np.float32) / 255.0
img_array = img_array * 2.0 - 1.0
img_tensor = torch.from_numpy(img_array).permute(2, 0, 1).unsqueeze(0)  # Shape: (1, 3, H, W)
img_tensor = img_tensor.to(device, dtype=vae.dtype)

# Encode the image to latents
with torch.no_grad():
    latent_dist = vae.encode(img_tensor)
    # Use mode (mean) for better reconstruction
    latents = latent_dist.latent_dist.mode()
    # Apply scale and shift (specific to Flux VAE)
    latents = latents * vae.config.scaling_factor + vae.config.shift_factor

# Decode the latents back to image
with torch.no_grad():
    # Reverse shift and scale
    latents_for_decode = latents / vae.config.scaling_factor + vae.config.shift_factor
    decoded = vae.decode(latents_for_decode).sample

# Postprocess decoded tensor to [0, 1] and convert to PIL image
decoded = (decoded + 1.0) / 2.0
decoded = torch.clamp(decoded, 0.0, 1.0)
decoded_img = (decoded[0].float().permute(1, 2, 0).cpu().numpy() * 255).astype(np.uint8)
output_image = Image.fromarray(decoded_img)

output_image.save("0-re.png")

# inference.py
import os
import hashlib
import argparse
import torch
from diffusers import AutoPipelineForText2Image
from safetensors.torch import load_file

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Inference script for Flux LoRA model")
    parser.add_argument("--model_id", type=str, default='black-forest-labs/FLUX.1-dev')
    parser.add_argument("--lora_name", type=str, required=True)
    parser.add_argument('--lora_scale', type=float, default=1.0, help='Scale for LoRA weights')
    parser.add_argument("--prompt", type=str, required=True)
    parser.add_argument("--steps", type=int, default=25)
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--guidance_scale", type=float, default=7.5)
    args = parser.parse_args()

    # Load the model
    model = AutoPipelineForText2Image.from_pretrained(args.model_id, torch_dtype=torch.bfloat16).to("cuda")
    # Load LoRA weights
    model.load_lora_weights(args.lora_name, weight_name="pytorch_lora_weights.safetensors")
    model.set_adapters("default_0", args.lora_scale) 

    # Inference
    with torch.no_grad():
        generator = torch.Generator(device=model.device).manual_seed(args.seed)
        output_image = model(prompt=args.prompt, num_inference_steps=args.steps, guidance_scale=args.guidance_scale, generator=generator).images[0]

    output_image.save('output.png')