Update app.py

app.py

@@ -7,7 +7,6 @@ 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")
 
@@ -30,151 +29,22 @@ 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_embeds=prompt_embeds,
-            pooled_prompt_embeds=pooled_prompt_embeds,
+            prompt=prompt,
+            negative_prompt=negative_prompt,
             guidance_scale=guidance_scale,
             num_inference_steps=num_inference_steps,
             width=width,
@@ -205,8 +75,8 @@ with gr.Blocks(css=css) as demo:
             prompt = gr.Text(
                 label="Prompt",
                 show_label=False,
-                max_lines=3,  # Increased to allow longer prompts to be visible
-                placeholder="Enter your prompt (supports long prompts now)",
+                max_lines=1,
+                placeholder="Enter your prompt (keep it under 60 words for best results)",
                 container=False,
             )
 
@@ -218,7 +88,7 @@ with gr.Blocks(css=css) as demo:
 
             negative_prompt = gr.Text(
                 label="Negative prompt",
-                max_lines=2,  # Also increased for consistency
+                max_lines=1,
                 placeholder="Enter a negative prompt",
                 value="nsfw, (low quality, worst quality:1.2), very displeasing, 3d, watermark, signature, ugly, poorly drawn"
             )