update: app.py

app.py

@@ -1,4 +1,3 @@
-import os
 import shutil
 
 import gradio as gr
@@ -7,7 +6,7 @@ import torch
 from diffusers import DiffusionPipeline
 
 import wandb
-from wandb_addons.diffusers import get_wandb_callback
+from wandb.integration.diffusers import autolog
 
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -36,21 +35,9 @@ def generate_image(
     guidance_rescale,
 ):
     if not (wandb_api_key is None or wandb_api_key == ""):
-        os.environ["WANDB_API_KEY"] = wandb_api_key
+        wandb.login(key=wandb_api_key, relogin=True)
         generator = torch.Generator(device="cuda").manual_seed(seed)
-        configs = {
-            "guidance_scale": guidance_scale,
-            "guidance_rescale": guidance_rescale,
-            "seed": seed,
-        }
-        callback = get_wandb_callback(
-            pipeline,
-            prompt=prompt,
-            negative_prompt=negative_prompt,
-            wandb_project=wandb_project,
-            num_inference_steps=num_inference_steps,
-            configs=configs,
-        )
+        autolog(init={"project": wandb_project})
         run_url = wandb.run.get_url()
         image = pipeline(
             prompt,
@@ -59,9 +46,10 @@ def generate_image(
             height=height,
             width=width,
             num_inference_steps=num_inference_steps,
-            callback=callback,
-            guidance_rescale=configs["guidance_rescale"],
+            guidance_scale=guidance_scale,
+            guidance_rescale=guidance_rescale,
         ).images[0]
+        wandb.finish()
         if torch.cuda.is_available():
             torch.cuda.empty_cache()
         shutil.rmtree("wandb")