Update app.py

app.py

@@ -1,13 +1,3 @@
-#!/usr/bin/env python
-#patch 2.0 ()
-# Permission is hereby granted, free of charge, to any person obtaining a copy
-# of this software and associated documentation files (the "Software"), to deal
-# in the Software without restriction, including without limitation the rights
-# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-# copies of the Software, and to permit persons to whom the Software is
-# furnished to do so, subject to the following conditions:
-#
-# ...
 import os
 import random
 import uuid
@@ -19,16 +9,9 @@ import spaces
 import torch
 from diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler
 
-#Load the HTML content
-#html_file_url = "https://prithivmlmods-hamster-static.static.hf.space/index.html"
-#html_content = f'<iframe src="{html_file_url}" style="width:100%; height:180px; border:none;"></iframe>'
-#html_file_url = "https://prithivmlmods-static-loading-theme.static.hf.space/index.html"
-
-#html_file_url = "https://prithivhamster.vercel.app/"
-#html_content = f'<iframe src="{html_file_url}" style="width:100%; height:400px; border:none"></iframe>'
-
 DESCRIPTIONx = """## STABLE HAMSTER 🐹
 
+
 """
 
 DESCRIPTIONy = """
@@ -39,7 +22,6 @@ DESCRIPTIONy = """
             </p>
 """
 
-
 css = '''
 .gradio-container{max-width: 560px !important}
 h1{text-align:center}
@@ -51,34 +33,17 @@ footer {
 examples = [
     "3d image, cute girl, in the style of Pixar --ar 1:2 --stylize 750, 4K resolution highlights, Sharp focus, octane render, ray tracing, Ultra-High-Definition, 8k, UHD, HDR, (Masterpiece:1.5), (best quality:1.5)",
     "Cold coffee in a cup bokeh --ar 85:128 --v 6.0 --style raw5, 4K",
-    "Vector illustration of a horse, vector graphic design with flat colors on an brown background in the style of vector art, using simple shapes and graphics with simple details, professionally designed as a tshirt logo ready for print on a white background. --ar 89:82 --v 6.0 --style raw",
+    "Vector illustration of a horse, vector graphic design with flat colors on a brown background in the style of vector art, using simple shapes and graphics with simple details, professionally designed as a tshirt logo ready for print on a white background. --ar 89:82 --v 6.0 --style raw",
     "Man in brown leather jacket posing for camera, in the style of sleek and stylized, clockpunk, subtle shades, exacting precision, ferrania p30  --ar 67:101 --v 5",
     "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 "
-    
 ]
 
-
-#examples = [
-  #  ["file/1.png", "3d image, cute girl, in the style of Pixar --ar 1:2 --stylize 750, 4K resolution highlights, Sharp focus, octane render, ray tracing, Ultra-High-Definition, 8k, UHD, HDR, (Masterpiece:1.5), (best quality:1.5)"],
-   # ["file/2.png", "Cold coffee in a cup bokeh --ar 85:128 --v 6.0 --style raw5, 4K"],
-    #["file/3.png", "Vector illustration of a horse, vector graphic design with flat colors on a brown background in the style of vector art, using simple shapes and graphics with simple details, professionally designed as a tshirt logo ready for print on a white background. --ar 89:82 --v 6.0 --style raw"],
-    #["file/4.png", "Man in brown leather jacket posing for the camera, in the style of sleek and stylized, clockpunk, subtle shades, exacting precision, ferrania p30  --ar 67:101 --v 5"],
-    #["file/5.png", "Commercial photography, giant burger, white lighting, studio light, 8k octane rendering, high resolution photography, insanely detailed, fine details, on a white isolated plain, 8k, commercial photography, stock photo, professional color grading, --v 4 --ar 9:16"]
-#]
-
-
-#Set an os.Getenv variable
-#set VAR_NAME=”VALUE”
-#Fetch an environment variable
-#echo %VAR_NAME%
-
-MODEL_ID = os.getenv("MODEL_VAL_PATH") #Use SDXL Model as "MODEL_REPO" --------->>> ”VALUE”. 
+MODEL_ID = os.getenv("MODEL_VAL_PATH") 
 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,
@@ -88,11 +53,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()
 
@@ -119,14 +82,24 @@ def generate(
     guidance_scale: float = 3,
     num_inference_steps: int = 25,
     randomize_seed: bool = False,
-    use_resolution_binning: bool = True, 
-    num_images: int = 1,  # Number of images to generate
+    use_resolution_binning: bool = True,
+    grid_size: str = "2x2",
     progress=gr.Progress(track_tqdm=True),
 ):
     seed = int(randomize_seed_fn(seed, randomize_seed))
     generator = torch.Generator(device=device).manual_seed(seed)
 
-    #Options 
+    grid_sizes = {
+        "2x1": (2, 1),
+        "1x2": (1, 2),
+        "2x2": (2, 2),
+        "2x3": (2, 3),
+        "3x2": (3, 2),
+        "1x1": (1, 1)
+    }
+    grid_size_x, grid_size_y = grid_sizes.get(grid_size, (2, 2))
+    num_images = grid_size_x * grid_size_y
+
     options = {
         "prompt": [prompt] * num_images,
         "negative_prompt": [negative_prompt] * num_images if use_negative_prompt else None,
@@ -138,11 +111,9 @@ def generate(
         "output_type": "pil",
     }
 
-    #VRAM usage Lesser
     if use_resolution_binning:
         options["use_resolution_binning"] = True
 
-    #Images potential batches
     images = []
     for i in range(0, num_images, BATCH_SIZE):
         batch_options = options.copy()
@@ -151,11 +122,17 @@ def generate(
             batch_options["negative_prompt"] = options["negative_prompt"][i:i+BATCH_SIZE]
         images.extend(pipe(**batch_options).images)
 
-    image_paths = [save_image(img) for img in images]
-    return image_paths, seed
-#Main gr.Block
+    torch.cuda.empty_cache()
+    grid_img = Image.new('RGB', (width * grid_size_x, height * grid_size_y))
+
+    for i, img in enumerate(images[:num_images]):
+        grid_img.paste(img, (i % grid_size_x * width, i // grid_size_x * height))
+
+    unique_name = save_image(grid_img)
+    return unique_name, seed
+
 with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
-    gr.Markdown(DESCRIPTIONx)  
+    gr.Markdown(DESCRIPTIONx)
 
     with gr.Group():
         with gr.Row():
@@ -168,14 +145,23 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
             )
             run_button = gr.Button("Run", scale=0)
         result = gr.Gallery(label="Result", columns=1, show_label=False) 
+        
+        with gr.Row(visible=True):
+            grid_size_selection = gr.Dropdown(
+                choices=["2x1", "1x2", "2x2", "2x3", "3x2", "1x1"],
+                value="1x1",
+                label="⚡Grid"
+            )
     with gr.Accordion("Advanced options", open=False, visible=False):
-        num_images = gr.Slider(
-            label="Number of Images",
-            minimum=1,
-            maximum=4,
-            step=1,
-            value=1,
-        )
+
+        with gr.Row():
+            num_images = gr.Slider(
+                label="Number of Images",
+                minimum=1,
+                maximum=4,
+                step=1,
+                value=1,
+            )
         with gr.Row():
             use_negative_prompt = gr.Checkbox(label="Use negative prompt", value=True)
             negative_prompt = gr.Text(
@@ -255,14 +241,12 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
             guidance_scale,
             num_inference_steps,
             randomize_seed,
-            num_images
+            grid_size_selection
         ],
         outputs=[result, seed],
         api_name="run",
     )   
 
-
-    
     gr.Markdown(DESCRIPTIONy)
 
     gr.Markdown("**Disclaimer:**")
@@ -271,7 +255,5 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
     gr.Markdown("**Note:**")
     gr.Markdown("⚠️ users are accountable for the content they generate and are responsible for ensuring it meets appropriate ethical standards.")
 
-    #gr.HTML(html_content)
-
 if __name__ == "__main__":
     demo.queue(max_size=40).launch()