Update app.py

app.py

@@ -21,6 +21,9 @@ import time
 import datetime
 from gradio import themes
 from image_gen_aux import UpscaleWithModel
+from ip_adapter import IPAdapterXL
+
+from huggingface_hub import snapshot_download
 
 torch.backends.cuda.matmul.allow_tf32 = False
 torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False
@@ -49,7 +52,7 @@ examples = [
 ]
 
 MODEL_OPTIONS = {
-    "REALVISXL V5.0 BF16": "Kim2091/RealVisXL_V5.0_BF16",
+    "REALVISXL V5.0 BF16": "ford442/RealVisXL_V5.0_BF16",
 }
 
 MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "4096"))
@@ -87,6 +90,16 @@ DEFAULT_STYLE_NAME = "Style Zero"
 STYLE_NAMES = list(styles.keys())
 HF_TOKEN = os.getenv("HF_TOKEN")
 
+## load IP Adapter
+repo_id = "ford442/SDXL-IP_ADAPTER"
+subfolder = "image_encoder"
+subfolder2 = "ip_adapter"
+local_repo_path = snapshot_download(repo_id=repo_id, repo_type="model")
+local_folder = os.path.join(local_repo_path, subfolder)
+local_folder2 = os.path.join(local_repo_path, subfolder2) # Path to the ip_adapter dir
+ip_ckpt = os.path.join(local_folder2, "ip-adapter_sdxl_vit-h.bin") # Correct path
+ip_model = IPAdapterXL(pipe, local_folder, ip_ckpt, device)
+
 #sampling_schedule = AysSchedules["StableDiffusionXLTimesteps"]
 upscaler = UpscaleWithModel.from_pretrained("Kim2091/ClearRealityV1").to(torch.device("cuda:0"))
 
@@ -99,7 +112,7 @@ def apply_style(style_name: str, positive: str, negative: str = "") -> Tuple[str
         negative = ""
     return p.replace("{prompt}", positive), n + negative
 
-def load_and_prepare_model(model_id):
+def load_and_prepare_model():
     #vae = AutoencoderKL.from_pretrained("ford442/sdxl-vae-bf16", safety_checker=None)
     vaeX = AutoencoderKL.from_pretrained("stabilityai/sdxl-vae", safety_checker=None,use_safetensors=False)
     #vae = AutoencoderKL.from_pretrained('cross-attention/asymmetric-autoencoder-kl-x-2',use_safetensors=False)
@@ -116,10 +129,10 @@ def load_and_prepare_model(model_id):
     #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',
+        'ford442/RealVisXL_V5.0_BF16',
        #'ford442/Juggernaut-XI-v11-fp32',
       #   'SG161222/RealVisXL_V5.0',
-        'John6666/uber-realistic-porn-merge-xl-urpmxl-v3-sdxl',
+        #'John6666/uber-realistic-porn-merge-xl-urpmxl-v3-sdxl',
         #torch_dtype=torch.bfloat16,
         add_watermarker=False,
        # custom_pipeline="lpw_stable_diffusion_xl",
@@ -202,7 +215,7 @@ def load_and_prepare_model(model_id):
     return pipe
     
 # Preload and compile both models
-models = {key: load_and_prepare_model(value) for key, value in MODEL_OPTIONS.items()}
+models = {key: load_and_prepare_model()}
 
 MAX_SEED = np.iinfo(np.int32).max
 
@@ -252,7 +265,6 @@ def uploadNote(prompt,num_inference_steps,guidance_scale,timestamp,denoise):
     
 @spaces.GPU(duration=40)
 def generate_30(
-    model_choice: str,
     prompt: str,
     negative_prompt: str = "",
     use_negative_prompt: bool = False,
@@ -264,57 +276,50 @@ def generate_30(
     num_inference_steps: int = 125,
     randomize_seed: bool = False,
     use_resolution_binning: bool = True, 
-    denoise: float = 0.3,
+    latent_file,  # Add latents file input
     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]
+    pipe = models
     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()
-    rv_image = pipe(**batch_options).images[0]
-    sd_image_path = f"rv50_A_{seed}.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()
-    torch.set_float32_matmul_precision("medium")
-    with torch.no_grad():
-        upscale = upscaler(rv_image, tiling=True, tile_width=256, tile_height=256)
-    downscale1 = upscale.resize((upscale.width // 4, upscale.height // 4), Image.LANCZOS)
-    downscale_path = f"rv50_upscale_{timestamp}.png"
-    downscale1.save(downscale_path,optimize=False,compress_level=0)
-    upload_to_ftp(downscale_path) 
-    image_paths = [save_image(downscale1)]
+    if latent_file:  # Check if a latent file is provided
+        #sd_image_a = torch.load(latent_file.name)  # Load the latent
+        sd_image_a = Image.open(latent_file.name)
+        timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
+        filename= f'rv_IP_{timestamp}.txt'
+        print("-- using image file --")
+        print('-- generating image --')
+        #with torch.no_grad():
+        sd_image = ip_model.generate(
+                pil_image=sd_image_a,
+                prompt=prompt,
+                num_samples=1,
+                num_inference_steps=num_inference_steps,
+                guidance_scale=guidance_scale,
+                seed=seed
+        )
+        sd_image[0].save(filename,optimize=False,compress_level=0)
+        upload_to_ftp(filename)
+        uploadNote(prompt,num_inference_steps,guidance_scale,timestamp)
+        image_paths = save_image(sd_image)
+        torch.setfloat32_matmul_precision("medium")
+        with torch.no_grad():
+            upscale = upscaler(sd_image, tiling=True, tile_width=256, tile_height=256)
+        downscale1 = upscale.resize((upscale.width // 4, upscale.height // 4), Image.LANCZOS)
+        downscale_path = f"rv50_upscale_{timestamp}.png"
+        downscale1.save(downscale_path,optimize=False,compress_level=0)
+        upload_to_ftp(downscale_path) 
+        image_paths = [save_image(downscale1)]     
+    else:
+        print('-- IMAGE REQUIRED --')
     return image_paths, seed
 
 @spaces.GPU(duration=60)
 def generate_60(
-    model_choice: str,
     prompt: str,
     negative_prompt: str = "",
     use_negative_prompt: bool = False,
@@ -326,57 +331,50 @@ def generate_60(
     num_inference_steps: int = 250,
     randomize_seed: bool = False,
     use_resolution_binning: bool = True, 
-    denoise: float = 0.3,
+    latent_file,  # Add latents file input
     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]
+    pipe = models
     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()
-    rv_image = pipe(**batch_options).images[0]
-    sd_image_path = f"rv50_A_{seed}.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()
-    torch.set_float32_matmul_precision("medium")
-    with torch.no_grad():
-        upscale = upscaler(rv_image, tiling=True, tile_width=256, tile_height=256)
-    downscale1 = upscale.resize((upscale.width // 4, upscale.height // 4), Image.LANCZOS)
-    downscale_path = f"rv50_upscale_{timestamp}.png"
-    downscale1.save(downscale_path,optimize=False,compress_level=0)
-    upload_to_ftp(downscale_path) 
-    image_paths = [save_image(downscale1)]
+    if latent_file:  # Check if a latent file is provided
+        #sd_image_a = torch.load(latent_file.name)  # Load the latent
+        sd_image_a = Image.open(latent_file.name)
+        timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
+        filename= f'rv_IP_{timestamp}.txt'
+        print("-- using image file --")
+        print('-- generating image --')
+        #with torch.no_grad():
+        sd_image = ip_model.generate(
+                pil_image=sd_image_a,
+                prompt=prompt,
+                num_samples=1,
+                num_inference_steps=num_inference_steps,
+                guidance_scale=guidance_scale,
+                seed=seed
+        )
+        sd_image[0].save(filename,optimize=False,compress_level=0)
+        upload_to_ftp(filename)
+        uploadNote(prompt,num_inference_steps,guidance_scale,timestamp)
+        image_paths = save_image(sd_image)
+        torch.setfloat32_matmul_precision("medium")
+        with torch.no_grad():
+            upscale = upscaler(sd_image, tiling=True, tile_width=256, tile_height=256)
+        downscale1 = upscale.resize((upscale.width // 4, upscale.height // 4), Image.LANCZOS)
+        downscale_path = f"rv50_upscale_{timestamp}.png"
+        downscale1.save(downscale_path,optimize=False,compress_level=0)
+        upload_to_ftp(downscale_path) 
+        image_paths = [save_image(downscale1)]     
+    else:
+        print('-- IMAGE REQUIRED --')
     return image_paths, seed
 
 @spaces.GPU(duration=90)
 def generate_90(
-    model_choice: str,
     prompt: str,
     negative_prompt: str = "",
     use_negative_prompt: bool = False,
@@ -388,52 +386,46 @@ def generate_90(
     num_inference_steps: int = 250,
     randomize_seed: bool = False,
     use_resolution_binning: bool = True, 
-    denoise: float = 0.3,
+    latent_file,  # Add latents file input
     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]
+    pipe = models
     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()
-    rv_image = pipe(**batch_options).images[0]
-    sd_image_path = f"rv50_A_{seed}.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()
-    torch.set_float32_matmul_precision("medium")
-    with torch.no_grad():
-        upscale = upscaler(rv_image, tiling=True, tile_width=256, tile_height=256)
-    downscale1 = upscale.resize((upscale.width // 4, upscale.height // 4), Image.LANCZOS)
-    downscale_path = f"rv50_upscale_{timestamp}.png"
-    downscale1.save(downscale_path,optimize=False,compress_level=0)
-    upload_to_ftp(downscale_path) 
-    image_paths = [save_image(downscale1)]
+    if latent_file:  # Check if a latent file is provided
+        #sd_image_a = torch.load(latent_file.name)  # Load the latent
+        sd_image_a = Image.open(latent_file.name)
+        timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
+        filename= f'rv_IP_{timestamp}.txt'
+        print("-- using image file --")
+        print('-- generating image --')
+        #with torch.no_grad():
+        sd_image = ip_model.generate(
+                pil_image=sd_image_a,
+                prompt=prompt,
+                num_samples=1,
+                num_inference_steps=num_inference_steps,
+                guidance_scale=guidance_scale,
+                seed=seed
+        )
+        sd_image[0].save(filename,optimize=False,compress_level=0)
+        upload_to_ftp(filename)
+        uploadNote(prompt,num_inference_steps,guidance_scale,timestamp)
+        image_paths = save_image(sd_image)
+        torch.setfloat32_matmul_precision("medium")
+        with torch.no_grad():
+            upscale = upscaler(sd_image, tiling=True, tile_width=256, tile_height=256)
+        downscale1 = upscale.resize((upscale.width // 4, upscale.height // 4), Image.LANCZOS)
+        downscale_path = f"rv50_upscale_{timestamp}.png"
+        downscale1.save(downscale_path,optimize=False,compress_level=0)
+        upload_to_ftp(downscale_path) 
+        image_paths = [save_image(downscale1)]     
+    else:
+        print('-- IMAGE REQUIRED --')
     return image_paths, seed
 
 def load_predefined_images1():
@@ -474,18 +466,13 @@ with gr.Blocks(theme=gr.themes.Origin(),css=css) as demo:
             placeholder="Enter your prompt",
             container=False,
         )
+        latent_file = gr.File(label="Image File (Required)")
         run_button_30 = gr.Button("Run 30 Seconds", scale=0)
         run_button_60 = gr.Button("Run 60 Seconds", scale=0)
         run_button_90 = gr.Button("Run 90 Seconds", scale=0)
     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,
             container=True,
@@ -512,13 +499,6 @@ with gr.Blocks(theme=gr.themes.Origin(),css=css) as demo:
             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)
         with gr.Row():
             width = gr.Slider(
@@ -548,7 +528,7 @@ with gr.Blocks(theme=gr.themes.Origin(),css=css) as demo:
                 minimum=10,
                 maximum=1000,
                 step=10,
-                value=180,
+                value=170,
             )
 
     gr.Examples(
@@ -571,7 +551,6 @@ 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,
@@ -582,7 +561,7 @@ with gr.Blocks(theme=gr.themes.Origin(),css=css) as demo:
             guidance_scale,
             num_inference_steps,
             randomize_seed, 
-            denoise
+            latent_file,
         ],
         outputs=[result, seed],
     )
@@ -594,7 +573,6 @@ 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,
@@ -605,7 +583,7 @@ with gr.Blocks(theme=gr.themes.Origin(),css=css) as demo:
             guidance_scale,
             num_inference_steps,
             randomize_seed, 
-            denoise
+            latent_file,
         ],
         outputs=[result, seed],
     )
@@ -617,7 +595,6 @@ 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,
@@ -628,7 +605,7 @@ with gr.Blocks(theme=gr.themes.Origin(),css=css) as demo:
             guidance_scale,
             num_inference_steps,
             randomize_seed, 
-            denoise
+            latent_file,
         ],
         outputs=[result, seed],
     )