update to diffusers code

app.py

@@ -4,8 +4,8 @@ import uuid
 import gradio as gr
 from PIL import Image  
 import torch
-from muse import PipelineMuse, MaskGiTUViT, VQGANModel
 from compel import Compel, ReturnedEmbeddingsType
+from diffusers import DiffusionPipeline
 
 # from swin_ir_2 import load_model, preprocesss_image, postprocess_image
 
@@ -23,24 +23,12 @@ def save_images(image_array):
     return paths
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
-# pipe = PipelineMuse.from_pretrained("openMUSE/muse-laiona6-uvit-clip-220k").to(device)
 
-pipe = PipelineMuse.from_pretrained(
-    transformer_path="valhalla/research-run",
-    text_encoder_path="openMUSE/clip-vit-large-patch14-text-enc",
-    vae_path="openMUSE/vqgan-f16-8192-laion",
+pipe = DiffusionPipeline.from_pretrained(
+    "amused/amused-512",
+    variant="fp16",
+    torch_dtype=torch.float16,
 ).to(device)
-pipe.transformer = MaskGiTUViT.from_pretrained("valhalla/research-run-finetuned-journeydb", subfolder="ema_model", revision="06bcd6ab6580a2ed3275ddfc17f463b8574457da").to(device)
-pipe.vae = VQGANModel.from_pretrained("valhalla/vqgan-finetune-512-2").to(device)
-pipe.tokenizer.pad_token_id = 49407
-
-# sr_model = load_model().to(device)
-
-if device == "cuda":
-    pipe.text_encoder.to(torch.float16)
-    pipe.transformer.to(torch.float16)
-    pipe.transformer.enable_xformers_memory_efficient_attention()
-
 
 compel = Compel(tokenizer=pipe.tokenizer, text_encoder=pipe.text_encoder, returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED, requires_pooled=True, truncate_long_prompts=False)
 
@@ -52,22 +40,13 @@ def infer(prompt, negative="", scale=10, progress=gr.Progress(track_tqdm=True)):
     conditioning, negative_conditioning = compel.pad_conditioning_tensors_to_same_length([conditioning, negative_conditioning])
 
     images = pipe(
-        prompt,
-        timesteps=16,
-        negative_text=negative,
         prompt_embeds=conditioning, 
-        pooled_embeds=pooled,
+        encoder_hidden_states=pooled,
         negative_prompt_embeds=negative_conditioning,
-        negative_pooled_embeds=negative_pooled,
+        negative_encoder_hidden_states=negative_pooled,
         guidance_scale=scale,
         num_images_per_prompt=4,
         temperature=(3, 1),
-        orig_size=(512, 512),
-        crop_coords=(0, 0),
-        aesthetic_score=6,
-        use_fp16=device == "cuda",
-        transformer_seq_len=1024,
-        use_tqdm=True,
     )
     print("Done Generating!")
     print("Num Images:", len(images))