Update pipeline_sdxs.py

pipeline_sdxs.py

@@ -1,372 +1,193 @@
+import gradio as gr
+import numpy as np
+import random
+import spaces 
 import torch
-from diffusers import DiffusionPipeline
-from diffusers.utils import BaseOutput
-from dataclasses import dataclass
-from typing import List, Union, Optional, Tuple, Any
+from diffusers import DiffusionPipeline, AutoencoderKL, UNet2DConditionModel, FlowMatchEulerDiscreteScheduler,AsymmetricAutoencoderKL
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from typing import Optional, Union, List, Tuple
 from PIL import Image
-import numpy as np
-from tqdm import tqdm
-
-@dataclass
-class SdxsPipelineOutput(BaseOutput):
-    images: Union[List[Image.Image], np.ndarray]
-    refined_prompt: Optional[Union[str, List[str]]] = None
-
-class SdxsPipeline(DiffusionPipeline):
-    # НОВОЕ: Константа для токена </think> в Qwen3
-    END_THINK_TOKEN_ID = 151668
-    
-    # Шаблон промпта по умолчанию
-    DEFAULT_REFINE_TEMPLATE = (
-        "You are a skilled text-to-image prompt engineer whose sole function is to transform the user's input into an aesthetically optimized, detailed, and visually descriptive three-sentence output. "
-        "**The primary subject (e.g., 'girl', 'dog', 'house') MUST be the main focus of the revised prompt and MUST be described in rich detail within the first sentence or two.** "
-        "If the input is short, elaborate the subject using diverse attributes (style, pose, expression, lighting/color palette/mood). **Descriptions must avoid cliches and include diverse options.** "
-        "If the input is long, concisely pack the core subject and essential details into the final three-sentence format without losing crucial information. "
-        "Output **only** the final revised prompt in **English**, with absolutely no commentary, thinking text, or surrounding quotes.\n"
-        "User input prompt: {prompt}"
-    )
-#User input prompt: {prompt}
-    def __init__(self, vae, text_encoder, tokenizer, unet, scheduler, max_length: int = 192):
-        super().__init__()
-        self.register_modules(
-            vae=vae, text_encoder=text_encoder, tokenizer=tokenizer, 
-            unet=unet, scheduler=scheduler
-        )
-        self.vae_scale_factor = 16
-        self.max_length = max_length
 
-    def encode_prompt(self, prompt=None, negative_prompt=None, device=None, dtype=None) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        device = device or self.device
-        dtype = dtype or next(self.unet.parameters()).dtype
-    
-        # Преобразуем в списки
-        if isinstance(prompt, str):
-            prompt = [prompt]
-        if isinstance(negative_prompt, str):
-            negative_prompt = [negative_prompt]
-    
-        # Если промпты не заданы, используем пустые эмбеддинги
-        if prompt is None and negative_prompt is None:
-            hidden_dim = 1024  # Размерность эмбеддинга
-            seq_len = self.max_length
-            batch_size = 1
-            # ИЗМЕНЕНО: Возвращаем три элемента: embeds, mask, pooled
-            empty_embeds = torch.zeros((batch_size, seq_len, hidden_dim), dtype=dtype, device=device)
-            empty_mask = torch.ones((batch_size, seq_len), dtype=torch.int64, device=device)
-            empty_pooled = torch.zeros((batch_size, hidden_dim), dtype=dtype, device=device)
-            return empty_embeds, empty_mask, empty_pooled
-    
-        # Токенизация с фиксированным max_length и padding="max_length"
-        def encode_texts(texts, max_length=self.max_length):
-            with torch.no_grad():
-                if isinstance(texts, str):
-                    texts = [texts]
+device = "cuda" if torch.cuda.is_available() else "cpu"
+dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+model_repo_id = "AiArtLab/sdxs-08b"
+
+pipe = DiffusionPipeline.from_pretrained(
+    model_repo_id,
+    torch_dtype=dtype,
+    trust_remote_code=True
+).to(device)
+
+# НОВОЕ: Инициализация Qwen3 для рефайнинга
+llm_model_id = "Qwen/Qwen3-0.6B"
+tokenizer = AutoTokenizer.from_pretrained(llm_model_id)
+llm_model = AutoModelForCausalLM.from_pretrained(llm_model_id, torch_dtype="auto", device_map="auto")
+
+MAX_SEED = np.iinfo(np.int32).max
+MIN_IMAGE_SIZE = 640
+MAX_IMAGE_SIZE = 1280
+STEP = 64
+
+# НОВОЕ: Настройки для LLM
+END_THINK_TOKEN_ID = 151668
+DEFAULT_REFINE_TEMPLATE = (
+    "You are a skilled text-to-image prompt engineer whose sole function is to transform the user's input into an aesthetically optimized, detailed, and visually descriptive three-sentence output. "
+    "**The primary subject (e.g., 'girl', 'dog', 'house') MUST be the main focus of the revised prompt and MUST be described in rich detail within the first sentence or two.** "
+    "Output **only** the final revised prompt in **English**, with absolutely no commentary, thinking text, or surrounding quotes.\n"
+    "User input prompt: {prompt}"
+)
+
+@spaces.GPU(duration=30)
+def infer(
+    prompt: str,
+    negative_prompt: str,
+    seed: int,
+    randomize_seed: bool,
+    width: int,
+    height: int,
+    guidance_scale: float,
+    num_inference_steps: int,
+    refine_prompt: bool,
+    progress=gr.Progress(track_tqdm=True),
+) -> Tuple[Image.Image, int, str]: # Возвращаем prompt в конце
+    
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    
+    # НОВОЕ: Логика улучшения промпта
+    if refine_prompt and prompt:
+        messages = [{"role": "user", "content": DEFAULT_REFINE_TEMPLATE.format(prompt=prompt)}]
+        text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, enable_thinking=True)
+        model_inputs = tokenizer([text], return_tensors="pt").to(llm_model.device)
         
-                for i, prompt_item in enumerate(texts):
-                    messages = [
-                        {"role": "user", "content": prompt_item},
-                    ]
-                    prompt_item = self.tokenizer.apply_chat_template(
-                        messages,
-                        tokenize=False,
-                        add_generation_prompt=True,
-                        enable_thinking=True,
-                    )
-                    texts[i] = prompt_item
-            
-                toks = self.tokenizer(
-                    texts,
-                    return_tensors="pt",
-                    padding="max_length",
-                    truncation=True,
-                    max_length=max_length
-                ).to(device)
-                outs = self.text_encoder(**toks, output_hidden_states=True, return_dict=True)
-                
-                # Токен-эмбеддинги (для Cross-Attention)
-                hidden = outs.hidden_states[-2] # Используем last hidden state -2???
-                # Маска внимания (для Cross-Attention)
-                attention_mask = toks["attention_mask"] 
-                
-                # Пулинг-эмбеддинг (для Class/Time Conditioning). Берем эмбеддинг последнего токена без padding.
-                sequence_lengths = attention_mask.sum(dim=1) - 1
-                batch_size = hidden.shape[0]
-                pooled = hidden[torch.arange(batch_size, device=hidden.device), sequence_lengths]
-                
-                # --- НОВАЯ ЛОГИКА: ОБЪЕДИНЕНИЕ ДЛЯ КРОСС-ВНИМАНИЯ ---
-                # 1. Расширяем пулинг-вектор до последовательности [B, 1, 1024]
-                pooled_expanded = pooled.unsqueeze(1) 
-                                
-                # 2. Объединяем последовательность токенов и пулинг-вектор
-                # !!! ИЗМЕНЕНИЕ ЗДЕСЬ !!!: Пулинг идет ПЕРВЫМ
-                # Теперь: [B, 1 + L, 1024]. Пулинг стал токеном в НАЧАЛЕ.
-                new_encoder_hidden_states = torch.cat([pooled_expanded, hidden], dim=1) 
-                                
-                # 3. Обновляем маску внимания для нового токена
-                # Маска внимания: [B, 1 + L]. Добавляем 1 в НАЧАЛО.
-                # torch.ones((batch_size, 1), device=device) создает маску [B, 1] со значениями 1.
-                new_attention_mask = torch.cat([torch.ones((batch_size, 1), device=device), attention_mask], dim=1)
-                
-                return new_encoder_hidden_states, new_attention_mask, pooled
-    
-        # Кодируем позитивные и негативные промпты
-        # ИСПРАВЛЕНИЕ: Теперь возвращаем (None, None, None), чтобы избежать UnboundLocalError
-        pos_result = encode_texts(prompt) if prompt is not None else (None, None, None)
-        neg_result = encode_texts(negative_prompt) if negative_prompt is not None else (None, None, None)
-    
-        pos_embeddings, pos_mask, pos_pooled = pos_result
-        neg_embeddings, neg_mask, neg_pooled = neg_result
-    
-        # Выравниваем размеры batch_size
-        batch_size = max(
-            pos_embeddings.shape[0] if pos_embeddings is not None else 0,
-            neg_embeddings.shape[0] if neg_embeddings is not None else 0
-        )
-    
-        # Повторяем эмбеддинги, маски и пулинг по batch_size
-        if pos_embeddings is not None and pos_embeddings.shape[0] < batch_size:
-            pos_embeddings = pos_embeddings.repeat(batch_size, 1, 1)
-            pos_mask = pos_mask.repeat(batch_size, 1)
-            pos_pooled = pos_pooled.repeat(batch_size, 1)
-            
-        # ИСПРАВЛЕНИЕ: Проверяем, существует ли neg_embeddings, прежде чем обращаться к его shape[0]
-        if neg_embeddings is not None and neg_embeddings.shape[0] < batch_size:
-            neg_embeddings = neg_embeddings.repeat(batch_size, 1, 1)
-            neg_mask = neg_mask.repeat(batch_size, 1)
-            neg_pooled = neg_pooled.repeat(batch_size, 1)
-    
-        # Конкатенируем для guidance (эмбеддинги и маски)
-        # Убеждаемся, что все три компонента существуют перед конкатенацией
-        if pos_embeddings is not None and neg_embeddings is not None:
-            text_embeddings = torch.cat([neg_embeddings, pos_embeddings], dim=0)
-            attention_mask = torch.cat([neg_mask, pos_mask], dim=0)
-            pooled_embeddings = torch.cat([neg_pooled, pos_pooled], dim=0)
-        elif pos_embeddings is not None:
-            text_embeddings = pos_embeddings
-            attention_mask = pos_mask
-            pooled_embeddings = pos_pooled
-        else: # Только neg_embeddings
-            text_embeddings = neg_embeddings
-            attention_mask = neg_mask
-            pooled_embeddings = neg_pooled
-    
-        # Возвращаем кортеж
-        return (
-            text_embeddings.to(device=device, dtype=dtype), 
-            attention_mask.to(device=device, dtype=torch.int64),
-            pooled_embeddings.to(device=device, dtype=dtype)
-        )
-
-
-    @torch.no_grad()
-    def generate_latents(
-        self,
-        text_embeddings,
-        attention_mask,
-        pooled_embeddings,
-        height: int = 1536,
-        width: int = 1280,
-        num_inference_steps: int = 40,
-        guidance_scale: float = 4.0,
-        latent_channels: int = 16,
-        batch_size: int = 1,
-        generator=None,
-    ):
-        device = self.device
-        dtype = next(self.unet.parameters()).dtype
+        generated_ids = llm_model.generate(**model_inputs, max_new_tokens=2048, do_sample=True, pad_token_id=tokenizer.eos_token_id)
+        output_ids = generated_ids[0][len(model_inputs.input_ids[0]):].tolist()
         
-        self.scheduler.set_timesteps(num_inference_steps, device=device)
-    
-        # Разделяем эмбеддинги и маски на условные и безусловные
-        if guidance_scale > 1:
-            neg_embeds, pos_embeds = text_embeddings.chunk(2)
-            neg_mask, pos_mask = attention_mask.chunk(2) 
-            neg_pooled, pos_pooled = pooled_embeddings.chunk(2)
-            
-            # Повторяем, если batch_size больше
-            if batch_size > pos_embeds.shape[0]:
-                pos_embeds = pos_embeds.repeat(batch_size, 1, 1)
-                neg_embeds = neg_embeds.repeat(batch_size, 1, 1)
-                pos_mask = pos_mask.repeat(batch_size, 1)
-                neg_mask = neg_mask.repeat(batch_size, 1)
-                pos_pooled = pos_pooled.repeat(batch_size, 1)
-                neg_pooled = neg_pooled.repeat(batch_size, 1)
-                
-            text_embeddings = torch.cat([neg_embeds, pos_embeds], dim=0)
-            unet_attention_mask = torch.cat([neg_mask, pos_mask], dim=0)
-            unet_pooled_embeddings = torch.cat([neg_pooled, pos_pooled], dim=0)
-        else:
-            text_embeddings = text_embeddings.repeat(batch_size, 1, 1)
-            unet_attention_mask = attention_mask.repeat(batch_size, 1)
-            unet_pooled_embeddings = pooled_embeddings.repeat(batch_size, 1)
-    
-        # Инициализация латентов
-        latent_shape = (
-            batch_size,
-            latent_channels,
-            height // self.vae_scale_factor,
-            width // self.vae_scale_factor
-        )
-        latents = torch.randn(latent_shape, device=device, dtype=dtype, generator=generator)
-    
-        # Процесс диффузии
-        for t in tqdm(self.scheduler.timesteps, desc="Генерация"):
-            latent_input = torch.cat([latents, latents], dim=0) if guidance_scale > 1 else latents
-            
-            noise_pred = self.unet(
-                latent_input, 
-                t, 
-                encoder_hidden_states=text_embeddings,
-                encoder_attention_mask=unet_attention_mask,
-                #added_cond_kwargs={'text_embeds': unet_pooled_embeddings}
-            ).sample
-    
-            if guidance_scale > 1:
-                noise_uncond, noise_text = noise_pred.chunk(2)
-                noise_pred = noise_uncond + guidance_scale * (noise_text - noise_uncond)
-    
-            latents = self.scheduler.step(noise_pred, t, latents).prev_sample
-    
-        return latents
-
+        try:
+            index = len(output_ids) - output_ids[::-1].index(END_THINK_TOKEN_ID)
+        except ValueError:
+            index = 0
+        
+        prompt = tokenizer.decode(output_ids[index:], skip_special_tokens=True).strip("\n").strip()
+
+    output = pipe(
+        prompt=prompt,
+        negative_prompt=negative_prompt,
+        guidance_scale=guidance_scale,
+        num_inference_steps=num_inference_steps,
+        width=width,
+        height=height,
+        seed=seed,
+    )
 
-    def decode_latents(self, latents, output_type="pil"):
-        """Декодирование латентов в изображения."""
-        latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor
-        with torch.no_grad():
-            images = self.vae.decode(latents).sample
-        images = (images / 2 + 0.5).clamp(0, 1)
+    image = output.images[0]
+    return image, seed, prompt # Возвращаем измененный промпт
+
+examples = [
+    "A frozen river, surrounded by snow-covered trees, reflects the clear blue sky, with a warm glow from the setting sun.",
+    "A young woman with striking blue eyes and pointed ears, adorned with a floral kimono and a tattoo. Her hair is styled in a braid, and she wears a pair of ears",
+    "A volcano explodes, creating a skull face shadow in embers with lightning illuminating the clouds.",
+    "There is a young male character standing against a vibrant, colorful graffiti wall. he is wearing a straw hat, a black jacket adorned with gold accents, and black shorts.",
+    "A man with dark hair and a beard is meticulously carving an intricate design on a piece of pottery. He is wearing a traditional scarf and a white shirt, and he is focused on his work.",
+    "girl, smiling, red eyes, blue hair, white shirt"
+]
+
+css = """
+#col-container {
+    margin: 0 auto;
+    max-width: 640px;
+}
+"""
+
+with gr.Blocks(css=css) as demo:
+    with gr.Column(elem_id="col-container"):
+        gr.Markdown(" # Simple Diffusion (sdxs-08b)")
+
+        with gr.Row():
+            prompt = gr.Text(
+                label="Prompt",
+                show_label=False,
+                max_lines=5,
+                placeholder="Enter your prompt",
+                container=False,
+            )
 
-        if output_type == "pil":
-            images = images.cpu().permute(0, 2, 3, 1).float().numpy()
-            images = (images * 255).round().astype("uint8")
-            return [Image.fromarray(image) for image in images]
-        return images.cpu().permute(0, 2, 3, 1).float().numpy()
+            run_button = gr.Button("Run", scale=0, variant="primary")
 
-    # ИЗМЕНЕНИЕ: Метод __call__ теперь корректно внутри класса SdxsPipeline
-    @torch.no_grad()
-    def __call__(
-        self,
-        prompt: Optional[Union[str, List[str]]] = None,
-        height: int = 1280,
-        width: int = 1024,
-        num_inference_steps: int = 40,
-        guidance_scale: float = 4.0,
-        latent_channels: int = 16,
-        output_type: str = "pil",
-        return_dict: bool = True,
-        batch_size: int = 1,
-        seed: Optional[int] = None,
-        negative_prompt: Optional[Union[str, List[str]]] = None,
-        text_embeddings: Optional[torch.FloatTensor] = None,
-        refine_prompt: bool = True,
-        refine_template: Optional[str] = None,
-    ):
-        device = self.device
-        generator = torch.Generator(device=device).manual_seed(seed) if seed is not None else None
+        result = gr.Image(label="Result", show_label=False)
 
-        refined_prompt_output = None
-        
-        # 1. ЛОГИКА УТОЧНЕНИЯ ПРОМПТА
-        if refine_prompt and prompt is not None and text_embeddings is None:
-            
-            is_str_input = isinstance(prompt, str)
-            original_prompts = [prompt] if is_str_input else prompt
-            
-            template = refine_template if refine_template is not None else self.DEFAULT_REFINE_TEMPLATE
-            
-            refined_list = []
+        with gr.Accordion("Advanced Settings", open=False):
+            # Изменено value на True
+            refine_prompt = gr.Checkbox(label="Refine Prompt with Qwen3", value=True)
             
-            for p in original_prompts:
-                # 1.1. Форматирование промпта по правилам Qwen
-                messages = [
-                    {"role": "user", "content": template.format(prompt=p)} # Шаблон с промптом пользователя внутри
-                ]
-                
-                # ИЗМЕНЕНИЕ: Используем chat_template для подготовки текста
-                text = self.tokenizer.apply_chat_template(
-                    messages,
-                    tokenize=False,
-                    add_generation_prompt=True,
-                    enable_thinking=True 
-                )
-                
-                model_inputs = self.tokenizer([text], return_tensors="pt", truncation=True).to(device)
+            negative_prompt = gr.Text(
+                label="Negative prompt",
+                max_lines=1,
+                placeholder="Enter a negative prompt",
+                value ="bad quality, low resolution"
+            )
 
-                try:
-                    # 1.2. Генерация текста (требует, чтобы self.text_encoder имел метод .generate())
-                    generated_ids = self.text_encoder.generate(
-                        **model_inputs,
-                        max_new_tokens=32768, # Ограничим, чтобы избежать слишком долгой генерации
-                        do_sample=True,
-                        pad_token_id=self.tokenizer.eos_token_id
-                    )
-                    
-                    # 1.3. Обрезка входного промпта
-                    # ИЗМЕНЕНИЕ: Обрезаем сгенерированные токены до тех, что были сгенерированы моделью
-                    output_ids = generated_ids[0][len(model_inputs.input_ids[0]):].tolist()
-                    
-                    # 1.4. Парсинг содержимого "мышления"
-                    try:
-                        # ИЗМЕНЕНИЕ: Ищем токен END_THINK_TOKEN_ID (151668) с конца
-                        # output_ids[::-1].index(151668) найдет индекс в обратном списке
-                        index = len(output_ids) - output_ids[::-1].index(self.END_THINK_TOKEN_ID)
-                    except ValueError:
-                        # Если токен </think> не найден, начинаем с начала
-                        index = 0
+            seed = gr.Slider(
+                label="Seed",
+                minimum=0,
+                maximum=MAX_SEED,
+                step=1,
+                value=0,
+            )
+            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+
+            with gr.Row():
+                width = gr.Slider(
+                    label="Width",
+                    minimum=MIN_IMAGE_SIZE,
+                    maximum=MAX_IMAGE_SIZE,
+                    step=STEP,
+                    value=1024,
+                )
 
-                    # ИЗМЕНЕНИЕ: Декодируем контент только после </think>
-                    refined_text = self.tokenizer.decode(output_ids[index:], skip_special_tokens=True).strip("\n")
-                    
-                    # 1.5. Добавление оригинального промпта в конец
-                    # Формат: refined_text + ", " + user_prompt
-                    final_refined_text = f"{refined_text.strip()}"#\n{p}"
-                    
-                except AttributeError:
-                    print("ВНИМАНИЕ: self.text_encoder не имеет метода .generate(). Уточнение промпта пропущено.")
-                    final_refined_text = p # Используем оригинальный промпт
-                except Exception as e:
-                    print(f"Ошибка при уточнении промпта: {e}. Используется оригинальный промпт.")
-                    final_refined_text = p
-                
-                refined_list.append(final_refined_text)
-                
-            # Обновление промпта и сохранение уточненного для вывода
-            prompt = refined_list[0] if is_str_input else refined_list
-            refined_prompt_output = prompt # Здесь уже список или строка
+                height = gr.Slider(
+                    label="Height",
+                    minimum=MIN_IMAGE_SIZE,
+                    maximum=MAX_IMAGE_SIZE,
+                    step=STEP,
+                    value=MAX_IMAGE_SIZE,
+                )
 
-        # 2. КОДИРОВАНИЕ ПРОМПТОВ (существующая логика)
-        if text_embeddings is None:
-            if prompt is None and negative_prompt is None:
-                raise ValueError("Необходимо указать prompt, negative_prompt или text_embeddings")
-            
-            # ВЫЗОВ СТАНДАРТНОЙ ФУНКЦИИ encode_prompt
-            text_embeddings, attention_mask, pooled_embeddings = self.encode_prompt(
-                prompt, negative_prompt, device=device, dtype=next(self.unet.parameters()).dtype
-            )
-        else:
-            raise NotImplementedError("Передача text_embeddings напрямую пока не поддерживает передачу маски и пулинга. Используйте prompt/negative_prompt.")
+            with gr.Row():
+                guidance_scale = gr.Slider(
+                    label="Guidance scale",
+                    minimum=0.0,
+                    maximum=10.0,
+                    step=0.5,
+                    value=4.0,
+                )
 
+                num_inference_steps = gr.Slider(
+                    label="Number of inference steps",
+                    minimum=1,
+                    maximum=50,
+                    step=1,
+                    value=40,
+                )
 
-        # 3. ГЕНЕРАЦИЯ ЛАТЕНТОВ (существующая логика)
-        latents = self.generate_latents(
-            text_embeddings=text_embeddings,
-            attention_mask=attention_mask,
-            pooled_embeddings=pooled_embeddings,
-            height=height,
-            width=width,
-            num_inference_steps=num_inference_steps,
-            guidance_scale=guidance_scale,
-            latent_channels=latent_channels,
-            batch_size=batch_size,
-            generator=generator
-        )
+        gr.Examples(examples=examples, inputs=[prompt])
+    
+    gr.on(
+        triggers=[run_button.click, prompt.submit],
+        fn=infer,
+        inputs=[
+            prompt,
+            negative_prompt,
+            seed,
+            randomize_seed,
+            width,
+            height,
+            guidance_scale,
+            num_inference_steps,
+            refine_prompt,
+        ],
+        outputs=[result, seed, prompt], # Добавлен prompt для обновления текста в интерфейсе
+    )
 
-        # 4. ДЕКОДИРОВАНИЕ (существующая логика)
-        images = self.decode_latents(latents, output_type=output_type)
-        
-        # 5. ВОЗВРАТ РЕЗУЛЬТАТА
-        if not return_dict:
-            return images
-        # ИЗМЕНЕНИЕ: Возвращаем уточненный промпт
-        return SdxsPipelineOutput(images=images, refined_prompt=refined_prompt_output)
+if __name__ == "__main__":
+    demo.launch()