import gradio as gr from PIL import Image import torch from diffusers import StableDiffusionPipeline # Load the diffusion pipeline from Hugging Face model_name = "Yaquv/rickthenpc" device = "cuda" if torch.cuda.is_available() else "cpu" try: pipe = StableDiffusionPipeline.from_pretrained(model_name) pipe = pipe.to(device) except Exception as e: print(f"Error loading the model: {e}") pipe = None def generate_image(prompt): """ Generates an image from the given prompt using the Hugging Face model. """ if pipe is None: raise ValueError("The model couldn't be loaded.") try: # Generate the image result = pipe(prompt) # Check that the result contains images if not hasattr(result, 'images') or len(result.images) == 0: raise ValueError("The model couldn't generate an image.") image = result.images[0] # Ensure the image is in PIL.Image format and convert to RGB if not isinstance(image, Image.Image): image = Image.fromarray(image) image = image.convert("RGB") return image except Exception as e: # Raise an exception for Gradio to handle raise ValueError(f"Error during image generation: {str(e)}") # Define the Gradio Interface iface = gr.Interface( fn=generate_image, inputs=gr.Textbox( label="Prompt", lines=2, placeholder="Enter your prompt here..." ), outputs=gr.Image( label="Generated Image", type="pil" # Ensure the output is a PIL Image ), title="Rick Generator", description="Enter a prompt to generate an image with the Rick Generator model." ) # Launch the Gradio app if __name__ == "__main__": iface.launch()