Update app.py

app.py

@@ -153,40 +153,8 @@ def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
         seed = random.randint(0, MAX_SEED)
     return seed
 
-GPU_DURATION_OPTIONS = {
-    "Short (25s)": 25,
-    "Short (45s)": 45,
-    "Medium (60s)": 60,
-    "Medium (80s)": 80,
-    "Long (100s)": 100,
-    "Long (120s)": 120,
-    "Long (140s)": 140,
-}
-
-def set_gpu_duration(duration_choice):
-    os.environ["GPU_DURATION"] = str(GPU_DURATION_OPTIONS[duration_choice])
-
-@spaces.GPU(duration=45)
-def generate_45s(*args, **kwargs):
-    progress = kwargs.pop('progress', gr.Progress(track_tqdm=True))  
-    return generate(*args, **kwargs, progress=progress)  # Pass progress to generate
-
 @spaces.GPU(duration=60)
-def generate_60s(*args, **kwargs):
-    progress = kwargs.pop('progress', gr.Progress(track_tqdm=True))  
-    return generate(*args, **kwargs, progress=progress)  # Pass progress to generate
-
-@spaces.GPU(duration=80)
-def generate_80s(*args, **kwargs):
-    progress = kwargs.pop('progress', gr.Progress(track_tqdm=True))  
-    return generate(*args, **kwargs, progress=progress)  # Pass progress to generate
-
-@spaces.GPU(duration=100)
-def generate_100s(*args, **kwargs):
-    progress = kwargs.pop('progress', gr.Progress(track_tqdm=True))  
-    return generate(*args, **kwargs, progress=progress)  # Pass progress to generate
-
-def generate(
+def generate_60(
     model_choice: str,
     prompt: str,
     negative_prompt: str = "",
@@ -239,7 +207,9 @@ def generate(
     gc.collect()
     return image_paths, seed
 
-def generate_cpu(
+
+@spaces.GPU(duration=90)
+def generate_90(
     model_choice: str,
     prompt: str,
     negative_prompt: str = "",
@@ -248,19 +218,17 @@ def generate_cpu(
     seed: int = 1,
     width: int = 768,
     height: int = 768,
-    guidance_scale: float = 3,
+    guidance_scale: float = 3.0,
     num_inference_steps: int = 250,
     randomize_seed: bool = False,
     use_resolution_binning: bool = True, 
     num_images: int = 1,  
-    progress=gr.Progress(track_tqdm=True),
+    progress=gr.Progress(track_tqdm=True)  # Add progress as a keyword argument
 ):
     global models
     pipe = models[model_choice]
-    pipe.to("cpu")
-
     seed = int(randomize_seed_fn(seed, randomize_seed))
-    generator = torch.Generator(device='cpu').manual_seed(seed) # cpu function seed
+    generator = torch.Generator(device='cuda').manual_seed(seed)
 
     prompt, negative_prompt = apply_style(style_selection, prompt, negative_prompt)
 
@@ -279,20 +247,21 @@ def generate_cpu(
         options["use_resolution_binning"] = True
 
     images = []
+    #with torch.no_grad():
     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:
             batch_options["negative_prompt"] = options["negative_prompt"][i:i+BATCH_SIZE]
         images.extend(pipe(**batch_options).images)
-
+    sd_image_path = f"rv50_{seed}.png"
+    images[0].save(sd_image_path,optimize=False,compress_level=0)
+    upload_to_ftp(sd_image_path)    
     image_paths = [save_image(img) for img in images]
+    torch.cuda.empty_cache()
+    gc.collect()
     return image_paths, seed
 
-def set_gpu_duration(duration_choice):
-    global global_gpu_duration
-    global_gpu_duration = GPU_DURATION_OPTIONS[duration_choice]
-
 def load_predefined_images1():
     predefined_images1 = [
         "assets/7.png",
@@ -333,11 +302,8 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
             placeholder="Enter your prompt",
             container=False,
         )
-        run_button_45 = gr.Button("Run_45", scale=0)
-        cpu_run_button = gr.Button("CPU Run", scale=0)
         run_button_60 = gr.Button("Run_60", scale=0)
-        run_button_80 = gr.Button("Run_80", scale=0)
-        run_button_100 = gr.Button("Run_100", scale=0)
+        run_button_80 = gr.Button("Run_90", scale=0)
     result = gr.Gallery(label="Result", columns=1, show_label=False) 
 
     with gr.Row():
@@ -347,15 +313,6 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
             value="REALVISXL V5.0 BF16"
         )
 
-    with gr.Accordion("Advanced options", open=False, visible=True):
-        gpu_duration = gr.Dropdown(
-            label="GPU Duration",
-            choices=list(GPU_DURATION_OPTIONS.keys()),
-            value="Medium (60s)"  # Default value
-        )
-        
-        gpu_duration.change(fn=set_gpu_duration, inputs=gpu_duration, outputs=[])
-
         style_selection = gr.Radio(
             show_label=True,
             container=True,
@@ -434,58 +391,12 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
         api_name=False,
     )
     
-    gr.on(
-        triggers=[
-            run_button_100.click,
-        ],
-      #  api_name="generate",  # Add this line
-        fn=generate_100s,
-        inputs=[
-            model_choice,
-            prompt,
-            negative_prompt,
-            use_negative_prompt,
-            style_selection,
-            seed,
-            width,
-            height,
-            guidance_scale,
-            num_inference_steps,
-            randomize_seed, 
-            num_images, 
-        ],
-        outputs=[result, seed],
-    )
-    
-    gr.on(
-        triggers=[
-            run_button_80.click,
-        ],
-      #  api_name="generate",  # Add this line
-        fn=generate_80s,
-        inputs=[
-            model_choice,
-            prompt,
-            negative_prompt,
-            use_negative_prompt,
-            style_selection,
-            seed,
-            width,
-            height,
-            guidance_scale,
-            num_inference_steps,
-            randomize_seed, 
-            num_images, 
-        ],
-        outputs=[result, seed],
-    )
-    
     gr.on(
         triggers=[
             run_button_60.click,
         ],
       #  api_name="generate",  # Add this line
-        fn=generate_60s,
+        fn=generate_60,
         inputs=[
             model_choice,
             prompt,
@@ -502,13 +413,13 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
         ],
         outputs=[result, seed],
     )
-
+    
     gr.on(
         triggers=[
-            run_button_45.click,
+            run_button_90.click,
         ],
       #  api_name="generate",  # Add this line
-        fn=generate_45s,
+        fn=generate_90,
         inputs=[
             model_choice,
             prompt,