Added refiner

app.py

@@ -1,4 +1,4 @@
-from diffusers import StableDiffusionXLPipeline, AutoencoderKL
+from diffusers import StableDiffusionXLPipeline, StableDiffusionXLImg2ImgPipeline, AutoencoderKL
 import torch
 import random
 #from controlnet_aux import OpenposeDetector
@@ -9,8 +9,8 @@ import gradio as gr
 model_base = "stabilityai/stable-diffusion-xl-base-1.0"
 
 #model_url = "https://huggingface.co/Krebzonide/Colossus_Project_XL/blob/main/colossusProjectXLSFW_v202BakedVAE.safetensors"
-#model_url = "https://huggingface.co/Krebzonide/Sevenof9_v3_sdxl/blob/main/nsfwSevenof9V3_nsfwSevenof9V3.safetensors"
-model_url = "https://civitai.com/api/download/models/182513?type=Model&format=SafeTensor&size=full&fp=fp16"
+model_url = "https://huggingface.co/Krebzonide/Sevenof9_v3_sdxl/blob/main/nsfwSevenof9V3_nsfwSevenof9V3.safetensors"
+
 vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
 
 #pipe = StableDiffusionXLPipeline.from_pretrained(
@@ -26,7 +26,15 @@ pipe = StableDiffusionXLPipeline.from_single_file(
     scheduler_type = "ddim",
     use_auth_token="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"
 )
-pipe = pipe.to("cuda")
+pipe.enable_model_cpu_offload()
+
+pipeRefiner = StableDiffusionXLImg2ImgPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-xl-refiner-1.0",
+    torch_dtype=torch.float16,
+    variant="fp16",
+    use_safetensors=True
+)
+pipeRefiner.enable_model_cpu_offload()
 
 css = """
 .btn-green {
@@ -49,11 +57,18 @@ def generate(prompt, neg_prompt, samp_steps, guide_scale, batch_size, seed, heig
         guidance_scale=guide_scale,
         #cross_attention_kwargs={"scale": lora_scale},
         num_images_per_prompt=batch_size,
-        height=height,
-        width=width,
+        height=height/2,
+        width=width/2,
         generator=torch.manual_seed(seed),
     ).images
-    return [(img, f"Image {i+1}") for i, img in enumerate(images)]
+    imagesRefined = pipeRefiner(
+        prompt,
+        image=images,
+        num_inference_steps=5,
+        height=height,
+        width=width
+    ).images
+    return [(img, f"Image {i+1}") for i, img in enumerate(imagesRefined)]
         
 
 with gr.Blocks(css=css) as demo:
@@ -67,8 +82,8 @@ with gr.Blocks(css=css) as demo:
             batch_size = gr.Slider(1, 6, value=1, step=1, label="Batch size")
         with gr.Row():
             seed = gr.Slider(label="Seed", value=-1, minimum=-1, maximum=999999, step=1)
-            height = gr.Slider(label="Height", value=1024, minumum=512, maximum=2048, step=8)
-            width = gr.Slider(label="Width", value=1024, minumum=512, maximum=2048, step=8)
+            height = gr.Slider(label="Height", value=1024, minimum=512, maximum=2048, step=16)
+            width = gr.Slider(label="Width", value=1024, minimum=512, maximum=2048, step=16)
         gallery = gr.Gallery(label="Generated images", height=800)
 
     submit_btn.click(generate, [prompt, negative_prompt, samp_steps, guide_scale, batch_size, seed, height, width], [gallery], queue=True)