Update app.py

app.py

@@ -3,13 +3,13 @@ import torch
 from diffusers import StableDiffusionPipeline
 
 # Load the model (replace with your actual model path)
-model_name = "alvdansen/midsommarcartoon"  # Example model
+model_name = "ZB-Tech/Text-to-Image"  # Specified model
 device = "cpu"  # Ensure we are using CPU
 
 print(f"Using device: {device}")
 
-# Load the Stable Diffusion model
 try:
+    # Load the Stable Diffusion model
     print("Loading model...")
     pipe = StableDiffusionPipeline.from_pretrained(model_name).to(device)
     print("Model loaded successfully.")
@@ -20,14 +20,12 @@ except Exception as e:
 def generate_image(prompt):
     try:
         print(f"Generating image for prompt: '{prompt}'")
-        # Generate image based on the prompt
         with torch.no_grad():
-            # Ensure this line correctly calls the model's generate method
-            image = pipe(prompt).images[0]  # Adjust this line if necessary based on your model's output
-        return image  # Return the generated image
+            image = pipe(prompt).images[0]  # Adjust this line based on your model's output structure
+        return image
     except Exception as e:
         print(f"Error generating image: {e}")
-        return None  # Return None or a placeholder image in case of error
+        return None
 
 # Gradio interface setup
 gr.Interface(