Update app.py

app.py

@@ -9,20 +9,38 @@ import spaces
 import torch
 from diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler
 
-
 DESCRIPTIONx = """## STABLE HAMSTER
 
+Drop your best results in the community: [rb.gy/klkbs7](http://rb.gy/klkbs7)
 """
 
 
-# Use environment variables for flexibility
+style_list = [
+    {
+        "name": "3840 x 2160",
+        "prompt": "hyper-realistic 8K image of {prompt}. ultra-detailed, lifelike, high-resolution, sharp, vibrant colors, photorealistic",
+        "negative_prompt": "cartoonish, low resolution, blurry, simplistic, abstract, deformed, ugly",
+    },
+    {
+        "name": "2560 x 1440",
+        "prompt": "hyper-realistic 4K image of {prompt}. ultra-detailed, lifelike, high-resolution, sharp, vibrant colors, photorealistic",
+        "negative_prompt": "cartoonish, low resolution, blurry, simplistic, abstract, deformed, ugly",
+    },
+    {
+        "name": "3D Model",
+        "prompt": "professional 3d model {prompt}. octane render, highly detailed, volumetric, dramatic lighting",
+        "negative_prompt": "ugly, deformed, noisy, low poly, blurry, painting",
+    },
+]
+
+#User -- Env -- .os -- Mode_Repo
+
 MODEL_ID = os.getenv("MODEL_REPO")
 MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "4096"))
 USE_TORCH_COMPILE = os.getenv("USE_TORCH_COMPILE", "0") == "1"
 ENABLE_CPU_OFFLOAD = os.getenv("ENABLE_CPU_OFFLOAD", "0") == "1"
-BATCH_SIZE = int(os.getenv("BATCH_SIZE", "1"))  # Allow generating multiple images at once
+BATCH_SIZE = int(os.getenv("BATCH_SIZE", "1"))  
 
-# Determine device and load model outside of function for efficiency
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 pipe = StableDiffusionXLPipeline.from_pretrained(
     MODEL_ID,
@@ -32,16 +50,26 @@ pipe = StableDiffusionXLPipeline.from_pretrained(
 ).to(device)
 pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
 
-# Torch compile for potential speedup (experimental)
+# potential speedup
 if USE_TORCH_COMPILE:
     pipe.compile()
 
-# CPU offloading for larger RAM capacity (experimental)
+# CPU offloading for Bigger RAM
 if ENABLE_CPU_OFFLOAD:
     pipe.enable_model_cpu_offload()
 
 MAX_SEED = np.iinfo(np.int32).max
 
+styles = {k["name"]: (k["prompt"], k["negative_prompt"]) for k in style_list}
+STYLE_NAMES = list(styles.keys())
+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
+
 def save_image(img):
     unique_name = str(uuid.uuid4()) + ".png"
     img.save(unique_name)
@@ -65,8 +93,10 @@ def generate(
     randomize_seed: bool = False,
     use_resolution_binning: bool = True, 
     num_images: int = 1,  # Number of images to generate
+    style: str = DEFAULT_STYLE_NAME,
     progress=gr.Progress(track_tqdm=True),
 ):
+    prompt, negative_prompt = apply_style(style, prompt, negative_prompt)
     seed = int(randomize_seed_fn(seed, randomize_seed))
     generator = torch.Generator(device=device).manual_seed(seed)
 
@@ -184,6 +214,14 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
                 step=1,
                 value=8,
             )
+        style_selection = gr.Radio(
+            show_label=True,
+            container=True,
+            interactive=True,
+            choices=STYLE_NAMES,
+            value=DEFAULT_STYLE_NAME,
+            label="Image Style",
+        )
 
     gr.Examples(
         examples=examples,
@@ -215,7 +253,8 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
             guidance_scale,
             num_inference_steps,
             randomize_seed,
-            num_images
+            num_images,
+            style_selection,
         ],
         outputs=[result, seed],
         api_name="run",