Upload 5 files

README.md

@@ -1,13 +1,14 @@
 ---
 title: CHATS
-emoji: 🚀
-colorFrom: blue
-colorTo: yellow
+emoji: 🖼
+colorFrom: purple
+colorTo: red
 sdk: gradio
-sdk_version: 5.31.0
+sdk_version: 5.25.2
 app_file: app.py
 pinned: false
 license: apache-2.0
+short_description: The demo for CHATS-SDXL text-to-image generation model
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,83 @@
+import torch
+import gradio as gr
+
+from pipeline import ChatsSDXLPipeline
+from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
+from transformers import CLIPFeatureExtractor
+from diffusers.utils import logging
+from PIL import Image
+
+logging.set_verbosity_error()
+
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+feature_extractor = CLIPFeatureExtractor.from_pretrained("openai/clip-vit-base-patch32")
+safety_checker = StableDiffusionSafetyChecker.from_pretrained("CompVis/stable-diffusion-safety-checker")
+
+# Load CHATS-SDXL pipeline
+pipe = ChatsSDXLPipeline.from_pretrained(
+        "AIDC-AI/CHATS",
+        safety_checker=safety_checker,
+        feature_extractor=feature_extractor,
+        torch_dtype=torch.float16
+)
+pipe.to(DEVICE)
+
+def generate(prompt, steps=50, guidance_scale=7.5, height=768, width=512):
+  output = pipe(
+      prompt=prompt,
+      num_inference_steps=steps,
+      guidance_scale=guidance_scale,
+      height=height,
+      width=width,
+      seed=0
+  )
+  image = output['images'][0]
+  image = Image.fromarray(image)
+  return image
+
+with gr.Blocks(title="🔥 CHATS-SDXL Demo") as demo:
+    gr.Markdown(
+        "## CHATS-SDXL Text-to-Image Demo\n\n"
+        "Enter your prompt and click **Generate Image**. All NSFW content will be automatically filtered."
+    )
+    with gr.Row():
+        prompt_input = gr.Textbox(
+            label="Prompt",
+            placeholder="Enter your description here...",
+            lines=2,
+        )
+    with gr.Row():
+        steps_slider = gr.Slider(
+            minimum=1, maximum=100, value=50, step=1,
+            label="Inference Steps"
+        )
+        scale_slider = gr.Slider(
+            minimum=1.0, maximum=14.0, value=5.0, step=0.1,
+            label="Guidance Scale"
+        )
+    with gr.Row():
+        height_slider = gr.Slider(
+            minimum=64, maximum=2048, value=1024, step=64, 
+            label="Image Height"
+        )
+        width_slider = gr.Slider(
+            minimum=64, maximum=2048, value=1024, step=64, 
+            label="Image Width"
+        )
+    generate_button = gr.Button("Generate Image")
+    gallery = gr.Gallery(
+        label="Generated Images",
+        show_label=False,
+        columns=2,     
+        elem_id="gallery"
+    )
+
+    generate_button.click(
+        fn=generate,
+        inputs=[prompt_input, steps_slider, scale_slider, height_slider, width_slider],
+        outputs=[gallery],
+    )
+
+if __name__ == "__main__":
+    demo.launch()

gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

pipeline.py

@@ -0,0 +1,189 @@
+#!/usr/bin/env python
+# coding=utf-8
+# Copyright (C) 2025 AIDC-AI
+#
+# Licensed under the Apache License, Version 2.0 (the "License"); 
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at 
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software 
+# distributed under the License is distributed on an "AS IS" BASIS, 
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 
+# See the License for the specific language governing permissions and 
+# limitations under the License.
+
+from typing import Optional, Union, List, Dict, Any
+
+import math
+import os
+import torch
+import torch.nn as nn
+from diffusers import DiffusionPipeline, EulerDiscreteScheduler, SchedulerMixin
+from diffusers.models import AutoencoderKL, UNet2DConditionModel
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.utils import logging
+from PIL import Image
+
+from transformers import CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer, CLIPFeatureExtractor
+from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
+
+def get_noise(
+    num_samples: int,
+    channel: int,
+    height: int,
+    width: int,
+    device: torch.device,
+    dtype: torch.dtype,
+    seed: int,
+):
+    return torch.randn(
+        num_samples,
+        channel,
+        # allow for packing
+        2 * math.ceil(height / 16),
+        2 * math.ceil(width / 16),
+        device=device,
+        dtype=dtype,
+        generator=torch.Generator(device=device).manual_seed(seed),
+    )
+
+class ChatsSDXLPipeline(DiffusionPipeline, ConfigMixin):
+
+    @register_to_config
+    def __init__(
+        self,
+        unet_win: nn.Module,
+        unet_lose: nn.Module,
+        text_encoder: CLIPTextModel,
+        text_encoder_two: CLIPTextModelWithProjection,
+        tokenizer: CLIPTokenizer,
+        tokenizer_two: CLIPTokenizer,
+        vae: AutoencoderKL,
+        scheduler: SchedulerMixin,
+        safety_checker: StableDiffusionSafetyChecker,
+        feature_extractor: CLIPFeatureExtractor
+    ):
+        super().__init__()
+
+        self.register_modules(
+            unet_win=unet_win,
+            unet_lose=unet_lose,
+            text_encoder=text_encoder,
+            text_encoder_two=text_encoder_two,
+            tokenizer=tokenizer,
+            tokenizer_two=tokenizer_two,
+            vae=vae,
+            scheduler=scheduler,
+            safety_checker=safety_checker,
+            feature_extractor=feature_extractor
+        )
+
+    
+    @classmethod
+    def from_pretrained(
+        cls,
+        pretrained_model_name_or_path: Union[str, os.PathLike],
+        **kwargs,
+    ) -> "ChatsSDXLPipeline":
+
+        return super().from_pretrained(pretrained_model_name_or_path, **kwargs)
+
+    def save_pretrained(self, save_directory: Union[str, os.PathLike]):
+        super().save_pretrained(save_directory)
+
+    @torch.no_grad()
+    def encode_text(self, tokenizers, text_encoders, prompt):
+      prompt_embeds_list = []
+
+      with torch.no_grad():
+          for tokenizer, text_encoder in zip(tokenizers, text_encoders):
+              text_inputs = tokenizer(prompt, padding="max_length", max_length=tokenizer.model_max_length, truncation=True, return_tensors="pt",)    
+              text_input_ids = text_inputs.input_ids
+              prompt_embeds = text_encoder(text_input_ids.to(self.unet_win.device), output_hidden_states=True)
+              pooled_prompt_embeds = prompt_embeds[0]
+              prompt_embeds = prompt_embeds.hidden_states[-2]
+              prompt_embeds_list.append(prompt_embeds)
+      
+      prompt_embeds = torch.concat(prompt_embeds_list, dim=-1)
+      prompt_embeds = prompt_embeds.to(dtype=text_encoders[-1].dtype, device=text_encoders[-1].device)
+
+      return prompt_embeds, pooled_prompt_embeds
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        prompt: Union[str, List[str]],
+        num_inference_steps: int = 50,
+        guidance_scale: float = 7.5,
+        latents: torch.FloatTensor = None,
+        height: int = 1024,
+        width: int = 1024,
+        seed: int = 0,
+        alpha: float=0.5
+    ):  
+        if isinstance(prompt, str):
+            prompt = [prompt]
+
+        device = self.unet_win.device
+
+        tokenizers = [self.tokenizer, self.tokenizer_two]
+        text_encoders = [self.text_encoder, self.text_encoder_two]
+
+        prompt_embeds, pooled_prompt_embeds = self.encode_text(tokenizers, text_encoders, prompt)
+        negative_prompt_embeds, negative_pooled_prompt_embeds = self.encode_text(tokenizers, text_encoders, "")
+
+        self.scheduler.set_timesteps(num_inference_steps, device=device)
+        timesteps = self.scheduler.timesteps
+
+        bs = len(prompt)
+        channel = self.vae.config.latent_channels
+        height = 16 * (height // 16)
+        width = 16 * (width // 16)
+
+        # prepare input
+        latents = get_noise(
+            bs,
+            channel,
+            height,
+            width,
+            device=device,
+            dtype=self.unet_win.dtype,
+            seed=seed,
+        )
+        latents = latents * self.scheduler.init_noise_sigma
+        
+        add_time_ids = torch.tensor([height, width, 0, 0, height, width], dtype=latents.dtype, device=device)[None, :].repeat(latents.size(0), 1)
+
+        for i, t in enumerate(timesteps):
+            latent_model_input = self.scheduler.scale_model_input(latents, t)
+
+            added_cond_kwargs_win = {"text_embeds": pooled_prompt_embeds, "time_ids": add_time_ids}
+            added_cond_kwargs_lose = {"text_embeds": pooled_prompt_embeds * (-alpha) + negative_pooled_prompt_embeds * (1. + alpha), "time_ids": add_time_ids}
+
+            pred_win = self.unet_win(latent_model_input, t, encoder_hidden_states=prompt_embeds, added_cond_kwargs=added_cond_kwargs_win, return_dict=False)[0]
+            pred_lose = self.unet_lose(latent_model_input, t, encoder_hidden_states=prompt_embeds * (-alpha) + negative_prompt_embeds * (1. + alpha), added_cond_kwargs=added_cond_kwargs_lose, return_dict=False)[0]
+
+            noise_pred = pred_win + guidance_scale * (pred_win - pred_lose)
+            latents = self.scheduler.step(noise_pred, t, latents, generator=None, return_dict=False)[0]
+
+        x = latents.float()
+
+        with torch.no_grad():
+            with torch.autocast(device_type=device.type, dtype=torch.float32):
+                if hasattr(self.vae.config, 'scaling_factor') and self.vae.config.scaling_factor is not None:
+                    x = x / self.vae.config.scaling_factor
+                if hasattr(self.vae.config, 'shift_factor') and self.vae.config.shift_factor is not None:
+                    x = x + self.vae.config.shift_factor
+                x = self.vae.decode(x, return_dict=False)[0]
+
+        # bring into PIL format and save
+        x = (x / 2 + 0.5).clamp(0, 1)
+        x = x.cpu().permute(0, 2, 3, 1).float().numpy()
+        images = (x * 255).round().astype("uint8")
+
+        clip_input = self.feature_extractor(images=images, return_tensors="pt").to(self.device)
+        filtered_images, has_nsfw_flags = self.safety_checker(images=images, clip_input=clip_input.pixel_values)
+
+        return {"images": filtered_images, "nsfw_flags": has_nsfw_flags}

requirements.txt

@@ -0,0 +1,22 @@
+transformers==4.44.2
+accelerate==0.31.0
+deepspeed==0.14.5
+numpy==1.24.3
+diffusers
+datasets
+requests
+fastapi
+scipy
+pandas
+xformers
+ftfy
+Jinja2
+bitsandbytes
+safetensors
+pyyaml
+pillow==10.3.0
+gradio
+--extra-index-url https://download.pytorch.org/whl/cu124
+torch==2.4.1
+torchvision==0.19.1
+torchaudio==2.4.1