Update app.py

app.py

@@ -8,15 +8,37 @@ from PIL import Image
 import spaces
 import torch
 from diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler
-
+from typing import Tuple  # Import Tuple from typing module
 
 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 = [
     {
@@ -29,34 +51,11 @@ 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}
@@ -69,37 +68,6 @@ 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)
@@ -123,19 +91,21 @@ 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] * 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,
+        "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),
         "generator": generator,
         "output_type": "pil",
     }
@@ -146,7 +116,7 @@ def generate(
 
     # Generate images potentially in batches
     images = []
-    for i in range(0, num_images, BATCH_SIZE):
+    for i in range(0, int(num_images), BATCH_SIZE):
         batch_options = options.copy()
         batch_options["prompt"] = options["prompt"][i:i+BATCH_SIZE]
         if "negative_prompt" in batch_options:
@@ -242,8 +212,6 @@ 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,
@@ -252,13 +220,11 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
             value=DEFAULT_STYLE_NAME,
             label="Image Style",
         )
-            
+
     gr.Examples(
         examples=examples,
         inputs=prompt,
-        outputs=[result, seed],
-        fn=generate,
-        cache_examples=CACHE_EXAMPLES,
+        cache_examples=False
     )
 
     use_negative_prompt.change(
@@ -279,14 +245,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
+            num_images,
+            style_selection 
         ],
         outputs=[result, seed],
         api_name="run",