Update app.py

app.py

@@ -36,23 +36,10 @@ sdxl_pipe = DiffusionPipeline.from_pretrained(
 )
 sdxl_pipe.to(device)
 
-# superprompt-v1
-tokenizer = T5Tokenizer.from_pretrained("roborovski/superprompt-v1")
-model = T5ForConditionalGeneration.from_pretrained(
-    "roborovski/superprompt-v1", device_map="auto", torch_dtype="auto"
-)
-model.to(device)
-
-# toggle visibility the enhanced prompt output
-def update_visibility(enhance_prompt):
-    return gr.update(visible=enhance_prompt)
-
-
 # Define the image generation function for the Arena tab
 @spaces.GPU(duration=80)
 def generate_arena_images(
     prompt,
-    enhance_prompt,
     negative_prompt,
     num_inference_steps,
     height,
@@ -67,23 +54,6 @@ def generate_arena_images(
     if seed == 0:
         seed = random.randint(1, 2**32 - 1)
 
-    if enhance_prompt:
-        transformers.set_seed(seed)
-
-        input_text = f"Expand the following prompt to add more detail: {prompt}"
-        input_ids = tokenizer(input_text, return_tensors="pt").input_ids.to(device)
-
-        outputs = model.generate(
-            input_ids,
-            max_new_tokens=512,
-            repetition_penalty=1.2,
-            do_sample=True,
-            temperature=0.7,
-            top_p=1,
-            top_k=50,
-        )
-        prompt = tokenizer.decode(outputs[0], skip_special_tokens=True)
-
     generator = torch.Generator().manual_seed(seed)
 
     # Generate images for both models
@@ -112,7 +82,7 @@ def generate_arena_images(
         generator,
     )
 
-    return images_1, images_2, prompt
+    return images_1, images_2
 
 
 # Helper function to generate images for a single model
@@ -155,7 +125,6 @@ def generate_single_image(
 @spaces.GPU(duration=80)
 def generate_individual_image(
     prompt,
-    enhance_prompt,
     negative_prompt,
     num_inference_steps,
     height,
@@ -169,23 +138,6 @@ def generate_individual_image(
     if seed == 0:
         seed = random.randint(1, 2**32 - 1)
 
-    if enhance_prompt:
-        transformers.set_seed(seed)
-
-        input_text = f"Expand the following prompt to add more detail: {prompt}"
-        input_ids = tokenizer(input_text, return_tensors="pt").input_ids.to(device)
-
-        outputs = model.generate(
-            input_ids,
-            max_new_tokens=512,
-            repetition_penalty=1.2,
-            do_sample=True,
-            temperature=0.7,
-            top_p=1,
-            top_k=50,
-        )
-        prompt = tokenizer.decode(outputs[0], skip_special_tokens=True)
-
     generator = torch.Generator().manual_seed(seed)
 
     output = generate_single_image(
@@ -201,14 +153,14 @@ def generate_individual_image(
         generator,
     )
 
-    return output, prompt
+    return output
 
 
 # Create the Gradio interface
 examples = [
-    ["A white car racing fast to the moon.", True],
-    ["A woman in a red dress singing on top of a building.", True],
-    ["An astronaut on mars in a futuristic cyborg suit.", True],
+    ["A white car racing fast to the moon."],
+    ["A woman in a red dress singing on top of a building."],
+    ["An astronaut on mars in a futuristic cyborg suit."],
 ]
 
 css = """
@@ -240,9 +192,6 @@ with gr.Blocks(css=css) as demo:
                         info="Describe the image you want",
                         placeholder="A cat...",
                     )
-                    enhance_prompt = gr.Checkbox(
-                        label="Prompt Enhancement with SuperPrompt-v1", value=True
-                    )
                     model_choice_1 = gr.Dropdown(
                         label="Stable Diffusion Model 1",
                         choices=["sd3 medium", "sd2.1", "sdxl"],
@@ -256,14 +205,6 @@ with gr.Blocks(css=css) as demo:
                     run_button = gr.Button("Run")
                 result_1 = gr.Gallery(label="Generated Images (Model 1)", elem_id="gallery_1")
                 result_2 = gr.Gallery(label="Generated Images (Model 2)", elem_id="gallery_2")
-                better_prompt = gr.Textbox(
-                    label="Enhanced Prompt",
-                    info="The output of your enhanced prompt used for the Image Generation",
-                    visible=True,
-                )
-                enhance_prompt.change(
-                    fn=update_visibility, inputs=enhance_prompt, outputs=better_prompt
-                )
             with gr.Accordion("Advanced options", open=False):
                 with gr.Row():
                     negative_prompt = gr.Textbox(
@@ -326,8 +267,8 @@ with gr.Blocks(css=css) as demo:
 
             gr.Examples(
                 examples=examples,
-                inputs=[prompt, enhance_prompt],
-                outputs=[result_1, result_2, better_prompt],
+                inputs=[prompt],
+                outputs=[result_1, result_2],
                 fn=generate_arena_images,
             )
 
@@ -339,7 +280,6 @@ with gr.Blocks(css=css) as demo:
                 fn=generate_arena_images,
                 inputs=[
                     prompt,
-                    enhance_prompt,
                     negative_prompt,
                     num_inference_steps,
                     width,
@@ -350,7 +290,7 @@ with gr.Blocks(css=css) as demo:
                     model_choice_1,
                     model_choice_2,
                 ],
-                outputs=[result_1, result_2, better_prompt],
+                outputs=[result_1, result_2],
             )
 
         with gr.TabItem("Individual"):
@@ -361,9 +301,6 @@ with gr.Blocks(css=css) as demo:
                         info="Describe the image you want",
                         placeholder="A cat...",
                     )
-                    enhance_prompt = gr.Checkbox(
-                        label="Prompt Enhancement with SuperPrompt-v1", value=True
-                    )
                     model_choice = gr.Dropdown(
                         label="Stable Diffusion Model",
                         choices=["sd3 medium", "sd2.1", "sdxl"],
@@ -371,14 +308,6 @@ with gr.Blocks(css=css) as demo:
                     )
                     run_button = gr.Button("Run")
                 result = gr.Gallery(label="Generated AI Images", elem_id="gallery")
-                better_prompt = gr.Textbox(
-                    label="Enhanced Prompt",
-                    info="The output of your enhanced prompt used for the Image Generation",
-                    visible=True,
-                )
-                enhance_prompt.change(
-                    fn=update_visibility, inputs=enhance_prompt, outputs=better_prompt
-                )
             with gr.Accordion("Advanced options", open=False):
                 with gr.Row():
                     negative_prompt = gr.Textbox(
@@ -441,8 +370,8 @@ with gr.Blocks(css=css) as demo:
 
             gr.Examples(
                 examples=examples,
-                inputs=[prompt, enhance_prompt],
-                outputs=[result, better_prompt],
+                inputs=[prompt],
+                outputs=[result],
                 fn=generate_individual_image,
             )
 
@@ -454,7 +383,6 @@ with gr.Blocks(css=css) as demo:
                 fn=generate_individual_image,
                 inputs=[
                     prompt,
-                    enhance_prompt,
                     negative_prompt,
                     num_inference_steps,
                     width,
@@ -464,7 +392,7 @@ with gr.Blocks(css=css) as demo:
                     num_images_per_prompt,
                     model_choice,
                 ],
-                outputs=[result, better_prompt],
+                outputs=[result],
             )
 
 demo.queue().launch(share=False)