Update app.py

app.py

@@ -12,18 +12,11 @@ import gradio as gr
 import numpy as np
 from PIL import Image
 import torch
-from diffusers import AutoencoderKL, StableDiffusionXLPipeline, UNet2DConditionModel
+from diffusers import AutoencoderKL, StableDiffusionXLPipeline
 from diffusers import EulerAncestralDiscreteScheduler
-#from diffusers import DPMSolverMultistepScheduler
-#from diffusers import AsymmetricAutoencoderKL
 from typing import Tuple
 import paramiko
-import gc
-import time
 import datetime
-#from diffusers.schedulers import AysSchedules
-from gradio import themes
-import gc
 
 torch.backends.cuda.matmul.allow_tf32 = False
 torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False
@@ -84,91 +77,31 @@ DEFAULT_STYLE_NAME = "Style Zero"
 STYLE_NAMES = list(styles.keys())
 HF_TOKEN = os.getenv("HF_TOKEN")
 
-#sampling_schedule = AysSchedules["StableDiffusionXLTimesteps"]
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
     
-def load_and_prepare_model(model_id):
+def load_and_prepare_model():
     model_dtypes = {"ford442/RealVisXL_V5.0_BF16": torch.bfloat16,}
     dtype = model_dtypes.get(model_id, torch.bfloat16)  # Default to float32 if not found
-    #vaeX = AutoencoderKL.from_pretrained("ford442/sdxl-vae-bf16", safety_checker=None)
     vaeXL = AutoencoderKL.from_pretrained("stabilityai/sdxl-vae", safety_checker=None, use_safetensors=False).to(device=device, dtype=torch.bfloat16)
-    #vae = AutoencoderKL.from_pretrained('cross-attention/asymmetric-autoencoder-kl-x-2',use_safetensors=False)
-    #vae = AutoencoderKL.from_single_file('https://huggingface.co/ford442/sdxl-vae-bf16/mySLR/myslrVAE_v10.safetensors')
-    #vaeX = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse",use_safetensors=True)
-    #vaeX = AutoencoderKL.from_pretrained('ford442/Juggernaut-XI-v11-fp32',subfolder='vae', safety_checker=None, use_safetensors=False)   #  ,use_safetensors=True  FAILS
-    #vaeX = AutoencoderKL.from_pretrained('ford442/RealVisXL_V5.0_FP64',subfolder='vae').to(torch.bfloat16)   #  ,use_safetensors=True  FAILS
-    #unetX = UNet2DConditionModel.from_pretrained('SG161222/RealVisXL_V5.0',subfolder='unet')  #  ,use_safetensors=True  FAILS
-  #  vae = AutoencoderKL.from_pretrained("BeastHF/MyBack_SDXL_Juggernaut_XL_VAE/MyBack_SDXL_Juggernaut_XL_VAE_V10(version_X).safetensors",safety_checker=None).to(torch.bfloat16)
-    #sched = EulerAncestralDiscreteScheduler.from_pretrained("SG161222/RealVisXL_V5.0", subfolder='scheduler',beta_schedule="scaled_linear", steps_offset=1,timestep_spacing="trailing"))
-    #sched = EulerAncestralDiscreteScheduler.from_pretrained("SG161222/RealVisXL_V5.0", subfolder='scheduler', steps_offset=1,timestep_spacing="trailing")
-    #sched = EulerAncestralDiscreteScheduler.from_pretrained('SG161222/RealVisXL_V5.0', subfolder='scheduler',beta_schedule="scaled_linear", beta_start=0.00085, beta_end=0.012, steps_offset=1,use_karras_sigmas=True)
     sched = EulerAncestralDiscreteScheduler.from_pretrained('ford442/RealVisXL_V5.0_BF16', subfolder='scheduler',beta_schedule="scaled_linear", beta_start=0.00085, beta_end=0.012, steps_offset=1,use_karras_sigmas=True)
-    #pipeX = StableDiffusionXLPipeline.from_pretrained("SG161222/RealVisXL_V5.0").to(torch.bfloat16)
-    #pipeX = StableDiffusionXLPipeline.from_pretrained("ford442/Juggernaut-XI-v11-fp32",use_safetensors=True)
-
     pipe = StableDiffusionXLPipeline.from_pretrained(
         'ford442/RealVisXL_V5.0_BF16',
-        #'John6666/pornworks-sexy-beauty-v04-sdxl',
-        #'John6666/uber-realistic-porn-merge-xl-urpmxl-v3-sdxl',
-       #'ford442/Juggernaut-XI-v11-fp32',
-      #   'SG161222/RealVisXL_V5.0',
         #torch_dtype=torch.bfloat16,
         add_watermarker=False,
-       # custom_pipeline="lpw_stable_diffusion_xl",
-        #use_safetensors=True,
-      #  use_auth_token=HF_TOKEN,
-     #  vae=AutoencoderKL.from_pretrained("BeastHF/MyBack_SDXL_Juggernaut_XL_VAE/MyBack_SDXL_Juggernaut_XL_VAE_V10(version_X).safetensors",repo_type='model',safety_checker=None),
-     #   vae=AutoencoderKL.from_pretrained("stabilityai/sdxl-vae",repo_type='model',safety_checker=None, torch_dtype=torch.float32),
-       # vae=AutoencoderKL.from_pretrained("ford442/sdxl-vae-bf16",repo_type='model',safety_checker=None),
-        #vae=vaeX.to(torch.bfloat16),
-        #unet=pipeX.unet,
-        #scheduler = sched,
-      #  scheduler = EulerAncestralDiscreteScheduler.from_config(pipeX.scheduler.config, beta_schedule="scaled_linear", beta_start=0.00085, beta_end=0.012, steps_offset=1)
-        #scheduler=EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config, beta_schedule="scaled_linear", beta_start=0.00085, beta_end=0.012, steps_offset =1)
     )
-    
-    #pipe.unet=UNet2DConditionModel.from_pretrained('ford442/RealVisXL_V5.0_FP64',subfolder='unet').to(torch.bfloat16)
-    #pipe.unet=UNet2DConditionModel.from_pretrained('SG161222/RealVisXL_V5.0',subfolder='unet').to(torch.bfloat16)
-    #pipe.vae = AsymmetricAutoencoderKL.from_pretrained('cross-attention/asymmetric-autoencoder-kl-x-2').to(torch.bfloat16)   #  ,use_safetensors=True  FAILS
-    #pipe.vae = AutoencoderKL.from_pretrained('ford442/Juggernaut-XI-v11-fp32',subfolder='vae')   #  ,use_safetensors=True  FAILS
-    #pipe.vae.to(torch.bfloat16)  
-    #sched = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config, beta_schedule="scaled_linear",use_karras_sigmas=True, algorithm_type="dpmsolver++")
-    #pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config, beta_schedule="scaled_linear", beta_start=0.00085, beta_end=0.012, steps_offset=1)
-    #pipe.scheduler = DPMSolverMultistepScheduler.from_pretrained('SG161222/RealVisXL_V5.0', subfolder='scheduler', algorithm_type='sde-dpmsolver++')
     pipe.vae = vaeXL #.to(torch.bfloat16)
-    #pipe.unet = unetX.to(torch.bfloat16)
     pipe.scheduler = sched
     pipe.vae.do_resize=False
-    #pipe.vae.do_rescale=False
-    #pipe.vae.do_convert_rgb=True
     pipe.vae.vae_scale_factor=8
-    #pipe.scheduler = sched
-    #pipe.vae=vae.to(torch.bfloat16)
-    #pipe.unet=pipeX.unet
-    #pipe.scheduler=EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config, beta_schedule="scaled_linear", beta_start=0.00085, beta_end=0.012, steps_offset=1)
-    #pipe.scheduler=EulerAncestralDiscreteScheduler.from_pretrained('ford442/RealVisXL_V5.0_BF16', subfolder='scheduler',beta_schedule="scaled_linear")
-
     pipe.to(device=device, dtype=torch.bfloat16)
-    #pipe.to(torch.bfloat16)
-    
-    #apply_hidiffusion(pipe)
-
-    #pipe.unet.set_default_attn_processor()
     pipe.vae.set_default_attn_processor()
-    
-    print(f'Pipeline: ')
-    #print(f'_optional_components: {pipe._optional_components}')
-    #print(f'watermark: {pipe.watermark}')
-    print(f'image_processor: {pipe.image_processor}')
-    #print(f'feature_extractor: {pipe.feature_extractor}')
     print(f'init noise scale: {pipe.scheduler.init_noise_sigma}')
-    #print(f'UNET: {pipe.unet}')
     pipe.watermark=None
     pipe.safety_checker=None  
     return pipe
     
 # Preload and compile both models
-models = {key: load_and_prepare_model(value) for key, value in MODEL_OPTIONS.items()}
+pipe =load_and_prepare_model()
 
 MAX_SEED = np.iinfo(np.int32).max
 
@@ -227,49 +160,33 @@ def uploadNote(prompt,num_inference_steps,guidance_scale,timestamp,denoise):
     
 @spaces.GPU(duration=30)
 def generate_30(
-    model_choice: str,
     prompt: str,
     negative_prompt: str = "",
     use_negative_prompt: bool = False,
     style_selection: str = "",
-    seed: int = 1,
     width: int = 768,
     height: int = 768,
     guidance_scale: float = 4,
     num_inference_steps: int = 125,
-    randomize_seed: bool = False,
     use_resolution_binning: bool = True, 
-    juggernaut: bool = False,
-    denoise: float = 0.3,
     progress=gr.Progress(track_tqdm=True)  # Add progress as a keyword argument
 ):
-    #torch.backends.cudnn.benchmark = False
-    #torch.cuda.empty_cache()
-    #gc.collect()
-    global models
-    pipe = models[model_choice]
-    #if juggernaut == True:
-    #    pipe.vae=vaeX
     seed = int(randomize_seed_fn(seed, randomize_seed))
     generator = torch.Generator(device='cuda').manual_seed(seed)
-    #prompt, negative_prompt = apply_style(style_selection, prompt, negative_prompt)
     options = {
         "prompt": [prompt],
         "negative_prompt": [negative_prompt],
         "negative_prompt_2": [neg_prompt_2],
-        "strength": denoise,
         "width": width,
         "height": height,
         "guidance_scale": guidance_scale,
         "num_inference_steps": num_inference_steps,
         "generator": generator,
-     #   "timesteps": sampling_schedule,
         "output_type": "pil",
     }
     if use_resolution_binning:
         options["use_resolution_binning"] = True
     images = []
-    pipe.scheduler.set_timesteps(num_inference_steps,device)
     timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
     uploadNote(prompt,num_inference_steps,guidance_scale,timestamp,denoise)
     batch_options = options.copy()
@@ -277,58 +194,39 @@ def generate_30(
     sd_image_path = f"rv50_B_{timestamp}.png"
     rv_image.save(sd_image_path,optimize=False,compress_level=0)
     upload_to_ftp(sd_image_path)    
-    #image_paths = save_image(rv_image)
-    #torch.cuda.empty_cache()
-    #gc.collect()
     unique_name = str(uuid.uuid4()) + ".png"  
     os.symlink(sd_image_path, unique_name)  
-    return [unique_name] , seed  
+    return [unique_name], seed
 
 @spaces.GPU(duration=60)
 def generate_60(
-    model_choice: str,
     prompt: str,
     negative_prompt: str = "",
     use_negative_prompt: bool = False,
     style_selection: str = "",
-    seed: int = 1,
     width: int = 768,
     height: int = 768,
     guidance_scale: float = 4,
-    num_inference_steps: int = 250,
-    randomize_seed: bool = False,
+    num_inference_steps: int = 125,
     use_resolution_binning: bool = True, 
-    juggernaut: bool = False,
-    denoise: float = 0.3,
     progress=gr.Progress(track_tqdm=True)  # Add progress as a keyword argument
 ):
-    #torch.backends.cudnn.benchmark = True
-    #torch.cuda.empty_cache()
-    #gc.collect()
-    global models
-    pipe = models[model_choice]
-    #if juggernaut == True:
-    #    pipe.vae=vaeX
     seed = int(randomize_seed_fn(seed, randomize_seed))
     generator = torch.Generator(device='cuda').manual_seed(seed)
-    #prompt, negative_prompt = apply_style(style_selection, prompt, negative_prompt)
     options = {
         "prompt": [prompt],
         "negative_prompt": [negative_prompt],
         "negative_prompt_2": [neg_prompt_2],
-        "strength": denoise,
         "width": width,
         "height": height,
         "guidance_scale": guidance_scale,
         "num_inference_steps": num_inference_steps,
         "generator": generator,
-     #   "timesteps": sampling_schedule,
         "output_type": "pil",
     }
     if use_resolution_binning:
         options["use_resolution_binning"] = True
     images = []
-    pipe.scheduler.set_timesteps(num_inference_steps,device)
     timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
     uploadNote(prompt,num_inference_steps,guidance_scale,timestamp,denoise)
     batch_options = options.copy()
@@ -336,58 +234,39 @@ def generate_60(
     sd_image_path = f"rv50_B_{timestamp}.png"
     rv_image.save(sd_image_path,optimize=False,compress_level=0)
     upload_to_ftp(sd_image_path)    
-    #image_paths = save_image(rv_image)
-    #torch.cuda.empty_cache()
-    #gc.collect()
     unique_name = str(uuid.uuid4()) + ".png"  
     os.symlink(sd_image_path, unique_name)  
-    return [unique_name] , seed  
+    return [unique_name], seed
     
 @spaces.GPU(duration=90)
 def generate_90(
-    model_choice: str,
     prompt: str,
     negative_prompt: str = "",
     use_negative_prompt: bool = False,
     style_selection: str = "",
-    seed: int = 1,
     width: int = 768,
     height: int = 768,
     guidance_scale: float = 4,
-    num_inference_steps: int = 250,
-    randomize_seed: bool = False,
+    num_inference_steps: int = 125,
     use_resolution_binning: bool = True, 
-    juggernaut: bool = False,
-    denoise: float = 0.3,
     progress=gr.Progress(track_tqdm=True)  # Add progress as a keyword argument
 ):
-    #torch.backends.cudnn.benchmark = True
-    #torch.cuda.empty_cache()
-    #gc.collect()
-    global models
-    pipe = models[model_choice]
-    #if juggernaut == True:
-    #    pipe.vae=vaeX
     seed = int(randomize_seed_fn(seed, randomize_seed))
     generator = torch.Generator(device='cuda').manual_seed(seed)
-    #prompt, negative_prompt = apply_style(style_selection, prompt, negative_prompt)
     options = {
         "prompt": [prompt],
         "negative_prompt": [negative_prompt],
         "negative_prompt_2": [neg_prompt_2],
-        "strength": denoise,
         "width": width,
         "height": height,
         "guidance_scale": guidance_scale,
         "num_inference_steps": num_inference_steps,
         "generator": generator,
-     #   "timesteps": sampling_schedule,
         "output_type": "pil",
     }
     if use_resolution_binning:
         options["use_resolution_binning"] = True
     images = []
-    pipe.scheduler.set_timesteps(num_inference_steps,device)
     timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
     uploadNote(prompt,num_inference_steps,guidance_scale,timestamp,denoise)
     batch_options = options.copy()
@@ -395,12 +274,9 @@ def generate_90(
     sd_image_path = f"rv50_B_{timestamp}.png"
     rv_image.save(sd_image_path,optimize=False,compress_level=0)
     upload_to_ftp(sd_image_path)    
-    #image_paths = save_image(rv_image)
-    #torch.cuda.empty_cache()
-    #gc.collect()
     unique_name = str(uuid.uuid4()) + ".png"  
     os.symlink(sd_image_path, unique_name)  
-    return [unique_name] , seed  
+    return [unique_name], seed
 
 def load_predefined_images1():
     predefined_images1 = [
@@ -446,11 +322,7 @@ with gr.Blocks(theme=gr.themes.Origin(),css=css) as demo:
     result = gr.Gallery(label="Result", columns=1, show_label=False) 
 
     with gr.Row():
-        model_choice = gr.Dropdown(
-            label="Model Selection🔻",
-            choices=list(MODEL_OPTIONS.keys()),
-            value="REALVISXL V5.0 BF16"
-        )
+
 
         style_selection = gr.Radio(
             show_label=True,
@@ -471,22 +343,6 @@ with gr.Blocks(theme=gr.themes.Origin(),css=css) as demo:
                     value="('deformed', 'distorted', 'disfigured':1.3),'not photorealistic':1.5, 'poorly drawn', 'bad anatomy', 'wrong anatomy', 'extra limb', 'missing limb', 'floating limbs', 'poorly drawn hands', 'poorly drawn feet', 'poorly drawn face':1.3, 'out of frame', 'extra limbs', 'bad anatomy', 'bad art', 'beginner',  'distorted face','amateur'",
                     visible=True,
                 )
-        seed = gr.Slider(
-            label="Seed",
-            minimum=0,
-            maximum=MAX_SEED,
-            step=1,
-            value=0,
-        )
-        denoise = gr.Slider(
-            label="Denoising Strength",
-            minimum=0.0,
-            maximum=1.0,
-            step=0.01,
-            value=0.3,
-        )
-        randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-        juggernaut = gr.Checkbox(label="Use Juggernaut VAE", value=False)
         with gr.Row():
             width = gr.Slider(
                 label="Width",
@@ -538,19 +394,14 @@ with gr.Blocks(theme=gr.themes.Origin(),css=css) as demo:
       #  api_name="generate",  # Add this line
         fn=generate_30,
         inputs=[
-            model_choice,
             prompt,
             negative_prompt,
             use_negative_prompt,
             style_selection,
-            seed,
             width,
             height,
             guidance_scale,
             num_inference_steps,
-            randomize_seed, 
-            juggernaut,
-            denoise
         ],
         outputs=[result, seed],
     )
@@ -562,19 +413,14 @@ with gr.Blocks(theme=gr.themes.Origin(),css=css) as demo:
       #  api_name="generate",  # Add this line
         fn=generate_60,
         inputs=[
-            model_choice,
             prompt,
             negative_prompt,
             use_negative_prompt,
             style_selection,
-            seed,
             width,
             height,
             guidance_scale,
             num_inference_steps,
-            randomize_seed, 
-            juggernaut,
-            denoise
         ],
         outputs=[result, seed],
     )
@@ -586,19 +432,14 @@ with gr.Blocks(theme=gr.themes.Origin(),css=css) as demo:
       #  api_name="generate",  # Add this line
         fn=generate_90,
         inputs=[
-            model_choice,
             prompt,
             negative_prompt,
             use_negative_prompt,
             style_selection,
-            seed,
             width,
             height,
             guidance_scale,
             num_inference_steps,
-            randomize_seed, 
-            juggernaut,
-            denoise
         ],
         outputs=[result, seed],
     )