Change to diffusion

app.py

@@ -1,24 +1,68 @@
+from diffusers import DiffusionPipeline
+from diffusers.utils import pt_to_pil
 import gradio as gr
-from transformers import pipeline
+import torch
 
-pipeline = pipeline(task="image-classification", model="julien-c/hotdog-not-hotdog")
+
+stage_1 = DiffusionPipeline.from_pretrained(
+    "DeepFloyd/IF-I-M-v1.0", variant="fp16", torch_dtype=torch.float16
+)
+stage_1.enable_xformers_memory_efficient_attention()  # remove line if torch.__version__ >= 2.0.0
+stage_1.enable_model_cpu_offload()
+stage_2 = DiffusionPipeline.from_pretrained(
+    "DeepFloyd/IF-II-M-v1.0",
+    text_encoder=None,
+    variant="fp16",
+    torch_dtype=torch.float16,
+)
+stage_2.enable_xformers_memory_efficient_attention()  # remove line if torch.__version__ >= 2.0.0
+stage_2.enable_model_cpu_offload()
+
+# stage 3
+safety_modules = {
+    "feature_extractor": stage_1.feature_extractor,
+    "safety_checker": stage_1.safety_checker,
+    "watermarker": stage_1.watermarker,
+}
+stage_3 = DiffusionPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-x4-upscaler",
+    **safety_modules,
+    torch_dtype=torch.float16
+)
+stage_3.enable_xformers_memory_efficient_attention()  # remove line if torch.__version__ >= 2.0.0
+stage_3.enable_model_cpu_offload()
 
 
 def predict(input_img):
-    predictions = pipeline(input_img)
-    return input_img, {p["label"]: p["score"] for p in predictions}
+    prompt = 'a photo of a kangaroo wearing an orange hoodie and blue sunglasses standing in front of the eiffel tower holding a sign that says "very deep learning"'
+
+    prompt_embeds, negative_embeds = stage_1.encode_prompt(prompt)
+    generator = torch.manual_seed(0)
+    image = stage_1(
+        prompt_embeds=prompt_embeds,
+        negative_prompt_embeds=negative_embeds,
+        generator=generator,
+        output_type="pt",
+    ).images
+    image = stage_2(
+        image=image,
+        prompt_embeds=prompt_embeds,
+        negative_prompt_embeds=negative_embeds,
+        generator=generator,
+        output_type="pt",
+    ).images
+    image = stage_3(
+        prompt=prompt, image=image, generator=generator, noise_level=100
+    ).images
+    return image
 
 
 gradio_app = gr.Interface(
-    predict,
-    inputs=gr.Image(
-        label="Select hot dog candidate", sources=["upload", "webcam"], type="pil"
-    ),
-    outputs=[
-        gr.Image(label="Processed Image"),
-        gr.Label(label="Result", num_top_classes=2),
-    ],
-    title="Hot Dog? Or Not?",
+    fn=predict,
+    inputs="text",
+    outputs="image",
+    title="Text to Image Generator",
+    description="Enter a text string to generate an image.",
 )
 
 if __name__ == "__main__":

requirements.txt

@@ -1,3 +1,8 @@
+accelerate
+diffusers
 gradio
+safetensors
+sentencepiece
 transformers
-torch
+torch
+xformers