Update app.py

app.py

@@ -8,14 +8,15 @@ from PIL import Image
 import spaces
 import torch
 from diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler
-from typing import Tuple
+
 
 DESCRIPTIONx = """## STABLE HAMSTER
 """
 
+
 # Use environment variables for flexibility
 MODEL_ID = os.getenv("MODEL_REPO")
-MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "2048"))
+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
@@ -40,34 +41,6 @@ if ENABLE_CPU_OFFLOAD:
 
 MAX_SEED = np.iinfo(np.int32).max
 
-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",
-    },
-]
-
-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)
@@ -91,21 +64,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 +87,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,14 +183,6 @@ 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,
@@ -234,7 +197,6 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
         api_name=False,
     )
 
-
     gr.on(
         triggers=[
             prompt.submit,
@@ -252,8 +214,7 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
             guidance_scale,
             num_inference_steps,
             randomize_seed,
-            num_images,
-            style_selection 
+            num_images
         ],
         outputs=[result, seed],
         api_name="run",