GPU version

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+checkpoint_000.pth filter=lfs diff=lfs merge=lfs -text
+checkpoint_*.pth filter=lfs diff=lfs merge=lfs -text

app.py

@@ -1,17 +1,34 @@
+import socketserver
+socketserver.TCPServer.allow_reuse_address = True
+
 import gradio as gr
-from optimum.intel.openvino import OVStableDiffusionPipeline
 
-pipe = OVStableDiffusionPipeline.from_pretrained("OpenVINO/stable-diffusion-2-1-quantized", compile=False)
-pipe.reshape(batch_size=1, height=512, width=512, num_images_per_prompt=1)
-pipe.compile()
+import torch
+
+from diffusers import StableDiffusionPipeline
+
+pipe = StableDiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1-base", torch_dtype=torch.float16)
+
+# load the patched VQ-VAE
+patched_decoder_ckpt = "checkpoint_000.pth"
+
+if patched_decoder_ckpt is not None:
+    sd2 = torch.load(patched_decoder_ckpt)['ldm_decoder']
+    #print("patching keys for first_stage_model: ", sd2.keys())
+    
+    msg = pipe.vae.load_state_dict(sd2, strict=False)
+    print(f"loaded LDM decoder state_dict with message\n{msg}")
+    print("you should check that the decoder keys are correctly matched")
+
+pipe = pipe.to("cuda")
 
 prompt = "sailing ship in storm by Rembrandt"
 
-def generate(name):
+def generate(prompt):
     output = pipe(prompt, num_inference_steps=50, output_type="pil")
     output.images[0].save("result.png")
     return output.images[0]
 
 iface = gr.Interface(fn=generate, inputs=[gr.Textbox(label="Prompt", value=prompt)], outputs=[gr.Image(type="pil")])
-iface.launch()
+iface.launch(server_name="0.0.0.0")