Update app.py

app.py

@@ -30,11 +30,9 @@ if not torch.cuda.is_available():
     DESCRIPTION += "\n<p>Running on CPU, This may not work on CPU.</p>"
 
 USE_TORCH_COMPILE = 0
-ENABLE_CPU_OFFLOAD = 0
+ENABLE_CPU_OFFLOAD = 1  # Enable CPU offloading to reduce GPU memory usage
 
 style_list = [
-    
-
     {
         "name": "3840 x 2160",
         "prompt": "hyper-realistic 8K image of {prompt} . ultra-detailed, lifelike, high-resolution, sharp, vibrant colors, photorealistic",
@@ -80,13 +78,11 @@ style_list = [
         "prompt": "professional 3d model {prompt} . octane render, highly detailed, volumetric, dramatic lighting",
         "negative_prompt": "ugly, deformed, noisy, low poly, blurry, painting",
     },
-
     {
         "name": "(No style)",
         "prompt": "{prompt}",
         "negative_prompt": "",
     },
-    
 ]
 
 styles = {k["name"]: (k["prompt"], k["negative_prompt"]) for k in style_list}
@@ -95,10 +91,21 @@ DEFAULT_STYLE_NAME = "3840 x 2160"
 
 def apply_style(style_name: str, positive: str, negative: str = "") -> Tuple[str, str]:
     p, n = styles.get(style_name, styles[DEFAULT_STYLE_NAME])
-    if not negative:
-        negative = ""
     return p.replace("{prompt}", positive), n + negative
 
+# Initialize models once during setup
+prior = StableCascadePriorPipeline.from_pretrained("stabilityai/stable-cascade-prior", variant="bf16", torch_dtype=torch.bfloat16)
+decoder = StableCascadeDecoderPipeline.from_pretrained("stabilityai/stable-cascade", variant="bf16", torch_dtype=torch.float16)
+
+if ENABLE_CPU_OFFLOAD:
+    prior.enable_model_cpu_offload()
+    decoder.enable_model_cpu_offload()
+else:
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    prior.to(device)
+    decoder.to(device)
+    print("Loaded on Device!")
+
 @spaces.GPU(enable_queue=True)
 def stab(
     prompt: str,
@@ -110,7 +117,7 @@ def stab(
     seed: int = 0,
     width: int = 1024,
     height: int = 1024,
-    guidance_scale: float = 3,
+    guidance_scale: float = 6,
     randomize_seed: bool = False,
     progress=gr.Progress(track_tqdm=True),
 ):
@@ -120,10 +127,6 @@ def stab(
         negative_prompt = ""
     prompt, negative_prompt = apply_style(style, prompt, negative_prompt)
 
-    prior = StableCascadePriorPipeline.from_pretrained("stabilityai/stable-cascade-prior", variant="bf16", torch_dtype=torch.bfloat16)
-    decoder = StableCascadeDecoderPipeline.from_pretrained("stabilityai/stable-cascade", variant="bf16", torch_dtype=torch.float16)
-
-    prior.enable_model_cpu_offload()
     prior_output = prior(
         prompt=prompt,
         height=height,
@@ -134,7 +137,6 @@ def stab(
         num_inference_steps=num_inference_steps
     )
 
-    decoder.enable_model_cpu_offload()
     images = decoder(
         image_embeddings=prior_output.image_embeddings.to(torch.float16),
         prompt=prompt,
@@ -149,10 +151,10 @@ def stab(
     return image_paths, seed
 
 examples = [
-    "3d image, cute girl, in the style of Pixar --ar 1:2 --stylize 750, 4K resolution highlights, Sharp focus, octane render, ray tracing, Ultra-High-Definition, 8k, UHD, HDR, (Masterpiece:1. 5), (best quality:1. 5)",
-    "(Pirate ship sailing into a bioluminescence sea with a galaxy in the sky), epic, 4k, ultra,  the space scene with planets and stars, in the style of ethereal escapism, richly colored skies, vibrant worlds --ar 8:5 ",
+    "3d image, cute girl, in the style of Pixar --ar 1:2 --stylize 750, 4K resolution highlights, Sharp focus, octane render, ray tracing, Ultra-High-Definition, 8k, UHD, HDR, (Masterpiece:1.5), (best quality:1.5)",
+    "(Pirate ship sailing into a bioluminescence sea with a galaxy in the sky), epic, 4k, ultra, the space scene with planets and stars, in the style of ethereal escapism, richly colored skies, vibrant worlds --ar 8:5",
     "Thin burger, realistic photo (without tomato or any other ingredient), smoky flavor, 4K resolution highlights every texture, providing an incredible and appetizing visual experience",
-    "A galaxy with blue water, a red star and many planets in one view, in the style of digital fantasy nubelas and cosmos, light black and violet, realistic nubelas paintings, james paick, steve henderson, ue5, cosmic horror --ar 8:5",
+    "A galaxy with blue water, a red star and many planets in one view, in the style of digital fantasy nebulae and cosmos, light black and violet, realistic nebulae paintings, james paick, steve henderson, ue5, cosmic horror --ar 8:5",
     "A dark night sky with thick, dense clouds and stars in the background. The main focus is on one of these large cloud formations that has been stylized to resemble an ancient dragon. There's no moon or other celestial bodies visible in the sky. This scene conveys mystery and magic, with the dark blue glow from distant galaxies adding depth and contrast to the night landscape. --ar 8:5 --v 5.2 --style raw"
 ]
 
@@ -174,7 +176,7 @@ with gr.Blocks(css=css, theme="xiaobaiyuan/theme_brief") as demo:
 
     with gr.Group():
         with gr.Row():
-            prompt = gr.Text(
+            prompt = gr.Textbox(
                 label="Prompt",
                 show_label=False,
                 max_lines=1,
@@ -185,7 +187,7 @@ with gr.Blocks(css=css, theme="xiaobaiyuan/theme_brief") as demo:
         result = gr.Gallery(label="Result", columns=1, preview=True)
     with gr.Accordion("Advanced options", open=False):
         use_negative_prompt = gr.Checkbox(label="Use negative prompt", value=True, visible=True)
-        negative_prompt = gr.Text(
+        negative_prompt = gr.Textbox(
             label="Negative prompt",
             max_lines=1,
             placeholder="Enter a negative prompt",