Update app.py

app.py

@@ -34,7 +34,14 @@ from torchvision import transforms
 from models.controlnet import ControlNetModel
 from models.unet_2d_condition import UNet2DConditionModel
 
+# VLM_NAME  = "Qwen/Qwen2.5-VL-3B-Instruct"
 
+# vlm_model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
+#     VLM_NAME,
+#     torch_dtype="auto",
+#     device_map="auto"   # immediately dispatches layers onto available GPUs
+# )
+# vlm_processor = AutoProcessor.from_pretrained(VLM_NAME)
 
 def _generate_vlm_prompt(
     vlm_model: Qwen2_5_VLForConditionalGeneration,
@@ -107,25 +114,25 @@ snapshot_download(
 
 
 snapshot_download(
-    repo_id="stabilityai/sd-turbo", 
-    local_dir="preset/models/sd-turbo"
+    repo_id="stabilityai/stable-diffusion-2-1-base", 
+    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/sd-turbo'
+pretrained_model_path = 'preset/models/stable-diffusion-2-1-base'
 seesr_model_path = 'preset/models/seesr'
 
 scheduler = DDIMScheduler.from_pretrained(pretrained_model_path, subfolder="scheduler")
 text_encoder = CLIPTextModel.from_pretrained(pretrained_model_path, subfolder="text_encoder")
 tokenizer = CLIPTokenizer.from_pretrained(pretrained_model_path, subfolder="tokenizer")
 vae = AutoencoderKL.from_pretrained(pretrained_model_path, subfolder="vae")
-# feature_extractor = CLIPImageProcessor.from_pretrained(f"{pretrained_model_path}/feature_extractor")
+feature_extractor = CLIPImageProcessor.from_pretrained(f"{pretrained_model_path}/feature_extractor")
 unet = UNet2DConditionModel.from_pretrained(seesr_model_path, subfolder="unet")
 controlnet = ControlNetModel.from_pretrained(seesr_model_path, subfolder="controlnet")
 
@@ -185,9 +192,9 @@ def magnify(
     user_prompt = "",
     positive_prompt = "clean, high-resolution, 8k, best quality, masterpiece",
     negative_prompt = "dotted, noise, blur, lowres, oversmooth, longbody, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality",
-    num_inference_steps = 2,
+    num_inference_steps = 50,
     scale_factor = 4,
-    cfg_scale = 1,
+    cfg_scale = 7.5,
     seed = 123,
     latent_tiled_size = 320,
     latent_tiled_overlap = 4,
@@ -288,15 +295,15 @@ with gr.Blocks(css=css, theme=theme) as demo:
                 input_image = gr.Image(type="pil", height=512)
                 run_button = gr.Button("🔎 Magnify 4x", variant="primary")
                 duration_time = gr.Text(label="duration time", value=60, visible=False)
-                with gr.Accordion("Options"):
+                with gr.Accordion("Options", visible=False):
                     user_prompt = gr.Textbox(label="User Prompt", value="")
                     positive_prompt = gr.Textbox(label="Positive Prompt", value="clean, high-resolution, 8k, best quality, masterpiece")
                     negative_prompt = gr.Textbox(
                         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 to 1.0 in sd-turbo)", minimum=1, maximum=10, value=1, step=0)
-                    num_inference_steps = gr.Slider(label="Inference Steps", minimum=2, maximum=100, value=2, step=1)
+                    cfg_scale = gr.Slider(label="Classifier Free Guidance Scale (Set to 1.0 in sd-turbo)", minimum=1, maximum=10, value=7.5, step=0)
+                    num_inference_steps = gr.Slider(label="Inference Steps", minimum=2, maximum=100, value=50, 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)
@@ -331,7 +338,7 @@ with gr.Blocks(css=css, theme=theme) as demo:
         inputs = [
             input_image,
         ]
-        run_button.click(fn=magnify, inputs=[input_image,user_prompt,positive_prompt,negative_prompt,num_inference_steps, scale_factor, cfg_scale] , outputs=[result_gallery])
+        run_button.click(fn=magnify, inputs=input_image, outputs=[result_gallery])
         input_image.upload(fn=preprocess_image,inputs=input_image, outputs=input_image)
 
 demo.launch(share=True)