Update app.py

app.py

@@ -19,20 +19,22 @@ model_id = "CompVis/stable-diffusion-v1-4"
 # register_free_crossattn_upblock2d(pip_freeu, b1=1.2, b2=1.4, s1=0.9, s2=0.2)
 # # -------- freeu block registration
 
-def infer(prompt, pip_sd, pip_freeu, seed, b1, b2, s1, s2):
-
- # # -------- freeu block registration
- # register_free_upblock2d(pip_freeu, b1=b1, b2=b2, s1=s1, s2=s1)
- # register_free_crossattn_upblock2d(pip_freeu, b1=b1, b2=b2, s1=s1, s2=s1)
- # # -------- freeu block registration
+def infer(prompt, pip, seed, b1, b2, s1, s2):
 
+ 
+ # register_free_upblock2d(pip, b1=1.0, b2=1.0, s1=1.0, s2=1.0)
+ # register_free_crossattn_upblock2d(pip, b1=1.0, b2=1.0, s1=1.0, s2=1.0)
+ 
  torch.manual_seed(seed)
  print("Generating SD:")
- sd_image = pip_sd(prompt).images[0] 
+ sd_image = pip(prompt).images[0] 
+
+ # register_free_upblock2d(pip, b1=b1, b2=b2, s1=s1, s2=s1)
+ # register_free_crossattn_upblock2d(pip, b1=b1, b2=b2, s1=s1, s2=s1)
 
  torch.manual_seed(seed)
  print("Generating FreeU:")
- freeu_image = pip_freeu(prompt).images[0] 
+ freeu_image = pip(prompt).images[0] 
 
  # First SD, then freeu
  images = [sd_image, freeu_image]
@@ -114,20 +116,20 @@ with block:
  else:
  model_id = "CompVis/stable-diffusion-v1-4"
  
- pip_sd = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
- pip_sd = pip_sd.to("cuda")
- 
- pip_freeu = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
- pip_freeu = pip_freeu.to("cuda")
+ pip = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
+ pip = pip_sd.to("cuda")
  
  with gr.Row():
- text = gr.Textbox(
- label="Enter your prompt",
- show_label=False,
- max_lines=1,
- placeholder="Enter your prompt",
- container=False,
- )
+ with gr.Column():
+ text = gr.Textbox(
+ label="Enter your prompt",
+ show_label=False,
+ max_lines=1,
+ placeholder="Enter your prompt",
+ container=False,
+ )
+ btn = gr.Button("Generate image", scale=0)
+ 
  seed = gr.Slider(label='seed',
  minimum=0,
  maximum=1000,
@@ -162,15 +164,14 @@ with block:
  
  with gr.Row():
  with gr.Group():
- btn = gr.Button("Generate image", scale=0)
+ # btn = gr.Button("Generate image", scale=0)
  with gr.Row():
  with gr.Column(min_width=256) as c1:
  image_1 = gr.Image(interactive=False)
  image_1_label = gr.Markdown("SD")
  
  with gr.Group():
- btn = gr.Button("Generate image", scale=0)
- 
+ # btn = gr.Button("Generate image", scale=0)
  with gr.Row():
  with gr.Column(min_width=256) as c2:
  image_2 = gr.Image(interactive=False)