Update app.py

app.py

@@ -2,45 +2,195 @@ import gradio as gr
 
 import torch
 from torch import autocast
-from diffusers import StableDiffusionPipeline
-from datasets import load_dataset
-from PIL import Image  
-import re
 
-model_id = "CompVis/stable-diffusion-v1-4"
+import gc
+import io
+import math
+import sys
+
+from PIL import Image, ImageOps
+import requests
+from torch import nn
+from torch.nn import functional as F
+from torchvision import transforms
+from torchvision.transforms import functional as TF
+from tqdm.notebook import tqdm
+
+import numpy as np
+
+from guided_diffusion.script_util import create_model_and_diffusion, model_and_diffusion_defaults, classifier_defaults, create_classifier
+
+from omegaconf import OmegaConf
+from ldm.util import instantiate_from_config
+
+from einops import rearrange
+from math import log2, sqrt
+
+import argparse
+import pickle
+
+import os
+
+from transformers import CLIPTokenizer, CLIPTextModel
+
+def fetch(url_or_path):
+    if str(url_or_path).startswith('http://') or str(url_or_path).startswith('https://'):
+        r = requests.get(url_or_path)
+        r.raise_for_status()
+        fd = io.BytesIO()
+        fd.write(r.content)
+        fd.seek(0)
+        return fd
+    return open(url_or_path, 'rb')
+
 device = "cuda"
 
-#If you are running this code locally, you need to either do a 'huggingface-cli login` or paste your User Access Token from here https://huggingface.co/settings/tokens into the use_auth_token field below. 
-pipe = StableDiffusionPipeline.from_pretrained(model_id, use_auth_token=False, revision="fp16", torch_dtype=torch.float16)
-pipe = pipe.to(device)
-#When running locally, you won`t have access to this, so you can remove this part
-word_list_dataset = load_dataset("stabilityai/word-list", data_files="list.txt", use_auth_token=True)
-word_list = word_list_dataset["train"]['text']
-
-def infer(prompt, samples, steps, scale, seed):
-    #When running locally you can also remove this filter
-    for filter in word_list:
-        if re.search(rf"\b{filter}\b", prompt):
-            raise gr.Error("Unsafe content found. Please try again with different prompts.")
-        
-    generator = torch.Generator(device=device).manual_seed(seed)
-    
-    #If you are running locally with CPU, you can remove the `with autocast("cuda")`
-    with autocast("cuda"):
-        images_list = pipe(
-            [prompt] * samples,
-            num_inference_steps=steps,
-            guidance_scale=scale,
-            generator=generator,
-        )
-    images = []
-    safe_image = Image.open(r"unsafe.png")
-    for i, image in enumerate(images_list["sample"]):
-        if(images_list["nsfw_content_detected"][i]):
-            images.append(safe_image)
-        else:
-            images.append(image)
-    return images
+#model_state_dict = torch.load('diffusion.pt', map_location='cpu')
+model_state_dict = torch.load(fetch('https://huggingface.co/Jack000/glid-3-xl-stable/resolve/main/default/diffusion-1.4.pt'), map_location='cpu')
+
+model_params = {
+    'attention_resolutions': '32,16,8',
+    'class_cond': False,
+    'diffusion_steps': 1000,
+    'rescale_timesteps': True,
+    'timestep_respacing': 'ddim100',
+    'image_size': 32,
+    'learn_sigma': False,
+    'noise_schedule': 'linear',
+    'num_channels': 320,
+    'num_heads': 8,
+    'num_res_blocks': 2,
+    'resblock_updown': False,
+    'use_fp16': True,
+    'use_scale_shift_norm': False,
+    'clip_embed_dim': None,
+    'image_condition': False,
+    'super_res_condition': False,
+}
+
+model_config = model_and_diffusion_defaults()
+model_config.update(model_params)
+
+# Load models
+model, diffusion = create_model_and_diffusion(**model_config)
+model.load_state_dict(model_state_dict, strict=True)
+model.requires_grad_(False).eval().to(device)
+
+if model_config['use_fp16']:
+    model.convert_to_fp16()
+else:
+    model.convert_to_fp32()
+
+def set_requires_grad(model, value):
+    for param in model.parameters():
+        param.requires_grad = value
+
+# vae
+kl_config = OmegaConf.load('kl.yaml')
+kl_sd = torch.load(fetch('https://huggingface.co/Jack000/glid-3-xl-stable/resolve/main/default/kl-1.4.pt'), map_location="cpu")
+
+ldm = instantiate_from_config(kl_config.model)
+ldm.load_state_dict(kl_sd, strict=True)
+
+ldm.to(device)
+ldm.eval()
+ldm.requires_grad_(False)
+set_requires_grad(ldm, False)
+
+# clip
+clip_version = 'openai/clip-vit-large-patch14'
+clip_tokenizer = CLIPTokenizer.from_pretrained(clip_version)
+clip_transformer = CLIPTextModel.from_pretrained(clip_version)
+clip_transformer.eval().requires_grad_(False).to(device)
+
+# classifier
+# load classifier
+classifier_config = classifier_defaults()
+classifier_config['classifier_width'] = 128
+classifier_config['classifier_depth'] = 4
+classifier_config['classifier_attention_resolutions'] = '64,32,16,8'
+
+classifier_photo = create_classifier(**classifier_config)
+classifier_photo.load_state_dict(
+    torch.load(fetch('https://huggingface.co/Jack000/glid-3-xl-stable/resolve/main/classifier_photo/model060000.pt'), map_location="cpu")
+)
+classifier_photo.to(device)
+classifier_photo.convert_to_fp16()
+classifier_photo.eval()
+
+classifier_art = create_classifier(**classifier_config)
+classifier_art.load_state_dict(
+    torch.load('https://huggingface.co/Jack000/glid-3-xl-stable/resolve/main/classifier_art/model110000.pt', map_location="cpu")
+)
+classifier_art.to(device)
+classifier_art.convert_to_fp16()
+classifier_art.eval()
+
+def infer(prompt, style, scale, classifier_scale, seed):
+    torch.manual_seed(seed)
+
+    # clip context
+    text = clip_tokenizer([prompt], truncation=True, max_length=77, return_length=True, return_overflowing_tokens=False, padding="max_length", return_tensors="pt")
+    text_blank = clip_tokenizer([''], truncation=True, max_length=77, return_length=True, return_overflowing_tokens=False, padding="max_length", return_tensors="pt")
+    text_tokens = text["input_ids"].to(device)
+    text_blank_tokens = text_blank["input_ids"].to(device)
+
+    text_emb = clip_transformer(input_ids=text_tokens).last_hidden_state
+    text_emb_blank = clip_transformer(input_ids=text_blank_tokens).last_hidden_state
+
+    kwargs = {
+        "context": torch.cat([text_emb, text_emb_blank], dim=0).half(),
+        "clip_embed": None,
+        "image_embed": None,
+    }
+
+    def model_fn(x_t, ts, **kwargs):
+        half = x_t[: len(x_t) // 2]
+        combined = torch.cat([half, half], dim=0)
+        model_out = model(combined, ts, **kwargs)
+        eps, rest = model_out[:, :3], model_out[:, 3:]
+        cond_eps, uncond_eps = torch.split(eps, len(eps) // 2, dim=0)
+        half_eps = uncond_eps + scale * (cond_eps - uncond_eps)
+        eps = torch.cat([half_eps, half_eps], dim=0)
+        return torch.cat([eps, rest], dim=1)
+
+    def cond_fn(x, t, context=None, clip_embed=None, image_embed=None):
+        with torch.enable_grad():
+            x_in = x[:x.shape[0]//2].detach().requires_grad_(True)
+            if style == 'photo':
+                logits = classifier_photo(x_in, t)
+            elif style == 'digital art':
+                logits = classifier_art(x_in, t)
+            else:
+                return 0
+
+            log_probs = F.log_softmax(logits, dim=-1)
+            selected = log_probs[range(len(logits)), torch.ones(x_in.shape[0], dtype=torch.long)]
+            return torch.autograd.grad(selected.sum(), x_in)[0] * classifier_scale
+
+    samples = diffusion.ddim_sample_loop_progressive(
+        model_fn,
+        (2, 4, 64, 64),
+        clip_denoised=False,
+        model_kwargs=kwargs,
+        cond_fn=cond_fn,
+        device=device,
+        progress=True,
+        init_image=None,
+        skip_timesteps=0,
+    )
+
+    for j, sample in enumerate(samples):
+        pass
+
+    emb = sample['pred_xstart'][0]
+    emb /= 0.18215
+    im = emb.unsqueeze(0)
+    im = ldm.decode(im)
+
+    im = TF.to_pil_image(im.squeeze(0).add(1).div(2).clamp(0, 1))
+
+    return [im]
     
 css = """
         .gradio-container {
@@ -97,8 +247,7 @@ css = """
             padding: 2px 8px;
             border-radius: 14px !important;
         }
-        #advanced-options {
-            display: none;
+        #advanced-options, #style-options {
             margin-bottom: 20px;
         }
         .footer {
@@ -213,19 +362,11 @@ with block:
                   <rect x="23" y="69" width="23" height="23" fill="black"></rect>
                 </svg>
                 <h1 style="font-weight: 900; margin-bottom: 7px;">
-                  Stable Diffusion Demo
+                  Classifier Guided Stable Diffusion
                 </h1>
               </div>
               <p style="margin-bottom: 10px; font-size: 94%">
-                Stable Diffusion is a state of the art text-to-image model that generates
-                images from text.<br>For faster generation and forthcoming API
-                access you can try
-                <a
-                  href="http://beta.dreamstudio.ai/"
-                  style="text-decoration: underline;"
-                  target="_blank"
-                  >DreamStudio Beta</a
-                >
+                a custom version of stable diffusion with classifier guidance
               </p>
             </div>
         """
@@ -252,13 +393,18 @@ with block:
             label="Generated images", show_label=False, elem_id="gallery"
         ).style(grid=[2], height="auto")
 
-        advanced_button = gr.Button("Advanced options", elem_id="advanced-btn")
+        #advanced_button = gr.Button("Advanced options", elem_id="advanced-btn")
 
+        with gr.Row(elem_id="style-options"):
+            style = gr.Radio(["none","photo","digital art","anime"], label="Image style")
         with gr.Row(elem_id="advanced-options"):
-            samples = gr.Slider(label="Images", minimum=1, maximum=4, value=4, step=1)
-            steps = gr.Slider(label="Steps", minimum=1, maximum=50, value=45, step=1)
+            #samples = gr.Slider(label="Images", minimum=1, maximum=4, value=4, step=1)
+            #steps = gr.Slider(label="Steps", minimum=1, maximum=50, value=45, step=1)
             scale = gr.Slider(
-                label="Guidance Scale", minimum=0, maximum=50, value=7.5, step=0.1
+                label="CFG Scale", minimum=0, maximum=50, value=7.5, step=0.1
+            )
+            classifier_scale = gr.Slider(
+                label="Classifier Scale", minimum=0, maximum=1000, value=100, step=1
             )
             seed = gr.Slider(
                 label="Seed",
@@ -268,22 +414,13 @@ with block:
                 randomize=True,
             )
 
-        ex = gr.Examples(examples=examples, fn=infer, inputs=[text, samples, steps, scale, seed], outputs=gallery, cache_examples=True)
+        ex = gr.Examples(examples=examples, fn=infer, inputs=[text, style, scale, classifier_scale, seed], outputs=gallery, cache_examples=True)
         ex.dataset.headers = [""]
 
         
-        text.submit(infer, inputs=[text, samples, steps, scale, seed], outputs=gallery)
-        btn.click(infer, inputs=[text, samples, steps, scale, seed], outputs=gallery)
-        advanced_button.click(
-            None,
-            [],
-            text,
-            _js="""
-            () => {
-                const options = document.querySelector("body > gradio-app").querySelector("#advanced-options");
-                options.style.display = ["none", ""].includes(options.style.display) ? "flex" : "none";
-            }""",
-        )
+        text.submit(infer, inputs=[text, style, scale, classifier_scale, seed], outputs=gallery)
+        btn.click(infer, inputs=[text, style, scale, classifier_scale, seed], outputs=gallery)
+
         gr.HTML(
             """
                 <div class="footer">