Update app.py

app.py

@@ -1,31 +1,67 @@
+import torch
+from diffusers import StableDiffusionPipeline
 import gradio as gr
-from llama_cpp import Llama
-import os
 
-# Load the GGUF model from the Hugging Face Space
-MODEL_PATH = "eric707/jibjab"
-
-# Initialize Llama model
-llm = Llama(model_path=MODEL_PATH)
+# ---------------------------------------------------------------------------
+# 1. Load the Stable Diffusion model from Hugging Face
+#    - We specify "runwayml/stable-diffusion-v1-5" as an example.
+#    - Use "revision='fp16'" and "torch_dtype=torch.float16" to use the half-precision weights.
+#    - .to('cuda') if GPU is available, else .to('cpu').
+# ---------------------------------------------------------------------------
+try:
+    pipe = StableDiffusionPipeline.from_pretrained(
+        "eric707/jibjab",
+        revision="fp16",
+        torch_dtype=torch.float16
+    ).to("cuda")
+    device = "cuda"
+except:
+    # If CUDA is not available, fall back to CPU (VERY slow for SD, but works in a pinch).
+    pipe = StableDiffusionPipeline.from_pretrained(
+        "runwayml/stable-diffusion-v1-5",
+        revision="fp16"
+        # If you're on CPU, you might remove the torch_dtype for better compatibility:
+        # torch_dtype=torch.float16 -> Not recommended on CPU
+    ).to("cpu")
+    device = "cpu"
 
+# ---------------------------------------------------------------------------
+# 2. Define a function to generate images given a prompt.
+#    - We'll keep things simple and only accept a single prompt string.
+#    - Feel free to modify the inference steps, guidance scale, image size, etc.
+# ---------------------------------------------------------------------------
 def generate_image(prompt):
-    """
-    Function to process input text and return a generated image.
-    """
-    response = llm(f"Generate an image of {prompt}")
-    return f"Model response: {response}"  # Replace this with actual image generation
+    # Lower the inference steps or guidance scale if you run out of memory
+    image = pipe(
+        prompt, 
+        num_inference_steps=30, 
+        guidance_scale=7.5
+    ).images[0]
+    return image
 
-# Create Gradio UI
+# ---------------------------------------------------------------------------
+# 3. Build the Gradio UI
+#    - We use a Textbox for user input,
+#      and an Image component for displaying the generated image.
+# ---------------------------------------------------------------------------
 with gr.Blocks() as demo:
-    gr.Markdown("# GGUF Image Generation")
+    gr.Markdown("## Stable Diffusion Image Generation")
+
     with gr.Row():
-        input_text = gr.Textbox(label="Enter Prompt")
-        submit_btn = gr.Button("Generate Image")
-    
-    output_text = gr.Textbox(label="Model Response")
-    
-    submit_btn.click(generate_image, inputs=input_text, outputs=output_text)
+        with gr.Column():
+            prompt_input = gr.Textbox(
+                label="Enter a prompt to generate an image",
+                placeholder="A photo of an astronaut riding a horse on Mars"
+            )
+            generate_button = gr.Button("Generate Image")
+        with gr.Column():
+            output_image = gr.Image(label="Generated Image")
+
+    generate_button.click(fn=generate_image, inputs=prompt_input, outputs=output_image)
 
-# Launch the Gradio app
+# ---------------------------------------------------------------------------
+# 4. Launch the Gradio app
+# ---------------------------------------------------------------------------
 if __name__ == "__main__":
+    # By default, .launch() will pick up the PORT from the environment if on HF Spaces
     demo.launch()