Update app.py

app.py

@@ -21,14 +21,11 @@ from transformers import (
     GemmaTokenizer
 )
 
+default_system_prompt = "You are an assistant designed to generate superior images with the superior degree of image-text alignment based on textual prompts or user prompts."
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model_repo_id = "Alpha-VLLM/Lumina-Image-2.0"
 transformer_repo_id = "benjamin-paine/Lumina-Image-2.0" # Temporarily fixed, change when main repo gets updated
-
-if torch.cuda.is_available():
-    torch_dtype = torch.bfloat16
-else:
-    torch_dtype = torch.float32
+torch_dtype = torch.float32
 
 ###
 transformer = Lumina2Transformer2DModel.from_pretrained(transformer_repo_id, subfolder="transformer")
@@ -60,6 +57,9 @@ def infer(
     height=1024,
     guidance_scale=4.0,
     num_inference_steps=30,
+    system_prompt=default_system_prompt,
+    cfg_normalization=True,
+    cfg_trunc_ratio=1.0,
     progress=gr.Progress(track_tqdm=True),
 ):
     if randomize_seed:
@@ -75,6 +75,9 @@ def infer(
         width=width,
         height=height,
         generator=generator,
+        system_prompt=system_prompt,
+        cfg_normalization=cfg_normalization,
+        cfg_trunc_ratio=cfg_trunc_ratio
     ).images[0]
 
     return image, seed
@@ -98,7 +101,7 @@ with gr.Blocks(css=css) as demo:
             prompt = gr.Text(
                 label="Prompt",
                 show_label=False,
-                max_lines=1,
+                max_lines=4,
                 placeholder="Enter your prompt",
                 container=False,
             )
@@ -108,21 +111,29 @@ with gr.Blocks(css=css) as demo:
         result = gr.Image(label="Result", show_label=False)
 
         with gr.Accordion("Advanced Settings", open=False):
-            negative_prompt = gr.Text(
-                label="Negative prompt",
-                max_lines=1,
-                placeholder="Enter a negative prompt",
-            )
+            with gr.Row():
+                system_prompt = gr.Text(
+                    label="System Prompt",
+                    max_lines=4,
+                    value=default_system_prompt
+                )
+                
+                negative_prompt = gr.Text(
+                    label="Negative prompt",
+                    max_lines=4,
+                    placeholder="Enter a negative prompt",
+                )
 
-            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():
+                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(
@@ -158,6 +169,22 @@ with gr.Blocks(css=css) as demo:
                     value=30, 
                 )
 
+            with gr.Row():
+                cfg_normalization = gr.Checkbox(
+                    label="CFG Normalization",
+                    value=True
+                )
+                cfg_trunc_ratio = gr.Slider(
+                    label="CFG Truncation Ratio",
+                    minimum=0.0,
+                    maximum=1.0,
+                    step=0.01,
+                    value=1.0
+                )
+
+            with gr.Row():
+                
+        
         gr.Examples(examples=examples, inputs=[prompt], outputs=[result, seed], fn=infer, cache_examples=True, cache_mode="lazy")
 
     gr.on(
@@ -172,6 +199,9 @@ with gr.Blocks(css=css) as demo:
             height,
             guidance_scale,
             num_inference_steps,
+            system_prompt,
+            cfg_normalization,
+            cfg_trunc_ratio,
         ],
         outputs=[result, seed],
     )