Men1scus commited on
Commit
127f197
·
1 Parent(s): 1cf84a1

Prioritize fidelity model in pipeline selection and update user prompts for clarity

Browse files
Files changed (1) hide show
  1. app.py +16 -16
app.py CHANGED
@@ -293,13 +293,13 @@ dit4sr_device = 'cuda:0'
293
 
294
  llava_agent = LLavaAgent("liuhaotian/llava-v1.5-13b", LLaVA_device, load_8bit=True, load_4bit=False)
295
 
296
- # Get the validation pipeline
297
- pipeline_dit4sr_q = load_dit4sr_q_pipeline(args, dit4sr_device)
298
 
299
- # Only load dit4sr_f if load_both_models is True
300
- pipeline_dit4sr_f = None
301
  if args.load_both_models == 'True':
302
- pipeline_dit4sr_f = load_dit4sr_f_pipeline(args, dit4sr_device)
303
 
304
  @spaces.GPU
305
  @torch.no_grad()
@@ -349,11 +349,11 @@ def process_sr(
349
 
350
  images = []
351
 
352
- # Choose pipeline based on model selection
353
- if model_choice == "dit4sr_f" and pipeline_dit4sr_f is not None:
354
- pipeline = pipeline_dit4sr_f
355
- else:
356
  pipeline = pipeline_dit4sr_q
 
 
357
 
358
  try:
359
  image = pipeline(
@@ -387,7 +387,7 @@ Intro= \
387
  if args.load_both_models == 'True':
388
  Prompt = \
389
  """
390
- First, select your preferred model (dit4sr_q or dit4sr_f). \\
391
  Then, click \"Run LLAVA\" to generate an initial prompt based on the input image. \\
392
  Modify the prompt for higher accuracy if needed. \\
393
  Finally, click \"Run DiT4SR\" to generate the SR result." \
@@ -397,7 +397,7 @@ else:
397
  """
398
  Click \"Run LLAVA\" to generate an initial prompt based on the input image. \\
399
  Modify the prompt for higher accuracy if needed. \\
400
- Finally, click \"Run DiT4SR\" to generate the SR result using dit4sr_q model." \
401
  """
402
 
403
  exaple_images = sorted(glob.glob('examples/*.png'))
@@ -414,16 +414,16 @@ with block:
414
  if args.load_both_models == 'True':
415
  model_choice = gr.Dropdown(
416
  label="Model Selection",
417
- choices=["dit4sr_q", "dit4sr_f"],
418
- value="dit4sr_q",
419
- info="Choose between dit4sr_q (default) and dit4sr_f models"
420
  )
421
  else:
422
  # Hidden component with default value when only one model is available
423
  model_choice = gr.Dropdown(
424
  label="Model Selection",
425
- choices=["dit4sr_q"],
426
- value="dit4sr_q",
427
  visible=False
428
  )
429
 
 
293
 
294
  llava_agent = LLavaAgent("liuhaotian/llava-v1.5-13b", LLaVA_device, load_8bit=True, load_4bit=False)
295
 
296
+ # Get the validation pipeline - prioritize dit4sr_f
297
+ pipeline_dit4sr_f = load_dit4sr_f_pipeline(args, dit4sr_device)
298
 
299
+ # Only load dit4sr_q if load_both_models is True
300
+ pipeline_dit4sr_q = None
301
  if args.load_both_models == 'True':
302
+ pipeline_dit4sr_q = load_dit4sr_q_pipeline(args, dit4sr_device)
303
 
304
  @spaces.GPU
305
  @torch.no_grad()
 
349
 
350
  images = []
351
 
352
+ # Choose pipeline based on model selection - prioritize dit4sr_f
353
+ if model_choice == "dit4sr_q" and pipeline_dit4sr_q is not None:
 
 
354
  pipeline = pipeline_dit4sr_q
355
+ else:
356
+ pipeline = pipeline_dit4sr_f
357
 
358
  try:
359
  image = pipeline(
 
387
  if args.load_both_models == 'True':
388
  Prompt = \
389
  """
390
+ First, select your preferred model (fidelity first or quality first). \\
391
  Then, click \"Run LLAVA\" to generate an initial prompt based on the input image. \\
392
  Modify the prompt for higher accuracy if needed. \\
393
  Finally, click \"Run DiT4SR\" to generate the SR result." \
 
397
  """
398
  Click \"Run LLAVA\" to generate an initial prompt based on the input image. \\
399
  Modify the prompt for higher accuracy if needed. \\
400
+ Finally, click \"Run DiT4SR\" to generate the SR result using fidelity first model." \
401
  """
402
 
403
  exaple_images = sorted(glob.glob('examples/*.png'))
 
414
  if args.load_both_models == 'True':
415
  model_choice = gr.Dropdown(
416
  label="Model Selection",
417
+ choices=[("Quality First", "dit4sr_q"), ("Fidelity First", "dit4sr_f")],
418
+ value="dit4sr_f",
419
+ info="Choose between Quality First and Fidelity First models"
420
  )
421
  else:
422
  # Hidden component with default value when only one model is available
423
  model_choice = gr.Dropdown(
424
  label="Model Selection",
425
+ choices=["dit4sr_f"],
426
+ value="dit4sr_f",
427
  visible=False
428
  )
429