Krebzonide commited on
Commit
08c17b3
·
1 Parent(s): 45bbd2a

Added half of model changing code

Browse files

I copy pasted some stuff on my phone during my 15 min break at work

Files changed (1) hide show
  1. app.py +18 -0
app.py CHANGED
@@ -54,6 +54,24 @@ def generate(prompt, neg_prompt, samp_steps, guide_scale, batch_size, seed, heig
54
  ).images
55
  return [(img, f"Image {i+1}") for i, img in enumerate(images)]
56
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
57
 
58
  with gr.Blocks(css=css) as demo:
59
  with gr.Column():
 
54
  ).images
55
  return [(img, f"Image {i+1}") for i, img in enumerate(images)]
56
 
57
+ def set_base_model(base_model_id):
58
+ del pipe
59
+ torch.cuda.empty_cache()
60
+ gc.collect()
61
+ pipe = load_pipe(base_model_id)
62
+ return pipe
63
+
64
+ def load_pipe(base_model_id):
65
+ model_id = CONTROLNET_MODEL_IDS[task_name]
66
+ pipe = DiffusionPipeline.from_pretrained(
67
+ base_model_id,
68
+ safety_checker=None,
69
+ require_safety_checker=False,
70
+ torch_dtype=torch.float16)
71
+ pipe.scheduler = UniPCMultistepScheduler.from_config(
72
+ pipe.scheduler.config)
73
+ pipe.to("cuda")
74
+ return pipe
75
 
76
  with gr.Blocks(css=css) as demo:
77
  with gr.Column():