update

README.md

@@ -15,4 +15,9 @@ pinned: false
 python run.py --model_id runwayml/stable-diffusion-v1-5
 python run.py --model_id stabilityai/stable-diffusion-2-1
 python run.py --model_id stabilityai/stable-diffusion-xl-base-1.0
+python run.py --model_id PixArt-alpha/PixArt-XL-2-512x512 
+python run.py --model_id PixArt-alpha/PixArt-XL-2-1024-MS 
+python run.py --model_id playgroundai/playground-v2-1024px-aesthetic
+python run.py --model_id kandinsky-community/kandinsky-3
+python run.py --model_id Lykon/dreamshaper-8
 ```

gradio_app.py

@@ -1,7 +1,60 @@
 import gradio as gr
+import os
+import csv
 
-def greet(name):
-    return "Hello " + name + "!!"
+root = 'saved'
+prompt_path = 'assets/ViLG-300.csv'
 
-iface = gr.Interface(fn=greet, inputs="text", outputs="text")
-iface.launch()
+
+def load_prompts(path):
+    if os.path.basename(path) == 'ViLG-300.csv':
+        def csv_to_dict(file_path):
+            result_dict = {}
+            with open(file_path, 'r', encoding='utf-8') as csv_file:
+                csv_reader = csv.DictReader(csv_file, delimiter=',')
+                for row in csv_reader:
+                    prompt = row['\ufeffPrompt']
+                    text = row['文本']
+                    category = row['类别']
+                    source = row['来源']
+                    result_dict[prompt] = {'prompt': prompt, 'text': text, 'category': category, 'source': source}
+            return result_dict
+        data = list(csv_to_dict(path).keys())
+    else:
+        return NotImplementedError
+    return data
+
+
+prompts = load_prompts(prompt_path)
+
+
+def load_images(methods, idx):
+    idx = int(idx)
+    prompt = prompts[idx].strip()
+    images = []
+    for method in methods:
+        image = os.path.join(root, method, f'{idx}.jpg')
+        images.append((image, method))
+    return prompt, images
+
+
+def load_methods():
+    methods = os.listdir(root)
+    return methods
+
+
+def main():
+    with gr.Blocks() as demo:
+        gr.Markdown("# Text to Image Models Comparison")
+        with gr.Row():
+            idx = gr.Number(value=0, label='Index')
+            prompt = gr.Textbox()
+        methods = gr.Dropdown(multiselect=True, choices=load_methods(), value=load_methods(), label='Methods')
+        gallery = gr.Gallery(show_label=False, object_fit='fill', height=600, columns=5)
+        idx.change(load_images, [methods, idx], [prompt, gallery])
+        methods.change(load_images, [methods, idx], [prompt, gallery])
+    demo.launch()
+
+
+if __name__ == '__main__':
+    main()

run.py

@@ -28,14 +28,19 @@ def main(
     model_id="runwayml/stable-diffusion-v1-5",
     prompt_path="assets/ViLG-300.csv",
     save_path=None,
+    dtype='fp16',
 ):
     if save_path is None:
         save_path = os.path.join('saved', model_id.replace('/', '_'))
         os.makedirs(save_path, exist_ok=True)
 
     prompts = load_prompts(prompt_path)
-    pipeline = AutoPipelineForText2Image.from_pretrained(model_id)
-    pipeline.to(device='cuda', dtype=torch.float16)
+    pipeline = AutoPipelineForText2Image.from_pretrained(
+        model_id, 
+        torch_dtype=torch.float32 if dtype == 'fp32' else torch.float16
+    )
+    pipeline.to(device='cuda')
+    pipeline.safety_checker = None
     for i, prompt in enumerate(prompts):
         print(f'{i}|{len(prompts)}: {prompt}')
         image = pipeline(prompt).images[0]