Update app.py

app.py

@@ -9,13 +9,11 @@ import spaces
 import torch
 from diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler
 
-DESCRIPTIONx = """## STABLE HAMSTER 🐹
-
-"""
+DESCRIPTIONx = """## STABLE HAMSTER 🐹"""
 
 css = '''
 .gradio-container {
-    max-width: 560px !important;
+    max-width: 1000px !important;
     margin: 0 auto !important;
 }
 h1 {
@@ -34,13 +32,12 @@ examples = [
     "Commercial photography, giant burger, white lighting, studio light, 8k octane rendering, high resolution photography, insanely detailed, fine details, on white isolated plain, 8k, commercial photography, stock photo, professional color grading, --v 4 --ar 9:16 "
 ]
 
-MODEL_ID = os.getenv("MODEL_VAL_PATH")  # uses SG161222/RealVisXL_V5.0_Lightning or SG161222/RealVisXL_V4.0_Lightning
+MODEL_ID = os.getenv("MODEL_VAL_PATH")  # SG161222/RealVisXL_V5.0_Lightning or SG161222/RealVisXL_V4.0_Lightning
 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"))
 
-# Load model outside of function
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 pipe = StableDiffusionXLPipeline.from_pretrained(
     MODEL_ID,
@@ -50,11 +47,9 @@ pipe = StableDiffusionXLPipeline.from_pretrained(
 ).to(device)
 pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
 
-# Compile speedup
 if USE_TORCH_COMPILE:
     pipe.compile()
 
-# Offloading capacity (RAM)
 if ENABLE_CPU_OFFLOAD:
     pipe.enable_model_cpu_offload()
 
@@ -82,13 +77,12 @@ def generate(
     num_inference_steps: int = 25,
     randomize_seed: bool = False,
     use_resolution_binning: bool = True,
-    num_images: int = 4,  # Number of images to generate
+    num_images: int = 4,
     progress=gr.Progress(track_tqdm=True),
 ):
     seed = int(randomize_seed_fn(seed, randomize_seed))
     generator = torch.Generator(device=device).manual_seed(seed)
 
-    # Options
     options = {
         "prompt": [prompt] * num_images,
         "negative_prompt": [negative_prompt] * num_images if use_negative_prompt else None,
@@ -103,7 +97,6 @@ def generate(
     if use_resolution_binning:
         options["use_resolution_binning"] = True
 
-    # Generate images in batches
     images = []
     for i in range(0, num_images, BATCH_SIZE):
         batch_options = options.copy()
@@ -117,6 +110,7 @@ def generate(
 
 with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
     gr.Markdown(DESCRIPTIONx)
+
     with gr.Row():
         with gr.Column(scale=3):
             with gr.Row():
@@ -128,72 +122,70 @@ with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
                 )
                 run_button = gr.Button("Run", scale=0, variant="primary")
 
-            # Changed from gr.Image to gr.Gallery to handle multiple images
             result = gr.Gallery(label="Result", columns=2, show_label=False)
-            
-    with gr.Accordion("Advanced options", open=False, visible=True):
-        num_images = gr.Slider(
-            label="Number of Images",
-            minimum=1,
-            maximum=4,
-            step=1,
-            value=4,
-        )
-        with gr.Row():
-            use_negative_prompt = gr.Checkbox(label="Use negative prompt", value=True)
-            negative_prompt = gr.Text(
-                label="Negative prompt",
-                max_lines=5,
-                lines=4,
-                placeholder="Enter a negative prompt",
-                value="(deformed, distorted, disfigured:1.3), poorly drawn, bad anatomy, wrong anatomy, extra limb, missing limb, floating limbs, (mutated hands and fingers:1.4), disconnected limbs, mutation, mutated, ugly, disgusting, blurry, amputation",
-                visible=True,
-            )
-        seed = gr.Slider(
-            label="Seed",
-            minimum=0,
-            maximum=MAX_SEED,
-            step=1,
-            value=0,
-        )
-        randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-        with gr.Row(visible=True):
-            width = gr.Slider(
-                label="Width",
-                minimum=512,
-                maximum=MAX_IMAGE_SIZE,
-                step=64,
-                value=1024,
-            )
-            height = gr.Slider(
-                label="Height",
-                minimum=512,
-                maximum=MAX_IMAGE_SIZE,
-                step=64,
-                value=1024,
-            )
-        with gr.Row():
-            guidance_scale = gr.Slider(
-                label="Guidance Scale",
-                minimum=0.1,
-                maximum=6,
-                step=0.1,
-                value=3.0,
-            )
-            num_inference_steps = gr.Slider(
-                label="Number of inference steps",
-                minimum=1,
-                maximum=25,
-                step=1,
-                value=23,
-            )
 
-    with gr.Column(scale=1):  # Examples on right
-        gr.Examples(
-            examples=examples,
-            inputs=prompt,
-            cache_examples=False,
-        )
+            with gr.Accordion("Advanced options", open=False, visible=True):
+                num_images = gr.Slider(
+                    label="Number of Images",
+                    minimum=1,
+                    maximum=4,
+                    step=1,
+                    value=4,
+                )
+                with gr.Row():
+                    use_negative_prompt = gr.Checkbox(label="Use negative prompt", value=True)
+                    negative_prompt = gr.Text(
+                        label="Negative prompt",
+                        max_lines=5,
+                        lines=4,
+                        placeholder="Enter a negative prompt",
+                        value="(deformed, distorted, disfigured:1.3), poorly drawn, bad anatomy, wrong anatomy, extra limb, missing limb, floating limbs, (mutated hands and fingers:1.4), disconnected limbs, mutation, mutated, ugly, disgusting, blurry, amputation",
+                        visible=True,
+                    )
+                seed = gr.Slider(
+                    label="Seed",
+                    minimum=0,
+                    maximum=MAX_SEED,
+                    step=1,
+                    value=0,
+                )
+                randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+                with gr.Row(visible=True):
+                    width = gr.Slider(
+                        label="Width",
+                        minimum=512,
+                        maximum=MAX_IMAGE_SIZE,
+                        step=64,
+                        value=1024,
+                    )
+                    height = gr.Slider(
+                        label="Height",
+                        minimum=512,
+                        maximum=MAX_IMAGE_SIZE,
+                        step=64,
+                        value=1024,
+                    )
+                with gr.Row():
+                    guidance_scale = gr.Slider(
+                        label="Guidance Scale",
+                        minimum=0.1,
+                        maximum=6,
+                        step=0.1,
+                        value=3.0,
+                    )
+                    num_inference_steps = gr.Slider(
+                        label="Number of inference steps",
+                        minimum=1,
+                        maximum=25,
+                        step=1,
+                        value=23,
+                    )
+        with gr.Column(scale=1):
+            gr.Examples(
+                examples=examples,
+                inputs=prompt,
+                cache_examples=False,
+            )
 
     use_negative_prompt.change(
         fn=lambda x: gr.update(visible=x),