File size: 3,046 Bytes
7b1a432
 
 
 
 
02a3a52
7b1a432
dfab3a9
02a3a52
dfab3a9
7b1a432
dfab3a9
 
509d782
dfab3a9
dc25832
02a3a52
 
 
 
 
 
 
 
 
 
 
 
7b1a432
dfab3a9
4a6aac0
7b1a432
dc25832
 
 
 
dfab3a9
dc25832
dfab3a9
 
dc25832
 
 
 
dfab3a9
 
dc25832
dfab3a9
b9c75da
dc25832
 
b9c75da
23d76b5
dfab3a9
7b1a432
 
23d76b5
 
 
 
7b1a432
23d76b5
 
 
 
 
 
 
7b1a432
dfab3a9
7b1a432
 
 
 
 
 
 
509d782
7b1a432
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
import gradio as gr
import torch
import os
from diffusers import StableDiffusion3Pipeline
from safetensors.torch import load_file
from spaces import GPU  # Remove if not in HF Space

# 1. Define model ID and HF_TOKEN (at the VERY beginning)
model_id = "stabilityai/stable-diffusion-3.5-large"  # Or your preferred model ID
hf_token = os.getenv("HF_TOKEN")  # For private models (set in HF Space settings)

# 2. Initialize pipeline (to None initially)
pipeline = None

# 3. Load Stable Diffusion and LoRA (before Gradio)
try:
    if hf_token: # check if the token exists, if not, then do not pass the token
        pipeline = StableDiffusion3Pipeline.from_pretrained(
            model_id,
            torch_dtype=torch.float16,
            cache_dir="./model_cache"  # For caching
        )
    else:
        pipeline = StableDiffusion3Pipeline.from_pretrained(
            model_id,
            torch_dtype=torch.float16,
            cache_dir="./model_cache"  # For caching
        )

    lora_filename = "lora_trained_model.safetensors"  # EXACT filename of your LoRA
    lora_path = os.path.join("./", lora_filename)

    if os.path.exists(lora_path):
        lora_weights = load_file(lora_path)
        text_encoder = pipeline.text_encoder
        text_encoder.load_state_dict(lora_weights, strict=False)
        print(f"LoRA loaded successfully from: {lora_path}")
    else:
        print(f"Error: LoRA file not found at: {lora_path}")
        exit()  # Stop if LoRA is not found

    print("Stable Diffusion model loaded successfully!")

except Exception as e:
    print(f"Error loading model or LoRA: {e}")
    exit()  # Stop if model loading fails

# 4. Image generation function (now decorated)
@GPU(duration=65)  # Only if in HF Space
def generate_image(prompt):
    global pipeline
    if pipeline is None:
        print("Error: Pipeline is None (model not loaded)")  # Log this specifically
        return "Error: Model not loaded!"

    try:
        print("Starting image generation...")  # Log before the image generation
        image = pipeline(prompt).images[0]
        print("Image generated successfully!")
        return image
    except Exception as e:
        error_message = f"Error during image generation: {type(e).__name__}: {e}"  # Include exception type
        print(f"Full Error Details:\n{error_message}")  # Print full details
        return error_message  # Return error message to Gradio
    except RuntimeError as re:
        error_message = f"Runtime Error during image generation: {type(re).__name__}: {re}"  # Include exception type
        print(f"Full Runtime Error Details:\n{error_message}")  # Print full details
        return error_message  # Return error message to Gradio

# 5. Gradio interface
with gr.Blocks() as demo:
    prompt_input = gr.Textbox(label="Prompt")
    image_output = gr.Image(label="Generated Image")
    generate_button = gr.Button("Generate")

    generate_button.click(
        fn=generate_image,
        inputs=prompt_input,
        outputs=image_output,
    )

demo.launch()