Update app.py

app.py

@@ -1,24 +1,32 @@
 import gradio as gr
 import numpy as np
 import random
-
 import spaces
 from diffusers import DiffusionPipeline
 import torch
 
+# Set device and model parameters
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model_repo_id = "stabilityai/stable-diffusion-3.5-large"
 torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
 
-if torch.cuda.is_available():
-    torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
-
+# Load the Stable Diffusion pipeline and move it to the appropriate device.
 pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
 pipe = pipe.to(device)
 
+# Maximum values as defined in your original code
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
 
+# Define a helper function to truncate text to a maximum of 77 tokens.
+def truncate_text(text, max_tokens=77):
+    if text.strip() == "":
+        return text
+    # Use the pipeline's tokenizer (CLIP tokenizer)
+    tokens = pipe.tokenizer(text, truncation=True, max_length=max_tokens, add_special_tokens=True)
+    truncated_text = pipe.tokenizer.decode(tokens["input_ids"], skip_special_tokens=True)
+    return truncated_text
+
 @spaces.GPU(duration=65)
 def infer(
     prompt,
@@ -31,11 +39,16 @@ def infer(
     num_inference_steps=40,
     progress=gr.Progress(track_tqdm=True),
 ):
+    # Optionally randomize seed
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
+    generator = torch.Generator(device=device).manual_seed(seed)
 
-    generator = torch.Generator().manual_seed(seed)
+    # Truncate both prompt and negative prompt to 77 tokens.
+    prompt = truncate_text(prompt, max_tokens=77)
+    negative_prompt = truncate_text(negative_prompt, max_tokens=77) if negative_prompt.strip() else ""
 
+    # Explicitly set pad_token_id to eos_token_id for open-end generation.
     image = pipe(
         prompt=prompt,
         negative_prompt=negative_prompt,
@@ -44,13 +57,14 @@ def infer(
         width=width,
         height=height,
         generator=generator,
+        pad_token_id=pipe.tokenizer.eos_token_id,
     ).images[0]
 
     return image, seed
 
-
+# Example prompt for testing
 examples = [
-        "A capybara wearing a suit holding a sign that reads Hello World",
+    "A capybara wearing a suit holding a sign that reads Hello World",
 ]
 
 css = """
@@ -63,7 +77,7 @@ css = """
 with gr.Blocks(css=css) as demo:
     with gr.Column(elem_id="col-container"):
         gr.Markdown(" # [Stable Diffusion 3.5 Large (8B)](https://huggingface.co/stabilityai/stable-diffusion-3.5-large)")
-        gr.Markdown("[Learn more](https://stability.ai/news/introducing-stable-diffusion-3-5) about the Stable Diffusion 3.5 series. Try on [Stability AI API](https://platform.stability.ai/docs/api-reference#tag/Generate/paths/~1v2beta~1stable-image~1generate~1sd3/post), or [download model](https://huggingface.co/stabilityai/stable-diffusion-3.5-large) to run locally with ComfyUI or diffusers.")
+        gr.Markdown("[Learn more](https://stability.ai/news/introducing-stable-diffusion-3.5) about the Stable Diffusion 3.5 series. Try on [Stability AI API](https://platform.stability.ai/docs/api-reference#tag/Generate/paths/~1v2beta~1stable-image~1generate~1sd3/post), or [download model](https://huggingface.co/stabilityai/stable-diffusion-3.5-large) to run locally with ComfyUI or diffusers.")
         with gr.Row():
             prompt = gr.Text(
                 label="Prompt",
@@ -72,11 +86,8 @@ with gr.Blocks(css=css) as demo:
                 placeholder="Enter your prompt",
                 container=False,
             )
-
             run_button = gr.Button("Run", scale=0, variant="primary")
-
         result = gr.Image(label="Result", show_label=False)
-
         with gr.Accordion("Advanced Settings", open=False):
             negative_prompt = gr.Text(
                 label="Negative prompt",
@@ -84,7 +95,6 @@ with gr.Blocks(css=css) as demo:
                 placeholder="Enter a negative prompt",
                 visible=False,
             )
-
             seed = gr.Slider(
                 label="Seed",
                 minimum=0,
@@ -92,9 +102,7 @@ with gr.Blocks(css=css) as demo:
                 step=1,
                 value=0,
             )
-
             randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-
             with gr.Row():
                 width = gr.Slider(
                     label="Width",
@@ -103,7 +111,6 @@ with gr.Blocks(css=css) as demo:
                     step=32,
                     value=1024, 
                 )
-
                 height = gr.Slider(
                     label="Height",
                     minimum=512,
@@ -111,7 +118,6 @@ with gr.Blocks(css=css) as demo:
                     step=32,
                     value=1024,
                 )
-
             with gr.Row():
                 guidance_scale = gr.Slider(
                     label="Guidance scale",
@@ -120,7 +126,6 @@ with gr.Blocks(css=css) as demo:
                     step=0.1,
                     value=4.5,
                 )
-
                 num_inference_steps = gr.Slider(
                     label="Number of inference steps",
                     minimum=1,
@@ -128,7 +133,6 @@ with gr.Blocks(css=css) as demo:
                     step=1,
                     value=40, 
                 )
-
         gr.Examples(examples=examples, inputs=[prompt], outputs=[result, seed], fn=infer, cache_examples=True, cache_mode="lazy")
     gr.on(
         triggers=[run_button.click, prompt.submit],
@@ -148,3 +152,4 @@ with gr.Blocks(css=css) as demo:
 
 if __name__ == "__main__":
     demo.launch()
+