Upload 2 files

Custom_Diffusion.py

@@ -0,0 +1,243 @@
+import os  
+import random  
+from pathlib import Path  
+
+import numpy as np  
+import safetensors  # 
+import torch
+import torch.nn.functional as F
+import transformers
+from accelerate import Accelerator
+from accelerate.logging import get_logger
+from accelerate.utils import ProjectConfiguration, set_seed
+from PIL import Image
+from torch.utils.data import Dataset
+from torchvision import transforms
+from tqdm.auto import tqdm
+from transformers import AutoTokenizer, CLIPTextModel
+from safetensors.torch import load_file
+
+import diffusers
+# from diffusers.pipelines import BlipDiffusionPipeline
+from diffusers import AutoencoderKL, DDPMScheduler, UNet2DConditionModel, DiffusionPipeline
+from diffusers.loaders import AttnProcsLayers
+from diffusers.models.attention_processor import CustomDiffusionAttnProcessor
+from diffusers.optimization import get_scheduler
+from diffusers.utils import load_image
+import streamlit as st
+
+import io
+
+import streamlit as st  # 用于创建交互式网页UI
+import io  # 处理文件流(后面用来生成下载按钮)
+
+# 设置页面标题和布局
+st.set_page_config(page_title="Fine-tuning style diffusion", layout="wide")
+
+st.title("Fine-tuning style diffusion 推理 Demo")
+
+st.write("支持 **A <new1> reference.(风格) + 文本*")
+
+st.write("只是训练了一个提示词 'A <new1> reference.'")
+
+st.write("即使用该提示词时以十二生肖为主要元素进行新年图片风格的生成，例如使用一下提示词")
+
+st.write("A <new1> reference. New Year image with a  rabbit as the main element, in a 2D or anime style, and a festive background")
+
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+dtype = torch.float16
+
+
+# ==========================
+# 模型加载（缓存）
+# ==========================
+
+@st.cache_resource
+def load_models():
+
+    model_path = "./stable-diffusion-v1-5"
+
+    tokenizer = AutoTokenizer.from_pretrained(model_path, subfolder="tokenizer")
+
+    text_encoder = CLIPTextModel.from_pretrained(
+        model_path,
+        subfolder="text_encoder",
+        torch_dtype=torch.float16 
+    ).to(device)
+
+    vae = AutoencoderKL.from_pretrained(
+        model_path,
+        subfolder="vae",
+        torch_dtype=torch.float16 
+    ).to(device)
+
+    unet = UNet2DConditionModel.from_pretrained(
+        model_path,
+        subfolder="unet",
+        torch_dtype=torch.float16  
+    ).to(device)
+
+    attn_path = "output/pytorch_custom_diffusion_weights.bin"
+
+    state_dict = torch.load(attn_path, map_location="cpu")
+    unet.load_attn_procs(state_dict)
+
+    token_path = "output/learned_embeds.safetensors"
+
+
+    try:
+
+        new_embed = torch.load(token_path)
+
+        token_id = tokenizer.convert_tokens_to_ids("<new1>")
+
+        text_encoder.get_input_embeddings().weight.data[token_id] = new_embed
+
+        print("Loaded <new1> token embedding")
+
+    except:
+        print("No trained <new1> token found")
+
+    scheduler = DDPMScheduler.from_pretrained(
+        model_path,
+        subfolder="scheduler"
+    )
+    
+    unet.enable_xformers_memory_efficient_attention()
+
+    return tokenizer, text_encoder, vae, unet, scheduler
+
+tokenizer, text_encoder, vae, unet, scheduler = load_models()
+
+
+prompt = st.text_input(
+    "Prompt",
+    "A <new1> reference."
+)
+
+# 调整参数
+steps = st.slider("Steps", 10, 320, 100)
+
+guidance = st.slider("Guidance", 1.0, 18.0, 6.0)
+
+
+# ==========================
+# 图像预处理
+# ==========================
+
+def preprocess(image):
+    # 调整图像，转换为tensor（张量）并归一化到[-1,1]
+    transform = transforms.Compose([
+        transforms.Resize((512,512)),
+        transforms.ToTensor(),
+        transforms.Normalize([0.5],[0.5])
+    ])
+    # 增加batch维度
+    return transform(image).unsqueeze(0)
+
+
+# ==========================
+# diffusion 推理
+# ==========================
+
+def generate(prompt):
+
+    with torch.no_grad():
+        # 文本向量化
+        text_input = tokenizer(
+            prompt,
+            padding="max_length",
+            max_length=tokenizer.model_max_length,
+            truncation=True,
+            return_tensors="pt"
+        ).to(device)
+        
+        text_emb = text_encoder(text_input.input_ids)[0]
+
+        # 无条件 embedding;
+        uncond_input = tokenizer(
+            "",
+            padding="max_length",
+            max_length=tokenizer.model_max_length,
+            return_tensors="pt"
+        ).to(device)
+
+        uncond_emb = text_encoder(uncond_input.input_ids)[0]
+
+
+        text_emb = torch.cat([uncond_emb, text_emb], dim=0)
+
+        # 初始化噪声潜变量
+        latents = torch.randn(
+            (1,4,64,64),
+            device=device,
+            dtype=torch.float16
+        )
+
+        # 设置diffusion时间步
+        scheduler.set_timesteps(steps)
+
+        # ----------------
+        # diffusion loop
+        # ----------------
+        # 采用
+        for t in scheduler.timesteps:
+            # 为什么要拼接两份
+            latent_model_input = torch.cat([latents]*2)
+
+            noise_pred = unet(
+                latent_model_input,
+                t,
+                encoder_hidden_states=text_emb
+            ).sample
+
+
+            noise_uncond, noise_text = noise_pred.chunk(2)
+
+            noise_pred = noise_uncond + guidance * (
+                noise_text - noise_uncond
+            )
+            # 调度程序/潜在的
+            latents = scheduler.step(
+                noise_pred,
+                t,
+                latents
+            ).prev_sample
+
+
+        # ----------------
+        # decode image；解码图像
+        # ----------------
+        # 解码生成图像；将latent解码成[0,1]的RGB图像
+        latents = latents / vae.config.scaling_factor
+
+        image = vae.decode(latents).sample
+
+        image = (image/2 + 0.5).clamp(0,1)
+        # 转成numpy数组，再用PIL转成可展示的图像
+        image = image.cpu().permute(0,2,3,1).numpy()[0]
+
+        image = (image*255).astype(np.uint8)
+
+        return Image.fromarray(image)
+
+
+if st.button("Generate"):
+
+    with st.spinner("Generating..."):
+
+        image = generate(prompt)
+
+    st.image(image,caption="Result",width=512)
+
+    buf = io.BytesIO()
+
+    image.save(buf,format="PNG")
+
+    st.download_button(
+        "Download",
+        buf.getvalue(),
+        "result.png"
+    )
+