Update app.py

Switch to diffusers testing

app.py

@@ -1,109 +1,109 @@
 import gradio as gr
 
-# import torch
+import torch
 
-# from diffusers import UniDiffuserPipeline
+from diffusers import UniDiffuserPipeline
 
 
-# device = 'cuda' if torch.cuda.is_available() else 'cpu'
-# model_id = "dg845/unidiffuser-diffusers"
-# # model_id = "dg845/unidiffuser-diffusers-v0"
-# pipeline = UniDiffuserPipeline.from_pretrained(
-#     model_id,
-# )
-# pipeline.to(device)
-
-
-# def convert_to_none(s):
-#     if s:
-#         return s
-#     else:
-#         return None
-
-
-# def set_mode(mode):
-#     if mode == "joint":
-#         pipeline.set_joint_mode()
-#     elif mode == "text2img":
-#         pipeline.set_text_to_image_mode()
-#     elif mode == "img2text":
-#         pipeline.set_image_text_mode()
-#     elif mode == "text":
-#         pipeline.set_text_mode()
-#     elif mode == "img":
-#         pipeline.set_image_mode()
-
-
-# def sample(mode, prompt, image, num_inference_steps, guidance_scale, seed):
-#     set_mode(mode)
-#     prompt = convert_to_none(prompt)
-#     image = convert_to_none(image)
-#     generator = torch.Generator(device=device).manual_seed(seed)
-#     output_sample = pipeline(
-#         prompt=prompt,
-#         image=image,
-#         num_inference_steps=num_inference_steps,
-#         guidance_scale=guidance_scale,
-#         generator=generator,
-#     )
-#     sample_image = None
-#     sample_text = ""
-#     if output_sample.images is not None:
-#         sample_image = output_sample.images[0]
-#     if output_sample.text is not None:
-#         sample_text = output_sample.text[0]
-#     return sample_image, sample_text
-    
-
-# iface = gr.Interface(
-#     fn=sample,
-#     inputs=[
-#         gr.Textbox(value="", label="Generation Task"),
-#         gr.Textbox(value="", label="Conditioning prompt"),
-#         gr.Image(value=None, label="Conditioning image", type="pil"),
-#         gr.Number(value=20, label="Num Inference Steps", precision=0),
-#         gr.Number(value=8.0, label="Guidance Scale"),
-#         gr.Number(value=0, label="Seed", precision=0),
-#     ],
-#     outputs=[
-#         gr.Image(label="Sample image"),
-#         gr.Textbox(label="Sample text"),
-#     ],
-# )
-# iface.launch()
-
-# from unidiffuser.sample_v0 import sample
-# from unidiffuser.sample_v0_test import sample
-# from unidiffuser.sample_v1 import sample
-from unidiffuser.sample_v1_test import sample
-
-
-def predict(mode, prompt, image, sample_steps, guidance_scale, seed):
-    output_images, output_text = sample(
-        mode, prompt, image, sample_steps=sample_steps, scale=guidance_scale, seed=seed,
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+model_id = "dg845/unidiffuser-diffusers"
+# model_id = "dg845/unidiffuser-diffusers-v0"
+pipeline = UniDiffuserPipeline.from_pretrained(
+    model_id,
+)
+pipeline.to(device)
+
+
+def convert_to_none(s):
+    if s:
+        return s
+    else:
+        return None
+
+
+def set_mode(mode):
+    if mode == "joint":
+        pipeline.set_joint_mode()
+    elif mode == "text2img":
+        pipeline.set_text_to_image_mode()
+    elif mode == "img2text":
+        pipeline.set_image_text_mode()
+    elif mode == "text":
+        pipeline.set_text_mode()
+    elif mode == "img":
+        pipeline.set_image_mode()
+
+
+def sample(mode, prompt, image, num_inference_steps, guidance_scale, seed):
+    set_mode(mode)
+    prompt = convert_to_none(prompt)
+    image = convert_to_none(image)
+    generator = torch.Generator(device=device).manual_seed(seed)
+    output_sample = pipeline(
+        prompt=prompt,
+        image=image,
+        num_inference_steps=num_inference_steps,
+        guidance_scale=guidance_scale,
+        generator=generator,
     )
     sample_image = None
     sample_text = ""
-    if output_images is not None:
-        sample_image = output_images[0]
-    if output_text is not None:
-        sample_text = output_text[0]
+    if output_sample.images is not None:
+        sample_image = output_sample.images[0]
+    if output_sample.text is not None:
+        sample_text = output_sample.text[0]
     return sample_image, sample_text
-
+    
 
 iface = gr.Interface(
-    fn=predict,
+    fn=sample,
     inputs=[
         gr.Textbox(value="", label="Generation Task"),
         gr.Textbox(value="", label="Conditioning prompt"),
-        gr.Image(value=None, label="Conditioning image", type="filepath"),
-        gr.Number(value=50, label="Num Inference Steps", precision=0),
-        gr.Number(value=7.0, label="Guidance Scale"),
-        gr.Number(value=1234, label="Seed", precision=0),
+        gr.Image(value=None, label="Conditioning image", type="pil"),
+        gr.Number(value=20, label="Num Inference Steps", precision=0),
+        gr.Number(value=8.0, label="Guidance Scale"),
+        gr.Number(value=0, label="Seed", precision=0),
     ],
     outputs=[
         gr.Image(label="Sample image"),
         gr.Textbox(label="Sample text"),
     ],
 )
-iface.launch()
+iface.launch()
+
+# from unidiffuser.sample_v0 import sample
+# from unidiffuser.sample_v0_test import sample
+# from unidiffuser.sample_v1 import sample
+# from unidiffuser.sample_v1_test import sample
+
+
+# def predict(mode, prompt, image, sample_steps, guidance_scale, seed):
+#     output_images, output_text = sample(
+#         mode, prompt, image, sample_steps=sample_steps, scale=guidance_scale, seed=seed,
+#     )
+#     sample_image = None
+#     sample_text = ""
+#     if output_images is not None:
+#         sample_image = output_images[0]
+#     if output_text is not None:
+#         sample_text = output_text[0]
+#     return sample_image, sample_text
+
+
+# iface = gr.Interface(
+#     fn=predict,
+#     inputs=[
+#         gr.Textbox(value="", label="Generation Task"),
+#         gr.Textbox(value="", label="Conditioning prompt"),
+#         gr.Image(value=None, label="Conditioning image", type="filepath"),
+#         gr.Number(value=50, label="Num Inference Steps", precision=0),
+#         gr.Number(value=7.0, label="Guidance Scale"),
+#         gr.Number(value=1234, label="Seed", precision=0),
+#     ],
+#     outputs=[
+#         gr.Image(label="Sample image"),
+#         gr.Textbox(label="Sample text"),
+#     ],
+# )
+# iface.launch()