Update app.py

app.py

@@ -8,37 +8,15 @@ from PIL import Image
 import spaces
 import torch
 from diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler
-from typing import Tuple  # Added this import
+
 
 DESCRIPTIONx = """## STABLE HAMSTER
+
 """
 
-# Use environment variables for flexibility
-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
 
-# 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,
-    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-    use_safetensors=True,
-    add_watermarker=False,
-).to(device)
-pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
 
-# Torch compile for potential speedup (experimental)
-if USE_TORCH_COMPILE:
-    pipe.compile()
 
-# CPU offloading for larger RAM capacity (experimental)
-if ENABLE_CPU_OFFLOAD:
-    pipe.enable_model_cpu_offload()
-
-MAX_SEED = np.iinfo(np.int32).max
 
 style_list = [
     {
@@ -51,11 +29,34 @@ style_list = [
         "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": "Photo",
+        "prompt": "cinematic photo {prompt}. 35mm photograph, film, bokeh, professional, 4k, highly detailed",
+        "negative_prompt": "drawing, painting, crayon, sketch, graphite, impressionist, noisy, blurry, soft, deformed, ugly",
+    },   
+
+    {
+        "name": "Cinematic",
+        "prompt": "cinematic still {prompt}. emotional, harmonious, vignette, highly detailed, high budget, bokeh, cinemascope, moody, epic, gorgeous, film grain, grainy",
+        "negative_prompt": "anime, cartoon, graphic, text, painting, crayon, graphite, abstract, glitch, deformed, mutated, ugly, disfigured",
+    },
+
+    {
+        "name": "Anime",
+        "prompt": "anime artwork {prompt}. anime style, key visual, vibrant, studio anime, highly detailed",
+        "negative_prompt": "photo, deformed, black and white, realism, disfigured, low contrast",
+    },
     {
         "name": "3D Model",
         "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}
@@ -68,6 +69,37 @@ def apply_style(style_name: str, positive: str, negative: str = "") -> Tuple[str
         negative = ""
     return p.replace("{prompt}", positive), n + negative
 
+
+
+
+
+# Use environment variables for flexibility
+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
+
+# 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,
+    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+    use_safetensors=True,
+    add_watermarker=False,
+).to(device)
+pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
+
+# Torch compile for potential speedup (experimental)
+if USE_TORCH_COMPILE:
+    pipe.compile()
+
+# CPU offloading for larger RAM capacity (experimental)
+if ENABLE_CPU_OFFLOAD:
+    pipe.enable_model_cpu_offload()
+
+MAX_SEED = np.iinfo(np.int32).max
+
 def save_image(img):
     unique_name = str(uuid.uuid4()) + ".png"
     img.save(unique_name)
@@ -91,21 +123,19 @@ 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)
 
     # Improved options handling
     options = {
-        "prompt": [prompt] * int(num_images),
-        "negative_prompt": [negative_prompt] * int(num_images) if use_negative_prompt else None,
-        "width": int(width),
-        "height": int(height),
-        "guidance_scale": float(guidance_scale),
-        "num_inference_steps": int(num_inference_steps),
+        "prompt": [prompt] * num_images,
+        "negative_prompt": [negative_prompt] * num_images if use_negative_prompt else None,
+        "width": width,
+        "height": height,
+        "guidance_scale": guidance_scale,
+        "num_inference_steps": num_inference_steps,
         "generator": generator,
         "output_type": "pil",
     }
@@ -116,7 +146,7 @@ def generate(
 
     # Generate images potentially in batches
     images = []
-    for i in range(0, int(num_images), BATCH_SIZE):
+    for i in range(0, num_images, BATCH_SIZE):
         batch_options = options.copy()
         batch_options["prompt"] = options["prompt"][i:i+BATCH_SIZE]
         if "negative_prompt" in batch_options:
@@ -212,6 +242,8 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
                 step=1,
                 value=8,
             )
+
+    with gr.Row(visible=True):
         style_selection = gr.Radio(
             show_label=True,
             container=True,
@@ -220,11 +252,13 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
             value=DEFAULT_STYLE_NAME,
             label="Image Style",
         )
-
+            
     gr.Examples(
         examples=examples,
         inputs=prompt,
-        cache_examples=False
+        outputs=[result, seed],
+        fn=generate,
+        cache_examples=CACHE_EXAMPLES,
     )
 
     use_negative_prompt.change(
@@ -245,14 +279,14 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
             prompt,
             negative_prompt,
             use_negative_prompt,
+            style_selection,
             seed,
             width,
             height,
             guidance_scale,
             num_inference_steps,
             randomize_seed,
-            num_images,
-            style_selection  # Added style_selection to inputs
+            num_images
         ],
         outputs=[result, seed],
         api_name="run",