Bakare - Unit 8

app.py

@@ -1,39 +1,88 @@
-import gradio as gr
-from transformers import VisionEncoderDecoderModel, ViTImageProcessor, AutoTokenizer
-from PIL import Image
 import torch
+from diffusers import StableDiffusionPipeline
+import gradio as gr
+
+# -------------------------------------------------------
+# 1. LOAD PRETRAINED TEXT-TO-IMAGE MODEL
+# -------------------------------------------------------
 
-# Load model
-model_name = "nlpconnect/vit-gpt2-image-captioning"
-model = VisionEncoderDecoderModel.from_pretrained(model_name)
-processor = ViTImageProcessor.from_pretrained(model_name)
-tokenizer = AutoTokenizer.from_pretrained(model_name)
+model_id = "runwayml/stable-diffusion-v1-5"
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
-model.to(device)
+dtype = torch.float16 if device == "cuda" else torch.float32
+
+pipe = StableDiffusionPipeline.from_pretrained(
+    model_id,
+    torch_dtype=dtype,
+    safety_checker=None,
+    use_safetensors=True
+)
+
+pipe = pipe.to(device)
+
+
+# -------------------------------------------------------
+# 2. CORE PREDICTION FUNCTION
+# -------------------------------------------------------
+
+def generate_image(prompt: str,
+                   num_inference_steps: int = 25,
+                   guidance_scale: float = 7.5):
 
-# Caption function
-def predict_caption(image):
-    if image is None:
-        return "Upload an image."
-    if image.mode != "RGB":
-        image = image.convert("RGB")
+    if not prompt or prompt.strip() == "":
+        prompt = "A friendly robot reading a book in a cozy library, digital art"
 
-    pixel_values = processor(images=image, return_tensors="pt").pixel_values.to(device)
+    if device == "cuda":
+        with torch.autocast(device_type="cuda"):
+            result = pipe(
+                prompt,
+                num_inference_steps=num_inference_steps,
+                guidance_scale=guidance_scale
+            )
+    else:
+        result = pipe(
+            prompt,
+            num_inference_steps=num_inference_steps,
+            guidance_scale=guidance_scale
+        )
 
-    with torch.no_grad():
-        output_ids = model.generate(pixel_values, max_length=32, num_beams=4)
+    return result.images[0]
+
+
+# -------------------------------------------------------
+# 3. GRADIO UI
+# -------------------------------------------------------
+
+prompt_input = gr.Textbox(
+    label="Enter your image prompt",
+    lines=2,
+    placeholder="e.g., 'A watercolor painting of a sunrise over mountains'"
+)
+
+steps_slider = gr.Slider(
+    minimum=10,
+    maximum=40,
+    value=25,
+    step=1,
+    label="Number of inference steps"
+)
+
+guidance_slider = gr.Slider(
+    minimum=1.0,
+    maximum=15.0,
+    value=7.5,
+    step=0.5,
+    label="Guidance scale"
+)
 
-    caption = tokenizer.decode(output_ids[0], skip_special_tokens=True)
-    return caption.strip()
+image_output = gr.Image(label="Generated image")
 
-# UI
 demo = gr.Interface(
-    fn=predict_caption,
-    inputs=gr.Image(type="pil", label="Upload Image"),
-    outputs=gr.Textbox(label="Caption"),
-    title="AI Image Captioning",
-    description="Upload an image to get an AI-generated caption."
+    fn=generate_image,
+    inputs=[prompt_input, steps_slider, guidance_slider],
+    outputs=image_output,
+    title="Multimodal Text-to-Image Generator",
+    description="Enter a prompt to generate an image using a pretrained text-to-image model."
 )
 
 if __name__ == "__main__":

requirements.txt

@@ -1,4 +1,6 @@
-gradio
-transformers
+gradio>=4.0.0
+diffusers>=0.30.0
+transformers>=4.40.0
+accelerate>=0.30.0
 torch
-Pillow
+safetensors