Update app.py

app.py

@@ -1,89 +1,145 @@
 import os
-import random
 import uuid
+import random
+from typing import Tuple
+
 import gradio as gr
 import numpy as np
 from PIL import Image
-import spaces
 import torch
+import spaces
 from diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler
-from typing import Tuple
 
+
+# -----------------------
+# Style handling
+# -----------------------
 def apply_style(style_name: str, positive: str, negative: str = "") -> Tuple[str, str]:
     styles = {
         "3840 x 2160": (
-            "hyper-realistic image of {prompt}. lifelike, authentic, natural colors, true-to-life details, landscape image, realistic lighting, immersive, highly detailed",
+            "hyper-realistic image of {prompt}. lifelike, authentic, natural colors, "
+            "true-to-life details, landscape image, realistic lighting, immersive, highly detailed",
             "unrealistic, low resolution, artificial, over-saturated, distorted, fake",
         ),
         "Style Zero": ("{prompt}", ""),
     }
     DEFAULT_STYLE_NAME = "3840 x 2160"
-
     p, n = styles.get(style_name, styles[DEFAULT_STYLE_NAME])
-    return p.replace("{prompt}", positive), n + negative
+    return p.replace("{prompt}", positive), (n + (" " + negative if negative else "")).strip()
+
 
-def load_and_prepare_model():
+# -----------------------
+# Model loader
+# -----------------------
+def load_and_prepare_model() -> StableDiffusionXLPipeline:
     model_id = "SG161222/RealVisXL_V5.0_Lightning"
+    use_cuda = torch.cuda.is_available()
+    dtype = torch.float16 if use_cuda else torch.float32
+    device = torch.device("cuda:0" if use_cuda else "cpu")
+
     pipe = StableDiffusionXLPipeline.from_pretrained(
         model_id,
-        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+        torch_dtype=dtype,
         use_safetensors=True,
         add_watermarker=False,
-    ).to(torch.device("cuda:0" if torch.cuda.is_available() else "cpu"))
+    )
+
+    # Use a stable, fast scheduler
     pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
+
+    # Memory/throughput optimizations
+    try:
+        pipe.enable_xformers_memory_efficient_attention()
+    except Exception:
+        # xformers may be missing on CPU or some build types — that's fine
+        pass
+
+    # Optional, safe perf knobs on CUDA
+    if use_cuda:
+        torch.backends.cuda.matmul.allow_tf32 = True
+        torch.set_grad_enabled(False)
+
+    pipe = pipe.to(device)
     return pipe
 
+
+# Global model (loaded once per Space instance)
 model = load_and_prepare_model()
 
+
+# -----------------------
+# Utils
+# -----------------------
 def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
     if randomize_seed:
         seed = random.randint(0, np.iinfo(np.int32).max)
-    return seed
+    return int(seed)
+
 
-def save_image(img):
-    unique_name = str(uuid.uuid4()) + ".png"
+def save_image(img: Image.Image) -> str:
+    # Save to the working dir so HF can expose it as an artifact
+    unique_name = f"{uuid.uuid4().hex}.png"
     img.save(unique_name)
     return unique_name
 
+
+# -----------------------
+# Generation
+# -----------------------
 @spaces.GPU(duration=60, enable_queue=True)
 def generate(
-        prompt: str,
-        seed: int = 1,
-        width: int = 1024,
-        height: int = 1024,
-        guidance_scale: float = 3,
-        num_inference_steps: int = 25,
-        randomize_seed: bool = False,
+    prompt: str,
+    seed: int = 1,
+    width: int = 1024,
+    height: int = 1024,
+    guidance_scale: float = 3.0,
+    num_inference_steps: int = 25,
+    randomize_seed: bool = False,
 ):
     global model
 
-    seed = int(randomize_seed_fn(seed, randomize_seed))
+    # Guardrails
+    if not prompt or not prompt.strip():
+        raise gr.Error("Please enter a prompt.")
+
+    # SDXL prefers dims divisible by 8
+    width = int(max(256, (width // 8) * 8))
+    height = int(max(256, (height // 8) * 8))
+
+    seed = randomize_seed_fn(seed, randomize_seed)
     generator = torch.Generator(device=model.device).manual_seed(seed)
 
     positive_prompt, negative_prompt = apply_style("3840 x 2160", prompt)
 
-    options = {
-        "prompt": [positive_prompt],
-        "negative_prompt": [negative_prompt],
-        "width": width,
-        "height": height,
-        "guidance_scale": guidance_scale,
-        "num_inference_steps": num_inference_steps,
-        "generator": generator,
-        "output_type": "pil",
-    }
-
-    images = model(**options).images
-    image_path = save_image(images[0])  # Saving the first generated image
+    # NOTE: pass strings (not one-element lists)
+    images = model(
+        prompt=positive_prompt,
+        negative_prompt=negative_prompt,
+        width=width,
+        height=height,
+        guidance_scale=float(guidance_scale),
+        num_inference_steps=int(num_inference_steps),
+        generator=generator,
+        output_type="pil",
+    ).images
+
+    image_path = save_image(images[0])
     return image_path
 
+
+# -----------------------
+# UI
+# -----------------------
 with gr.Blocks(theme="soft") as demo:
-    # Block for "SNAPSCRIBE" centered at the top
     with gr.Row():
         with gr.Column(scale=12, elem_id="title_block"):
-            gr.Markdown("<h1 style='text-align:center; color:white; font-weight:bold; text-decoration:underline;'>SNAPSCRIBE</h1>")
-            gr.Markdown("<h2 style='text-align:center; color:white; font-weight:bold; text-decoration:underline;'>Developed using RealVisXL_V5.0_Lightning model with ❤ by Aklavya</h2>")
-    
+            gr.Markdown(
+                "<h1 style='text-align:center; color:white; font-weight:bold; text-decoration:underline;'>SNAPSCRIBE</h1>"
+            )
+            gr.Markdown(
+                "<h2 style='text-align:center; color:white; font-weight:bold; text-decoration:underline;'>Developed using RealVisXL_V5.0_Lightning model with ❤ by Aklavya</h2>"
+            )
+
     with gr.Row():
         with gr.Column(scale=3):
             prompt = gr.Textbox(
@@ -91,9 +147,15 @@ with gr.Blocks(theme="soft") as demo:
                 placeholder="Describe the image you want to create",
                 lines=2,
             )
-            run_button = gr.Button("Generate Image")
+            seed = gr.Number(value=1, label="Seed", precision=0)
+            randomize_seed = gr.Checkbox(value=True, label="Randomize Seed")
+            width = gr.Slider(512, 1536, value=1024, step=8, label="Width")
+            height = gr.Slider(512, 1536, value=1024, step=8, label="Height")
+            guidance_scale = gr.Slider(1.0, 10.0, value=3.0, step=0.5, label="Guidance Scale")
+            steps = gr.Slider(10, 35, value=25, step=1, label="Inference Steps")
+
+            run_button = gr.Button("Generate Image", variant="primary")
 
-            # Example prompts box
             example_prompts_text = (
                 "Dew-covered spider web in morning sunlight, with blurred greenery\n"
                 "--------------------------------------------\n"
@@ -104,9 +166,9 @@ with gr.Blocks(theme="soft") as demo:
                 "Autumn forest with golden leaves, sunlight through trees, and a breeze"
             )
 
-            example_prompts = gr.Textbox(
+            gr.Textbox(
                 value=example_prompts_text,
-                lines=5,
+                lines=8,
                 label="Sample Inputs",
                 interactive=False,
             )
@@ -116,13 +178,14 @@ with gr.Blocks(theme="soft") as demo:
                 label="Generated Image",
                 type="filepath",
                 elem_id="output_image",
-                height=650  # Increased the height by 100%
             )
 
     run_button.click(
         fn=generate,
-        inputs=[prompt],
+        inputs=[prompt, seed, width, height, guidance_scale, steps, randomize_seed],
         outputs=[result_image],
+        api_name="generate",
     )
 
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()