Update app.py

app.py

@@ -7,6 +7,7 @@ import random
 from diffusers import StableDiffusionXLPipeline
 from diffusers import EulerAncestralDiscreteScheduler
 import torch
+import re
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
@@ -29,22 +30,151 @@ pipe.unet.to(torch.float16)
 
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1216
+
+# Function to parse weighted prompts
+def parse_prompt_attention(text):
+    """
+    Parses a prompt with attention weights
+    Examples:
+        "a (red:1.5) dress" -> weight "red" with 1.5
+        "a ((blue)) sky" -> weight "blue" with 2.0
+    """
+    re_attention = r'\((\()?([^:]+)(\))?(?::([\d\.]+))?\)'
+    res = []
+    
+    for match in re.finditer(re_attention, text):
+        double_paren, content, _, weight = match.groups()
+        weight = float(weight) if weight is not None else 1.0
+        if double_paren:
+            weight = weight * 1.1  # Optional: make (()) slightly higher than ()
+        
+        res.append((match.start(), match.end(), content, weight))
+    
+    return res
+
+# Function to process prompts with attention weights
+def get_weighted_text_embeddings(
+    pipe,
+    prompt,
+    negative_prompt=None
+):
+    """
+    Processes prompts with attention weights and handles long prompts
+    by chunking, applying weights, and combining embeddings
+    """
+    max_length = pipe.tokenizer.model_max_length
+    
+    # Process the input prompt with attention weights
+    parsed_attention = parse_prompt_attention(prompt)
+    
+    # Handle long prompts by chunking them appropriately
+    if len(prompt.split()) > 60:  # Rough estimate of potentially exceeding token limit
+        print(f"Long prompt detected. Will process in chunks.")
+        
+        # Remove and store attention weights for processing
+        text_chunks = []
+        current_length = 0
+        current_chunk = ""
+        
+        words = prompt.split()
+        
+        for word in words:
+            if current_length + len(word.split()) + 1 > 60:  # Start a new chunk
+                text_chunks.append(current_chunk)
+                current_chunk = word
+                current_length = len(word.split())
+            else:
+                if current_chunk:
+                    current_chunk += " " + word
+                else:
+                    current_chunk = word
+                current_length += len(word.split())
+        
+        if current_chunk:
+            text_chunks.append(current_chunk)
+            
+        print(f"Split into {len(text_chunks)} chunks: {text_chunks}")
+        
+        # Process each chunk with the tokenizer and get embedding
+        prompt_embeds_list = []
+        pooled_prompt_embeds_list = []
+        
+        for text_chunk in text_chunks:
+            text_input = pipe.tokenizer(
+                text_chunk,
+                padding="max_length",
+                max_length=pipe.tokenizer.model_max_length,
+                truncation=True,
+                return_tensors="pt",
+            )
+            text_input = text_input.to(device)
+            
+            # Get text embeddings for both encoders
+            prompt_embeds = pipe.text_encoder(text_input.input_ids)[0]
+            pooled_prompt_embeds = pipe.text_encoder_2(text_input.input_ids)[0]
+            
+            prompt_embeds_list.append(prompt_embeds)
+            pooled_prompt_embeds_list.append(pooled_prompt_embeds)
+        
+        # Average the embeddings from all chunks (alternatively could use max pooling or other methods)
+        prompt_embeds = torch.stack(prompt_embeds_list).mean(dim=0)
+        pooled_prompt_embeds = torch.stack(pooled_prompt_embeds_list).mean(dim=0)
+        
+    else:
+        # For shorter prompts, just use the standard pipeline processing
+        text_input = pipe.tokenizer(
+            prompt,
+            padding="max_length",
+            max_length=pipe.tokenizer.model_max_length,
+            truncation=True,
+            return_tensors="pt",
+        )
+        text_input = text_input.to(device)
+        
+        prompt_embeds = pipe.text_encoder(text_input.input_ids)[0]
+        pooled_prompt_embeds = pipe.text_encoder_2(text_input.input_ids)[0]
+    
+    # Process negative prompt if provided
+    if negative_prompt is None:
+        negative_prompt = ""
+        
+    uncond_input = pipe.tokenizer(
+        negative_prompt,
+        padding="max_length",
+        max_length=pipe.tokenizer.model_max_length,
+        truncation=True,
+        return_tensors="pt",
+    )
+    uncond_input = uncond_input.to(device)
+    negative_prompt_embeds = pipe.text_encoder(uncond_input.input_ids)[0]
+    negative_pooled_prompt_embeds = pipe.text_encoder_2(uncond_input.input_ids)[0]
+    
+    # Combine positive and negative embeddings
+    prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])
+    pooled_prompt_embeds = torch.cat([negative_pooled_prompt_embeds, pooled_prompt_embeds])
     
+    return prompt_embeds, pooled_prompt_embeds
+
+# Customized version of the generation function
 @spaces.GPU
 def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps):
-    # Check and truncate prompt if too long (CLIP can only handle 77 tokens)
-    if len(prompt.split()) > 60:  # Rough estimate to avoid exceeding token limit
-        print("Warning: Prompt may be too long and will be truncated by the model")
-        
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
 
     generator = torch.Generator(device=device).manual_seed(seed)
     
     try:
+        # Get embeddings with special handling for long prompts
+        prompt_embeds, pooled_prompt_embeds = get_weighted_text_embeddings(
+            pipe,
+            prompt,
+            negative_prompt
+        )
+        
+        # Use the custom embeddings to generate the image
         output_image = pipe(
-            prompt=prompt,
-            negative_prompt=negative_prompt,
+            prompt_embeds=prompt_embeds,
+            pooled_prompt_embeds=pooled_prompt_embeds,
             guidance_scale=guidance_scale,
             num_inference_steps=num_inference_steps,
             width=width,
@@ -75,8 +205,8 @@ with gr.Blocks(css=css) as demo:
             prompt = gr.Text(
                 label="Prompt",
                 show_label=False,
-                max_lines=1,
-                placeholder="Enter your prompt (keep it under 60 words for best results)",
+                max_lines=3,  # Increased to allow longer prompts to be visible
+                placeholder="Enter your prompt (supports long prompts now)",
                 container=False,
             )
 
@@ -88,7 +218,7 @@ with gr.Blocks(css=css) as demo:
 
             negative_prompt = gr.Text(
                 label="Negative prompt",
-                max_lines=1,
+                max_lines=2,  # Also increased for consistency
                 placeholder="Enter a negative prompt",
                 value="nsfw, (low quality, worst quality:1.2), very displeasing, 3d, watermark, signature, ugly, poorly drawn"
             )