Update gradio_seesr.py

gradio_seesr.py

@@ -1,3 +1,4 @@
+import spaces
 import gradio as gr
 import os
 import sys
@@ -8,9 +9,29 @@ import numpy as np
 from PIL import Image
 
 import torch
+
+print(f'torch version:{torch.__version__}')
+
+# import subprocess
+# import importlib, site, sys
+
+# # Re-discover all .pth/.egg-link files
+# for sitedir in site.getsitepackages():
+#     site.addsitedir(sitedir)
+
+# # Clear caches so importlib will pick up new modules
+# importlib.invalidate_caches()
+
+# def sh(cmd): subprocess.check_call(cmd, shell=True)
+
+# sh("pip install -U xformers --index-url https://download.pytorch.org/whl/cu126")
+
+# # tell Python to re-scan site-packages now that the egg-link exists
+# import importlib, site; site.addsitedir(site.getsitepackages()[0]); importlib.invalidate_caches()
+
 import torch.utils.checkpoint
 from pytorch_lightning import seed_everything
-from diffusers import AutoencoderKL, DDPMScheduler
+from diffusers import AutoencoderKL, DDIMScheduler
 from diffusers.utils import check_min_version
 from diffusers.utils.import_utils import is_xformers_available
 from transformers import CLIPTextModel, CLIPTokenizer, CLIPImageProcessor
@@ -36,7 +57,7 @@ ram_transforms = transforms.Compose([
         ])
 
 snapshot_download(
-    repo_id="CSWRY/SeeSR", 
+    repo_id="alexnasa/SEESR", 
     local_dir="preset/models"
 )
 
@@ -46,6 +67,11 @@ snapshot_download(
     local_dir="preset/models/stable-diffusion-2-1-base"
 )
 
+snapshot_download(
+    repo_id="xinyu1205/recognize_anything_model", 
+    local_dir="preset/models/"
+)
+
 
 # Load scheduler, tokenizer and models.
 pretrained_model_path = 'preset/models/stable-diffusion-2-1-base'
@@ -65,15 +91,14 @@ text_encoder.requires_grad_(False)
 unet.requires_grad_(False)
 controlnet.requires_grad_(False)
 
-if is_xformers_available():
-    unet.enable_xformers_memory_efficient_attention()
-    controlnet.enable_xformers_memory_efficient_attention()
-else:
-    raise ValueError("xformers is not available. Make sure it is installed correctly")
+# unet.to("cuda")
+# controlnet.to("cuda")
+# unet.enable_xformers_memory_efficient_attention()
+# controlnet.enable_xformers_memory_efficient_attention()
 
 # Get the validation pipeline
 validation_pipeline = StableDiffusionControlNetPipeline(
-    vae=vae, text_encoder=text_encoder, tokenizer=tokenizer, feature_extractor=feature_extractor,
+    vae=vae, text_encoder=text_encoder, tokenizer=tokenizer, feature_extractor=None,
     unet=unet, controlnet=controlnet, scheduler=scheduler, safety_checker=None, requires_safety_checker=False,
 )
 
@@ -96,10 +121,12 @@ tag_model = ram(pretrained='preset/models/ram_swin_large_14m.pth',
 tag_model.eval()
 tag_model.to(device, dtype=weight_dtype)
 
-@torch.no_grad()
+@spaces.GPU()
 def process(
     input_image: Image.Image,
     user_prompt: str,
+    use_KDS: bool,
+    num_particles: int,
     positive_prompt: str,
     negative_prompt: str,
     num_inference_steps: int,
@@ -150,7 +177,8 @@ def process(
                     height=height, width=width,
                     guidance_scale=cfg_scale,  conditioning_scale=1,
                     start_point='lr', start_steps=999,ram_encoder_hidden_states=ram_encoder_hidden_states,
-                    latent_tiled_size=latent_tiled_size, latent_tiled_overlap=latent_tiled_overlap
+                    latent_tiled_size=latent_tiled_size, latent_tiled_overlap=latent_tiled_overlap, use_KDS=use_KDS,
+                    num_particles=num_particles
                 ).images[0]
 
             if True:  # alpha<1.0:
@@ -181,8 +209,10 @@ with block:
         gr.Markdown(MARKDOWN)
     with gr.Row():
         with gr.Column():
-            input_image = gr.Image(source="upload", type="pil")
-            run_button = gr.Button(label="Run")
+            input_image = gr.Image(type="pil")
+            num_particles = gr.Slider(label="Num of Partickes", minimum=1, maximum=16, step=1, value=4)
+            use_KDS = gr.Checkbox(label="Use Kernel Density Steering")
+            run_button = gr.Button("Run")
             with gr.Accordion("Options", open=True):
                 user_prompt = gr.Textbox(label="User Prompt", value="")
                 positive_prompt = gr.Textbox(label="Positive Prompt", value="clean, high-resolution, 8k, best quality, masterpiece")
@@ -190,19 +220,21 @@ with block:
                     label="Negative Prompt",
                     value="dotted, noise, blur, lowres, oversmooth, longbody, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality"
                 )
-                cfg_scale = gr.Slider(label="Classifier Free Guidance Scale (Set a value larger than 1 to enable it!)", minimum=0.1, maximum=10.0, value=5.5, step=0.1)
-                num_inference_steps = gr.Slider(label="Inference Steps", minimum=10, maximum=100, value=50, step=1)
+                cfg_scale = gr.Slider(label="Classifier Free Guidance Scale (Set to 1.0 in sd-turbo)", minimum=1, maximum=1, value=1, step=0)
+                num_inference_steps = gr.Slider(label="Inference Steps", minimum=2, maximum=8, value=2, step=1)
                 seed = gr.Slider(label="Seed", minimum=-1, maximum=2147483647, step=1, value=231)
                 sample_times = gr.Slider(label="Sample Times", minimum=1, maximum=10, step=1, value=1)
                 latent_tiled_size = gr.Slider(label="Diffusion Tile Size", minimum=128, maximum=480, value=320, step=1)
                 latent_tiled_overlap = gr.Slider(label="Diffusion Tile Overlap", minimum=4, maximum=16, value=4, step=1)
                 scale_factor = gr.Number(label="SR Scale", value=4)
         with gr.Column():
-            result_gallery = gr.Gallery(label="Output", show_label=False, elem_id="gallery").style(grid=2, height="auto")
+            result_gallery = gr.Gallery(label="Output", show_label=False, elem_id="gallery")
 
     inputs = [
         input_image,
         user_prompt,
+        use_KDS,
+        num_particles,
         positive_prompt,
         negative_prompt,
         num_inference_steps,
@@ -215,5 +247,5 @@ with block:
     ]
     run_button.click(fn=process, inputs=inputs, outputs=[result_gallery])
 
-block.launch()
+block.launch(share=True)