Upload app.py

app.py

@@ -59,10 +59,10 @@ def change_base_model(repo_id: str, cn_on: bool):
             #progress(0, desc=f"Loading model: {repo_id} / Loading ControlNet: {controlnet_model_union_repo}")
             print(f"Loading model: {repo_id} / Loading ControlNet: {controlnet_model_union_repo}")
             clear_cache()
-            controlnet_union = FluxControlNetModel.from_pretrained(controlnet_model_union_repo, torch_dtype=dtype).to(device)
-            controlnet = FluxMultiControlNetModel([controlnet_union]).to(device)
-            pipe = FluxControlNetPipeline.from_pretrained(repo_id, controlnet=controlnet, torch_dtype=dtype, vae=taef1).to(device)
-            pipe.flux_pipe_call_that_returns_an_iterable_of_images = flux_pipe_call_that_returns_an_iterable_of_images.__get__(pipe)
+            controlnet_union = FluxControlNetModel.from_pretrained(controlnet_model_union_repo, torch_dtype=dtype)
+            controlnet = FluxMultiControlNetModel([controlnet_union])
+            pipe = FluxControlNetPipeline.from_pretrained(repo_id, controlnet=controlnet, torch_dtype=dtype)
+            #pipe.flux_pipe_call_that_returns_an_iterable_of_images = flux_pipe_call_that_returns_an_iterable_of_images.__get__(pipe)
             last_model = repo_id
             last_cn_on = cn_on
             #progress(1, desc=f"Model loaded: {repo_id} / ControlNet Loaded: {controlnet_model_union_repo}")
@@ -71,7 +71,7 @@ def change_base_model(repo_id: str, cn_on: bool):
             #progress(0, desc=f"Loading model: {repo_id}")
             print(f"Loading model: {repo_id}")
             clear_cache()
-            pipe = DiffusionPipeline.from_pretrained(repo_id, torch_dtype=dtype, vae=taef1).to(device)
+            pipe = DiffusionPipeline.from_pretrained(repo_id, torch_dtype=dtype, vae=taef1)
             pipe.flux_pipe_call_that_returns_an_iterable_of_images = flux_pipe_call_that_returns_an_iterable_of_images.__get__(pipe)
             last_model = repo_id
             last_cn_on = cn_on
@@ -154,7 +154,7 @@ def generate_image(prompt_mash, steps, seed, cfg_scale, width, height, lora_scal
                 progress(0, desc="Start Inference with ControlNet.")
                 if controlnet is not None: controlnet.to("cuda")
                 if controlnet_union is not None: controlnet_union.to("cuda")
-                for img in pipe.flux_pipe_call_that_returns_an_iterable_of_images(
+                for img in pipe(
                     prompt=prompt_mash,
                     control_image=images,
                     control_mode=modes,
@@ -165,9 +165,7 @@ def generate_image(prompt_mash, steps, seed, cfg_scale, width, height, lora_scal
                     controlnet_conditioning_scale=scales,
                     generator=generator,
                     joint_attention_kwargs={"scale": lora_scale},
-                    output_type="pil",
-                    good_vae=good_vae,
-                ):
+                ).images:
                     yield img
     except Exception as e:
         print(e)